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