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