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