poulpy/poulpy-schemes/benches/circuit_bootstrapping.rs
poulpy-schemes/benches/circuit_bootstrapping.rsRaw
1use std::hint::black_box;
2
3use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
4use poulpy_backend::{FFT64Avx, FFT64Ref, FFT64Spqlios};
5use poulpy_core::layouts::{
6 AutomorphismKeyLayout, Dsize, GGSW, GGSWLayout, GLWESecret, LWE, LWELayout, LWESecret, TensorKeyLayout,
7 prepared::PrepareAlloc,
8};
9use poulpy_hal::{
10 api::{
11 ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc,
12 SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism,
13 VecZnxAutomorphismInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAutomorphismInplace,
14 VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy,
15 VecZnxDftAddInplace, VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxFillUniform,
16 VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace,
17 VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRotateInplaceTmpBytes, VecZnxRshInplace, VecZnxSub,
18 VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc,
19 VmpPrepare, ZnAddNormal, ZnFillUniform, ZnNormalizeInplace,
20 },
21 layouts::{Backend, Module, ScratchOwned},
22 oep::{
23 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeMatZnxImpl, TakeScalarZnxImpl, TakeSliceImpl,
24 TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl,
25 },
26 source::Source,
27};
28use poulpy_schemes::tfhe::{
29 blind_rotation::{
30 BlincRotationExecute, BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk,
31 BlindRotationKeyInfos, BlindRotationKeyLayout, BlindRotationKeyPrepared, CGGI,
32 },
33 circuit_bootstrapping::{
34 CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
35 CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
36 },
37};
38
39pub fn benc_circuit_bootstrapping<B: Backend, BRA: BlindRotationAlgo>(c: &mut Criterion, label: &str)
40where
41 Module<B>: ModuleNew<B>
42 + VecZnxFillUniform
43 + VecZnxAddNormal
44 + VecZnxNormalizeInplace<B>
45 + VecZnxDftBytesOf
46 + VecZnxBigNormalize<B>
47 + VecZnxDftApply<B>
48 + SvpApplyDftToDftInplace<B>
49 + VecZnxIdftApplyConsume<B>
50 + VecZnxNormalizeTmpBytes
51 + VecZnxSubInplace
52 + VecZnxAddInplace
53 + VecZnxNormalize<B>
54 + VecZnxSub
55 + VecZnxAddScalarInplace
56 + VecZnxAutomorphism
57 + VecZnxSwitchRing
58 + VecZnxBigBytesOf
59 + VecZnxIdftApplyTmpA<B>
60 + SvpApplyDftToDft<B>
61 + VecZnxBigAddInplace<B>
62 + VecZnxBigAddSmallInplace<B>
63 + VecZnxBigAlloc<B>
64 + VecZnxDftAlloc<B>
65 + VecZnxBigNormalizeTmpBytes
66 + VmpPMatAlloc<B>
67 + VmpPrepare<B>
68 + SvpPrepare<B>
69 + SvpPPolAlloc<B>
70 + VmpApplyDftToDftTmpBytes
71 + VmpApplyDftToDft<B>
72 + VmpApplyDftToDftAdd<B>
73 + SvpPPolBytesOf
74 + VecZnxRotateInplace<B>
75 + VecZnxBigAutomorphismInplace<B>
76 + VecZnxRshInplace<B>
77 + VecZnxDftCopy<B>
78 + VecZnxNegateInplace
79 + VecZnxCopy
80 + VecZnxAutomorphismInplace<B>
81 + VecZnxBigSubSmallNegateInplace<B>
82 + VecZnxRotateInplaceTmpBytes
83 + VecZnxBigBytesOf
84 + VecZnxDftAddInplace<B>
85 + VecZnxRotate
86 + ZnFillUniform
87 + ZnAddNormal
88 + ZnNormalizeInplace<B>,
89 ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
90 B: Backend
91 + ScratchOwnedAllocImpl<B>
92 + ScratchOwnedBorrowImpl<B>
93 + TakeVecZnxDftImpl<B>
94 + ScratchAvailableImpl<B>
95 + TakeVecZnxImpl<B>
96 + TakeScalarZnxImpl<B>
97 + TakeSvpPPolImpl<B>
98 + TakeVecZnxBigImpl<B>
99 + TakeVecZnxDftSliceImpl<B>
100 + TakeMatZnxImpl<B>
101 + TakeVecZnxSliceImpl<B>
102 + TakeSliceImpl<B>,
103 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
104 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
105 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
106{
107 let group_name: String = format!("circuit_bootstrapping::{label}");
108
109 let mut group = c.benchmark_group(group_name);
110
111 struct Params {
112 name: String,
113 extension_factor: usize,
114 k_pt: usize,
115 block_size: usize,
116 lwe_infos: LWELayout,
117 ggsw_infos: GGSWLayout,
118 cbt_infos: CircuitBootstrappingKeyLayout,
119 }
120
121 fn runner<B: Backend, BRA: BlindRotationAlgo>(params: &Params) -> impl FnMut()
122 where
123 Module<B>: ModuleNew<B>
124 + VecZnxFillUniform
125 + VecZnxAddNormal
126 + VecZnxNormalizeInplace<B>
127 + VecZnxDftBytesOf
128 + VecZnxBigNormalize<B>
129 + VecZnxDftApply<B>
130 + SvpApplyDftToDftInplace<B>
131 + VecZnxIdftApplyConsume<B>
132 + VecZnxNormalizeTmpBytes
133 + VecZnxSubInplace
134 + VecZnxAddInplace
135 + VecZnxNormalize<B>
136 + VecZnxSub
137 + VecZnxAddScalarInplace
138 + VecZnxAutomorphism
139 + VecZnxSwitchRing
140 + VecZnxBigBytesOf
141 + VecZnxIdftApplyTmpA<B>
142 + SvpApplyDftToDft<B>
143 + VecZnxBigAddInplace<B>
144 + VecZnxBigAddSmallInplace<B>
145 + VecZnxBigAlloc<B>
146 + VecZnxDftAlloc<B>
147 + VecZnxBigNormalizeTmpBytes
148 + VmpPMatAlloc<B>
149 + VmpPrepare<B>
150 + SvpPrepare<B>
151 + SvpPPolAlloc<B>
152 + VmpApplyDftToDftTmpBytes
153 + VmpApplyDftToDft<B>
154 + VmpApplyDftToDftAdd<B>
155 + SvpPPolBytesOf
156 + VecZnxRotateInplace<B>
157 + VecZnxBigAutomorphismInplace<B>
158 + VecZnxRshInplace<B>
159 + VecZnxDftCopy<B>
160 + VecZnxNegateInplace
161 + VecZnxCopy
162 + VecZnxAutomorphismInplace<B>
163 + VecZnxBigSubSmallNegateInplace<B>
164 + VecZnxRotateInplaceTmpBytes
165 + VecZnxBigBytesOf
166 + VecZnxDftAddInplace<B>
167 + VecZnxRotate
168 + ZnFillUniform
169 + ZnAddNormal
170 + ZnNormalizeInplace<B>,
171 B: Backend
172 + ScratchOwnedAllocImpl<B>
173 + ScratchOwnedBorrowImpl<B>
174 + TakeVecZnxDftImpl<B>
175 + ScratchAvailableImpl<B>
176 + TakeVecZnxImpl<B>
177 + TakeScalarZnxImpl<B>
178 + TakeSvpPPolImpl<B>
179 + TakeVecZnxBigImpl<B>
180 + TakeVecZnxDftSliceImpl<B>
181 + TakeMatZnxImpl<B>
182 + TakeVecZnxSliceImpl<B>
183 + TakeSliceImpl<B>,
184 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
185 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
186 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
187 {
188 // Scratch space (4MB)
189 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 22);
190
191 let n_glwe: poulpy_core::layouts::RingDegree = params.cbt_infos.layout_brk.n_glwe();
192 let n_lwe: poulpy_core::layouts::RingDegree = params.cbt_infos.layout_brk.n_lwe();
193 let rank: poulpy_core::layouts::Rank = params.cbt_infos.layout_brk.rank;
194
195 let module: Module<B> = Module::<B>::new(n_glwe.as_u32() as u64);
196
197 let mut source_xs: Source = Source::new([1u8; 32]);
198 let mut source_xa: Source = Source::new([1u8; 32]);
199 let mut source_xe: Source = Source::new([1u8; 32]);
200
201 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
202 sk_lwe.fill_binary_block(params.block_size, &mut source_xs);
203 sk_lwe.fill_zero();
204
205 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe, rank);
206 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
207
208 let ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(¶ms.lwe_infos);
209
210 // Circuit bootstrapping evaluation key
211 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
212 &module,
213 &sk_lwe,
214 &sk_glwe,
215 ¶ms.cbt_infos,
216 &mut source_xa,
217 &mut source_xe,
218 scratch.borrow(),
219 );
220
221 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(¶ms.ggsw_infos);
222 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(&module, scratch.borrow());
223
224 move || {
225 cbt_prepared.execute_to_constant(
226 &module,
227 &mut res,
228 &ct_lwe,
229 params.k_pt,
230 params.extension_factor,
231 scratch.borrow(),
232 );
233 black_box(());
234 }
235 }
236
237 for params in [Params {
238 name: String::from("1-bit"),
239 extension_factor: 1,
240 k_pt: 1,
241 lwe_infos: LWELayout {
242 n: 574_u32.into(),
243 k: 13_u32.into(),
244 base2k: 13_u32.into(),
245 },
246 block_size: 7,
247 ggsw_infos: GGSWLayout {
248 n: 1024_u32.into(),
249 base2k: 13_u32.into(),
250 k: 26_u32.into(),
251 dnum: 2_u32.into(),
252 dsize: 1_u32.into(),
253 rank: 2_u32.into(),
254 },
255 cbt_infos: CircuitBootstrappingKeyLayout {
256 layout_brk: BlindRotationKeyLayout {
257 n_glwe: 1024_u32.into(),
258 n_lwe: 574_u32.into(),
259 base2k: 13_u32.into(),
260 k: 52_u32.into(),
261 dnum: 3_u32.into(),
262 rank: 2_u32.into(),
263 },
264 layout_atk: AutomorphismKeyLayout {
265 n: 1024_u32.into(),
266 base2k: 13_u32.into(),
267 k: 52_u32.into(),
268 dnum: 3_u32.into(),
269 dsize: Dsize(1),
270 rank: 2_u32.into(),
271 },
272 layout_tsk: TensorKeyLayout {
273 n: 1024_u32.into(),
274 base2k: 13_u32.into(),
275 k: 52_u32.into(),
276 dnum: 3_u32.into(),
277 dsize: Dsize(1),
278 rank: 2_u32.into(),
279 },
280 },
281 }] {
282 let id: BenchmarkId = BenchmarkId::from_parameter(params.name.clone());
283 let mut runner = runner::<B, BRA>(¶ms);
284 group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
285 }
286
287 group.finish();
288}
289
290fn bench_circuit_bootstrapping_cpu_ref_fft64(c: &mut Criterion) {
291 benc_circuit_bootstrapping::<FFT64Ref, CGGI>(c, "fft64_ref");
292}
293
294fn bench_circuit_bootstrapping_cpu_avx_fft64(c: &mut Criterion) {
295 benc_circuit_bootstrapping::<FFT64Avx, CGGI>(c, "fft64_avx");
296}
297
298fn bench_circuit_bootstrapping_cpu_spqlios_fft64(c: &mut Criterion) {
299 benc_circuit_bootstrapping::<FFT64Spqlios, CGGI>(c, "fft64_spqlios");
300}
301
302criterion_group!(
303 benches,
304 bench_circuit_bootstrapping_cpu_ref_fft64,
305 bench_circuit_bootstrapping_cpu_avx_fft64,
306 bench_circuit_bootstrapping_cpu_spqlios_fft64,
307);
308
309criterion_main!(benches);