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