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, GLWERotate, 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 + GLWERotate<BE>,
120 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
121 Scratch<BE>: ScratchTakeCore<BE>,
122{
123 fn circuit_bootstrapping_execute_tmp_bytes<R, A>(
124 &self,
125 block_size: usize,
126 extension_factor: usize,
127 res_infos: &R,
128 cbt_infos: &A,
129 ) -> usize
130 where
131 R: GGSWInfos,
132 A: CircuitBootstrappingKeyInfos,
133 {
134 self.blind_rotation_execute_tmp_bytes(
135 block_size,
136 extension_factor,
137 res_infos,
138 &cbt_infos.brk_infos(),
139 )
140 .max(self.glwe_trace_tmp_bytes(res_infos, res_infos, &cbt_infos.atk_infos()))
141 .max(self.ggsw_from_gglwe_tmp_bytes(res_infos, &cbt_infos.tsk_infos()))
142 }
143
144 fn circuit_bootstrapping_execute_to_constant<R, L, D>(
145 &self,
146 res: &mut R,
147 lwe: &L,
148 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
149 log_domain: usize,
150 extension_factor: usize,
151 scratch: &mut Scratch<BE>,
152 ) where
153 R: GGSWToMut + GGSWInfos,
154 L: LWEToRef + LWEInfos,
155 D: DataRef,
156 {
157 circuit_bootstrap_core(
158 false,
159 self,
160 0,
161 res,
162 lwe,
163 log_domain,
164 extension_factor,
165 key,
166 scratch,
167 );
168 }
169
170 fn circuit_bootstrapping_execute_to_exponent<R, L, D>(
171 &self,
172 log_gap_out: usize,
173 res: &mut R,
174 lwe: &L,
175 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
176 log_domain: usize,
177 extension_factor: usize,
178 scratch: &mut Scratch<BE>,
179 ) where
180 R: GGSWToMut + GGSWInfos,
181 L: LWEToRef + LWEInfos,
182 D: DataRef,
183 {
184 circuit_bootstrap_core(
185 true,
186 self,
187 log_gap_out,
188 res,
189 lwe,
190 log_domain,
191 extension_factor,
192 key,
193 scratch,
194 );
195 }
196}
197
198#[allow(clippy::too_many_arguments)]
199pub fn circuit_bootstrap_core<R, L, D, M, BRA: BlindRotationAlgo, BE: Backend>(
200 to_exponent: bool,
201 module: &M,
202 log_gap_out: usize,
203 res: &mut R,
204 lwe: &L,
205 log_domain: usize,
206 extension_factor: usize,
207 key: &CircuitBootstrappingKeyPrepared<D, BRA, BE>,
208 scratch: &mut Scratch<BE>,
209) where
210 R: GGSWToMut,
211 L: LWEToRef,
212 D: DataRef,
213 M: ModuleN
214 + LookupTableFactory
215 + BlindRotationExecute<BRA, BE>
216 + GLWETrace<BE>
217 + GLWEPacking<BE>
218 + GGSWFromGGLWE<BE>
219 + GLWESecretPreparedFactory<BE>
220 + GLWEDecrypt<BE>
221 + GLWERotate<BE>
222 + ModuleLogN,
223 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
224 Scratch<BE>: ScratchTakeCore<BE>,
225{
226 let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
227 let lwe: &LWE<&[u8]> = &lwe.to_ref();
228
229 assert_eq!(res.n(), key.brk.n());
230 assert_eq!(lwe.base2k(), key.brk.base2k());
231 assert_eq!(res.base2k(), key.brk.base2k());
232
233 let n: usize = res.n().into();
234 let base2k: usize = res.base2k().into();
235 let dnum: usize = res.dnum().into();
236 let rank: usize = res.rank().into();
237 let k: usize = res.k().into();
238
239 let alpha: usize = dnum.next_power_of_two();
240
241 let mut f: Vec<i64> = vec![0i64; (1 << log_domain) * alpha];
242
243 if to_exponent {
244 (0..dnum).for_each(|i| {
245 f[i] = 1 << (base2k * (dnum - 1 - i));
246 });
247 } else {
248 (0..1 << log_domain).for_each(|j| {
249 (0..dnum).for_each(|i| {
250 f[j * alpha + i] = j as i64 * (1 << (base2k * (dnum - 1 - i)));
251 });
252 });
253 }
254
255 let lut_infos: LookUpTableLayout = LookUpTableLayout {
256 n: module.n().into(),
257 extension_factor,
258 k: (base2k * dnum).into(),
259 base2k: base2k.into(),
260 };
261
262 // Lut precision, basically must be able to hold the decomposition power basis of the GGSW
263 let mut lut: LookupTable = LookupTable::alloc(&lut_infos);
264 lut.set(module, &f, base2k * dnum);
265
266 if to_exponent {
267 lut.set_rotation_direction(LookUpTableRotationDirection::Right);
268 }
269
270 // TODO: separate GGSW k from output of blind rotation k
271 let (mut res_glwe, scratch_1) = scratch.take_glwe(res);
272
273 let gglwe_infos: GGLWELayout = GGLWELayout {
274 n: n.into(),
275 base2k: base2k.into(),
276 k: k.into(),
277 dnum: dnum.into(),
278 dsize: Dsize(1),
279 rank_in: rank.max(1).into(),
280 rank_out: rank.into(),
281 };
282
283 let (mut tmp_gglwe, scratch_2) = scratch_1.take_gglwe(&gglwe_infos);
284
285 key.brk.execute(module, &mut res_glwe, lwe, &lut, scratch_2);
286
287 let gap: usize = 2 * lut.drift / lut.extension_factor();
288
289 let log_gap_in: usize = (usize::BITS - (gap * alpha - 1).leading_zeros()) as _;
290
291 (0..dnum).for_each(|i| {
292 let mut tmp_glwe: GLWE<&mut [u8]> = tmp_gglwe.at_mut(i, 0);
293
294 if to_exponent {
295 // Isolates i-th LUT and moves coefficients according to requested gap.
296 post_process(
297 module,
298 &mut tmp_glwe,
299 &res_glwe,
300 log_gap_in,
301 log_gap_out,
302 log_domain,
303 &key.atk,
304 scratch_2,
305 );
306 } else {
307 tmp_glwe.trace(module, 0, module.log_n(), &res_glwe, &key.atk, scratch_2);
308 }
309
310 // let sk_glwe: &poulpy_core::layouts::GLWESecret<&[u8]> = &sk_glwe.to_ref();
311 // let sk_glwe_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, sk_glwe.rank());
312 // let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&res_glwe);
313 // res_glwe.decrypt(module, &mut pt, &sk_glwe_prepared, scratch_2);
314 // println!("pt[{i}]: {}", pt);
315
316 if i < dnum {
317 module.glwe_rotate_inplace(-(gap as i64), &mut res_glwe, scratch_2);
318 }
319 });
320
321 // Expands GGLWE to GGSW using GGLWE(s^2)
322 res.from_gglwe(module, &tmp_gglwe, &key.tsk, scratch_2);
323}
324
325#[allow(clippy::too_many_arguments)]
326fn post_process<R, A, M, BE: Backend>(
327 module: &M,
328 res: &mut R,
329 a: &A,
330 log_gap_in: usize,
331 log_gap_out: usize,
332 log_domain: usize,
333 auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, BE>>,
334 scratch: &mut Scratch<BE>,
335) where
336 R: GLWEToMut,
337 A: GLWEToRef,
338 M: ModuleLogN + GLWETrace<BE> + GLWEPacking<BE> + GLWERotate<BE>,
339 Scratch<BE>: ScratchTakeCore<BE>,
340{
341 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
342 let a: &GLWE<&[u8]> = &a.to_ref();
343
344 let log_n: usize = module.log_n();
345
346 let mut cts: HashMap<usize, &mut GLWE<Vec<u8>>> = HashMap::new();
347
348 // First partial trace, vanishes all coefficients which are not multiples of gap_in
349 // [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]
350 res.trace(
351 module,
352 module.log_n() - log_gap_in + 1,
353 log_n,
354 a,
355 auto_keys,
356 scratch,
357 );
358
359 // TODO: optimize with packing and final partial trace
360 // If gap_out < gap_in, then we need to repack, i.e. reduce the cap between coefficients.
361 if log_gap_in != log_gap_out {
362 let steps: usize = 1 << log_domain;
363
364 // TODO: from Scratch
365 let mut cts_vec: Vec<GLWE<Vec<u8>>> = Vec::new();
366
367 for i in 0..steps {
368 if i != 0 {
369 module.glwe_rotate_inplace(-(1 << log_gap_in), res, scratch);
370 }
371 cts_vec.push(res.to_owned_deep());
372 }
373
374 for (i, ct) in cts_vec.iter_mut().enumerate().take(steps) {
375 cts.insert(i * (1 << log_gap_out), ct);
376 }
377
378 module.glwe_pack(&mut cts, log_gap_out, auto_keys, scratch);
379
380 let packed: &mut GLWE<Vec<u8>> = cts.remove(&0).unwrap();
381 res.trace(
382 module,
383 log_n - log_gap_out,
384 log_n,
385 packed,
386 auto_keys,
387 scratch,
388 );
389 }
390}