poulpy/poulpy-schemes/src/bin_fhe/blind_rotation/algorithms/cggi/algorithm.rs
poulpy-schemes/src/bin_fhe/blind_rotation/algorithms/cggi/algorithm.rsRaw
1use itertools::izip;
2use poulpy_hal::{
3 api::{
4 ModuleN, ScratchTakeBasic, SvpApplyDftToDft, VecZnxBigAddSmallInplace, VecZnxBigBytesOf, VecZnxBigNormalize,
5 VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftSubInplace,
6 VecZnxDftZero, VecZnxIdftApply, VecZnxIdftApplyTmpBytes, VecZnxRotate, VmpApplyDftToDft, VmpApplyDftToDftTmpBytes,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, Scratch, SvpPPol, VecZnx, ZnxZero},
9};
10
11use poulpy_core::{
12 Distribution, GLWEAdd, GLWEExternalProduct, GLWEMulXpMinusOne, GLWENormalize, ScratchTakeCore,
13 layouts::{GGSWInfos, GLWE, GLWEInfos, GLWEToMut, LWE, LWEInfos, LWEToRef},
14};
15
16use crate::bin_fhe::blind_rotation::{
17 BlindRotationExecute, BlindRotationKeyInfos, BlindRotationKeyPrepared, CGGI, LookupTable, mod_switch_2n,
18};
19
20impl<BE: Backend> BlindRotationExecute<CGGI, BE> for Module<BE>
21where
22 Self: VecZnxDftBytesOf
23 + VecZnxBigBytesOf
24 + VmpApplyDftToDftTmpBytes
25 + VecZnxBigNormalizeTmpBytes
26 + VecZnxIdftApplyTmpBytes
27 + GLWEExternalProduct<BE>
28 + ModuleN
29 + VecZnxRotate
30 + VecZnxDftApply<BE>
31 + VecZnxDftZero<BE>
32 + VmpApplyDftToDft<BE>
33 + SvpApplyDftToDft<BE>
34 + VecZnxDftAddInplace<BE>
35 + VecZnxDftSubInplace<BE>
36 + VecZnxIdftApply<BE>
37 + VecZnxBigAddSmallInplace<BE>
38 + VecZnxBigNormalize<BE>
39 + VecZnxCopy
40 + GLWEMulXpMinusOne<BE>
41 + GLWEAdd
42 + GLWENormalize<BE>,
43 Scratch<BE>: ScratchTakeCore<BE>,
44{
45 fn blind_rotation_execute_tmp_bytes<G, B>(
46 &self,
47 block_size: usize,
48 extension_factor: usize,
49 glwe_infos: &G,
50 brk_infos: &B,
51 ) -> usize
52 where
53 G: GLWEInfos,
54 B: GGSWInfos,
55 {
56 let brk_size: usize = brk_infos.size();
57
58 if block_size > 1 {
59 let cols: usize = (brk_infos.rank() + 1).into();
60 let dnum: usize = brk_infos.dnum().into();
61 let acc_dft: usize = self.bytes_of_vec_znx_dft(cols, dnum) * extension_factor;
62 let acc_big: usize = self.bytes_of_vec_znx_big(1, brk_size);
63 let vmp_res: usize = self.bytes_of_vec_znx_dft(cols, brk_size) * extension_factor;
64 let vmp_xai: usize = self.bytes_of_vec_znx_dft(1, brk_size);
65 let acc_dft_add: usize = vmp_res;
66 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(brk_size, dnum, dnum, 2, 2, brk_size); // GGSW product: (1 x 2) x (2 x 2)
67 let acc: usize = if extension_factor > 1 {
68 VecZnx::bytes_of(self.n(), cols, glwe_infos.size()) * extension_factor
69 } else {
70 0
71 };
72
73 acc + acc_dft
74 + acc_dft_add
75 + vmp_res
76 + vmp_xai
77 + (vmp
78 | (acc_big
79 + (self
80 .vec_znx_big_normalize_tmp_bytes()
81 .max(self.vec_znx_idft_apply_tmp_bytes()))))
82 } else {
83 GLWE::bytes_of_from_infos(glwe_infos) + GLWE::external_product_tmp_bytes(self, glwe_infos, glwe_infos, brk_infos)
84 }
85 }
86
87 fn blind_rotation_execute<DR, DL, DB>(
88 &self,
89 res: &mut GLWE<DR>,
90 lwe: &LWE<DL>,
91 lut: &LookupTable,
92 brk: &BlindRotationKeyPrepared<DB, CGGI, BE>,
93 scratch: &mut Scratch<BE>,
94 ) where
95 DR: DataMut,
96 DL: DataRef,
97 DB: DataRef,
98 {
99 match brk.dist {
100 Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {
101 if lut.extension_factor() > 1 {
102 execute_block_binary_extended(self, res, lwe, lut, brk, scratch)
103 } else if brk.block_size() > 1 {
104 execute_block_binary(self, res, lwe, lut, brk, scratch);
105 } else {
106 execute_standard(self, res, lwe, lut, brk, scratch);
107 }
108 }
109 _ => panic!("invalid CGGI distribution (have you prepared the key?)"),
110 }
111 }
112}
113
114fn execute_block_binary_extended<DataRes, DataIn, DataBrk, M, BE: Backend>(
115 module: &M,
116 res: &mut GLWE<DataRes>,
117 lwe: &LWE<DataIn>,
118 lut: &LookupTable,
119 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, BE>,
120 scratch: &mut Scratch<BE>,
121) where
122 DataRes: DataMut,
123 DataIn: DataRef,
124 DataBrk: DataRef,
125 M: VecZnxDftBytesOf
126 + ModuleN
127 + VecZnxRotate
128 + VecZnxDftApply<BE>
129 + VecZnxDftZero<BE>
130 + VmpApplyDftToDft<BE>
131 + SvpApplyDftToDft<BE>
132 + VecZnxDftAddInplace<BE>
133 + VecZnxDftSubInplace<BE>
134 + VecZnxIdftApply<BE>
135 + VecZnxBigAddSmallInplace<BE>
136 + VecZnxBigNormalize<BE>
137 + VecZnxCopy
138 + VecZnxBigBytesOf,
139 Scratch<BE>: ScratchTakeCore<BE>,
140{
141 let n_glwe: usize = brk.n_glwe().into();
142 let extension_factor: usize = lut.extension_factor();
143 let base2k: usize = res.base2k().into();
144 let dnum: usize = brk.dnum().into();
145 let cols: usize = (res.rank() + 1).into();
146
147 let (mut acc, scratch_1) = scratch.take_vec_znx_slice(extension_factor, n_glwe, cols, res.size());
148 let (mut acc_dft, scratch_2) = scratch_1.take_vec_znx_dft_slice(module, extension_factor, cols, dnum);
149 let (mut vmp_res, scratch_3) = scratch_2.take_vec_znx_dft_slice(module, extension_factor, cols, brk.size());
150 let (mut acc_add_dft, scratch_4) = scratch_3.take_vec_znx_dft_slice(module, extension_factor, cols, brk.size());
151 let (mut vmp_xai, scratch_5) = scratch_4.take_vec_znx_dft(module, 1, brk.size());
152
153 (0..extension_factor).for_each(|i| {
154 acc[i].zero();
155 });
156
157 let x_pow_a: &Vec<SvpPPol<Vec<u8>, BE>>;
158 if let Some(b) = &brk.x_pow_a {
159 x_pow_a = b
160 } else {
161 panic!("invalid key: x_pow_a has not been initialized")
162 }
163
164 let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).as_usize()]; // TODO: from scratch space
165 let lwe_ref: LWE<&[u8]> = lwe.to_ref();
166
167 let two_n: usize = 2 * n_glwe;
168 let two_n_ext: usize = 2 * lut.domain_size();
169
170 mod_switch_2n(two_n_ext, &mut lwe_2n, &lwe_ref, lut.rotation_direction());
171
172 let a: &[i64] = &lwe_2n[1..];
173 let b_pos: usize = ((lwe_2n[0] + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
174
175 let b_hi: usize = b_pos / extension_factor;
176 let b_lo: usize = b_pos & (extension_factor - 1);
177
178 for (i, j) in (0..b_lo).zip(extension_factor - b_lo..extension_factor) {
179 module.vec_znx_rotate(b_hi as i64 + 1, &mut acc[i], 0, &lut.data[j], 0);
180 }
181 for (i, j) in (b_lo..extension_factor).zip(0..extension_factor - b_lo) {
182 module.vec_znx_rotate(b_hi as i64, &mut acc[i], 0, &lut.data[j], 0);
183 }
184
185 let block_size: usize = brk.block_size();
186
187 izip!(a.chunks_exact(block_size), brk.data.chunks_exact(block_size)).for_each(|(ai, ski)| {
188 for i in 0..extension_factor {
189 for j in 0..cols {
190 module.vec_znx_dft_apply(1, 0, &mut acc_dft[i], j, &acc[i], j);
191 module.vec_znx_dft_zero(&mut acc_add_dft[i], j)
192 }
193 }
194
195 // TODO: first & last iterations can be optimized
196 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
197 let ai_pos: usize = ((aii + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
198 let ai_hi: usize = ai_pos / extension_factor;
199 let ai_lo: usize = ai_pos & (extension_factor - 1);
200
201 // vmp_res = DFT(acc) * BRK[i]
202 (0..extension_factor).for_each(|i| {
203 module.vmp_apply_dft_to_dft(&mut vmp_res[i], &acc_dft[i], skii.data(), scratch_5);
204 });
205
206 // Trivial case: no rotation between polynomials, we can directly multiply with (X^{-ai} - 1)
207 if ai_lo == 0 {
208 // Sets acc_add_dft[i] = (acc[i] * sk) * X^{-ai} - (acc[i] * sk)
209 if ai_hi != 0 {
210 // DFT X^{-ai}
211 (0..extension_factor).for_each(|j| {
212 (0..cols).for_each(|i| {
213 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], i);
214 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j], i, &vmp_xai, 0);
215 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[j], i, &vmp_res[j], i);
216 });
217 });
218 }
219
220 // Non trivial case: rotation between polynomials
221 // In this case we can't directly multiply with (X^{-ai} - 1) because of the
222 // ring homomorphism R^{N} -> prod R^{N/extension_factor}, so we split the
223 // computation in two steps: acc_add_dft = (acc * sk) * (-1) + (acc * sk) * X^{-ai}
224 } else {
225 // Sets acc_add_dft[0..ai_lo] += (acc[extension_factor - ai_lo..extension_factor] * sk) * X^{-ai+1}
226 if (ai_hi + 1) & (two_n - 1) != 0 {
227 for (i, j) in (0..ai_lo).zip(extension_factor - ai_lo..extension_factor) {
228 (0..cols).for_each(|k| {
229 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi + 1], 0, &vmp_res[j], k);
230 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
231 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
232 });
233 }
234 }
235
236 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
237 if ai_hi != 0 {
238 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
239 for (i, j) in (ai_lo..extension_factor).zip(0..extension_factor - ai_lo) {
240 (0..cols).for_each(|k| {
241 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], k);
242 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
243 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
244 });
245 }
246 }
247 }
248 });
249
250 {
251 let (mut acc_add_big, scratch7) = scratch_5.take_vec_znx_big(module, 1, brk.size());
252
253 (0..extension_factor).for_each(|j| {
254 (0..cols).for_each(|i| {
255 module.vec_znx_idft_apply(&mut acc_add_big, 0, &acc_add_dft[j], i, scratch7);
256 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &acc[j], i);
257 module.vec_znx_big_normalize(&mut acc[j], base2k, 0, i, &acc_add_big, base2k, 0, scratch7);
258 });
259 });
260 }
261 });
262
263 (0..cols).for_each(|i| {
264 module.vec_znx_copy(res.data_mut(), i, &acc[0], i);
265 });
266}
267
268fn execute_block_binary<DataRes, DataIn, DataBrk, M, BE: Backend>(
269 module: &M,
270 res: &mut GLWE<DataRes>,
271 lwe: &LWE<DataIn>,
272 lut: &LookupTable,
273 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, BE>,
274 scratch: &mut Scratch<BE>,
275) where
276 DataRes: DataMut,
277 DataIn: DataRef,
278 DataBrk: DataRef,
279 M: VecZnxDftBytesOf
280 + ModuleN
281 + VecZnxRotate
282 + VecZnxDftApply<BE>
283 + VecZnxDftZero<BE>
284 + VmpApplyDftToDft<BE>
285 + SvpApplyDftToDft<BE>
286 + VecZnxDftAddInplace<BE>
287 + VecZnxDftSubInplace<BE>
288 + VecZnxIdftApply<BE>
289 + VecZnxBigAddSmallInplace<BE>
290 + VecZnxBigNormalize<BE>
291 + VecZnxCopy
292 + VecZnxBigBytesOf,
293 Scratch<BE>: ScratchTakeCore<BE>,
294{
295 let n_glwe: usize = brk.n_glwe().into();
296 let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).into()]; // TODO: from scratch space
297 let mut out_mut: GLWE<&mut [u8]> = res.to_mut();
298 let lwe_ref: LWE<&[u8]> = lwe.to_ref();
299 let two_n: usize = n_glwe << 1;
300 let base2k: usize = brk.base2k().into();
301 let dnum: usize = brk.dnum().into();
302
303 let cols: usize = (out_mut.rank() + 1).into();
304
305 mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe_ref, lut.rotation_direction());
306
307 let a: &[i64] = &lwe_2n[1..];
308 let b: i64 = lwe_2n[0];
309
310 out_mut.data_mut().zero();
311
312 // Initialize out to X^{b} * LUT(X)
313 module.vec_znx_rotate(b, out_mut.data_mut(), 0, &lut.data[0], 0);
314
315 let block_size: usize = brk.block_size();
316
317 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
318
319 let (mut acc_dft, scratch_1) = scratch.take_vec_znx_dft(module, cols, dnum);
320 let (mut vmp_res, scratch_2) = scratch_1.take_vec_znx_dft(module, cols, brk.size());
321 let (mut acc_add_dft, scratch_3) = scratch_2.take_vec_znx_dft(module, cols, brk.size());
322 let (mut vmp_xai, scratch_4) = scratch_3.take_vec_znx_dft(module, 1, brk.size());
323
324 let x_pow_a: &Vec<SvpPPol<Vec<u8>, BE>>;
325 if let Some(b) = &brk.x_pow_a {
326 x_pow_a = b
327 } else {
328 panic!("invalid key: x_pow_a has not been initialized")
329 }
330
331 izip!(a.chunks_exact(block_size), brk.data.chunks_exact(block_size)).for_each(|(ai, ski)| {
332 for j in 0..cols {
333 module.vec_znx_dft_apply(1, 0, &mut acc_dft, j, out_mut.data_mut(), j);
334 module.vec_znx_dft_zero(&mut acc_add_dft, j)
335 }
336
337 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
338 let ai_pos: usize = ((aii + two_n as i64) & (two_n - 1) as i64) as usize;
339
340 // vmp_res = DFT(acc) * BRK[i]
341 module.vmp_apply_dft_to_dft(&mut vmp_res, &acc_dft, skii.data(), scratch_4);
342
343 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
344 (0..cols).for_each(|i| {
345 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_pos], 0, &vmp_res, i);
346 module.vec_znx_dft_add_inplace(&mut acc_add_dft, i, &vmp_xai, 0);
347 module.vec_znx_dft_sub_inplace(&mut acc_add_dft, i, &vmp_res, i);
348 });
349 });
350
351 {
352 let (mut acc_add_big, scratch_5) = scratch_4.take_vec_znx_big(module, 1, brk.size());
353
354 (0..cols).for_each(|i| {
355 module.vec_znx_idft_apply(&mut acc_add_big, 0, &acc_add_dft, i, scratch_5);
356 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, out_mut.data_mut(), i);
357 module.vec_znx_big_normalize(out_mut.data_mut(), base2k, 0, i, &acc_add_big, base2k, 0, scratch_5);
358 });
359 }
360 });
361}
362
363fn execute_standard<DataRes, DataIn, DataBrk, M, BE: Backend>(
364 module: &M,
365 res: &mut GLWE<DataRes>,
366 lwe: &LWE<DataIn>,
367 lut: &LookupTable,
368 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, BE>,
369 scratch: &mut Scratch<BE>,
370) where
371 DataRes: DataMut,
372 DataIn: DataRef,
373 DataBrk: DataRef,
374 M: VecZnxRotate + GLWEExternalProduct<BE> + GLWEMulXpMinusOne<BE> + GLWEAdd + GLWENormalize<BE>,
375 Scratch<BE>: ScratchTakeCore<BE>,
376{
377 #[cfg(debug_assertions)]
378 {
379 assert_eq!(res.n(), brk.n(), "res.n(): {} != brk.n(): {}", res.n(), brk.n());
380 assert_eq!(
381 lut.domain_size(),
382 brk.n_glwe().as_usize(),
383 "lut.n(): {} != brk.n(): {}",
384 lut.domain_size(),
385 brk.n_glwe().as_usize()
386 );
387 assert_eq!(
388 res.rank(),
389 brk.rank(),
390 "res.rank(): {} != brk.rank(): {}",
391 res.rank(),
392 brk.rank()
393 );
394 assert_eq!(
395 lwe.n(),
396 brk.n_lwe(),
397 "lwe.n(): {} != brk.data.len(): {}",
398 lwe.n(),
399 brk.n_lwe()
400 );
401 }
402
403 let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).into()]; // TODO: from scratch space
404 let mut out_mut: GLWE<&mut [u8]> = res.to_mut();
405 let lwe_ref: LWE<&[u8]> = lwe.to_ref();
406
407 mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe_ref, lut.rotation_direction());
408
409 let a: &[i64] = &lwe_2n[1..];
410 let b: i64 = lwe_2n[0];
411
412 out_mut.data_mut().zero();
413
414 // Initialize out to X^{b} * LUT(X)
415 module.vec_znx_rotate(b, out_mut.data_mut(), 0, &lut.data[0], 0);
416
417 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
418 let (mut acc_tmp, scratch_1) = scratch.take_glwe(&out_mut);
419
420 // TODO: see if faster by skipping normalization in external product and keeping acc in big coeffs
421 // TODO: first iteration can be optimized to be a gglwe product
422 izip!(a.iter(), brk.data.iter()).for_each(|(ai, ski)| {
423 // acc_tmp = sk[i] * acc
424 acc_tmp.external_product(module, &out_mut, ski, scratch_1);
425
426 // acc_tmp = (sk[i] * acc) * (X^{ai} - 1)
427 module.glwe_mul_xp_minus_one_inplace(*ai, &mut acc_tmp, scratch_1);
428
429 // acc = acc + (sk[i] * acc) * (X^{ai} - 1)
430 module.glwe_add_inplace(&mut out_mut, &acc_tmp);
431 });
432
433 // We can normalize only at the end because we add normalized values in [-2^{base2k-1}, 2^{base2k-1}]
434 // on top of each others, thus ~ 2^{63-base2k} additions are supported before overflow.
435 module.glwe_normalize_inplace(&mut out_mut, scratch_1);
436}