poulpy/poulpy-schemes/src/tfhe/circuit_bootstrapping/tests/circuit_bootstrapping.rs
poulpy-schemes/src/tfhe/circuit_bootstrapping/tests/circuit_bootstrapping.rsRaw
1use std::time::Instant;
2
3use poulpy_hal::{
4 api::{
5 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf,
6 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
7 VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAutomorphismInplace, VecZnxBigBytesOf,
8 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftAddInplace,
9 VecZnxDftAlloc, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxFillUniform, VecZnxIdftApplyConsume,
10 VecZnxIdftApplyTmpA, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate,
11 VecZnxRotateInplace, VecZnxRotateInplaceTmpBytes, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
12 VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare, ZnAddNormal, ZnFillUniform,
13 ZnNormalizeInplace,
14 },
15 layouts::{Backend, Module, ScalarZnx, ScratchOwned, ZnxView, ZnxViewMut},
16 oep::{
17 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeMatZnxImpl, TakeScalarZnxImpl, TakeSliceImpl,
18 TakeSvpPPolImpl, TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl,
19 },
20 source::Source,
21};
22
23use crate::tfhe::{
24 blind_rotation::{
25 BlincRotationExecute, BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk,
26 BlindRotationKeyLayout, BlindRotationKeyPrepared,
27 },
28 circuit_bootstrapping::{
29 CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
30 CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
31 },
32};
33
34use poulpy_core::layouts::{
35 AutomorphismKeyLayout, Dsize, GGSWLayout, LWELayout, TensorKeyLayout, prepared::PrepareAlloc,
36};
37
38use poulpy_core::layouts::{
39 GGSW, GLWE, GLWEPlaintext, GLWESecret, LWE, LWEPlaintext, LWESecret,
40 prepared::{GGSWPrepared, GLWESecretPrepared},
41};
42
43pub fn test_circuit_bootstrapping_to_exponent<B, BRA: BlindRotationAlgo>(module: &Module<B>)
44where
45 Module<B>: VecZnxFillUniform
46 + VecZnxAddNormal
47 + VecZnxNormalizeInplace<B>
48 + VecZnxDftBytesOf
49 + VecZnxBigNormalize<B>
50 + VecZnxDftApply<B>
51 + SvpApplyDftToDftInplace<B>
52 + VecZnxIdftApplyConsume<B>
53 + VecZnxNormalizeTmpBytes
54 + VecZnxSubInplace
55 + VecZnxAddInplace
56 + VecZnxNormalize<B>
57 + VecZnxSub
58 + VecZnxAddScalarInplace
59 + VecZnxAutomorphism
60 + VecZnxSwitchRing
61 + VecZnxBigBytesOf
62 + VecZnxIdftApplyTmpA<B>
63 + SvpApplyDftToDft<B>
64 + VecZnxBigAddInplace<B>
65 + VecZnxBigAddSmallInplace<B>
66 + VecZnxBigAlloc<B>
67 + VecZnxDftAlloc<B>
68 + VecZnxBigNormalizeTmpBytes
69 + VmpPMatAlloc<B>
70 + VmpPrepare<B>
71 + SvpPrepare<B>
72 + SvpPPolAlloc<B>
73 + VmpApplyDftToDftTmpBytes
74 + VmpApplyDftToDft<B>
75 + VmpApplyDftToDftAdd<B>
76 + SvpPPolBytesOf
77 + VecZnxRotateInplace<B>
78 + VecZnxBigAutomorphismInplace<B>
79 + VecZnxRshInplace<B>
80 + VecZnxDftCopy<B>
81 + VecZnxNegateInplace
82 + VecZnxCopy
83 + VecZnxAutomorphismInplace<B>
84 + VecZnxBigSubSmallNegateInplace<B>
85 + VecZnxRotateInplaceTmpBytes
86 + VecZnxBigBytesOf
87 + VecZnxDftAddInplace<B>
88 + VecZnxRotate
89 + ZnFillUniform
90 + ZnAddNormal
91 + ZnNormalizeInplace<B>,
92 B: Backend
93 + ScratchOwnedAllocImpl<B>
94 + ScratchOwnedBorrowImpl<B>
95 + TakeVecZnxDftImpl<B>
96 + ScratchAvailableImpl<B>
97 + TakeVecZnxImpl<B>
98 + TakeScalarZnxImpl<B>
99 + TakeSvpPPolImpl<B>
100 + TakeVecZnxBigImpl<B>
101 + TakeVecZnxDftSliceImpl<B>
102 + TakeMatZnxImpl<B>
103 + TakeVecZnxSliceImpl<B>
104 + TakeSliceImpl<B>,
105 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
106 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
107 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
108{
109 let n_glwe: usize = module.n();
110 let base2k: usize = 17;
111 let extension_factor: usize = 1;
112 let rank: usize = 1;
113
114 let n_lwe: usize = 77;
115 let k_lwe_pt: usize = 4;
116 let k_lwe_ct: usize = 22;
117 let block_size: usize = 7;
118
119 let k_brk: usize = 5 * base2k;
120 let rows_brk: usize = 4;
121
122 let k_atk: usize = 5 * base2k;
123 let rows_atk: usize = 4;
124
125 let k_tsk: usize = 5 * base2k;
126 let rows_tsk: usize = 4;
127
128 let k_ggsw_res: usize = 4 * base2k;
129 let rows_ggsw_res: usize = 2;
130
131 let lwe_infos: LWELayout = LWELayout {
132 n: n_lwe.into(),
133 k: k_lwe_ct.into(),
134 base2k: base2k.into(),
135 };
136
137 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
138 layout_brk: BlindRotationKeyLayout {
139 n_glwe: n_glwe.into(),
140 n_lwe: n_lwe.into(),
141 base2k: base2k.into(),
142 k: k_brk.into(),
143 dnum: rows_brk.into(),
144 rank: rank.into(),
145 },
146 layout_atk: AutomorphismKeyLayout {
147 n: n_glwe.into(),
148 base2k: base2k.into(),
149 k: k_atk.into(),
150 dnum: rows_atk.into(),
151 rank: rank.into(),
152 dsize: Dsize(1),
153 },
154 layout_tsk: TensorKeyLayout {
155 n: n_glwe.into(),
156 base2k: base2k.into(),
157 k: k_tsk.into(),
158 dnum: rows_tsk.into(),
159 dsize: Dsize(1),
160 rank: rank.into(),
161 },
162 };
163
164 let ggsw_infos: GGSWLayout = GGSWLayout {
165 n: n_glwe.into(),
166 base2k: base2k.into(),
167 k: k_ggsw_res.into(),
168 dnum: rows_ggsw_res.into(),
169 dsize: Dsize(1),
170 rank: rank.into(),
171 };
172
173 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
174
175 let mut source_xs: Source = Source::new([1u8; 32]);
176 let mut source_xa: Source = Source::new([1u8; 32]);
177 let mut source_xe: Source = Source::new([1u8; 32]);
178
179 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
180 sk_lwe.fill_binary_block(block_size, &mut source_xs);
181
182 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
183 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
184
185 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
186
187 let data: i64 = 1;
188
189 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
190 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
191
192 println!("pt_lwe: {pt_lwe}");
193
194 let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
195 ct_lwe.encrypt_sk(module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
196
197 let now: Instant = Instant::now();
198 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
199 module,
200 &sk_lwe,
201 &sk_glwe,
202 &cbt_infos,
203 &mut source_xa,
204 &mut source_xe,
205 scratch.borrow(),
206 );
207 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
208
209 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
210
211 let log_gap_out = 1;
212
213 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
214
215 let now: Instant = Instant::now();
216 cbt_prepared.execute_to_exponent(
217 module,
218 log_gap_out,
219 &mut res,
220 &ct_lwe,
221 k_lwe_pt,
222 extension_factor,
223 scratch.borrow(),
224 );
225 println!("CBT: {} ms", now.elapsed().as_millis());
226
227 // X^{data * 2^log_gap_out}
228 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
229 pt_ggsw.at_mut(0, 0)[0] = 1;
230 module.vec_znx_rotate_inplace(
231 data * (1 << log_gap_out),
232 &mut pt_ggsw.as_vec_znx_mut(),
233 0,
234 scratch.borrow(),
235 );
236
237 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
238
239 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
240 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
241 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - 2);
242
243 ct_glwe.encrypt_sk(
244 module,
245 &pt_glwe,
246 &sk_glwe_prepared,
247 &mut source_xa,
248 &mut source_xe,
249 scratch.borrow(),
250 );
251
252 let res_prepared: GGSWPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
253
254 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
255
256 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
257 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
258
259 // Parameters are set such that the first limb should be noiseless.
260 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
261 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
262 assert_eq!(pt_res.data.at(0, 0), pt_want);
263}
264
265pub fn test_circuit_bootstrapping_to_constant<B, BRA: BlindRotationAlgo>(module: &Module<B>)
266where
267 Module<B>: VecZnxFillUniform
268 + VecZnxAddNormal
269 + VecZnxNormalizeInplace<B>
270 + VecZnxDftBytesOf
271 + VecZnxBigNormalize<B>
272 + VecZnxDftApply<B>
273 + SvpApplyDftToDftInplace<B>
274 + VecZnxIdftApplyConsume<B>
275 + VecZnxNormalizeTmpBytes
276 + VecZnxSubInplace
277 + VecZnxAddInplace
278 + VecZnxNormalize<B>
279 + VecZnxSub
280 + VecZnxAddScalarInplace
281 + VecZnxAutomorphism
282 + VecZnxSwitchRing
283 + VecZnxBigBytesOf
284 + VecZnxIdftApplyTmpA<B>
285 + SvpApplyDftToDft<B>
286 + VecZnxBigAddInplace<B>
287 + VecZnxBigAddSmallInplace<B>
288 + VecZnxBigAlloc<B>
289 + VecZnxDftAlloc<B>
290 + VecZnxBigNormalizeTmpBytes
291 + VmpPMatAlloc<B>
292 + VmpPrepare<B>
293 + SvpPrepare<B>
294 + SvpPPolAlloc<B>
295 + VmpApplyDftToDftTmpBytes
296 + VmpApplyDftToDft<B>
297 + VmpApplyDftToDftAdd<B>
298 + SvpPPolBytesOf
299 + VecZnxRotateInplace<B>
300 + VecZnxBigAutomorphismInplace<B>
301 + VecZnxRotateInplaceTmpBytes
302 + VecZnxRshInplace<B>
303 + VecZnxDftCopy<B>
304 + VecZnxNegateInplace
305 + VecZnxCopy
306 + VecZnxAutomorphismInplace<B>
307 + VecZnxBigSubSmallNegateInplace<B>
308 + VecZnxBigBytesOf
309 + VecZnxDftAddInplace<B>
310 + VecZnxRotate
311 + ZnFillUniform
312 + ZnAddNormal
313 + ZnNormalizeInplace<B>,
314 B: Backend
315 + ScratchOwnedAllocImpl<B>
316 + ScratchOwnedBorrowImpl<B>
317 + TakeVecZnxDftImpl<B>
318 + ScratchAvailableImpl<B>
319 + TakeVecZnxImpl<B>
320 + TakeScalarZnxImpl<B>
321 + TakeSvpPPolImpl<B>
322 + TakeVecZnxBigImpl<B>
323 + TakeVecZnxDftSliceImpl<B>
324 + TakeMatZnxImpl<B>
325 + TakeVecZnxSliceImpl<B>
326 + TakeSliceImpl<B>,
327 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
328 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
329 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
330{
331 let n_glwe: usize = module.n();
332 let base2k: usize = 14;
333 let extension_factor: usize = 1;
334 let rank: usize = 2;
335
336 let n_lwe: usize = 77;
337 let k_lwe_pt: usize = 1;
338 let k_lwe_ct: usize = 13;
339 let block_size: usize = 7;
340
341 let k_brk: usize = 5 * base2k;
342 let rows_brk: usize = 3;
343
344 let k_atk: usize = 5 * base2k;
345 let rows_atk: usize = 4;
346
347 let k_tsk: usize = 5 * base2k;
348 let rows_tsk: usize = 4;
349
350 let k_ggsw_res: usize = 4 * base2k;
351 let rows_ggsw_res: usize = 3;
352
353 let lwe_infos: LWELayout = LWELayout {
354 n: n_lwe.into(),
355 k: k_lwe_ct.into(),
356 base2k: base2k.into(),
357 };
358
359 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
360 layout_brk: BlindRotationKeyLayout {
361 n_glwe: n_glwe.into(),
362 n_lwe: n_lwe.into(),
363 base2k: base2k.into(),
364 k: k_brk.into(),
365 dnum: rows_brk.into(),
366 rank: rank.into(),
367 },
368 layout_atk: AutomorphismKeyLayout {
369 n: n_glwe.into(),
370 base2k: base2k.into(),
371 k: k_atk.into(),
372 dnum: rows_atk.into(),
373 rank: rank.into(),
374 dsize: Dsize(1),
375 },
376 layout_tsk: TensorKeyLayout {
377 n: n_glwe.into(),
378 base2k: base2k.into(),
379 k: k_tsk.into(),
380 dnum: rows_tsk.into(),
381 dsize: Dsize(1),
382 rank: rank.into(),
383 },
384 };
385
386 let ggsw_infos: GGSWLayout = GGSWLayout {
387 n: n_glwe.into(),
388 base2k: base2k.into(),
389 k: k_ggsw_res.into(),
390 dnum: rows_ggsw_res.into(),
391 dsize: Dsize(1),
392 rank: rank.into(),
393 };
394
395 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
396
397 let mut source_xs: Source = Source::new([1u8; 32]);
398 let mut source_xa: Source = Source::new([1u8; 32]);
399 let mut source_xe: Source = Source::new([1u8; 32]);
400
401 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
402 sk_lwe.fill_binary_block(block_size, &mut source_xs);
403
404 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
405 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
406
407 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
408
409 let data: i64 = 1;
410
411 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
412 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
413
414 println!("pt_lwe: {pt_lwe}");
415
416 let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
417 ct_lwe.encrypt_sk(module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
418
419 let now: Instant = Instant::now();
420 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
421 module,
422 &sk_lwe,
423 &sk_glwe,
424 &cbt_infos,
425 &mut source_xa,
426 &mut source_xe,
427 scratch.borrow(),
428 );
429 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
430
431 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
432
433 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
434
435 let now: Instant = Instant::now();
436 cbt_prepared.execute_to_constant(
437 module,
438 &mut res,
439 &ct_lwe,
440 k_lwe_pt,
441 extension_factor,
442 scratch.borrow(),
443 );
444 println!("CBT: {} ms", now.elapsed().as_millis());
445
446 // X^{data * 2^log_gap_out}
447 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
448 pt_ggsw.at_mut(0, 0)[0] = data;
449
450 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
451
452 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
453 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
454 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - k_lwe_pt - 1);
455
456 ct_glwe.encrypt_sk(
457 module,
458 &pt_glwe,
459 &sk_glwe_prepared,
460 &mut source_xa,
461 &mut source_xe,
462 scratch.borrow(),
463 );
464
465 let res_prepared: GGSWPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
466
467 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
468
469 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
470 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
471
472 // Parameters are set such that the first limb should be noiseless.
473 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
474 pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
475 assert_eq!(pt_res.data.at(0, 0), pt_want);
476}