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