poulpy/schemes/src/tfhe/circuit_bootstrapping/tests/circuit_bootstrapping.rs
schemes/src/tfhe/circuit_bootstrapping/tests/circuit_bootstrapping.rsRaw
1use std::time::Instant;
2
3use backend::hal::{
4 api::{
5 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpPPolAlloc, SvpPrepare, VecZnxAddNormal, VecZnxAddScalarInplace,
6 VecZnxAutomorphism, VecZnxFillUniform, VecZnxNormalizeInplace, VecZnxRotateInplace, VecZnxSwithcDegree, VmpPMatAlloc,
7 VmpPMatPrepare, ZnxView, ZnxViewMut,
8 },
9 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeMatZnxImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxDftSliceImpl, TakeVecZnxImpl, TakeVecZnxSliceImpl,
13 },
14};
15use sampling::source::Source;
16
17use crate::tfhe::{
18 blind_rotation::{
19 BlincRotationExecute, BlindRotationAlgo, BlindRotationKey, BlindRotationKeyAlloc, BlindRotationKeyEncryptSk,
20 BlindRotationKeyPrepared,
21 },
22 circuit_bootstrapping::{
23 CircuitBootstrapFamily, CircuitBootstrappingKey, CircuitBootstrappingKeyEncryptSk, CircuitBootstrappingKeyPrepared,
24 CirtuitBootstrappingExecute,
25 },
26};
27
28use core::{
29 layouts::prepared::PrepareAlloc,
30 trait_families::{
31 GGLWEAutomorphismKeyEncryptSkFamily, GGLWETensorKeyEncryptSkFamily, GGSWAssertNoiseFamily, GGSWEncryptSkFamily,
32 GLWEDecryptFamily,
33 },
34};
35
36use core::layouts::{
37 GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, LWECiphertext, LWEPlaintext, LWESecret,
38 prepared::{GGSWCiphertextPrepared, GLWESecretPrepared},
39};
40
41pub fn test_circuit_bootstrapping_to_exponent<B: Backend, BRA: BlindRotationAlgo>(module: &Module<B>)
42where
43 Module<B>: VecZnxFillUniform
44 + VecZnxAddNormal
45 + VecZnxNormalizeInplace<B>
46 + GGSWEncryptSkFamily<B>
47 + VecZnxAddScalarInplace
48 + GGLWEAutomorphismKeyEncryptSkFamily<B>
49 + VecZnxAutomorphism
50 + VecZnxSwithcDegree
51 + GGLWETensorKeyEncryptSkFamily<B>
52 + CircuitBootstrapFamily<B>
53 + GLWEDecryptFamily<B>
54 + GGSWAssertNoiseFamily<B>
55 + VmpPMatAlloc<B>
56 + VmpPMatPrepare<B>
57 + SvpPrepare<B>
58 + SvpPPolAlloc<B>,
59 B: ScratchOwnedAllocImpl<B>
60 + ScratchOwnedBorrowImpl<B>
61 + TakeVecZnxDftImpl<B>
62 + ScratchAvailableImpl<B>
63 + TakeVecZnxImpl<B>
64 + TakeScalarZnxImpl<B>
65 + TakeSvpPPolImpl<B>
66 + TakeVecZnxBigImpl<B>
67 + TakeVecZnxDftSliceImpl<B>
68 + TakeMatZnxImpl<B>
69 + TakeVecZnxSliceImpl<B>,
70 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
71 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
72 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
73{
74 let n: usize = module.n();
75 let basek: usize = 17;
76 let extension_factor: usize = 1;
77 let rank: usize = 1;
78 let sigma: f64 = 3.2;
79
80 let n_lwe: usize = 77;
81 let k_lwe_pt: usize = 4;
82 let k_lwe_ct: usize = 22;
83 let block_size: usize = 7;
84
85 let k_brk: usize = 5 * basek;
86 let rows_brk: usize = 4;
87
88 let k_trace: usize = 5 * basek;
89 let rows_trace: usize = 4;
90
91 let k_tsk: usize = 5 * basek;
92 let rows_tsk: usize = 4;
93
94 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
95
96 let mut source_xs: Source = Source::new([1u8; 32]);
97 let mut source_xa: Source = Source::new([1u8; 32]);
98 let mut source_xe: Source = Source::new([1u8; 32]);
99
100 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
101 sk_lwe.fill_binary_block(block_size, &mut source_xs);
102
103 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
104 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
105
106 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
107
108 let data: i64 = 1;
109
110 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
111 pt_lwe.encode_i64(data, k_lwe_pt + 2);
112
113 println!("pt_lwe: {}", pt_lwe);
114
115 let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
116 ct_lwe.encrypt_sk(
117 module,
118 &pt_lwe,
119 &sk_lwe,
120 &mut source_xa,
121 &mut source_xe,
122 sigma,
123 );
124
125 let now: Instant = Instant::now();
126 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
127 module,
128 basek,
129 &sk_lwe,
130 &sk_glwe,
131 k_brk,
132 rows_brk,
133 k_trace,
134 rows_trace,
135 k_tsk,
136 rows_tsk,
137 &mut source_xa,
138 &mut source_xe,
139 sigma,
140 scratch.borrow(),
141 );
142 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
143
144 let k_ggsw_res: usize = 4 * basek;
145 let rows_ggsw_res: usize = 2;
146
147 let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw_res, rows_ggsw_res, 1, rank);
148
149 let log_gap_out = 1;
150
151 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
152
153 let now: Instant = Instant::now();
154 cbt_prepared.execute_to_exponent(
155 module,
156 log_gap_out,
157 &mut res,
158 &ct_lwe,
159 k_lwe_pt,
160 extension_factor,
161 scratch.borrow(),
162 );
163 println!("CBT: {} ms", now.elapsed().as_millis());
164
165 // X^{data * 2^log_gap_out}
166 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
167 pt_ggsw.at_mut(0, 0)[0] = 1;
168 module.vec_znx_rotate_inplace(data * (1 << log_gap_out), &mut pt_ggsw.as_vec_znx_mut(), 0);
169
170 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
171
172 let k_glwe: usize = k_ggsw_res;
173
174 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe, rank);
175 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, basek);
176 pt_glwe.data.at_mut(0, 0)[0] = 1 << (basek - 2);
177
178 ct_glwe.encrypt_sk(
179 module,
180 &pt_glwe,
181 &sk_glwe_prepared,
182 &mut source_xa,
183 &mut source_xe,
184 sigma,
185 scratch.borrow(),
186 );
187
188 let res_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
189
190 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
191
192 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe);
193 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
194
195 // Parameters are set such that the first limb should be noiseless.
196 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
197 pt_want[data as usize * (1 << log_gap_out)] = pt_glwe.data.at(0, 0)[0];
198 assert_eq!(pt_res.data.at(0, 0), pt_want);
199}
200
201pub fn test_circuit_bootstrapping_to_constant<B: Backend, BRA: BlindRotationAlgo>(module: &Module<B>)
202where
203 Module<B>: VecZnxFillUniform
204 + VecZnxAddNormal
205 + VecZnxNormalizeInplace<B>
206 + GGSWEncryptSkFamily<B>
207 + VecZnxAddScalarInplace
208 + GGLWEAutomorphismKeyEncryptSkFamily<B>
209 + VecZnxAutomorphism
210 + VecZnxSwithcDegree
211 + GGLWETensorKeyEncryptSkFamily<B>
212 + CircuitBootstrapFamily<B>
213 + GLWEDecryptFamily<B>
214 + GGSWAssertNoiseFamily<B>
215 + VmpPMatAlloc<B>
216 + VmpPMatPrepare<B>
217 + SvpPrepare<B>
218 + SvpPPolAlloc<B>,
219 B: ScratchOwnedAllocImpl<B>
220 + ScratchOwnedBorrowImpl<B>
221 + TakeVecZnxDftImpl<B>
222 + ScratchAvailableImpl<B>
223 + TakeVecZnxImpl<B>
224 + TakeScalarZnxImpl<B>
225 + TakeSvpPPolImpl<B>
226 + TakeVecZnxBigImpl<B>
227 + TakeVecZnxDftSliceImpl<B>
228 + TakeMatZnxImpl<B>
229 + TakeVecZnxSliceImpl<B>,
230 BlindRotationKey<Vec<u8>, BRA>: PrepareAlloc<B, BlindRotationKeyPrepared<Vec<u8>, BRA, B>>,
231 BlindRotationKeyPrepared<Vec<u8>, BRA, B>: BlincRotationExecute<B>,
232 BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyAlloc + BlindRotationKeyEncryptSk<B>,
233{
234 let n = module.n();
235 let basek: usize = 14;
236 let extension_factor: usize = 1;
237 let rank: usize = 2;
238 let sigma: f64 = 3.2;
239
240 let n_lwe: usize = 77;
241 let k_lwe_pt: usize = 1;
242 let k_lwe_ct: usize = 13;
243 let block_size: usize = 7;
244
245 let k_brk: usize = 5 * basek;
246 let rows_brk: usize = 3;
247
248 let k_trace: usize = 5 * basek;
249 let rows_trace: usize = 4;
250
251 let k_tsk: usize = 5 * basek;
252 let rows_tsk: usize = 4;
253
254 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 23);
255
256 let mut source_xs: Source = Source::new([1u8; 32]);
257 let mut source_xa: Source = Source::new([1u8; 32]);
258 let mut source_xe: Source = Source::new([1u8; 32]);
259
260 let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
261 sk_lwe.fill_binary_block(block_size, &mut source_xs);
262
263 let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
264 sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
265
266 let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_glwe.prepare_alloc(module, scratch.borrow());
267
268 let data: i64 = 1;
269
270 let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
271 pt_lwe.encode_i64(data, k_lwe_pt + 2);
272
273 println!("pt_lwe: {}", pt_lwe);
274
275 let mut ct_lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
276 ct_lwe.encrypt_sk(
277 module,
278 &pt_lwe,
279 &sk_lwe,
280 &mut source_xa,
281 &mut source_xe,
282 sigma,
283 );
284
285 let now: Instant = Instant::now();
286 let cbt_key: CircuitBootstrappingKey<Vec<u8>, BRA> = CircuitBootstrappingKey::encrypt_sk(
287 module,
288 basek,
289 &sk_lwe,
290 &sk_glwe,
291 k_brk,
292 rows_brk,
293 k_trace,
294 rows_trace,
295 k_tsk,
296 rows_tsk,
297 &mut source_xa,
298 &mut source_xe,
299 sigma,
300 scratch.borrow(),
301 );
302 println!("CBT-KGEN: {} ms", now.elapsed().as_millis());
303
304 let k_ggsw_res: usize = 4 * basek;
305 let rows_ggsw_res: usize = 3;
306
307 let mut res: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ggsw_res, rows_ggsw_res, 1, rank);
308
309 let cbt_prepared: CircuitBootstrappingKeyPrepared<Vec<u8>, BRA, B> = cbt_key.prepare_alloc(module, scratch.borrow());
310
311 let now: Instant = Instant::now();
312 cbt_prepared.execute_to_constant(
313 module,
314 &mut res,
315 &ct_lwe,
316 k_lwe_pt,
317 extension_factor,
318 scratch.borrow(),
319 );
320 println!("CBT: {} ms", now.elapsed().as_millis());
321
322 // X^{data * 2^log_gap_out}
323 let mut pt_ggsw: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
324 pt_ggsw.at_mut(0, 0)[0] = data;
325
326 res.print_noise(module, &sk_glwe_prepared, &pt_ggsw);
327
328 let k_glwe: usize = k_ggsw_res;
329
330 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe, rank);
331 let mut pt_glwe: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, basek);
332 pt_glwe.data.at_mut(0, 0)[0] = 1 << (basek - k_lwe_pt - 1);
333
334 ct_glwe.encrypt_sk(
335 module,
336 &pt_glwe,
337 &sk_glwe_prepared,
338 &mut source_xa,
339 &mut source_xe,
340 sigma,
341 scratch.borrow(),
342 );
343
344 let res_prepared: GGSWCiphertextPrepared<Vec<u8>, B> = res.prepare_alloc(module, scratch.borrow());
345
346 ct_glwe.external_product_inplace(module, &res_prepared, scratch.borrow());
347
348 let mut pt_res: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe);
349 ct_glwe.decrypt(module, &mut pt_res, &sk_glwe_prepared, scratch.borrow());
350
351 // Parameters are set such that the first limb should be noiseless.
352 let mut pt_want: Vec<i64> = vec![0i64; module.n()];
353 pt_want[0] = pt_glwe.data.at(0, 0)[0] * data;
354 assert_eq!(pt_res.data.at(0, 0), pt_want);
355}