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, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
6 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
7 VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigAutomorphismInplace,
8 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallBInplace, VecZnxCopy, VecZnxDftAddInplace,
9 VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftCopy, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
10 VecZnxDftToVecZnxBigTmpA, VecZnxFillUniform, VecZnxNegateInplace, VecZnxNormalize, VecZnxNormalizeInplace,
11 VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubABInplace,
12 VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare, ZnxView, ZnxViewMut,
13 },
14 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
15 oep::{
16 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeMatZnxImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
17 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl,
18 },
19 source::Source,
20};
21
22use crate::tfhe::{
23 blind_rotation::{
24 BlincRotationExecute, BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk,
25 BlindRotationKeyPrepared,
26 },
27 circuit_bootstrapping::{
28 CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyPrepared, CirtuitBootstrappingExecute,
29 },
30};
31
32use poulpy_core::layouts::prepared::PrepareAlloc;
33
34use poulpy_core::layouts::{
35 GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, LWECiphertext, LWEPlaintext, LWESecret,
36 prepared::{GGSWCiphertextPrepared, GLWESecretPrepared},
37};
38
39pub fn test_circuit_bootstrapping_to_exponent<B, BRA: BlindRotationAlgo>(module: &Module<B>)
40where
41 Module<B>: VecZnxFillUniform
42 + VecZnxAddNormal
43 + VecZnxNormalizeInplace<B>
44 + VecZnxDftAllocBytes
45 + VecZnxBigNormalize<B>
46 + VecZnxDftFromVecZnx<B>
47 + SvpApplyInplace<B>
48 + VecZnxDftToVecZnxBigConsume<B>
49 + VecZnxNormalizeTmpBytes
50 + VecZnxSubABInplace
51 + VecZnxAddInplace
52 + VecZnxNormalize<B>
53 + VecZnxSub
54 + VecZnxAddScalarInplace
55 + VecZnxAutomorphism
56 + VecZnxSwithcDegree
57 + VecZnxBigAllocBytes
58 + VecZnxDftToVecZnxBigTmpA<B>
59 + SvpApply<B>
60 + VecZnxBigAddInplace<B>
61 + VecZnxBigAddSmallInplace<B>
62 + VecZnxBigAlloc<B>
63 + VecZnxDftAlloc<B>
64 + VecZnxBigNormalizeTmpBytes
65 + VmpPMatAlloc<B>
66 + VmpPrepare<B>
67 + SvpPrepare<B>
68 + SvpPPolAlloc<B>
69 + VmpApplyTmpBytes
70 + VmpApply<B>
71 + VmpApplyAdd<B>
72 + SvpPPolAllocBytes
73 + VecZnxRotateInplace
74 + VecZnxBigAutomorphismInplace<B>
75 + VecZnxRshInplace
76 + VecZnxDftCopy<B>
77 + VecZnxNegateInplace
78 + VecZnxCopy
79 + VecZnxAutomorphismInplace
80 + VecZnxBigSubSmallBInplace<B>
81 + VecZnxBigAllocBytes
82 + VecZnxDftAddInplace<B>
83 + VecZnxRotate,
84 B: Backend
85 + ScratchOwnedAllocImpl<B>
86 + ScratchOwnedBorrowImpl<B>
87 + TakeVecZnxDftImpl<B>
88 + ScratchAvailableImpl<B>
89 + TakeVecZnxImpl<B>
90 + TakeScalarZnxImpl<B>
91 + TakeSvpPPolImpl<B>
92 + TakeVecZnxBigImpl<B>
93 + TakeVecZnxDftSliceImpl<B>
94 + TakeMatZnxImpl<B>
95 + TakeVecZnxSliceImpl<B>,
96 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
97 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
98 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
99{
100 let n: usize = module.n();
101 let basek: usize = 17;
102 let extension_factor: usize = 1;
103 let rank: usize = 1;
104 let sigma: f64 = 3.2;
105
106 let n_lwe: usize = 77;
107 let k_lwe_pt: usize = 4;
108 let k_lwe_ct: usize = 22;
109 let block_size: usize = 7;
110
111 let k_brk: usize = 5 * basek;
112 let rows_brk: usize = 4;
113
114 let k_trace: usize = 5 * basek;
115 let rows_trace: usize = 4;
116
117 let k_tsk: usize = 5 * basek;
118 let rows_tsk: usize = 4;
119
120 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
121
122 let mut source_xs: Source = Source::new([1u8; 32]);
123 let mut source_xa: Source = Source::new([1u8; 32]);
124 let mut source_xe: Source = Source::new([1u8; 32]);
125
126 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
127 sk_lwe.fill_binary_block(block_size, &mut source_xs);
128
129 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
130 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
131
132 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
133
134 let data: i64 = 1;
135
136 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
137 pt_lwe.encode_i64(data, k_lwe_pt + 1);
138
139 println!("pt_lwe: {}", pt_lwe);
140
141 let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
142 ct_lwe.encrypt_sk(
143 module,
144 &pt_lwe,
145 &sk_lwe,
146 &mut source_xa,
147 &mut source_xe,
148 sigma,
149 );
150
151 let now: Instant = Instant::now();
152 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
153 module,
154 basek,
155 &sk_lwe,
156 &sk_glwe,
157 k_brk,
158 rows_brk,
159 k_trace,
160 rows_trace,
161 k_tsk,
162 rows_tsk,
163 &mut source_xa,
164 &mut source_xe,
165 sigma,
166 scratch.borrow(),
167 );
168 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
169
170 let k_ggsw_res: usize = 4 * basek;
171 let rows_ggsw_res: usize = 2;
172
173 let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw_res, rows_ggsw_res, 1, rank);
174
175 let log_gap_out = 1;
176
177 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
178
179 let now: Instant = Instant::now();
180 cbt_prepared.execute_to_exponent(
181 module,
182 log_gap_out,
183 &mut res,
184 &ct_lwe,
185 k_lwe_pt,
186 extension_factor,
187 scratch.borrow(),
188 );
189 println!("CBT: {} ms", now.elapsed().as_millis());
190
191 // X^{data * 2^log_gap_out}
192 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
193 pt_ggsw.at_mut(0, 0)[0] = 1;
194 module.vec_znx_rotate_inplace(data * (1 << log_gap_out), &mut pt_ggsw.as_vec_znx_mut(), 0);
195
196 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
197
198 let k_glwe: usize = k_ggsw_res;
199
200 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe, rank);
201 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, basek);
202 pt_glwe.data.at_mut(0, 0)[0] = 1 << (basek - 2);
203
204 ct_glwe.encrypt_sk(
205 module,
206 &pt_glwe,
207 &sk_glwe_prepared,
208 &mut source_xa,
209 &mut source_xe,
210 sigma,
211 scratch.borrow(),
212 );
213
214 let res_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
215
216 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
217
218 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe);
219 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
220
221 // Parameters are set such that the first limb should be noiseless.
222 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
223 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
224 assert_eq!(pt_res.data.at(0, 0), pt_want);
225}
226
227pub fn test_circuit_bootstrapping_to_constant<B, BRA: BlindRotationAlgo>(module: &Module<B>)
228where
229 Module<B>: VecZnxFillUniform
230 + VecZnxAddNormal
231 + VecZnxNormalizeInplace<B>
232 + VecZnxDftAllocBytes
233 + VecZnxBigNormalize<B>
234 + VecZnxDftFromVecZnx<B>
235 + SvpApplyInplace<B>
236 + VecZnxDftToVecZnxBigConsume<B>
237 + VecZnxNormalizeTmpBytes
238 + VecZnxSubABInplace
239 + VecZnxAddInplace
240 + VecZnxNormalize<B>
241 + VecZnxSub
242 + VecZnxAddScalarInplace
243 + VecZnxAutomorphism
244 + VecZnxSwithcDegree
245 + VecZnxBigAllocBytes
246 + VecZnxDftToVecZnxBigTmpA<B>
247 + SvpApply<B>
248 + VecZnxBigAddInplace<B>
249 + VecZnxBigAddSmallInplace<B>
250 + VecZnxBigAlloc<B>
251 + VecZnxDftAlloc<B>
252 + VecZnxBigNormalizeTmpBytes
253 + VmpPMatAlloc<B>
254 + VmpPrepare<B>
255 + SvpPrepare<B>
256 + SvpPPolAlloc<B>
257 + VmpApplyTmpBytes
258 + VmpApply<B>
259 + VmpApplyAdd<B>
260 + SvpPPolAllocBytes
261 + VecZnxRotateInplace
262 + VecZnxBigAutomorphismInplace<B>
263 + VecZnxRshInplace
264 + VecZnxDftCopy<B>
265 + VecZnxNegateInplace
266 + VecZnxCopy
267 + VecZnxAutomorphismInplace
268 + VecZnxBigSubSmallBInplace<B>
269 + VecZnxBigAllocBytes
270 + VecZnxDftAddInplace<B>
271 + VecZnxRotate,
272 B: Backend
273 + ScratchOwnedAllocImpl<B>
274 + ScratchOwnedBorrowImpl<B>
275 + TakeVecZnxDftImpl<B>
276 + ScratchAvailableImpl<B>
277 + TakeVecZnxImpl<B>
278 + TakeScalarZnxImpl<B>
279 + TakeSvpPPolImpl<B>
280 + TakeVecZnxBigImpl<B>
281 + TakeVecZnxDftSliceImpl<B>
282 + TakeMatZnxImpl<B>
283 + TakeVecZnxSliceImpl<B>,
284 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
285 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
286 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
287{
288 let n = module.n();
289 let basek: usize = 14;
290 let extension_factor: usize = 1;
291 let rank: usize = 2;
292 let sigma: f64 = 3.2;
293
294 let n_lwe: usize = 77;
295 let k_lwe_pt: usize = 1;
296 let k_lwe_ct: usize = 13;
297 let block_size: usize = 7;
298
299 let k_brk: usize = 5 * basek;
300 let rows_brk: usize = 3;
301
302 let k_trace: usize = 5 * basek;
303 let rows_trace: usize = 4;
304
305 let k_tsk: usize = 5 * basek;
306 let rows_tsk: usize = 4;
307
308 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
309
310 let mut source_xs: Source = Source::new([1u8; 32]);
311 let mut source_xa: Source = Source::new([1u8; 32]);
312 let mut source_xe: Source = Source::new([1u8; 32]);
313
314 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
315 sk_lwe.fill_binary_block(block_size, &mut source_xs);
316
317 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
318 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
319
320 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
321
322 let data: i64 = 1;
323
324 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
325 pt_lwe.encode_i64(data, k_lwe_pt + 1);
326
327 println!("pt_lwe: {}", pt_lwe);
328
329 let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
330 ct_lwe.encrypt_sk(
331 module,
332 &pt_lwe,
333 &sk_lwe,
334 &mut source_xa,
335 &mut source_xe,
336 sigma,
337 );
338
339 let now: Instant = Instant::now();
340 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
341 module,
342 basek,
343 &sk_lwe,
344 &sk_glwe,
345 k_brk,
346 rows_brk,
347 k_trace,
348 rows_trace,
349 k_tsk,
350 rows_tsk,
351 &mut source_xa,
352 &mut source_xe,
353 sigma,
354 scratch.borrow(),
355 );
356 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
357
358 let k_ggsw_res: usize = 4 * basek;
359 let rows_ggsw_res: usize = 3;
360
361 let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw_res, rows_ggsw_res, 1, rank);
362
363 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
364
365 let now: Instant = Instant::now();
366 cbt_prepared.execute_to_constant(
367 module,
368 &mut res,
369 &ct_lwe,
370 k_lwe_pt,
371 extension_factor,
372 scratch.borrow(),
373 );
374 println!("CBT: {} ms", now.elapsed().as_millis());
375
376 // X^{data * 2^log_gap_out}
377 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
378 pt_ggsw.at_mut(0, 0)[0] = data;
379
380 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
381
382 let k_glwe: usize = k_ggsw_res;
383
384 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe, rank);
385 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, basek);
386 pt_glwe.data.at_mut(0, 0)[0] = 1 << (basek - k_lwe_pt - 1);
387
388 ct_glwe.encrypt_sk(
389 module,
390 &pt_glwe,
391 &sk_glwe_prepared,
392 &mut source_xa,
393 &mut source_xe,
394 sigma,
395 scratch.borrow(),
396 );
397
398 let res_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
399
400 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
401
402 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe);
403 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
404
405 // Parameters are set such that the first limb should be noiseless.
406 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
407 pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
408 assert_eq!(pt_res.data.at(0, 0), pt_want);
409}