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