poulpy/poulpy-schemes/src/tfhe/blind_rotation/algorithms/cggi/algorithm.rs
poulpy-schemes/src/tfhe/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::tfhe::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!(
188 a.chunks_exact(block_size),
189 brk.data.chunks_exact(block_size)
190 )
191 .for_each(|(ai, ski)| {
192 (0..extension_factor).for_each(|i| {
193 (0..cols).for_each(|j| {
194 module.vec_znx_dft_apply(1, 0, &mut acc_dft[i], j, &acc[i], j);
195 });
196 module.vec_znx_dft_zero(&mut acc_add_dft[i])
197 });
198
199 // TODO: first & last iterations can be optimized
200 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
201 let ai_pos: usize = ((aii + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
202 let ai_hi: usize = ai_pos / extension_factor;
203 let ai_lo: usize = ai_pos & (extension_factor - 1);
204
205 // vmp_res = DFT(acc) * BRK[i]
206 (0..extension_factor).for_each(|i| {
207 module.vmp_apply_dft_to_dft(&mut vmp_res[i], &acc_dft[i], skii.data(), scratch_5);
208 });
209
210 // Trivial case: no rotation between polynomials, we can directly multiply with (X^{-ai} - 1)
211 if ai_lo == 0 {
212 // Sets acc_add_dft[i] = (acc[i] * sk) * X^{-ai} - (acc[i] * sk)
213 if ai_hi != 0 {
214 // DFT X^{-ai}
215 (0..extension_factor).for_each(|j| {
216 (0..cols).for_each(|i| {
217 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], i);
218 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j], i, &vmp_xai, 0);
219 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[j], i, &vmp_res[j], i);
220 });
221 });
222 }
223
224 // Non trivial case: rotation between polynomials
225 // In this case we can't directly multiply with (X^{-ai} - 1) because of the
226 // ring homomorphism R^{N} -> prod R^{N/extension_factor}, so we split the
227 // computation in two steps: acc_add_dft = (acc * sk) * (-1) + (acc * sk) * X^{-ai}
228 } else {
229 // Sets acc_add_dft[0..ai_lo] += (acc[extension_factor - ai_lo..extension_factor] * sk) * X^{-ai+1}
230 if (ai_hi + 1) & (two_n - 1) != 0 {
231 for (i, j) in (0..ai_lo).zip(extension_factor - ai_lo..extension_factor) {
232 (0..cols).for_each(|k| {
233 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi + 1], 0, &vmp_res[j], k);
234 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
235 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
236 });
237 }
238 }
239
240 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
241 if ai_hi != 0 {
242 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
243 for (i, j) in (ai_lo..extension_factor).zip(0..extension_factor - ai_lo) {
244 (0..cols).for_each(|k| {
245 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], k);
246 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
247 module.vec_znx_dft_sub_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
248 });
249 }
250 }
251 }
252 });
253
254 {
255 let (mut acc_add_big, scratch7) = scratch_5.take_vec_znx_big(module, 1, brk.size());
256
257 (0..extension_factor).for_each(|j| {
258 (0..cols).for_each(|i| {
259 module.vec_znx_idft_apply(&mut acc_add_big, 0, &acc_add_dft[j], i, scratch7);
260 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &acc[j], i);
261 module.vec_znx_big_normalize(base2k, &mut acc[j], i, base2k, &acc_add_big, 0, scratch7);
262 });
263 });
264 }
265 });
266
267 (0..cols).for_each(|i| {
268 module.vec_znx_copy(res.data_mut(), i, &acc[0], i);
269 });
270}
271
272fn execute_block_binary<DataRes, DataIn, DataBrk, M, BE: Backend>(
273 module: &M,
274 res: &mut GLWE<DataRes>,
275 lwe: &LWE<DataIn>,
276 lut: &LookupTable,
277 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, BE>,
278 scratch: &mut Scratch<BE>,
279) where
280 DataRes: DataMut,
281 DataIn: DataRef,
282 DataBrk: DataRef,
283 M: VecZnxDftBytesOf
284 + ModuleN
285 + VecZnxRotate
286 + VecZnxDftApply<BE>
287 + VecZnxDftZero<BE>
288 + VmpApplyDftToDft<BE>
289 + SvpApplyDftToDft<BE>
290 + VecZnxDftAddInplace<BE>
291 + VecZnxDftSubInplace<BE>
292 + VecZnxIdftApply<BE>
293 + VecZnxBigAddSmallInplace<BE>
294 + VecZnxBigNormalize<BE>
295 + VecZnxCopy
296 + VecZnxBigBytesOf,
297 Scratch<BE>: ScratchTakeCore<BE>,
298{
299 let n_glwe: usize = brk.n_glwe().into();
300 let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).into()]; // TODO: from scratch space
301 let mut out_mut: GLWE<&mut [u8]> = res.to_mut();
302 let lwe_ref: LWE<&[u8]> = lwe.to_ref();
303 let two_n: usize = n_glwe << 1;
304 let base2k: usize = brk.base2k().into();
305 let dnum: usize = brk.dnum().into();
306
307 let cols: usize = (out_mut.rank() + 1).into();
308
309 mod_switch_2n(
310 2 * lut.domain_size(),
311 &mut lwe_2n,
312 &lwe_ref,
313 lut.rotation_direction(),
314 );
315
316 let a: &[i64] = &lwe_2n[1..];
317 let b: i64 = lwe_2n[0];
318
319 out_mut.data_mut().zero();
320
321 // Initialize out to X^{b} * LUT(X)
322 module.vec_znx_rotate(b, out_mut.data_mut(), 0, &lut.data[0], 0);
323
324 let block_size: usize = brk.block_size();
325
326 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
327
328 let (mut acc_dft, scratch_1) = scratch.take_vec_znx_dft(module, cols, dnum);
329 let (mut vmp_res, scratch_2) = scratch_1.take_vec_znx_dft(module, cols, brk.size());
330 let (mut acc_add_dft, scratch_3) = scratch_2.take_vec_znx_dft(module, cols, brk.size());
331 let (mut vmp_xai, scratch_4) = scratch_3.take_vec_znx_dft(module, 1, brk.size());
332
333 let x_pow_a: &Vec<SvpPPol<Vec<u8>, BE>>;
334 if let Some(b) = &brk.x_pow_a {
335 x_pow_a = b
336 } else {
337 panic!("invalid key: x_pow_a has not been initialized")
338 }
339
340 izip!(
341 a.chunks_exact(block_size),
342 brk.data.chunks_exact(block_size)
343 )
344 .for_each(|(ai, ski)| {
345 (0..cols).for_each(|j| {
346 module.vec_znx_dft_apply(1, 0, &mut acc_dft, j, out_mut.data_mut(), j);
347 });
348
349 module.vec_znx_dft_zero(&mut acc_add_dft);
350
351 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
352 let ai_pos: usize = ((aii + two_n as i64) & (two_n - 1) as i64) as usize;
353
354 // vmp_res = DFT(acc) * BRK[i]
355 module.vmp_apply_dft_to_dft(&mut vmp_res, &acc_dft, skii.data(), scratch_4);
356
357 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
358 (0..cols).for_each(|i| {
359 module.svp_apply_dft_to_dft(&mut vmp_xai, 0, &x_pow_a[ai_pos], 0, &vmp_res, i);
360 module.vec_znx_dft_add_inplace(&mut acc_add_dft, i, &vmp_xai, 0);
361 module.vec_znx_dft_sub_inplace(&mut acc_add_dft, i, &vmp_res, i);
362 });
363 });
364
365 {
366 let (mut acc_add_big, scratch_5) = scratch_4.take_vec_znx_big(module, 1, brk.size());
367
368 (0..cols).for_each(|i| {
369 module.vec_znx_idft_apply(&mut acc_add_big, 0, &acc_add_dft, i, scratch_5);
370 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, out_mut.data_mut(), i);
371 module.vec_znx_big_normalize(
372 base2k,
373 out_mut.data_mut(),
374 i,
375 base2k,
376 &acc_add_big,
377 0,
378 scratch_5,
379 );
380 });
381 }
382 });
383}
384
385fn execute_standard<DataRes, DataIn, DataBrk, M, BE: Backend>(
386 module: &M,
387 res: &mut GLWE<DataRes>,
388 lwe: &LWE<DataIn>,
389 lut: &LookupTable,
390 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, BE>,
391 scratch: &mut Scratch<BE>,
392) where
393 DataRes: DataMut,
394 DataIn: DataRef,
395 DataBrk: DataRef,
396 M: VecZnxRotate + GLWEExternalProduct<BE> + GLWEMulXpMinusOne<BE> + GLWEAdd + GLWENormalize<BE>,
397 Scratch<BE>: ScratchTakeCore<BE>,
398{
399 #[cfg(debug_assertions)]
400 {
401 assert_eq!(
402 res.n(),
403 brk.n(),
404 "res.n(): {} != brk.n(): {}",
405 res.n(),
406 brk.n()
407 );
408 assert_eq!(
409 lut.domain_size(),
410 brk.n_glwe().as_usize(),
411 "lut.n(): {} != brk.n(): {}",
412 lut.domain_size(),
413 brk.n_glwe().as_usize()
414 );
415 assert_eq!(
416 res.rank(),
417 brk.rank(),
418 "res.rank(): {} != brk.rank(): {}",
419 res.rank(),
420 brk.rank()
421 );
422 assert_eq!(
423 lwe.n(),
424 brk.n_lwe(),
425 "lwe.n(): {} != brk.data.len(): {}",
426 lwe.n(),
427 brk.n_lwe()
428 );
429 }
430
431 let mut lwe_2n: Vec<i64> = vec![0i64; (lwe.n() + 1).into()]; // TODO: from scratch space
432 let mut out_mut: GLWE<&mut [u8]> = res.to_mut();
433 let lwe_ref: LWE<&[u8]> = lwe.to_ref();
434
435 mod_switch_2n(
436 2 * lut.domain_size(),
437 &mut lwe_2n,
438 &lwe_ref,
439 lut.rotation_direction(),
440 );
441
442 let a: &[i64] = &lwe_2n[1..];
443 let b: i64 = lwe_2n[0];
444
445 out_mut.data_mut().zero();
446
447 // Initialize out to X^{b} * LUT(X)
448 module.vec_znx_rotate(b, out_mut.data_mut(), 0, &lut.data[0], 0);
449
450 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
451 let (mut acc_tmp, scratch_1) = scratch.take_glwe(&out_mut);
452
453 // TODO: see if faster by skipping normalization in external product and keeping acc in big coeffs
454 // TODO: first iteration can be optimized to be a gglwe product
455 izip!(a.iter(), brk.data.iter()).for_each(|(ai, ski)| {
456 // acc_tmp = sk[i] * acc
457 acc_tmp.external_product(module, &out_mut, ski, scratch_1);
458
459 // acc_tmp = (sk[i] * acc) * (X^{ai} - 1)
460 module.glwe_mul_xp_minus_one_inplace(*ai, &mut acc_tmp, scratch_1);
461
462 // acc = acc + (sk[i] * acc) * (X^{ai} - 1)
463 module.glwe_add_inplace(&mut out_mut, &acc_tmp);
464 });
465
466 // We can normalize only at the end because we add normalized values in [-2^{base2k-1}, 2^{base2k-1}]
467 // on top of each others, thus ~ 2^{63-base2k} additions are supported before overflow.
468 module.glwe_normalize_inplace(&mut out_mut, scratch_1);
469}