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