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