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