poulpy/bin_fhe/circuit_bootstrapping/circuit_bootstrapping.rs
bin_fhe/circuit_bootstrapping/circuit_bootstrapping.rsRaw
1use std::{collections::HashMap, time::Instant, usize};
2
3use backend::hal::{
4 api::{
5 ScratchAvailable, TakeMatZnx, TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, TakeVecZnxDftSlice,
6 TakeVecZnxSlice, VecZnxAddInplace, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
7 VecZnxBigAutomorphismInplace, VecZnxBigSubSmallBInplace, VecZnxCopy, VecZnxDftCopy, VecZnxDftToVecZnxBigTmpA,
8 VecZnxNegateInplace, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
9 VecZnxSubABInplace, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
10 },
11 layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
12 oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
13};
14use sampling::source::Source;
15
16use crate::{
17 BlindRotationKeyCGGI, BlindRotationKeyCGGIExec, BlindRotationKeyCGGIExecLayoutFamily, CCGIBlindRotationFamily, GLWEOps,
18 Infos, LookUpTable, LookUpTableRotationDirection, TakeGGLWE, TakeGLWECt, cggi_blind_rotate,
19 layouts::{
20 AutomorphismKey, GGSWCiphertext, GLWECiphertext, GLWESecret, GLWETensorKey, LWECiphertext, LWESecret,
21 prepared::{AutomorphismKeyExec, GLWESecretExec, GLWETensorKeyExec},
22 },
23};
24
25use crate::trait_families::{
26 AutomorphismKeyEncryptSkFamily, GGSWEncryptSkFamily, GLWESecretExecModuleFamily, GLWETensorKeyEncryptSkFamily,
27 GLWETraceModuleFamily,
28};
29
30pub struct CircuitBootstrappingKeyCGGI<D: Data> {
31 pub(crate) brk: BlindRotationKeyCGGI<D>,
32 pub(crate) tsk: GLWETensorKey<Vec<u8>>,
33 pub(crate) atk: HashMap<i64, AutomorphismKey<Vec<u8>>>,
34}
35
36impl CircuitBootstrappingKeyCGGI<Vec<u8>> {
37 pub fn generate<DLwe, DGlwe, B: Backend>(
38 module: &Module<B>,
39 basek: usize,
40 sk_lwe: &LWESecret<DLwe>,
41 sk_glwe: &GLWESecret<DGlwe>,
42 k_brk: usize,
43 rows_brk: usize,
44 k_trace: usize,
45 rows_trace: usize,
46 k_tsk: usize,
47 rows_tsk: usize,
48 source_xa: &mut Source,
49 source_xe: &mut Source,
50 sigma: f64,
51 scratch: &mut Scratch<B>,
52 ) -> Self
53 where
54 Module<B>: GGSWEncryptSkFamily<B>
55 + GLWESecretExecModuleFamily<B>
56 + VecZnxAddScalarInplace
57 + AutomorphismKeyEncryptSkFamily<B>
58 + VecZnxAutomorphism
59 + VecZnxSwithcDegree
60 + GLWETensorKeyEncryptSkFamily<B>,
61 DLwe: DataRef,
62 DGlwe: DataRef,
63 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeSvpPPol<B> + TakeVecZnxBig<B>,
64 {
65 let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>>> = HashMap::new();
66 let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(&module);
67 gal_els.iter().for_each(|gal_el| {
68 let mut key: AutomorphismKey<Vec<u8>> =
69 AutomorphismKey::alloc(sk_glwe.n(), basek, k_trace, rows_trace, 1, sk_glwe.rank());
70 key.encrypt_sk(
71 &module, *gal_el, &sk_glwe, source_xa, source_xe, sigma, scratch,
72 );
73 auto_keys.insert(*gal_el, key);
74 });
75
76 let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
77
78 let mut brk: BlindRotationKeyCGGI<Vec<u8>> = BlindRotationKeyCGGI::alloc(
79 sk_glwe.n(),
80 sk_lwe.n(),
81 basek,
82 k_brk,
83 rows_brk,
84 sk_glwe.rank(),
85 );
86 brk.generate_from_sk(
87 module,
88 &sk_glwe_exec,
89 sk_lwe,
90 source_xa,
91 source_xe,
92 sigma,
93 scratch,
94 );
95
96 let mut tsk: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc(sk_glwe.n(), basek, k_tsk, rows_tsk, 1, sk_glwe.rank());
97 tsk.encrypt_sk(module, &sk_glwe, source_xa, source_xe, sigma, scratch);
98
99 Self {
100 brk,
101 atk: auto_keys,
102 tsk,
103 }
104 }
105}
106
107pub struct CircuitBootstrappingKeyCGGIExec<D: Data, B: Backend> {
108 pub(crate) brk: BlindRotationKeyCGGIExec<D, B>,
109 pub(crate) tsk: GLWETensorKeyExec<Vec<u8>, B>,
110 pub(crate) atk: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
111}
112
113impl<B: Backend> CircuitBootstrappingKeyCGGIExec<Vec<u8>, B> {
114 pub fn from<DataOther: DataRef>(
115 module: &Module<B>,
116 other: &CircuitBootstrappingKeyCGGI<DataOther>,
117 scratch: &mut Scratch<B>,
118 ) -> CircuitBootstrappingKeyCGGIExec<Vec<u8>, B>
119 where
120 Module<B>: BlindRotationKeyCGGIExecLayoutFamily<B> + VmpPMatAlloc<B> + VmpPMatPrepare<B>,
121 {
122 let brk: BlindRotationKeyCGGIExec<Vec<u8>, B> = BlindRotationKeyCGGIExec::from(module, &other.brk, scratch);
123 let tsk: GLWETensorKeyExec<Vec<u8>, B> = GLWETensorKeyExec::from(module, &other.tsk, scratch);
124 let mut atk: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>> = HashMap::new();
125 for (key, value) in &other.atk {
126 atk.insert(*key, AutomorphismKeyExec::from(module, value, scratch));
127 }
128 CircuitBootstrappingKeyCGGIExec { brk, tsk, atk }
129 }
130}
131
132pub trait CGGICircuitBootstrapFamily<B: Backend> = VecZnxRotateInplace
133 + VecZnxNormalizeInplace<B>
134 + VecZnxNormalizeTmpBytes
135 + CCGIBlindRotationFamily<B>
136 + VecZnxSwithcDegree
137 + VecZnxBigAutomorphismInplace<B>
138 + VecZnxRshInplace
139 + VecZnxDftCopy<B>
140 + VecZnxDftToVecZnxBigTmpA<B>
141 + VecZnxSub
142 + VecZnxAddInplace
143 + VecZnxNegateInplace
144 + VecZnxCopy
145 + VecZnxSubABInplace
146 + GLWETraceModuleFamily<B>
147 + VecZnxRotateInplace
148 + VecZnxAutomorphismInplace
149 + VecZnxBigSubSmallBInplace<B>;
150
151pub fn circuit_bootstrap_to_constant_cggi<DRes, DLwe, DBrk, B: Backend>(
152 module: &Module<B>,
153 res: &mut GGSWCiphertext<DRes>,
154 lwe: &LWECiphertext<DLwe>,
155 log_domain: usize,
156 extension_factor: usize,
157 key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
158 scratch: &mut Scratch<B>,
159) where
160 DRes: DataMut,
161 DLwe: DataRef,
162 DBrk: DataRef,
163 Module<B>: CGGICircuitBootstrapFamily<B>,
164 B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
165 Scratch<B>: TakeVecZnx
166 + TakeVecZnxDftSlice<B>
167 + TakeVecZnxBig<B>
168 + TakeVecZnxDft<B>
169 + TakeMatZnx
170 + ScratchAvailable
171 + TakeVecZnxSlice,
172{
173 circuit_bootstrap_core_cggi(
174 false,
175 module,
176 0,
177 res,
178 lwe,
179 log_domain,
180 extension_factor,
181 key,
182 scratch,
183 );
184}
185
186pub fn circuit_bootstrap_to_exponent_cggi<DRes, DLwe, DBrk, B: Backend>(
187 module: &Module<B>,
188 log_gap_out: usize,
189 res: &mut GGSWCiphertext<DRes>,
190 lwe: &LWECiphertext<DLwe>,
191 log_domain: usize,
192 extension_factor: usize,
193 key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
194 scratch: &mut Scratch<B>,
195) where
196 DRes: DataMut,
197 DLwe: DataRef,
198 DBrk: DataRef,
199 Module<B>: CGGICircuitBootstrapFamily<B>,
200 B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
201 Scratch<B>: TakeVecZnx
202 + TakeVecZnxDftSlice<B>
203 + TakeVecZnxBig<B>
204 + TakeVecZnxDft<B>
205 + TakeMatZnx
206 + ScratchAvailable
207 + TakeVecZnxSlice,
208{
209 circuit_bootstrap_core_cggi(
210 true,
211 module,
212 log_gap_out,
213 res,
214 lwe,
215 log_domain,
216 extension_factor,
217 key,
218 scratch,
219 );
220}
221
222pub fn circuit_bootstrap_core_cggi<DRes, DLwe, DBrk, B: Backend>(
223 to_exponent: bool,
224 module: &Module<B>,
225 log_gap_out: usize,
226 res: &mut GGSWCiphertext<DRes>,
227 lwe: &LWECiphertext<DLwe>,
228 log_domain: usize,
229 extension_factor: usize,
230 key: &CircuitBootstrappingKeyCGGIExec<DBrk, B>,
231 scratch: &mut Scratch<B>,
232) where
233 DRes: DataMut,
234 DLwe: DataRef,
235 DBrk: DataRef,
236 Module<B>: CGGICircuitBootstrapFamily<B>,
237 B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
238 Scratch<B>: TakeGGLWE<B>
239 + TakeVecZnxDftSlice<B>
240 + TakeVecZnxBig<B>
241 + TakeVecZnxDft<B>
242 + TakeVecZnx
243 + ScratchAvailable
244 + TakeVecZnxSlice,
245{
246 #[cfg(debug_assertions)]
247 {
248 use crate::Infos;
249
250 assert_eq!(res.n(), key.brk.n());
251 assert_eq!(lwe.basek(), key.brk.basek());
252 assert_eq!(res.basek(), key.brk.basek());
253 }
254
255 let n: usize = res.n();
256 let basek: usize = res.basek();
257 let rows: usize = res.rows();
258 let rank: usize = res.rank();
259 let k: usize = res.k();
260
261 let alpha: usize = rows.next_power_of_two();
262
263 let mut f: Vec<i64> = vec![0i64; (1 << log_domain) * alpha];
264
265 if to_exponent {
266 (0..rows).for_each(|i| {
267 f[i] = 1 << (basek * (rows - 1 - i));
268 });
269 } else {
270 (0..1 << log_domain).for_each(|j| {
271 (0..rows).for_each(|i| {
272 f[j * alpha + i] = j as i64 * (1 << (basek * (rows - 1 - i)));
273 });
274 });
275 }
276
277 // Lut precision, basically must be able to hold the decomposition power basis of the GGSW
278 let mut lut: LookUpTable = LookUpTable::alloc(n, basek, basek * rows, extension_factor);
279 lut.set(module, &f, basek * rows);
280
281 if to_exponent {
282 lut.set_rotation_direction(LookUpTableRotationDirection::Right);
283 }
284
285 // TODO: separate GGSW k from output of blind rotation k
286 let (mut res_glwe, scratch1) = scratch.take_glwe_ct(n, basek, k, rank);
287 let (mut tmp_gglwe, scratch2) = scratch1.take_gglwe(n, basek, k, rows, 1, rank, rank);
288
289 let now: Instant = Instant::now();
290 cggi_blind_rotate(module, &mut res_glwe, &lwe, &lut, &key.brk, scratch2);
291 println!("cggi_blind_rotate: {} ms", now.elapsed().as_millis());
292
293 let gap: usize = 2 * lut.drift / lut.extension_factor();
294
295 let log_gap_in: usize = (usize::BITS - (gap * alpha - 1).leading_zeros()) as _;
296
297 (0..rows).for_each(|i| {
298 let mut tmp_glwe: GLWECiphertext<&mut [u8]> = tmp_gglwe.at_mut(i, 0);
299
300 if to_exponent {
301 let now: Instant = Instant::now();
302
303 // Isolates i-th LUT and moves coefficients according to requested gap.
304 post_process(
305 module,
306 &mut tmp_glwe,
307 &res_glwe,
308 log_gap_in,
309 log_gap_out,
310 log_domain,
311 &key.atk,
312 scratch2,
313 );
314 println!("post_process: {} ms", now.elapsed().as_millis());
315 } else {
316 tmp_glwe.trace(module, 0, module.log_n(), &res_glwe, &key.atk, scratch2);
317 }
318
319 if i < rows {
320 res_glwe.rotate_inplace(module, -(gap as i64));
321 }
322 });
323
324 // Expands GGLWE to GGSW using GGLWE(s^2)
325 res.from_gglwe(module, &tmp_gglwe, &key.tsk, scratch2);
326}
327
328fn post_process<DataRes, DataA, B: Backend>(
329 module: &Module<B>,
330 res: &mut GLWECiphertext<DataRes>,
331 a: &GLWECiphertext<DataA>,
332 log_gap_in: usize,
333 log_gap_out: usize,
334 log_domain: usize,
335 auto_keys: &HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
336 scratch: &mut Scratch<B>,
337) where
338 DataRes: DataMut,
339 DataA: DataRef,
340 Module<B>: CGGICircuitBootstrapFamily<B>,
341 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
342{
343 let log_n: usize = module.log_n();
344
345 let mut cts: HashMap<usize, GLWECiphertext<Vec<u8>>> = HashMap::new();
346
347 // First partial trace, vanishes all coefficients which are not multiples of gap_in
348 // [1, 1, 1, 1, 0, 0, 0, ..., 0, 0, -1, -1, -1, -1] -> [1, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0, 0, 0]
349 res.trace(
350 module,
351 module.log_n() - log_gap_in as usize + 1,
352 log_n,
353 &a,
354 auto_keys,
355 scratch,
356 );
357
358 // TODO: optimize with packing and final partial trace
359 // If gap_out < gap_in, then we need to repack, i.e. reduce the cap between coefficients.
360 if log_gap_in != log_gap_out {
361 let steps: i32 = 1 << log_domain;
362 (0..steps).for_each(|i| {
363 if i != 0 {
364 res.rotate_inplace(module, -(1 << log_gap_in));
365 }
366 cts.insert(i as usize * (1 << log_gap_out), res.clone());
367 });
368
369 let now: Instant = Instant::now();
370 pack(module, &mut cts, log_gap_out, auto_keys, scratch);
371 println!("pack: {} ms", now.elapsed().as_millis());
372 let packed: GLWECiphertext<Vec<u8>> = cts.remove(&0).unwrap();
373 res.trace(
374 module,
375 log_n - log_gap_out,
376 log_n,
377 &packed,
378 auto_keys,
379 scratch,
380 );
381 }
382}
383
384pub fn pack<D: DataMut, B: Backend>(
385 module: &Module<B>,
386 cts: &mut HashMap<usize, GLWECiphertext<D>>,
387 log_gap_out: usize,
388 auto_keys: &HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>>,
389 scratch: &mut Scratch<B>,
390) where
391 Module<B>: CGGICircuitBootstrapFamily<B>,
392 Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
393{
394 let log_n: usize = module.log_n();
395
396 let basek: usize = cts.get(&0).unwrap().basek();
397 let k: usize = cts.get(&0).unwrap().k();
398 let rank: usize = cts.get(&0).unwrap().rank();
399
400 (0..log_n - log_gap_out).for_each(|i| {
401 let now: Instant = Instant::now();
402
403 let t = 16.min(1 << (log_n - 1 - i));
404
405 let auto_key: &AutomorphismKeyExec<Vec<u8>, B>;
406 if i == 0 {
407 auto_key = auto_keys.get(&-1).unwrap()
408 } else {
409 auto_key = auto_keys.get(&module.galois_element(1 << (i - 1))).unwrap();
410 }
411
412 (0..t).for_each(|j| {
413 let mut a: Option<GLWECiphertext<D>> = cts.remove(&j);
414 let mut b: Option<GLWECiphertext<D>> = cts.remove(&(j + t));
415
416 combine(
417 module,
418 basek,
419 k,
420 rank,
421 a.as_mut(),
422 b.as_mut(),
423 i,
424 auto_key,
425 scratch,
426 );
427
428 if let Some(a) = a {
429 cts.insert(j, a);
430 } else if let Some(b) = b {
431 cts.insert(j, b);
432 }
433 });
434
435 println!("combine: {} us", now.elapsed().as_micros());
436 });
437}
438
439fn combine<A: DataMut, D: DataMut, DataAK: DataRef, B: Backend>(
440 module: &Module<B>,
441 basek: usize,
442 k: usize,
443 rank: usize,
444 a: Option<&mut GLWECiphertext<A>>,
445 b: Option<&mut GLWECiphertext<D>>,
446 i: usize,
447 auto_key: &AutomorphismKeyExec<DataAK, B>,
448 scratch: &mut Scratch<B>,
449) where
450 Module<B>: CGGICircuitBootstrapFamily<B>,
451 Scratch<B>: TakeVecZnx + TakeVecZnxDft<B> + ScratchAvailable,
452{
453 // Goal is to evaluate: a = a + b*X^t + phi(a - b*X^t))
454 // We also use the identity: AUTO(a * X^t, g) = -X^t * AUTO(a, g)
455 // where t = 2^(log_n - i - 1) and g = 5^{2^(i - 1)}
456 // Different cases for wether a and/or b are zero.
457 //
458 // Implicite RSH without modulus switch, introduces extra I(X) * Q/2 on decryption.
459 // Necessary so that the scaling of the plaintext remains constant.
460 // It however is ok to do so here because coefficients are eventually
461 // either mapped to garbage or twice their value which vanishes I(X)
462 // since 2*(I(X) * Q/2) = I(X) * Q = 0 mod Q.
463 if let Some(a) = a {
464 let n: usize = a.n();
465 let log_n: usize = (u64::BITS - (n - 1).leading_zeros()) as _;
466 let t: i64 = 1 << (log_n - i - 1);
467
468 if let Some(b) = b {
469 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
470
471 // a = a * X^-t
472 a.rotate_inplace(module, -t);
473
474 // tmp_b = a * X^-t - b
475 tmp_b.sub(module, a, b);
476 tmp_b.rsh(module, 1);
477
478 // a = a * X^-t + b
479 a.add_inplace(module, b);
480 a.rsh(module, 1);
481
482 tmp_b.normalize_inplace(module, scratch_1);
483
484 // tmp_b = phi(a * X^-t - b)
485 tmp_b.automorphism_inplace(module, auto_key, scratch_1);
486
487 // a = a * X^-t + b - phi(a * X^-t - b)
488 a.sub_inplace_ab(module, &tmp_b);
489 a.normalize_inplace(module, scratch_1);
490
491 // a = a + b * X^t - phi(a * X^-t - b) * X^t
492 // = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
493 // = a + b * X^t + phi(a - b * X^t)
494 a.rotate_inplace(module, t);
495 } else {
496 a.rsh(module, 1);
497 // a = a + phi(a)
498 a.automorphism_add_inplace(module, auto_key, scratch);
499 }
500 } else {
501 if let Some(b) = b {
502 let n: usize = b.n();
503 let log_n: usize = (u64::BITS - (n - 1).leading_zeros()) as _;
504 let t: i64 = 1 << (log_n - i - 1);
505
506 let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(n, basek, k, rank);
507 tmp_b.rotate(module, t, b);
508 tmp_b.rsh(module, 1);
509
510 // a = (b* X^t - phi(b* X^t))
511 b.automorphism_sub_ba(module, &tmp_b, auto_key, scratch_1);
512 }
513 }
514}