poulpy/poulpy-schemes/src/tfhe/blind_rotation/cggi_algo.rs
poulpy-schemes/src/tfhe/blind_rotation/cggi_algo.rsRaw
1use itertools::izip;
2use poulpy_hal::{
3 api::{
4 ScratchAvailable, SvpApply, SvpPPolAllocBytes, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, TakeVecZnxDftSlice,
5 TakeVecZnxSlice, VecZnxAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize,
6 VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAdd, VecZnxDftAddInplace, VecZnxDftAllocBytes, VecZnxDftFromVecZnx,
7 VecZnxDftSubABInplace, VecZnxDftToVecZnxBig, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpBytes, VecZnxDftZero,
8 VecZnxMulXpMinusOneInplace, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate,
9 VecZnxSubABInplace, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, ZnxView, ZnxZero,
10 },
11 layouts::{Backend, DataMut, DataRef, Module, Scratch, SvpPPol, VecZnx},
12};
13
14use poulpy_core::{
15 Distribution, GLWEOperations, TakeGLWECt,
16 layouts::{GLWECiphertext, GLWECiphertextToMut, Infos, LWECiphertext, LWECiphertextToRef},
17};
18
19use crate::tfhe::blind_rotation::{
20 BlincRotationExecute, BlindRotationKeyPrepared, CGGI, LookUpTable, LookUpTableRotationDirection,
21};
22
23#[allow(clippy::too_many_arguments)]
24pub fn cggi_blind_rotate_scratch_space<B: Backend>(
25 module: &Module<B>,
26 n: usize,
27 block_size: usize,
28 extension_factor: usize,
29 basek: usize,
30 k_res: usize,
31 k_brk: usize,
32 rows: usize,
33 rank: usize,
34) -> usize
35where
36 Module<B>: VecZnxDftAllocBytes
37 + VmpApplyTmpBytes
38 + VecZnxNormalizeTmpBytes
39 + VecZnxBigAllocBytes
40 + VecZnxDftToVecZnxBigTmpBytes
41 + VecZnxBigNormalizeTmpBytes,
42{
43 let brk_size: usize = k_brk.div_ceil(basek);
44
45 if block_size > 1 {
46 let cols: usize = rank + 1;
47 let acc_dft: usize = module.vec_znx_dft_alloc_bytes(n, cols, rows) * extension_factor;
48 let acc_big: usize = module.vec_znx_big_alloc_bytes(n, 1, brk_size);
49 let vmp_res: usize = module.vec_znx_dft_alloc_bytes(n, cols, brk_size) * extension_factor;
50 let vmp_xai: usize = module.vec_znx_dft_alloc_bytes(n, 1, brk_size);
51 let acc_dft_add: usize = vmp_res;
52 let vmp: usize = module.vmp_apply_tmp_bytes(n, brk_size, rows, rows, 2, 2, brk_size); // GGSW product: (1 x 2) x (2 x 2)
53 let acc: usize = if extension_factor > 1 {
54 VecZnx::alloc_bytes(n, cols, k_res.div_ceil(basek)) * extension_factor
55 } else {
56 0
57 };
58
59 acc + acc_dft
60 + acc_dft_add
61 + vmp_res
62 + vmp_xai
63 + (vmp | (acc_big + (module.vec_znx_big_normalize_tmp_bytes(n) | module.vec_znx_dft_to_vec_znx_big_tmp_bytes(n))))
64 } else {
65 GLWECiphertext::bytes_of(n, basek, k_res, rank)
66 + GLWECiphertext::external_product_scratch_space(module, n, basek, k_res, k_res, k_brk, 1, rank)
67 }
68}
69
70impl<D: DataRef, B: Backend> BlincRotationExecute<B> for BlindRotationKeyPrepared<D, CGGI, B>
71where
72 Module<B>: VecZnxBigAllocBytes
73 + VecZnxDftAllocBytes
74 + SvpPPolAllocBytes
75 + VmpApplyTmpBytes
76 + VecZnxBigNormalizeTmpBytes
77 + VecZnxDftToVecZnxBigTmpBytes
78 + VecZnxDftToVecZnxBig<B>
79 + VecZnxDftAdd<B>
80 + VecZnxDftAddInplace<B>
81 + VecZnxDftFromVecZnx<B>
82 + VecZnxDftZero<B>
83 + SvpApply<B>
84 + VecZnxDftSubABInplace<B>
85 + VecZnxBigAddSmallInplace<B>
86 + VecZnxRotate
87 + VecZnxAddInplace
88 + VecZnxSubABInplace
89 + VecZnxNormalize<B>
90 + VecZnxNormalizeInplace<B>
91 + VecZnxCopy
92 + VecZnxMulXpMinusOneInplace
93 + VmpApply<B>
94 + VmpApplyAdd<B>
95 + VecZnxDftToVecZnxBigConsume<B>
96 + VecZnxBigNormalize<B>
97 + VecZnxNormalizeTmpBytes,
98 Scratch<B>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnxSlice + TakeVecZnx + ScratchAvailable,
99{
100 fn execute<DR: DataMut, DI: DataRef>(
101 &self,
102 module: &Module<B>,
103 res: &mut GLWECiphertext<DR>,
104 lwe: &LWECiphertext<DI>,
105 lut: &LookUpTable,
106 scratch: &mut Scratch<B>,
107 ) {
108 match self.dist {
109 Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) | Distribution::ZERO => {
110 if lut.extension_factor() > 1 {
111 execute_block_binary_extended(module, res, lwe, lut, self, scratch)
112 } else if self.block_size() > 1 {
113 execute_block_binary(module, res, lwe, lut, self, scratch);
114 } else {
115 execute_standard(module, res, lwe, lut, self, scratch);
116 }
117 }
118 _ => panic!("invalid CGGI distribution"),
119 }
120 }
121}
122
123fn execute_block_binary_extended<DataRes, DataIn, DataBrk, B: Backend>(
124 module: &Module<B>,
125 res: &mut GLWECiphertext<DataRes>,
126 lwe: &LWECiphertext<DataIn>,
127 lut: &LookUpTable,
128 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, B>,
129 scratch: &mut Scratch<B>,
130) where
131 DataRes: DataMut,
132 DataIn: DataRef,
133 DataBrk: DataRef,
134 Module<B>: VecZnxBigAllocBytes
135 + VecZnxDftAllocBytes
136 + SvpPPolAllocBytes
137 + VmpApplyTmpBytes
138 + VecZnxBigNormalizeTmpBytes
139 + VecZnxDftToVecZnxBigTmpBytes
140 + VecZnxDftToVecZnxBig<B>
141 + VecZnxDftAdd<B>
142 + VecZnxDftAddInplace<B>
143 + VecZnxDftFromVecZnx<B>
144 + VecZnxDftZero<B>
145 + SvpApply<B>
146 + VecZnxDftSubABInplace<B>
147 + VecZnxBigAddSmallInplace<B>
148 + VecZnxRotate
149 + VecZnxAddInplace
150 + VecZnxSubABInplace
151 + VecZnxNormalize<B>
152 + VecZnxNormalizeInplace<B>
153 + VecZnxCopy
154 + VecZnxMulXpMinusOneInplace
155 + VecZnxBigNormalize<B>
156 + VmpApply<B>,
157 Scratch<B>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnxSlice + ScratchAvailable + TakeVecZnx,
158{
159 let n_glwe: usize = brk.n();
160 let extension_factor: usize = lut.extension_factor();
161 let basek: usize = res.basek();
162 let rows: usize = brk.rows();
163 let cols: usize = res.rank() + 1;
164
165 let (mut acc, scratch1) = scratch.take_vec_znx_slice(extension_factor, n_glwe, cols, res.size());
166 let (mut acc_dft, scratch2) = scratch1.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, rows);
167 let (mut vmp_res, scratch3) = scratch2.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, brk.size());
168 let (mut acc_add_dft, scratch4) = scratch3.take_vec_znx_dft_slice(extension_factor, n_glwe, cols, brk.size());
169 let (mut vmp_xai, scratch5) = scratch4.take_vec_znx_dft(n_glwe, 1, brk.size());
170
171 (0..extension_factor).for_each(|i| {
172 acc[i].zero();
173 });
174
175 let x_pow_a: &Vec<SvpPPol<Vec<u8>, B>>;
176 if let Some(b) = &brk.x_pow_a {
177 x_pow_a = b
178 } else {
179 panic!("invalid key: x_pow_a has not been initialized")
180 }
181
182 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
183 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
184
185 let two_n: usize = 2 * n_glwe;
186 let two_n_ext: usize = 2 * lut.domain_size();
187
188 mod_switch_2n(two_n_ext, &mut lwe_2n, &lwe_ref, lut.rotation_direction());
189
190 let a: &[i64] = &lwe_2n[1..];
191 let b_pos: usize = ((lwe_2n[0] + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
192
193 let b_hi: usize = b_pos / extension_factor;
194 let b_lo: usize = b_pos & (extension_factor - 1);
195
196 for (i, j) in (0..b_lo).zip(extension_factor - b_lo..extension_factor) {
197 module.vec_znx_rotate(b_hi as i64 + 1, &mut acc[i], 0, &lut.data[j], 0);
198 }
199 for (i, j) in (b_lo..extension_factor).zip(0..extension_factor - b_lo) {
200 module.vec_znx_rotate(b_hi as i64, &mut acc[i], 0, &lut.data[j], 0);
201 }
202
203 let block_size: usize = brk.block_size();
204
205 izip!(
206 a.chunks_exact(block_size),
207 brk.data.chunks_exact(block_size)
208 )
209 .for_each(|(ai, ski)| {
210 (0..extension_factor).for_each(|i| {
211 (0..cols).for_each(|j| {
212 module.vec_znx_dft_from_vec_znx(1, 0, &mut acc_dft[i], j, &acc[i], j);
213 });
214 module.vec_znx_dft_zero(&mut acc_add_dft[i])
215 });
216
217 // TODO: first & last iterations can be optimized
218 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
219 let ai_pos: usize = ((aii + two_n_ext as i64) & (two_n_ext - 1) as i64) as usize;
220 let ai_hi: usize = ai_pos / extension_factor;
221 let ai_lo: usize = ai_pos & (extension_factor - 1);
222
223 // vmp_res = DFT(acc) * BRK[i]
224 (0..extension_factor).for_each(|i| {
225 module.vmp_apply(&mut vmp_res[i], &acc_dft[i], skii.data(), scratch5);
226 });
227
228 // Trivial case: no rotation between polynomials, we can directly multiply with (X^{-ai} - 1)
229 if ai_lo == 0 {
230 // Sets acc_add_dft[i] = (acc[i] * sk) * X^{-ai} - (acc[i] * sk)
231 if ai_hi != 0 {
232 // DFT X^{-ai}
233 (0..extension_factor).for_each(|j| {
234 (0..cols).for_each(|i| {
235 module.svp_apply(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], i);
236 module.vec_znx_dft_add_inplace(&mut acc_add_dft[j], i, &vmp_xai, 0);
237 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[j], i, &vmp_res[j], i);
238 });
239 });
240 }
241
242 // Non trivial case: rotation between polynomials
243 // In this case we can't directly multiply with (X^{-ai} - 1) because of the
244 // ring homomorphism R^{N} -> prod R^{N/extension_factor}, so we split the
245 // computation in two steps: acc_add_dft = (acc * sk) * (-1) + (acc * sk) * X^{-ai}
246 } else {
247 // Sets acc_add_dft[0..ai_lo] += (acc[extension_factor - ai_lo..extension_factor] * sk) * X^{-ai+1}
248 if (ai_hi + 1) & (two_n - 1) != 0 {
249 for (i, j) in (0..ai_lo).zip(extension_factor - ai_lo..extension_factor) {
250 (0..cols).for_each(|k| {
251 module.svp_apply(&mut vmp_xai, 0, &x_pow_a[ai_hi + 1], 0, &vmp_res[j], k);
252 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
253 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
254 });
255 }
256 }
257
258 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
259 if ai_hi != 0 {
260 // Sets acc_add_dft[ai_lo..extension_factor] += (acc[0..extension_factor - ai_lo] * sk) * X^{-ai}
261 for (i, j) in (ai_lo..extension_factor).zip(0..extension_factor - ai_lo) {
262 (0..cols).for_each(|k| {
263 module.svp_apply(&mut vmp_xai, 0, &x_pow_a[ai_hi], 0, &vmp_res[j], k);
264 module.vec_znx_dft_add_inplace(&mut acc_add_dft[i], k, &vmp_xai, 0);
265 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft[i], k, &vmp_res[i], k);
266 });
267 }
268 }
269 }
270 });
271
272 {
273 let (mut acc_add_big, scratch7) = scratch5.take_vec_znx_big(n_glwe, 1, brk.size());
274
275 (0..extension_factor).for_each(|j| {
276 (0..cols).for_each(|i| {
277 module.vec_znx_dft_to_vec_znx_big(&mut acc_add_big, 0, &acc_add_dft[j], i, scratch7);
278 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &acc[j], i);
279 module.vec_znx_big_normalize(basek, &mut acc[j], i, &acc_add_big, 0, scratch7);
280 });
281 });
282 }
283 });
284
285 (0..cols).for_each(|i| {
286 module.vec_znx_copy(&mut res.data, i, &acc[0], i);
287 });
288}
289
290fn execute_block_binary<DataRes, DataIn, DataBrk, B: Backend>(
291 module: &Module<B>,
292 res: &mut GLWECiphertext<DataRes>,
293 lwe: &LWECiphertext<DataIn>,
294 lut: &LookUpTable,
295 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, B>,
296 scratch: &mut Scratch<B>,
297) where
298 DataRes: DataMut,
299 DataIn: DataRef,
300 DataBrk: DataRef,
301 Module<B>: VecZnxBigAllocBytes
302 + VecZnxDftAllocBytes
303 + SvpPPolAllocBytes
304 + VmpApplyTmpBytes
305 + VecZnxBigNormalizeTmpBytes
306 + VecZnxDftToVecZnxBigTmpBytes
307 + VecZnxDftToVecZnxBig<B>
308 + VecZnxDftAdd<B>
309 + VecZnxDftAddInplace<B>
310 + VecZnxDftFromVecZnx<B>
311 + VecZnxDftZero<B>
312 + SvpApply<B>
313 + VecZnxDftSubABInplace<B>
314 + VecZnxBigAddSmallInplace<B>
315 + VecZnxRotate
316 + VecZnxAddInplace
317 + VecZnxSubABInplace
318 + VecZnxNormalize<B>
319 + VecZnxNormalizeInplace<B>
320 + VecZnxCopy
321 + VecZnxMulXpMinusOneInplace
322 + VmpApply<B>
323 + VecZnxBigNormalize<B>,
324 Scratch<B>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnxSlice + ScratchAvailable + TakeVecZnx,
325{
326 let n_glwe: usize = brk.n();
327 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
328 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
329 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
330 let two_n: usize = n_glwe << 1;
331 let basek: usize = brk.basek();
332 let rows: usize = brk.rows();
333
334 let cols: usize = out_mut.rank() + 1;
335
336 mod_switch_2n(
337 2 * lut.domain_size(),
338 &mut lwe_2n,
339 &lwe_ref,
340 lut.rotation_direction(),
341 );
342
343 let a: &[i64] = &lwe_2n[1..];
344 let b: i64 = lwe_2n[0];
345
346 out_mut.data.zero();
347
348 // Initialize out to X^{b} * LUT(X)
349 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
350
351 let block_size: usize = brk.block_size();
352
353 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
354
355 let (mut acc_dft, scratch1) = scratch.take_vec_znx_dft(n_glwe, cols, rows);
356 let (mut vmp_res, scratch2) = scratch1.take_vec_znx_dft(n_glwe, cols, brk.size());
357 let (mut acc_add_dft, scratch3) = scratch2.take_vec_znx_dft(n_glwe, cols, brk.size());
358 let (mut vmp_xai, scratch4) = scratch3.take_vec_znx_dft(n_glwe, 1, brk.size());
359
360 let x_pow_a: &Vec<SvpPPol<Vec<u8>, B>>;
361 if let Some(b) = &brk.x_pow_a {
362 x_pow_a = b
363 } else {
364 panic!("invalid key: x_pow_a has not been initialized")
365 }
366
367 izip!(
368 a.chunks_exact(block_size),
369 brk.data.chunks_exact(block_size)
370 )
371 .for_each(|(ai, ski)| {
372 (0..cols).for_each(|j| {
373 module.vec_znx_dft_from_vec_znx(1, 0, &mut acc_dft, j, &out_mut.data, j);
374 });
375
376 module.vec_znx_dft_zero(&mut acc_add_dft);
377
378 izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
379 let ai_pos: usize = ((aii + two_n as i64) & (two_n - 1) as i64) as usize;
380
381 // vmp_res = DFT(acc) * BRK[i]
382 module.vmp_apply(&mut vmp_res, &acc_dft, skii.data(), scratch4);
383
384 // DFT(X^ai -1) * (DFT(acc) * BRK[i])
385 (0..cols).for_each(|i| {
386 module.svp_apply(&mut vmp_xai, 0, &x_pow_a[ai_pos], 0, &vmp_res, i);
387 module.vec_znx_dft_add_inplace(&mut acc_add_dft, i, &vmp_xai, 0);
388 module.vec_znx_dft_sub_ab_inplace(&mut acc_add_dft, i, &vmp_res, i);
389 });
390 });
391
392 {
393 let (mut acc_add_big, scratch5) = scratch4.take_vec_znx_big(n_glwe, 1, brk.size());
394
395 (0..cols).for_each(|i| {
396 module.vec_znx_dft_to_vec_znx_big(&mut acc_add_big, 0, &acc_add_dft, i, scratch5);
397 module.vec_znx_big_add_small_inplace(&mut acc_add_big, 0, &out_mut.data, i);
398 module.vec_znx_big_normalize(basek, &mut out_mut.data, i, &acc_add_big, 0, scratch5);
399 });
400 }
401 });
402}
403
404fn execute_standard<DataRes, DataIn, DataBrk, B: Backend>(
405 module: &Module<B>,
406 res: &mut GLWECiphertext<DataRes>,
407 lwe: &LWECiphertext<DataIn>,
408 lut: &LookUpTable,
409 brk: &BlindRotationKeyPrepared<DataBrk, CGGI, B>,
410 scratch: &mut Scratch<B>,
411) where
412 DataRes: DataMut,
413 DataIn: DataRef,
414 DataBrk: DataRef,
415 Module<B>: VecZnxBigAllocBytes
416 + VecZnxDftAllocBytes
417 + SvpPPolAllocBytes
418 + VmpApplyTmpBytes
419 + VecZnxBigNormalizeTmpBytes
420 + VecZnxDftToVecZnxBigTmpBytes
421 + VecZnxDftToVecZnxBig<B>
422 + VecZnxDftAdd<B>
423 + VecZnxDftAddInplace<B>
424 + VecZnxDftFromVecZnx<B>
425 + VecZnxDftZero<B>
426 + SvpApply<B>
427 + VecZnxDftSubABInplace<B>
428 + VecZnxBigAddSmallInplace<B>
429 + VecZnxRotate
430 + VecZnxAddInplace
431 + VecZnxSubABInplace
432 + VecZnxNormalize<B>
433 + VecZnxNormalizeInplace<B>
434 + VecZnxCopy
435 + VecZnxMulXpMinusOneInplace
436 + VmpApply<B>
437 + VmpApplyAdd<B>
438 + VecZnxDftToVecZnxBigConsume<B>
439 + VecZnxBigNormalize<B>
440 + VecZnxNormalizeTmpBytes,
441 Scratch<B>: TakeVecZnxDftSlice<B> + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnxSlice + ScratchAvailable + TakeVecZnx,
442{
443 #[cfg(debug_assertions)]
444 {
445 assert_eq!(
446 res.n(),
447 brk.n(),
448 "res.n(): {} != brk.n(): {}",
449 res.n(),
450 brk.n()
451 );
452 assert_eq!(
453 lut.domain_size(),
454 brk.n(),
455 "lut.n(): {} != brk.n(): {}",
456 lut.domain_size(),
457 brk.n()
458 );
459 assert_eq!(
460 res.rank(),
461 brk.rank(),
462 "res.rank(): {} != brk.rank(): {}",
463 res.rank(),
464 brk.rank()
465 );
466 assert_eq!(
467 lwe.n(),
468 brk.data.len(),
469 "lwe.n(): {} != brk.data.len(): {}",
470 lwe.n(),
471 brk.data.len()
472 );
473 }
474
475 let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
476 let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
477 let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
478 let basek: usize = brk.basek();
479
480 mod_switch_2n(
481 2 * lut.domain_size(),
482 &mut lwe_2n,
483 &lwe_ref,
484 lut.rotation_direction(),
485 );
486
487 let a: &[i64] = &lwe_2n[1..];
488 let b: i64 = lwe_2n[0];
489
490 out_mut.data.zero();
491
492 // Initialize out to X^{b} * LUT(X)
493 module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
494
495 // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
496 let (mut acc_tmp, scratch1) = scratch.take_glwe_ct(out_mut.n(), basek, out_mut.k(), out_mut.rank());
497
498 // TODO: see if faster by skipping normalization in external product and keeping acc in big coeffs
499 // TODO: first iteration can be optimized to be a gglwe product
500 izip!(a.iter(), brk.data.iter()).for_each(|(ai, ski)| {
501 // acc_tmp = sk[i] * acc
502 acc_tmp.external_product(module, &out_mut, ski, scratch1);
503
504 // acc_tmp = (sk[i] * acc) * (X^{ai} - 1)
505 acc_tmp.mul_xp_minus_one_inplace(module, *ai);
506
507 // acc = acc + (sk[i] * acc) * (X^{ai} - 1)
508 out_mut.add_inplace(module, &acc_tmp);
509 });
510
511 // We can normalize only at the end because we add normalized values in [-2^{basek-1}, 2^{basek-1}]
512 // on top of each others, thus ~ 2^{63-basek} additions are supported before overflow.
513 out_mut.normalize_inplace(module, scratch1);
514}
515
516pub fn mod_switch_2n(n: usize, res: &mut [i64], lwe: &LWECiphertext<&[u8]>, rot_dir: LookUpTableRotationDirection) {
517 let basek: usize = lwe.basek();
518
519 let log2n: usize = usize::BITS as usize - (n - 1).leading_zeros() as usize + 1;
520
521 res.copy_from_slice(lwe.data().at(0, 0));
522
523 match rot_dir {
524 LookUpTableRotationDirection::Left => {
525 res.iter_mut().for_each(|x| *x = -*x);
526 }
527 LookUpTableRotationDirection::Right => {}
528 }
529
530 if basek > log2n {
531 let diff: usize = basek - (log2n - 1); // additional -1 because we map to [-N/2, N/2) instead of [0, N)
532 res.iter_mut().for_each(|x| {
533 *x = div_round_by_pow2(x, diff);
534 })
535 } else {
536 let rem: usize = basek - (log2n % basek);
537 let size: usize = log2n.div_ceil(basek);
538 (1..size).for_each(|i| {
539 if i == size - 1 && rem != basek {
540 let k_rem: usize = basek - rem;
541 izip!(lwe.data().at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
542 *y = (*y << k_rem) + (x >> rem);
543 });
544 } else {
545 izip!(lwe.data().at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
546 *y = (*y << basek) + x;
547 });
548 }
549 })
550 }
551}
552
553#[inline(always)]
554fn div_round_by_pow2(x: &i64, k: usize) -> i64 {
555 (x + (1 << (k - 1))) >> k
556}