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::{ModuleLogN, ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow},
5 layouts::{Backend, DataRef, Module, Scratch, ScratchOwned, ToOwnedDeep},
6};
7
8use poulpy_core::{
9 GGSWFromGGLWE, GLWEDecrypt, GLWEPacking, GLWETrace, ScratchTakeCore,
10 layouts::{
11 Dsize, GGLWELayout, GGSWInfos, GGSWToMut, GLWEInfos, GLWESecretPreparedFactory, GLWEToMut, GLWEToRef, LWEInfos, LWEToRef,
12 },
13};
14
15use poulpy_core::layouts::{GGSW, GLWE, LWE, prepared::GLWEAutomorphismKeyPrepared};
16
17use crate::tfhe::{
18 blind_rotation::{
19 BlindRotationAlgo, BlindRotationExecute, LookUpTableLayout, LookUpTableRotationDirection, LookupTable, LookupTableFactory,
20 },
21 circuit_bootstrapping::{CircuitBootstrappingKeyInfos, CircuitBootstrappingKeyPrepared},
22};
23
24pub trait CirtuitBootstrappingExecute<BRA: BlindRotationAlgo, BE: Backend> {
25 fn circuit_bootstrapping_execute_tmp_bytes<R, A>(
26 &self,
27 block_size: usize,
28 extension_factor: usize,
29 res_infos: &R,
30 cbt_infos: &A,
31 ) -> usize
32 where
33 R: GGSWInfos,
34 A: CircuitBootstrappingKeyInfos;
35
36 fn circuit_bootstrapping_execute_to_constant<R, L, D>(
37 &self,
38 res: &mut R,
39 lwe: &L,
40 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
41 log_domain: usize,
42 extension_factor: usize,
43 scratch: &mut Scratch<BE>,
44 ) where
45 R: GGSWToMut + GGSWInfos,
46 L: LWEToRef + LWEInfos,
47 D: DataRef;
48
49 #[allow(clippy::too_many_arguments)]
50 fn circuit_bootstrapping_execute_to_exponent<R, L, D>(
51 &self,
52 log_gap_out: usize,
53 res: &mut R,
54 lwe: &L,
55 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
56 log_domain: usize,
57 extension_factor: usize,
58 scratch: &mut Scratch<BE>,
59 ) where
60 R: GGSWToMut + GGSWInfos,
61 L: LWEToRef + LWEInfos,
62 D: DataRef;
63}
64
65impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> CircuitBootstrappingKeyPrepared<D, BRA, BE> {
66 pub fn execute_to_constant<M, L, R>(
67 &self,
68 module: &M,
69 res: &mut R,
70 lwe: &L,
71 log_domain: usize,
72 extension_factor: usize,
73 scratch: &mut Scratch<BE>,
74 ) where
75 M: CirtuitBootstrappingExecute<BRA, BE>,
76 R: GGSWToMut + GGSWInfos,
77 L: LWEToRef + LWEInfos,
78 {
79 module.circuit_bootstrapping_execute_to_constant(res, lwe, self, log_domain, extension_factor, scratch);
80 }
81
82 #[allow(clippy::too_many_arguments)]
83 pub fn execute_to_exponent<R, L, M>(
84 &self,
85 module: &M,
86 log_gap_out: usize,
87 res: &mut R,
88 lwe: &L,
89 log_domain: usize,
90 extension_factor: usize,
91 scratch: &mut Scratch<BE>,
92 ) where
93 M: CirtuitBootstrappingExecute<BRA, BE>,
94 R: GGSWToMut + GGSWInfos,
95 L: LWEToRef + LWEInfos,
96 {
97 module.circuit_bootstrapping_execute_to_exponent(
98 log_gap_out,
99 res,
100 lwe,
101 self,
102 log_domain,
103 extension_factor,
104 scratch,
105 );
106 }
107}
108
109impl<BRA: BlindRotationAlgo, BE: Backend> CirtuitBootstrappingExecute<BRA, BE> for Module<BE>
110where
111 Self: ModuleN
112 + LookupTableFactory
113 + BlindRotationExecute<BRA, BE>
114 + GLWETrace<BE>
115 + GLWEPacking<BE>
116 + GGSWFromGGLWE<BE>
117 + GLWESecretPreparedFactory<BE>
118 + GLWEDecrypt<BE>,
119 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
120 Scratch<BE>: ScratchTakeCore<BE>,
121{
122 fn circuit_bootstrapping_execute_tmp_bytes<R, A>(
123 &self,
124 block_size: usize,
125 extension_factor: usize,
126 res_infos: &R,
127 cbt_infos: &A,
128 ) -> usize
129 where
130 R: GGSWInfos,
131 A: CircuitBootstrappingKeyInfos,
132 {
133 self.blind_rotation_execute_tmp_bytes(
134 block_size,
135 extension_factor,
136 res_infos,
137 &cbt_infos.brk_infos(),
138 )
139 .max(self.glwe_trace_tmp_bytes(res_infos, res_infos, &cbt_infos.atk_infos()))
140 .max(self.ggsw_from_gglwe_tmp_bytes(res_infos, &cbt_infos.tsk_infos()))
141 }
142
143 fn circuit_bootstrapping_execute_to_constant<R, L, D>(
144 &self,
145 res: &mut R,
146 lwe: &L,
147 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
148 log_domain: usize,
149 extension_factor: usize,
150 scratch: &mut Scratch<BE>,
151 ) where
152 R: GGSWToMut + GGSWInfos,
153 L: LWEToRef + LWEInfos,
154 D: DataRef,
155 {
156 circuit_bootstrap_core(
157 false,
158 self,
159 0,
160 res,
161 lwe,
162 log_domain,
163 extension_factor,
164 key,
165 scratch,
166 );
167 }
168
169 fn circuit_bootstrapping_execute_to_exponent<R, L, D>(
170 &self,
171 log_gap_out: usize,
172 res: &mut R,
173 lwe: &L,
174 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
175 log_domain: usize,
176 extension_factor: usize,
177 scratch: &mut Scratch<BE>,
178 ) where
179 R: GGSWToMut + GGSWInfos,
180 L: LWEToRef + LWEInfos,
181 D: DataRef,
182 {
183 circuit_bootstrap_core(
184 true,
185 self,
186 log_gap_out,
187 res,
188 lwe,
189 log_domain,
190 extension_factor,
191 key,
192 scratch,
193 );
194 }
195}
196
197#[allow(clippy::too_many_arguments)]
198pub fn circuit_bootstrap_core<R, L, D, M, BRA: BlindRotationAlgo, BE: Backend>(
199 to_exponent: bool,
200 module: &M,
201 log_gap_out: usize,
202 res: &mut R,
203 lwe: &L,
204 log_domain: usize,
205 extension_factor: usize,
206 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
207 scratch: &mut Scratch<BE>,
208) where
209 R: GGSWToMut,
210 L: LWEToRef,
211 D: DataRef,
212 M: ModuleN
213 + LookupTableFactory
214 + BlindRotationExecute<BRA, BE>
215 + GLWETrace<BE>
216 + GLWEPacking<BE>
217 + GGSWFromGGLWE<BE>
218 + GLWESecretPreparedFactory<BE>
219 + GLWEDecrypt<BE>,
220 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
221 Scratch<BE>: ScratchTakeCore<BE>,
222{
223 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
224 let lwe: &LWE<&[u8]> = &lwe.to_ref();
225
226 assert_eq!(res.n(), key.brk.n());
227 assert_eq!(lwe.base2k(), key.brk.base2k());
228 assert_eq!(res.base2k(), key.brk.base2k());
229
230 let n: usize = res.n().into();
231 let base2k: usize = res.base2k().into();
232 let dnum: usize = res.dnum().into();
233 let rank: usize = res.rank().into();
234 let k: usize = res.k().into();
235
236 let alpha: usize = dnum.next_power_of_two();
237
238 let mut f: Vec<i64> = vec![0i64; (1 << log_domain) * alpha];
239
240 if to_exponent {
241 (0..dnum).for_each(|i| {
242 f[i] = 1 << (base2k * (dnum - 1 - i));
243 });
244 } else {
245 (0..1 << log_domain).for_each(|j| {
246 (0..dnum).for_each(|i| {
247 f[j * alpha + i] = j as i64 * (1 << (base2k * (dnum - 1 - i)));
248 });
249 });
250 }
251
252 let lut_infos: LookUpTableLayout = LookUpTableLayout {
253 n: module.n().into(),
254 extension_factor,
255 k: (base2k * dnum).into(),
256 base2k: base2k.into(),
257 };
258
259 // Lut precision, basically must be able to hold the decomposition power basis of the GGSW
260 let mut lut: LookupTable = LookupTable::alloc(&lut_infos);
261 lut.set(module, &f, base2k * dnum);
262
263 if to_exponent {
264 lut.set_rotation_direction(LookUpTableRotationDirection::Right);
265 }
266
267 // TODO: separate GGSW k from output of blind rotation k
268 let (mut res_glwe, scratch_1) = scratch.take_glwe(res);
269
270 let gglwe_infos: GGLWELayout = GGLWELayout {
271 n: n.into(),
272 base2k: base2k.into(),
273 k: k.into(),
274 dnum: dnum.into(),
275 dsize: Dsize(1),
276 rank_in: rank.max(1).into(),
277 rank_out: rank.into(),
278 };
279
280 let (mut tmp_gglwe, scratch_2) = scratch_1.take_gglwe(&gglwe_infos);
281
282 key.brk.execute(module, &mut res_glwe, lwe, &lut, scratch_2);
283
284 let gap: usize = 2 * lut.drift / lut.extension_factor();
285
286 let log_gap_in: usize = (usize::BITS - (gap * alpha - 1).leading_zeros()) as _;
287
288 (0..dnum).for_each(|i| {
289 let mut tmp_glwe: GLWE<&mut [u8]> = tmp_gglwe.at_mut(i, 0);
290
291 if to_exponent {
292 // Isolates i-th LUT and moves coefficients according to requested gap.
293 post_process(
294 module,
295 &mut tmp_glwe,
296 &res_glwe,
297 log_gap_in,
298 log_gap_out,
299 log_domain,
300 &key.atk,
301 scratch_2,
302 );
303 } else {
304 tmp_glwe.trace(module, 0, module.log_n(), &res_glwe, &key.atk, scratch_2);
305 }
306
307 // let sk_glwe: &poulpy_core::layouts::GLWESecret<&[u8]> = &sk_glwe.to_ref();
308 // let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_glwe.rank());
309 // let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&res_glwe);
310 // res_glwe.decrypt(module, &mut pt, &sk_glwe_prepared, scratch_2);
311 // println!("pt[{i}]: {}", pt);
312
313 if i < dnum {
314 module.glwe_rotate_inplace(-(gap as i64), &mut res_glwe, scratch_2);
315 }
316 });
317
318 // Expands GGLWE to GGSW using GGLWE(s^2)
319 res.from_gglwe(module, &tmp_gglwe, &key.tsk, scratch_2);
320}
321
322#[allow(clippy::too_many_arguments)]
323fn post_process<R, A, M, BE: Backend>(
324 module: &M,
325 res: &mut R,
326 a: &A,
327 log_gap_in: usize,
328 log_gap_out: usize,
329 log_domain: usize,
330 auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, BE>>,
331 scratch: &mut Scratch<BE>,
332) where
333 R: GLWEToMut,
334 A: GLWEToRef,
335 M: ModuleLogN + GLWETrace<BE> + GLWEPacking<BE>,
336 Scratch<BE>: ScratchTakeCore<BE>,
337{
338 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
339 let a: &GLWE<&[u8]> = &a.to_ref();
340
341 let log_n: usize = module.log_n();
342
343 let mut cts: HashMap<usize, &mut GLWE<Vec<u8>>> = HashMap::new();
344
345 // First partial trace, vanishes all coefficients which are not multiples of gap_in
346 // [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]
347 res.trace(
348 module,
349 module.log_n() - log_gap_in + 1,
350 log_n,
351 a,
352 auto_keys,
353 scratch,
354 );
355
356 // TODO: optimize with packing and final partial trace
357 // If gap_out < gap_in, then we need to repack, i.e. reduce the cap between coefficients.
358 if log_gap_in != log_gap_out {
359 let steps: usize = 1 << log_domain;
360
361 // TODO: from Scratch
362 let mut cts_vec: Vec<GLWE<Vec<u8>>> = Vec::new();
363
364 for i in 0..steps {
365 if i != 0 {
366 module.glwe_rotate_inplace(-(1 << log_gap_in), res, scratch);
367 }
368 cts_vec.push(res.to_owned_deep());
369 }
370
371 for (i, ct) in cts_vec.iter_mut().enumerate().take(steps) {
372 cts.insert(i * (1 << log_gap_out), ct);
373 }
374
375 module.glwe_pack(&mut cts, log_gap_out, auto_keys, scratch);
376
377 let packed: &mut GLWE<Vec<u8>> = cts.remove(&0).unwrap();
378 res.trace(
379 module,
380 log_n - log_gap_out,
381 log_n,
382 packed,
383 auto_keys,
384 scratch,
385 );
386 }
387}