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