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