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