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