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