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