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