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