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