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::{ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxRotateInplace},
5 layouts::{Backend, ScalarZnx, Scratch, ScratchOwned, ZnxView, ZnxViewMut},
6 source::Source,
7};
8
9use crate::tfhe::{
10 blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, BlindRotationKeyLayout},
11 circuit_bootstrapping::{
12 CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyLayout,
13 CircuitBootstrappingKeyPrepared, CircuitBootstrappingKeyPreparedFactory, CirtuitBootstrappingExecute,
14 },
15};
16
17use poulpy_core::{
18 GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWEExternalProduct, LWEEncryptSk, ScratchTakeCore,
19 layouts::{
20 Dsize, GGLWEToGGSWKeyLayout, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout, GLWESecretPreparedFactory,
21 LWELayout,
22 },
23};
24
25use poulpy_core::layouts::{
26 GGSW, GLWE, GLWEPlaintext, GLWESecret, LWE, LWEPlaintext, LWESecret,
27 prepared::{GGSWPrepared, GLWESecretPrepared},
28};
29
30pub fn test_circuit_bootstrapping_to_exponent<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
31where
32 M: ModuleN
33 + GLWESecretPreparedFactory<BE>
34 + GLWEExternalProduct<BE>
35 + GLWEDecrypt<BE>
36 + LWEEncryptSk<BE>
37 + CircuitBootstrappingKeyEncryptSk<BRA, BE>
38 + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
39 + CirtuitBootstrappingExecute<BRA, BE>
40 + GGSWPreparedFactory<BE>
41 + GGSWNoise<BE>
42 + GLWEEncryptSk<BE>
43 + VecZnxRotateInplace<BE>,
44 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
45 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
46 Scratch<BE>: ScratchTakeCore<BE>,
47{
48 let n_glwe: usize = module.n();
49 let base2k: usize = 17;
50 let extension_factor: usize = 1;
51 let rank: usize = 1;
52
53 let n_lwe: usize = 77;
54 let k_lwe_pt: usize = 4;
55 let k_lwe_ct: usize = 22;
56 let block_size: usize = 7;
57
58 let k_brk: usize = 5 * base2k;
59 let rows_brk: usize = 4;
60
61 let k_atk: usize = 5 * base2k;
62 let rows_atk: usize = 4;
63
64 let k_tsk: usize = 5 * base2k;
65 let rows_tsk: usize = 4;
66
67 let k_ggsw_res: usize = 4 * base2k;
68 let rows_ggsw_res: usize = 2;
69
70 let lwe_infos: LWELayout = LWELayout {
71 n: n_lwe.into(),
72 k: k_lwe_ct.into(),
73 base2k: base2k.into(),
74 };
75
76 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
77 layout_brk: BlindRotationKeyLayout {
78 n_glwe: n_glwe.into(),
79 n_lwe: n_lwe.into(),
80 base2k: base2k.into(),
81 k: k_brk.into(),
82 dnum: rows_brk.into(),
83 rank: rank.into(),
84 },
85 layout_atk: GLWEAutomorphismKeyLayout {
86 n: n_glwe.into(),
87 base2k: base2k.into(),
88 k: k_atk.into(),
89 dnum: rows_atk.into(),
90 rank: rank.into(),
91 dsize: Dsize(1),
92 },
93 layout_tsk: GGLWEToGGSWKeyLayout {
94 n: n_glwe.into(),
95 base2k: base2k.into(),
96 k: k_tsk.into(),
97 dnum: rows_tsk.into(),
98 dsize: Dsize(1),
99 rank: rank.into(),
100 },
101 };
102
103 let ggsw_infos: GGSWLayout = GGSWLayout {
104 n: n_glwe.into(),
105 base2k: base2k.into(),
106 k: k_ggsw_res.into(),
107 dnum: rows_ggsw_res.into(),
108 dsize: Dsize(1),
109 rank: rank.into(),
110 };
111
112 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
113
114 let mut source_xs: Source = Source::new([1u8; 32]);
115 let mut source_xa: Source = Source::new([1u8; 32]);
116 let mut source_xe: Source = Source::new([1u8; 32]);
117
118 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
119 sk_lwe.fill_binary_block(block_size, &mut source_xs);
120
121 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
122 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
123
124 let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
125 sk_glwe_prepared.prepare(module, &sk_glwe);
126
127 let data: i64 = 1;
128
129 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
130 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
131
132 println!("pt_lwe: {pt_lwe}");
133
134 let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
135 ct_lwe.encrypt_sk(
136 module,
137 &pt_lwe,
138 &sk_lwe,
139 &mut source_xa,
140 &mut source_xe,
141 scratch.borrow(),
142 );
143
144 let now: Instant = Instant::now();
145 let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
146 println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
147
148 let now: Instant = Instant::now();
149 cbt_key.encrypt_sk(
150 module,
151 &sk_lwe,
152 &sk_glwe,
153 &mut source_xa,
154 &mut source_xe,
155 scratch.borrow(),
156 );
157 println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
158
159 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
160
161 let log_gap_out = 1;
162
163 let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
164 CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
165 cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
166
167 let now: Instant = Instant::now();
168 cbt_prepared.execute_to_exponent(
169 module,
170 log_gap_out,
171 &mut res,
172 &ct_lwe,
173 k_lwe_pt,
174 extension_factor,
175 scratch.borrow(),
176 );
177 println!("CBT: {} ms", now.elapsed().as_millis());
178
179 // X^{data * 2^log_gap_out}
180 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
181 pt_ggsw.at_mut(0, 0)[0] = 1;
182 module.vec_znx_rotate_inplace(
183 data * (1 << log_gap_out),
184 &mut pt_ggsw.as_vec_znx_mut(),
185 0,
186 scratch.borrow(),
187 );
188
189 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
190
191 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
192 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
193 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - 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 mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
205 res_prepared.prepare(module, &res, scratch.borrow());
206
207 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
208
209 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
210 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
211
212 // Parameters are set such that the first limb should be noiseless.
213 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
214 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
215 assert_eq!(pt_res.data.at(0, 0), pt_want);
216}
217
218pub fn test_circuit_bootstrapping_to_constant<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
219where
220 M: ModuleN
221 + GLWESecretPreparedFactory<BE>
222 + GLWEExternalProduct<BE>
223 + GLWEDecrypt<BE>
224 + LWEEncryptSk<BE>
225 + CircuitBootstrappingKeyEncryptSk<BRA, BE>
226 + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
227 + CirtuitBootstrappingExecute<BRA, BE>
228 + GGSWPreparedFactory<BE>
229 + GGSWNoise<BE>
230 + GLWEEncryptSk<BE>
231 + VecZnxRotateInplace<BE>,
232 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
233 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
234 Scratch<BE>: ScratchTakeCore<BE>,
235{
236 let n_glwe: usize = module.n();
237 let base2k: usize = 14;
238 let extension_factor: usize = 1;
239 let rank: usize = 1;
240
241 let n_lwe: usize = 77;
242 let k_lwe_pt: usize = 1;
243 let k_lwe_ct: usize = 13;
244 let block_size: usize = 7;
245
246 let k_brk: usize = 5 * base2k;
247 let rows_brk: usize = 3;
248
249 let k_atk: usize = 5 * base2k;
250 let rows_atk: usize = 4;
251
252 let k_tsk: usize = 5 * base2k;
253 let rows_tsk: usize = 4;
254
255 let k_ggsw_res: usize = 4 * base2k;
256 let rows_ggsw_res: usize = 3;
257
258 let lwe_infos: LWELayout = LWELayout {
259 n: n_lwe.into(),
260 k: k_lwe_ct.into(),
261 base2k: base2k.into(),
262 };
263
264 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
265 layout_brk: BlindRotationKeyLayout {
266 n_glwe: n_glwe.into(),
267 n_lwe: n_lwe.into(),
268 base2k: base2k.into(),
269 k: k_brk.into(),
270 dnum: rows_brk.into(),
271 rank: rank.into(),
272 },
273 layout_atk: GLWEAutomorphismKeyLayout {
274 n: n_glwe.into(),
275 base2k: base2k.into(),
276 k: k_atk.into(),
277 dnum: rows_atk.into(),
278 rank: rank.into(),
279 dsize: Dsize(1),
280 },
281 layout_tsk: GGLWEToGGSWKeyLayout {
282 n: n_glwe.into(),
283 base2k: base2k.into(),
284 k: k_tsk.into(),
285 dnum: rows_tsk.into(),
286 dsize: Dsize(1),
287 rank: rank.into(),
288 },
289 };
290
291 let ggsw_infos: GGSWLayout = GGSWLayout {
292 n: n_glwe.into(),
293 base2k: base2k.into(),
294 k: k_ggsw_res.into(),
295 dnum: rows_ggsw_res.into(),
296 dsize: Dsize(1),
297 rank: rank.into(),
298 };
299
300 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
301
302 let mut source_xs: Source = Source::new([1u8; 32]);
303 let mut source_xa: Source = Source::new([1u8; 32]);
304 let mut source_xe: Source = Source::new([1u8; 32]);
305
306 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
307 sk_lwe.fill_binary_block(block_size, &mut source_xs);
308
309 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
310 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
311
312 let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
313 sk_glwe_prepared.prepare(module, &sk_glwe);
314
315 let data: i64 = 1;
316
317 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k.into(), k_lwe_pt.into());
318 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
319
320 println!("pt_lwe: {pt_lwe}");
321
322 let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
323 ct_lwe.encrypt_sk(
324 module,
325 &pt_lwe,
326 &sk_lwe,
327 &mut source_xa,
328 &mut source_xe,
329 scratch.borrow(),
330 );
331
332 let now: Instant = Instant::now();
333 let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
334 println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
335
336 let now: Instant = Instant::now();
337 cbt_key.encrypt_sk(
338 module,
339 &sk_lwe,
340 &sk_glwe,
341 &mut source_xa,
342 &mut source_xe,
343 scratch.borrow(),
344 );
345 println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
346
347 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
348
349 let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
350 CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
351 cbt_prepared.prepare(module, &cbt_key, 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_glwe, 1);
366 pt_ggsw.at_mut(0, 0)[0] = data;
367
368 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
369
370 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
371 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
372 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - k_lwe_pt - 1);
373
374 ct_glwe.encrypt_sk(
375 module,
376 &pt_glwe,
377 &sk_glwe_prepared,
378 &mut source_xa,
379 &mut source_xe,
380 scratch.borrow(),
381 );
382
383 let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
384 res_prepared.prepare(module, &res, 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_from_infos(&ggsw_infos);
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}