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