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