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, GGSWInfos, GGSWLayout, GGSWPreparedFactory, GLWEAutomorphismKeyLayout, GLWEInfos,
21 GLWESecretPreparedFactory, 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 for row in 0..res.dnum().as_usize() {
190 for col in 0..res.rank().as_usize() + 1 {
191 println!(
192 "row:{row} col:{col} -> {}",
193 res.noise(
194 module,
195 row,
196 col,
197 &pt_ggsw,
198 &sk_glwe_prepared,
199 scratch.borrow()
200 )
201 .std()
202 .log2()
203 )
204 }
205 }
206 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
207 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
208 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k - 2);
209
210 ct_glwe.encrypt_sk(
211 module,
212 &pt_glwe,
213 &sk_glwe_prepared,
214 &mut source_xa,
215 &mut source_xe,
216 scratch.borrow(),
217 );
218
219 let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
220 res_prepared.prepare(module, &res, scratch.borrow());
221
222 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
223
224 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
225 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
226
227 // Parameters are set such that the first limb should be noiseless.
228 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
229 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
230 assert_eq!(pt_res.data.at(0, 0), pt_want);
231}
232
233pub fn test_circuit_bootstrapping_to_constant<BE: Backend, M, BRA: BlindRotationAlgo>(module: &M)
234where
235 M: ModuleN
236 + GLWESecretPreparedFactory<BE>
237 + GLWEExternalProduct<BE>
238 + GLWEDecrypt<BE>
239 + LWEEncryptSk<BE>
240 + CircuitBootstrappingKeyEncryptSk<BRA, BE>
241 + CircuitBootstrappingKeyPreparedFactory<BRA, BE>
242 + CirtuitBootstrappingExecute<BRA, BE>
243 + GGSWPreparedFactory<BE>
244 + GGSWNoise<BE>
245 + GLWEEncryptSk<BE>
246 + VecZnxRotateInplace<BE>,
247 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>, // TODO find a way to remove this bound or move it to CBT KEY
248 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
249 Scratch<BE>: ScratchTakeCore<BE>,
250{
251 let n_glwe: usize = module.n();
252 let base2k_res: usize = 14;
253 let base2k_brk: usize = 13;
254 let base2k_tsk: usize = 12;
255 let base2k_atk: usize = 11;
256 let extension_factor: usize = 1;
257 let rank: usize = 1;
258
259 let n_lwe: usize = 77;
260 let k_lwe_pt: usize = 1;
261 let k_lwe_ct: usize = 13;
262 let block_size: usize = 7;
263
264 let k_ggsw_res: usize = 4 * base2k_res;
265 let rows_ggsw_res: usize = 3;
266
267 let k_brk: usize = k_ggsw_res + base2k_brk;
268 let rows_brk: usize = 3;
269
270 let k_atk: usize = k_ggsw_res + base2k_tsk;
271 let rows_atk: usize = 4;
272
273 let k_tsk: usize = k_ggsw_res + base2k_atk;
274 let rows_tsk: usize = 4;
275
276 let lwe_infos: LWELayout = LWELayout {
277 n: n_lwe.into(),
278 k: k_lwe_ct.into(),
279 base2k: base2k_brk.into(),
280 };
281
282 let cbt_infos: CircuitBootstrappingKeyLayout = CircuitBootstrappingKeyLayout {
283 layout_brk: BlindRotationKeyLayout {
284 n_glwe: n_glwe.into(),
285 n_lwe: n_lwe.into(),
286 base2k: base2k_brk.into(),
287 k: k_brk.into(),
288 dnum: rows_brk.into(),
289 rank: rank.into(),
290 },
291 layout_atk: GLWEAutomorphismKeyLayout {
292 n: n_glwe.into(),
293 base2k: base2k_atk.into(),
294 k: k_atk.into(),
295 dnum: rows_atk.into(),
296 rank: rank.into(),
297 dsize: Dsize(1),
298 },
299 layout_tsk: GGLWEToGGSWKeyLayout {
300 n: n_glwe.into(),
301 base2k: base2k_tsk.into(),
302 k: k_tsk.into(),
303 dnum: rows_tsk.into(),
304 dsize: Dsize(1),
305 rank: rank.into(),
306 },
307 };
308
309 let ggsw_infos: GGSWLayout = GGSWLayout {
310 n: n_glwe.into(),
311 base2k: base2k_res.into(),
312 k: k_ggsw_res.into(),
313 dnum: rows_ggsw_res.into(),
314 dsize: Dsize(1),
315 rank: rank.into(),
316 };
317
318 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(1 << 23);
319
320 let mut source_xs: Source = Source::new([1u8; 32]);
321 let mut source_xa: Source = Source::new([1u8; 32]);
322 let mut source_xe: Source = Source::new([1u8; 32]);
323
324 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe.into());
325 sk_lwe.fill_binary_block(block_size, &mut source_xs);
326
327 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n_glwe.into(), rank.into());
328 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
329
330 let mut sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
331 sk_glwe_prepared.prepare(module, &sk_glwe);
332
333 let data: i64 = 1;
334
335 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(base2k_res.into(), k_lwe_pt.into());
336 pt_lwe.encode_i64(data, (k_lwe_pt + 1).into());
337
338 println!("pt_lwe: {pt_lwe}");
339
340 let mut ct_lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(&lwe_infos);
341 ct_lwe.encrypt_sk(
342 module,
343 &pt_lwe,
344 &sk_lwe,
345 &mut source_xa,
346 &mut source_xe,
347 scratch.borrow(),
348 );
349
350 let now: Instant = Instant::now();
351 let mut cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::alloc_from_infos(&cbt_infos);
352 println!("CBT-ALLOC: {} ms", now.elapsed().as_millis());
353
354 let now: Instant = Instant::now();
355 cbt_key.encrypt_sk(
356 module,
357 &sk_lwe,
358 &sk_glwe,
359 &mut source_xa,
360 &mut source_xe,
361 scratch.borrow(),
362 );
363 println!("CBT-ENCRYPT: {} ms", now.elapsed().as_millis());
364
365 let mut res: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_infos);
366
367 let mut cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, BE> =
368 CircuitBootstrappingKeyPrepared::alloc_from_infos(module, &cbt_infos);
369 cbt_prepared.prepare(module, &cbt_key, scratch.borrow());
370
371 let now: Instant = Instant::now();
372 cbt_prepared.execute_to_constant(
373 module,
374 &mut res,
375 &ct_lwe,
376 k_lwe_pt,
377 extension_factor,
378 scratch.borrow(),
379 );
380 println!("CBT: {} ms", now.elapsed().as_millis());
381
382 // X^{data * 2^log_gap_out}
383 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n_glwe, 1);
384 pt_ggsw.at_mut(0, 0)[0] = data;
385
386 for row in 0..res.dnum().as_usize() {
387 for col in 0..res.rank().as_usize() + 1 {
388 println!(
389 "row:{row} col:{col} -> {}",
390 res.noise(
391 module,
392 row,
393 col,
394 &pt_ggsw,
395 &sk_glwe_prepared,
396 scratch.borrow()
397 )
398 .std()
399 .log2()
400 )
401 }
402 }
403
404 let mut ct_glwe: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&ggsw_infos);
405 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
406 pt_glwe.data.at_mut(0, 0)[0] = 1 << (base2k_res - k_lwe_pt - 1);
407
408 ct_glwe.encrypt_sk(
409 module,
410 &pt_glwe,
411 &sk_glwe_prepared,
412 &mut source_xa,
413 &mut source_xe,
414 scratch.borrow(),
415 );
416
417 let mut res_prepared: GGSWPrepared<Vec<u8>, BE> = GGSWPrepared::alloc_from_infos(module, &res);
418 res_prepared.prepare(module, &res, scratch.borrow());
419
420 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
421
422 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&ggsw_infos);
423 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
424
425 // Parameters are set such that the first limb should be noiseless.
426 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
427 pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
428 println!("pt_res: {pt_res}");
429 assert_eq!(pt_res.data.at(0, 0), pt_want);
430}