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