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