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