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