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, SvpPPolAllocBytes,
6 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
7 VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigAutomorphismInplace,
8 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftAddInplace,
9 VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftApply, 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, TakeSvpPPolImpl,
18 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 Digits, GGLWEAutomorphismKeyLayout, GGLWETensorKeyLayout, GGSWCiphertextLayout, LWECiphertextLayout, prepared::PrepareAlloc,
36};
37
38use poulpy_core::layouts::{
39 GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, LWECiphertext, LWEPlaintext, LWESecret,
40 prepared::{GGSWCiphertextPrepared, 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 + VecZnxDftAllocBytes
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 + VecZnxBigAllocBytes
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 + SvpPPolAllocBytes
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 + VecZnxBigAllocBytes
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 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
105 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
106 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
107{
108 let n_glwe: usize = module.n();
109 let base2k: usize = 17;
110 let extension_factor: usize = 1;
111 let rank: usize = 1;
112
113 let n_lwe: usize = 77;
114 let k_lwe_pt: usize = 4;
115 let k_lwe_ct: usize = 22;
116 let block_size: usize = 7;
117
118 let k_brk: usize = 5 * base2k;
119 let rows_brk: usize = 4;
120
121 let k_atk: usize = 5 * base2k;
122 let rows_atk: usize = 4;
123
124 let k_tsk: usize = 5 * base2k;
125 let rows_tsk: usize = 4;
126
127 let k_ggsw_res: usize = 4 * base2k;
128 let rows_ggsw_res: usize = 2;
129
130 let lwe_infos: LWECiphertextLayout = LWECiphertextLayout {
131 n: n_lwe.into(),
132 k: k_lwe_ct.into(),
133 base2k: base2k.into(),
134 };
135
136 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
137 layout_brk: BlindRotationKeyLayout {
138 n_glwe: n_glwe.into(),
139 n_lwe: n_lwe.into(),
140 base2k: base2k.into(),
141 k: k_brk.into(),
142 rows: rows_brk.into(),
143 rank: rank.into(),
144 },
145 layout_atk: GGLWEAutomorphismKeyLayout {
146 n: n_glwe.into(),
147 base2k: base2k.into(),
148 k: k_atk.into(),
149 rows: rows_atk.into(),
150 rank: rank.into(),
151 digits: Digits(1),
152 },
153 layout_tsk: GGLWETensorKeyLayout {
154 n: n_glwe.into(),
155 base2k: base2k.into(),
156 k: k_tsk.into(),
157 rows: rows_tsk.into(),
158 digits: Digits(1),
159 rank: rank.into(),
160 },
161 };
162
163 let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
164 n: n_glwe.into(),
165 base2k: base2k.into(),
166 k: k_ggsw_res.into(),
167 rows: rows_ggsw_res.into(),
168 digits: Digits(1),
169 rank: rank.into(),
170 };
171
172 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
173
174 let mut source_xs: Source = Source::new([1u8; 32]);
175 let mut source_xa: Source = Source::new([1u8; 32]);
176 let mut source_xe: Source = Source::new([1u8; 32]);
177
178 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
179 sk_lwe.fill_binary_block(block_size, &mut source_xs);
180
181 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n_glwe.into(), rank.into());
182 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
183
184 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
185
186 let data: i64 = 1;
187
188 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc_with(base2k.into(), k_lwe_pt.into());
189 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
190
191 println!("pt_lwe: {pt_lwe}");
192
193 let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&lwe_infos);
194 ct_lwe.encrypt_sk(module, &pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe);
195
196 let now: Instant = Instant::now();
197 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
198 module,
199 &sk_lwe,
200 &sk_glwe,
201 &cbt_infos,
202 &mut source_xa,
203 &mut source_xe,
204 scratch.borrow(),
205 );
206 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
207
208 let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_infos);
209
210 let log_gap_out = 1;
211
212 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
213
214 let now: Instant = Instant::now();
215 cbt_prepared.execute_to_exponent(
216 module,
217 log_gap_out,
218 &mut res,
219 &ct_lwe,
220 k_lwe_pt,
221 extension_factor,
222 scratch.borrow(),
223 );
224 println!("CBT: {} ms", now.elapsed().as_millis());
225
226 // X^{data * 2^log_gap_out}
227 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
228 pt_ggsw.at_mut(0, 0)[0] = 1;
229 module.vec_znx_rotate_inplace(
230 data * (1 << log_gap_out),
231 &mut pt_ggsw.as_vec_znx_mut(),
232 0,
233 scratch.borrow(),
234 );
235
236 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
237
238 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&ggsw_infos);
239 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&ggsw_infos);
240 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - 2);
241
242 ct_glwe.encrypt_sk(
243 module,
244 &pt_glwe,
245 &sk_glwe_prepared,
246 &mut source_xa,
247 &mut source_xe,
248 scratch.borrow(),
249 );
250
251 let res_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
252
253 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
254
255 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&ggsw_infos);
256 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
257
258 // Parameters are set such that the first limb should be noiseless.
259 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
260 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
261 assert_eq!(pt_res.data.at(0, 0), pt_want);
262}
263
264pub fn test_circuit_bootstrapping_to_constant<B, BRA: BlindRotationAlgo>(module: &Module<B>)
265where
266 Module<B>: VecZnxFillUniform
267 + VecZnxAddNormal
268 + VecZnxNormalizeInplace<B>
269 + VecZnxDftAllocBytes
270 + VecZnxBigNormalize<B>
271 + VecZnxDftApply<B>
272 + SvpApplyDftToDftInplace<B>
273 + VecZnxIdftApplyConsume<B>
274 + VecZnxNormalizeTmpBytes
275 + VecZnxSubInplace
276 + VecZnxAddInplace
277 + VecZnxNormalize<B>
278 + VecZnxSub
279 + VecZnxAddScalarInplace
280 + VecZnxAutomorphism
281 + VecZnxSwitchRing
282 + VecZnxBigAllocBytes
283 + VecZnxIdftApplyTmpA<B>
284 + SvpApplyDftToDft<B>
285 + VecZnxBigAddInplace<B>
286 + VecZnxBigAddSmallInplace<B>
287 + VecZnxBigAlloc<B>
288 + VecZnxDftAlloc<B>
289 + VecZnxBigNormalizeTmpBytes
290 + VmpPMatAlloc<B>
291 + VmpPrepare<B>
292 + SvpPrepare<B>
293 + SvpPPolAlloc<B>
294 + VmpApplyDftToDftTmpBytes
295 + VmpApplyDftToDft<B>
296 + VmpApplyDftToDftAdd<B>
297 + SvpPPolAllocBytes
298 + VecZnxRotateInplace<B>
299 + VecZnxBigAutomorphismInplace<B>
300 + VecZnxRotateInplaceTmpBytes
301 + VecZnxRshInplace<B>
302 + VecZnxDftCopy<B>
303 + VecZnxNegateInplace
304 + VecZnxCopy
305 + VecZnxAutomorphismInplace<B>
306 + VecZnxBigSubSmallNegateInplace<B>
307 + VecZnxBigAllocBytes
308 + VecZnxDftAddInplace<B>
309 + VecZnxRotate
310 + ZnFillUniform
311 + ZnAddNormal
312 + ZnNormalizeInplace<B>,
313 B: Backend
314 + ScratchOwnedAllocImpl<B>
315 + ScratchOwnedBorrowImpl<B>
316 + TakeVecZnxDftImpl<B>
317 + ScratchAvailableImpl<B>
318 + TakeVecZnxImpl<B>
319 + TakeScalarZnxImpl<B>
320 + TakeSvpPPolImpl<B>
321 + TakeVecZnxBigImpl<B>
322 + TakeVecZnxDftSliceImpl<B>
323 + TakeMatZnxImpl<B>
324 + TakeVecZnxSliceImpl<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: LWECiphertextLayout = LWECiphertextLayout {
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 rows: rows_brk.into(),
364 rank: rank.into(),
365 },
366 layout_atk: GGLWEAutomorphismKeyLayout {
367 n: n_glwe.into(),
368 base2k: base2k.into(),
369 k: k_atk.into(),
370 rows: rows_atk.into(),
371 rank: rank.into(),
372 digits: Digits(1),
373 },
374 layout_tsk: GGLWETensorKeyLayout {
375 n: n_glwe.into(),
376 base2k: base2k.into(),
377 k: k_tsk.into(),
378 rows: rows_tsk.into(),
379 digits: Digits(1),
380 rank: rank.into(),
381 },
382 };
383
384 let ggsw_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
385 n: n_glwe.into(),
386 base2k: base2k.into(),
387 k: k_ggsw_res.into(),
388 rows: rows_ggsw_res.into(),
389 digits: Digits(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_with(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_with(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: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(&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: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&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: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&ggsw_infos);
451 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&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: GGSWCiphertextPrepared<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(&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}