poulpy/backend/src/implementation/cpu_spqlios/vec_znx.rs
backend/src/implementation/cpu_spqlios/vec_znx.rsRaw
1use itertools::izip;
2use rand_distr::Normal;
3use sampling::source::Source;
4
5use crate::{
6 hal::{
7 api::{
8 TakeSlice, TakeVecZnx, VecZnxAddDistF64, VecZnxCopy, VecZnxFillDistF64, VecZnxNormalizeTmpBytes, VecZnxRotate,
9 VecZnxRotateInplace, VecZnxSwithcDegree, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
10 },
11 layouts::{Backend, Module, ScalarZnx, ScalarZnxToRef, Scratch, VecZnx, VecZnxToMut, VecZnxToRef},
12 oep::{
13 VecZnxAddDistF64Impl, VecZnxAddImpl, VecZnxAddInplaceImpl, VecZnxAddNormalImpl, VecZnxAddScalarInplaceImpl,
14 VecZnxAutomorphismImpl, VecZnxAutomorphismInplaceImpl, VecZnxCopyImpl, VecZnxFillDistF64Impl, VecZnxFillNormalImpl,
15 VecZnxFillUniformImpl, VecZnxLshInplaceImpl, VecZnxMergeImpl, VecZnxMulXpMinusOneImpl,
16 VecZnxMulXpMinusOneInplaceImpl, VecZnxNegateImpl, VecZnxNegateInplaceImpl, VecZnxNormalizeImpl,
17 VecZnxNormalizeInplaceImpl, VecZnxNormalizeTmpBytesImpl, VecZnxRotateImpl, VecZnxRotateInplaceImpl,
18 VecZnxRshInplaceImpl, VecZnxSplitImpl, VecZnxSubABInplaceImpl, VecZnxSubBAInplaceImpl, VecZnxSubImpl,
19 VecZnxSubScalarInplaceImpl, VecZnxSwithcDegreeImpl,
20 },
21 },
22 implementation::cpu_spqlios::{
23 CPUAVX,
24 ffi::{module::module_info_t, vec_znx, znx},
25 },
26};
27
28unsafe impl<B: Backend> VecZnxNormalizeTmpBytesImpl<B> for B
29where
30 B: CPUAVX,
31{
32 fn vec_znx_normalize_tmp_bytes_impl(module: &Module<B>, n: usize) -> usize {
33 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(module.ptr() as *const module_info_t, n as u64) as usize }
34 }
35}
36
37unsafe impl<B: Backend + CPUAVX> VecZnxNormalizeImpl<B> for B {
38 fn vec_znx_normalize_impl<R, A>(
39 module: &Module<B>,
40 basek: usize,
41 res: &mut R,
42 res_col: usize,
43 a: &A,
44 a_col: usize,
45 scratch: &mut Scratch<B>,
46 ) where
47 R: VecZnxToMut,
48 A: VecZnxToRef,
49 {
50 let a: VecZnx<&[u8]> = a.to_ref();
51 let mut res: VecZnx<&mut [u8]> = res.to_mut();
52
53 #[cfg(debug_assertions)]
54 {
55 assert_eq!(res.n(), a.n());
56 }
57
58 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes(a.n()));
59
60 unsafe {
61 vec_znx::vec_znx_normalize_base2k(
62 module.ptr() as *const module_info_t,
63 a.n() as u64,
64 basek as u64,
65 res.at_mut_ptr(res_col, 0),
66 res.size() as u64,
67 res.sl() as u64,
68 a.at_ptr(a_col, 0),
69 a.size() as u64,
70 a.sl() as u64,
71 tmp_bytes.as_mut_ptr(),
72 );
73 }
74 }
75}
76
77unsafe impl<B: Backend + CPUAVX> VecZnxNormalizeInplaceImpl<B> for B {
78 fn vec_znx_normalize_inplace_impl<A>(module: &Module<B>, basek: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
79 where
80 A: VecZnxToMut,
81 {
82 let mut a: VecZnx<&mut [u8]> = a.to_mut();
83
84 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes(a.n()));
85
86 unsafe {
87 vec_znx::vec_znx_normalize_base2k(
88 module.ptr() as *const module_info_t,
89 a.n() as u64,
90 basek as u64,
91 a.at_mut_ptr(a_col, 0),
92 a.size() as u64,
93 a.sl() as u64,
94 a.at_ptr(a_col, 0),
95 a.size() as u64,
96 a.sl() as u64,
97 tmp_bytes.as_mut_ptr(),
98 );
99 }
100 }
101}
102
103unsafe impl<B: Backend + CPUAVX> VecZnxAddImpl<B> for B {
104 fn vec_znx_add_impl<R, A, C>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
105 where
106 R: VecZnxToMut,
107 A: VecZnxToRef,
108 C: VecZnxToRef,
109 {
110 let a: VecZnx<&[u8]> = a.to_ref();
111 let b: VecZnx<&[u8]> = b.to_ref();
112 let mut res: VecZnx<&mut [u8]> = res.to_mut();
113
114 #[cfg(debug_assertions)]
115 {
116 assert_eq!(a.n(), res.n());
117 assert_eq!(b.n(), res.n());
118 assert_ne!(a.as_ptr(), b.as_ptr());
119 }
120 unsafe {
121 vec_znx::vec_znx_add(
122 module.ptr() as *const module_info_t,
123 res.at_mut_ptr(res_col, 0),
124 res.size() as u64,
125 res.sl() as u64,
126 a.at_ptr(a_col, 0),
127 a.size() as u64,
128 a.sl() as u64,
129 b.at_ptr(b_col, 0),
130 b.size() as u64,
131 b.sl() as u64,
132 )
133 }
134 }
135}
136
137unsafe impl<B: Backend + CPUAVX> VecZnxAddInplaceImpl<B> for B {
138 fn vec_znx_add_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
139 where
140 R: VecZnxToMut,
141 A: VecZnxToRef,
142 {
143 let a: VecZnx<&[u8]> = a.to_ref();
144 let mut res: VecZnx<&mut [u8]> = res.to_mut();
145
146 #[cfg(debug_assertions)]
147 {
148 assert_eq!(a.n(), res.n());
149 }
150 unsafe {
151 vec_znx::vec_znx_add(
152 module.ptr() as *const module_info_t,
153 res.at_mut_ptr(res_col, 0),
154 res.size() as u64,
155 res.sl() as u64,
156 a.at_ptr(a_col, 0),
157 a.size() as u64,
158 a.sl() as u64,
159 res.at_ptr(res_col, 0),
160 res.size() as u64,
161 res.sl() as u64,
162 )
163 }
164 }
165}
166
167unsafe impl<B: Backend + CPUAVX> VecZnxAddScalarInplaceImpl<B> for B {
168 fn vec_znx_add_scalar_inplace_impl<R, A>(
169 module: &Module<B>,
170 res: &mut R,
171 res_col: usize,
172 res_limb: usize,
173 a: &A,
174 a_col: usize,
175 ) where
176 R: VecZnxToMut,
177 A: ScalarZnxToRef,
178 {
179 let mut res: VecZnx<&mut [u8]> = res.to_mut();
180 let a: ScalarZnx<&[u8]> = a.to_ref();
181
182 #[cfg(debug_assertions)]
183 {
184 assert_eq!(a.n(), res.n());
185 }
186
187 unsafe {
188 vec_znx::vec_znx_add(
189 module.ptr() as *const module_info_t,
190 res.at_mut_ptr(res_col, res_limb),
191 1_u64,
192 res.sl() as u64,
193 a.at_ptr(a_col, 0),
194 a.size() as u64,
195 a.sl() as u64,
196 res.at_ptr(res_col, res_limb),
197 1_u64,
198 res.sl() as u64,
199 )
200 }
201 }
202}
203
204unsafe impl<B: Backend + CPUAVX> VecZnxSubImpl<B> for B {
205 fn vec_znx_sub_impl<R, A, C>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
206 where
207 R: VecZnxToMut,
208 A: VecZnxToRef,
209 C: VecZnxToRef,
210 {
211 let a: VecZnx<&[u8]> = a.to_ref();
212 let b: VecZnx<&[u8]> = b.to_ref();
213 let mut res: VecZnx<&mut [u8]> = res.to_mut();
214
215 #[cfg(debug_assertions)]
216 {
217 assert_eq!(a.n(), res.n());
218 assert_eq!(b.n(), res.n());
219 assert_ne!(a.as_ptr(), b.as_ptr());
220 }
221 unsafe {
222 vec_znx::vec_znx_sub(
223 module.ptr() as *const module_info_t,
224 res.at_mut_ptr(res_col, 0),
225 res.size() as u64,
226 res.sl() as u64,
227 a.at_ptr(a_col, 0),
228 a.size() as u64,
229 a.sl() as u64,
230 b.at_ptr(b_col, 0),
231 b.size() as u64,
232 b.sl() as u64,
233 )
234 }
235 }
236}
237
238unsafe impl<B: Backend + CPUAVX> VecZnxSubABInplaceImpl<B> for B {
239 fn vec_znx_sub_ab_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
240 where
241 R: VecZnxToMut,
242 A: VecZnxToRef,
243 {
244 let a: VecZnx<&[u8]> = a.to_ref();
245 let mut res: VecZnx<&mut [u8]> = res.to_mut();
246 #[cfg(debug_assertions)]
247 {
248 assert_eq!(a.n(), res.n());
249 }
250 unsafe {
251 vec_znx::vec_znx_sub(
252 module.ptr() as *const module_info_t,
253 res.at_mut_ptr(res_col, 0),
254 res.size() as u64,
255 res.sl() as u64,
256 res.at_ptr(res_col, 0),
257 res.size() as u64,
258 res.sl() as u64,
259 a.at_ptr(a_col, 0),
260 a.size() as u64,
261 a.sl() as u64,
262 )
263 }
264 }
265}
266
267unsafe impl<B: Backend + CPUAVX> VecZnxSubBAInplaceImpl<B> for B {
268 fn vec_znx_sub_ba_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
269 where
270 R: VecZnxToMut,
271 A: VecZnxToRef,
272 {
273 let a: VecZnx<&[u8]> = a.to_ref();
274 let mut res: VecZnx<&mut [u8]> = res.to_mut();
275 #[cfg(debug_assertions)]
276 {
277 assert_eq!(a.n(), res.n());
278 }
279 unsafe {
280 vec_znx::vec_znx_sub(
281 module.ptr() as *const module_info_t,
282 res.at_mut_ptr(res_col, 0),
283 res.size() as u64,
284 res.sl() as u64,
285 a.at_ptr(a_col, 0),
286 a.size() as u64,
287 a.sl() as u64,
288 res.at_ptr(res_col, 0),
289 res.size() as u64,
290 res.sl() as u64,
291 )
292 }
293 }
294}
295
296unsafe impl<B: Backend + CPUAVX> VecZnxSubScalarInplaceImpl<B> for B {
297 fn vec_znx_sub_scalar_inplace_impl<R, A>(
298 module: &Module<B>,
299 res: &mut R,
300 res_col: usize,
301 res_limb: usize,
302 a: &A,
303 a_col: usize,
304 ) where
305 R: VecZnxToMut,
306 A: ScalarZnxToRef,
307 {
308 let mut res: VecZnx<&mut [u8]> = res.to_mut();
309 let a: ScalarZnx<&[u8]> = a.to_ref();
310
311 #[cfg(debug_assertions)]
312 {
313 assert_eq!(a.n(), res.n());
314 }
315
316 unsafe {
317 vec_znx::vec_znx_sub(
318 module.ptr() as *const module_info_t,
319 res.at_mut_ptr(res_col, res_limb),
320 1_u64,
321 res.sl() as u64,
322 a.at_ptr(a_col, 0),
323 a.size() as u64,
324 a.sl() as u64,
325 res.at_ptr(res_col, res_limb),
326 1_u64,
327 res.sl() as u64,
328 )
329 }
330 }
331}
332
333unsafe impl<B: Backend + CPUAVX> VecZnxNegateImpl<B> for B {
334 fn vec_znx_negate_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
335 where
336 R: VecZnxToMut,
337 A: VecZnxToRef,
338 {
339 let a: VecZnx<&[u8]> = a.to_ref();
340 let mut res: VecZnx<&mut [u8]> = res.to_mut();
341 #[cfg(debug_assertions)]
342 {
343 assert_eq!(a.n(), res.n());
344 }
345 unsafe {
346 vec_znx::vec_znx_negate(
347 module.ptr() as *const module_info_t,
348 res.at_mut_ptr(res_col, 0),
349 res.size() as u64,
350 res.sl() as u64,
351 a.at_ptr(a_col, 0),
352 a.size() as u64,
353 a.sl() as u64,
354 )
355 }
356 }
357}
358
359unsafe impl<B: Backend + CPUAVX> VecZnxNegateInplaceImpl<B> for B {
360 fn vec_znx_negate_inplace_impl<A>(module: &Module<B>, a: &mut A, a_col: usize)
361 where
362 A: VecZnxToMut,
363 {
364 let mut a: VecZnx<&mut [u8]> = a.to_mut();
365 unsafe {
366 vec_znx::vec_znx_negate(
367 module.ptr() as *const module_info_t,
368 a.at_mut_ptr(a_col, 0),
369 a.size() as u64,
370 a.sl() as u64,
371 a.at_ptr(a_col, 0),
372 a.size() as u64,
373 a.sl() as u64,
374 )
375 }
376 }
377}
378
379unsafe impl<B: Backend + CPUAVX> VecZnxLshInplaceImpl<B> for B {
380 fn vec_znx_lsh_inplace_impl<A>(_module: &Module<B>, basek: usize, k: usize, a: &mut A)
381 where
382 A: VecZnxToMut,
383 {
384 vec_znx_lsh_inplace_ref(basek, k, a)
385 }
386}
387
388pub fn vec_znx_lsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
389where
390 A: VecZnxToMut,
391{
392 let mut a: VecZnx<&mut [u8]> = a.to_mut();
393
394 let n: usize = a.n();
395 let cols: usize = a.cols();
396 let size: usize = a.size();
397 let steps: usize = k / basek;
398
399 a.raw_mut().rotate_left(n * steps * cols);
400 (0..cols).for_each(|i| {
401 (size - steps..size).for_each(|j| {
402 a.zero_at(i, j);
403 })
404 });
405
406 let k_rem: usize = k % basek;
407
408 if k_rem != 0 {
409 let shift: usize = i64::BITS as usize - k_rem;
410 (0..cols).for_each(|i| {
411 (0..steps).for_each(|j| {
412 a.at_mut(i, j).iter_mut().for_each(|xi| {
413 *xi <<= shift;
414 });
415 });
416 });
417 }
418}
419
420unsafe impl<B: Backend + CPUAVX> VecZnxRshInplaceImpl<B> for B {
421 fn vec_znx_rsh_inplace_impl<A>(_module: &Module<B>, basek: usize, k: usize, a: &mut A)
422 where
423 A: VecZnxToMut,
424 {
425 vec_znx_rsh_inplace_ref(basek, k, a)
426 }
427}
428
429pub fn vec_znx_rsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
430where
431 A: VecZnxToMut,
432{
433 let mut a: VecZnx<&mut [u8]> = a.to_mut();
434 let n: usize = a.n();
435 let cols: usize = a.cols();
436 let size: usize = a.size();
437 let steps: usize = k / basek;
438
439 a.raw_mut().rotate_right(n * steps * cols);
440 (0..cols).for_each(|i| {
441 (0..steps).for_each(|j| {
442 a.zero_at(i, j);
443 })
444 });
445
446 let k_rem: usize = k % basek;
447
448 if k_rem != 0 {
449 let mut carry: Vec<i64> = vec![0i64; n]; // ALLOC (but small so OK)
450 let shift: usize = i64::BITS as usize - k_rem;
451 (0..cols).for_each(|i| {
452 carry.fill(0);
453 (steps..size).for_each(|j| {
454 izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
455 *xi += *ci << basek;
456 *ci = (*xi << shift) >> shift;
457 *xi = (*xi - *ci) >> k_rem;
458 });
459 });
460 })
461 }
462}
463
464unsafe impl<B: Backend + CPUAVX> VecZnxRotateImpl<B> for B {
465 fn vec_znx_rotate_impl<R, A>(_module: &Module<B>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
466 where
467 R: VecZnxToMut,
468 A: VecZnxToRef,
469 {
470 let a: VecZnx<&[u8]> = a.to_ref();
471 let mut res: VecZnx<&mut [u8]> = res.to_mut();
472 #[cfg(debug_assertions)]
473 {
474 assert_eq!(res.n(), a.n());
475 }
476 unsafe {
477 (0..a.size()).for_each(|j| {
478 znx::znx_rotate_i64(
479 a.n() as u64,
480 k,
481 res.at_mut_ptr(res_col, j),
482 a.at_ptr(a_col, j),
483 );
484 });
485 }
486 }
487}
488
489unsafe impl<B: Backend + CPUAVX> VecZnxRotateInplaceImpl<B> for B {
490 fn vec_znx_rotate_inplace_impl<A>(_module: &Module<B>, k: i64, a: &mut A, a_col: usize)
491 where
492 A: VecZnxToMut,
493 {
494 let mut a: VecZnx<&mut [u8]> = a.to_mut();
495 unsafe {
496 (0..a.size()).for_each(|j| {
497 znx::znx_rotate_inplace_i64(a.n() as u64, k, a.at_mut_ptr(a_col, j));
498 });
499 }
500 }
501}
502
503unsafe impl<B: Backend + CPUAVX> VecZnxAutomorphismImpl<B> for B {
504 fn vec_znx_automorphism_impl<R, A>(module: &Module<B>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
505 where
506 R: VecZnxToMut,
507 A: VecZnxToRef,
508 {
509 let a: VecZnx<&[u8]> = a.to_ref();
510 let mut res: VecZnx<&mut [u8]> = res.to_mut();
511 #[cfg(debug_assertions)]
512 {
513 assert_eq!(a.n(), res.n());
514 }
515 unsafe {
516 vec_znx::vec_znx_automorphism(
517 module.ptr() as *const module_info_t,
518 k,
519 res.at_mut_ptr(res_col, 0),
520 res.size() as u64,
521 res.sl() as u64,
522 a.at_ptr(a_col, 0),
523 a.size() as u64,
524 a.sl() as u64,
525 )
526 }
527 }
528}
529
530unsafe impl<B: Backend + CPUAVX> VecZnxAutomorphismInplaceImpl<B> for B {
531 fn vec_znx_automorphism_inplace_impl<A>(module: &Module<B>, k: i64, a: &mut A, a_col: usize)
532 where
533 A: VecZnxToMut,
534 {
535 let mut a: VecZnx<&mut [u8]> = a.to_mut();
536 #[cfg(debug_assertions)]
537 {
538 assert!(
539 k & 1 != 0,
540 "invalid galois element: must be odd but is {}",
541 k
542 );
543 }
544 unsafe {
545 vec_znx::vec_znx_automorphism(
546 module.ptr() as *const module_info_t,
547 k,
548 a.at_mut_ptr(a_col, 0),
549 a.size() as u64,
550 a.sl() as u64,
551 a.at_ptr(a_col, 0),
552 a.size() as u64,
553 a.sl() as u64,
554 )
555 }
556 }
557}
558
559unsafe impl<B: Backend + CPUAVX> VecZnxMulXpMinusOneImpl<B> for B {
560 fn vec_znx_mul_xp_minus_one_impl<R, A>(module: &Module<B>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
561 where
562 R: VecZnxToMut,
563 A: VecZnxToRef,
564 {
565 let a: VecZnx<&[u8]> = a.to_ref();
566 let mut res: VecZnx<&mut [u8]> = res.to_mut();
567 #[cfg(debug_assertions)]
568 {
569 assert_eq!(a.n(), res.n());
570 assert_eq!(res.n(), res.n());
571 }
572 unsafe {
573 vec_znx::vec_znx_mul_xp_minus_one(
574 module.ptr() as *const module_info_t,
575 p,
576 res.at_mut_ptr(res_col, 0),
577 res.size() as u64,
578 res.sl() as u64,
579 a.at_ptr(a_col, 0),
580 a.size() as u64,
581 a.sl() as u64,
582 )
583 }
584 }
585}
586
587unsafe impl<B: Backend + CPUAVX> VecZnxMulXpMinusOneInplaceImpl<B> for B {
588 fn vec_znx_mul_xp_minus_one_inplace_impl<R>(module: &Module<B>, p: i64, res: &mut R, res_col: usize)
589 where
590 R: VecZnxToMut,
591 {
592 let mut res: VecZnx<&mut [u8]> = res.to_mut();
593 #[cfg(debug_assertions)]
594 {
595 assert_eq!(res.n(), res.n());
596 }
597 unsafe {
598 vec_znx::vec_znx_mul_xp_minus_one(
599 module.ptr() as *const module_info_t,
600 p,
601 res.at_mut_ptr(res_col, 0),
602 res.size() as u64,
603 res.sl() as u64,
604 res.at_ptr(res_col, 0),
605 res.size() as u64,
606 res.sl() as u64,
607 )
608 }
609 }
610}
611
612unsafe impl<B: Backend + CPUAVX> VecZnxSplitImpl<B> for B {
613 fn vec_znx_split_impl<R, A>(module: &Module<B>, res: &mut [R], res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<B>)
614 where
615 R: VecZnxToMut,
616 A: VecZnxToRef,
617 {
618 vec_znx_split_ref(module, res, res_col, a, a_col, scratch)
619 }
620}
621
622pub fn vec_znx_split_ref<R, A, B>(
623 module: &Module<B>,
624 res: &mut [R],
625 res_col: usize,
626 a: &A,
627 a_col: usize,
628 scratch: &mut Scratch<B>,
629) where
630 B: Backend + CPUAVX,
631 R: VecZnxToMut,
632 A: VecZnxToRef,
633{
634 let a: VecZnx<&[u8]> = a.to_ref();
635
636 let (n_in, n_out) = (a.n(), res[0].to_mut().n());
637
638 let (mut buf, _) = scratch.take_vec_znx(n_in.max(n_out), 1, a.size());
639
640 debug_assert!(
641 n_out < n_in,
642 "invalid a: output ring degree should be smaller"
643 );
644 res[1..].iter_mut().for_each(|bi| {
645 debug_assert_eq!(
646 bi.to_mut().n(),
647 n_out,
648 "invalid input a: all VecZnx must have the same degree"
649 )
650 });
651
652 res.iter_mut().enumerate().for_each(|(i, bi)| {
653 if i == 0 {
654 module.vec_znx_switch_degree(bi, res_col, &a, a_col);
655 module.vec_znx_rotate(-1, &mut buf, 0, &a, a_col);
656 } else {
657 module.vec_znx_switch_degree(bi, res_col, &buf, a_col);
658 module.vec_znx_rotate_inplace(-1, &mut buf, a_col);
659 }
660 })
661}
662
663unsafe impl<B: Backend + CPUAVX> VecZnxMergeImpl<B> for B {
664 fn vec_znx_merge_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &[A], a_col: usize)
665 where
666 R: VecZnxToMut,
667 A: VecZnxToRef,
668 {
669 vec_znx_merge_ref(module, res, res_col, a, a_col)
670 }
671}
672
673pub fn vec_znx_merge_ref<R, A, B>(module: &Module<B>, res: &mut R, res_col: usize, a: &[A], a_col: usize)
674where
675 B: Backend + CPUAVX,
676 R: VecZnxToMut,
677 A: VecZnxToRef,
678{
679 let mut res: VecZnx<&mut [u8]> = res.to_mut();
680
681 let (n_in, n_out) = (res.n(), a[0].to_ref().n());
682
683 debug_assert!(
684 n_out < n_in,
685 "invalid a: output ring degree should be smaller"
686 );
687 a[1..].iter().for_each(|ai| {
688 debug_assert_eq!(
689 ai.to_ref().n(),
690 n_out,
691 "invalid input a: all VecZnx must have the same degree"
692 )
693 });
694
695 a.iter().for_each(|ai| {
696 module.vec_znx_switch_degree(&mut res, res_col, ai, a_col);
697 module.vec_znx_rotate_inplace(-1, &mut res, res_col);
698 });
699
700 module.vec_znx_rotate_inplace(a.len() as i64, &mut res, res_col);
701}
702
703unsafe impl<B: Backend + CPUAVX> VecZnxSwithcDegreeImpl<B> for B {
704 fn vec_znx_switch_degree_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
705 where
706 R: VecZnxToMut,
707 A: VecZnxToRef,
708 {
709 vec_znx_switch_degree_ref(module, res, res_col, a, a_col)
710 }
711}
712
713pub fn vec_znx_switch_degree_ref<R, A, B>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
714where
715 B: Backend + CPUAVX,
716 R: VecZnxToMut,
717 A: VecZnxToRef,
718{
719 let a: VecZnx<&[u8]> = a.to_ref();
720 let mut res: VecZnx<&mut [u8]> = res.to_mut();
721
722 let (n_in, n_out) = (a.n(), res.n());
723
724 if n_in == n_out {
725 module.vec_znx_copy(&mut res, res_col, &a, a_col);
726 return;
727 }
728
729 let (gap_in, gap_out): (usize, usize);
730 if n_in > n_out {
731 (gap_in, gap_out) = (n_in / n_out, 1)
732 } else {
733 (gap_in, gap_out) = (1, n_out / n_in);
734 res.zero();
735 }
736
737 let size: usize = a.size().min(res.size());
738
739 (0..size).for_each(|i| {
740 izip!(
741 a.at(a_col, i).iter().step_by(gap_in),
742 res.at_mut(res_col, i).iter_mut().step_by(gap_out)
743 )
744 .for_each(|(x_in, x_out)| *x_out = *x_in);
745 });
746}
747
748unsafe impl<B: Backend + CPUAVX> VecZnxCopyImpl<B> for B {
749 fn vec_znx_copy_impl<R, A>(_module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
750 where
751 R: VecZnxToMut,
752 A: VecZnxToRef,
753 {
754 vec_znx_copy_ref(res, res_col, a, a_col)
755 }
756}
757
758pub fn vec_znx_copy_ref<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
759where
760 R: VecZnxToMut,
761 A: VecZnxToRef,
762{
763 let mut res_mut: VecZnx<&mut [u8]> = res.to_mut();
764 let a_ref: VecZnx<&[u8]> = a.to_ref();
765
766 let min_size: usize = res_mut.size().min(a_ref.size());
767
768 (0..min_size).for_each(|j| {
769 res_mut
770 .at_mut(res_col, j)
771 .copy_from_slice(a_ref.at(a_col, j));
772 });
773 (min_size..res_mut.size()).for_each(|j| {
774 res_mut.zero_at(res_col, j);
775 })
776}
777
778unsafe impl<B: Backend + CPUAVX> VecZnxFillUniformImpl<B> for B {
779 fn vec_znx_fill_uniform_impl<R>(_module: &Module<B>, basek: usize, res: &mut R, res_col: usize, k: usize, source: &mut Source)
780 where
781 R: VecZnxToMut,
782 {
783 let mut a: VecZnx<&mut [u8]> = res.to_mut();
784 let base2k: u64 = 1 << basek;
785 let mask: u64 = base2k - 1;
786 let base2k_half: i64 = (base2k >> 1) as i64;
787 (0..k.div_ceil(basek)).for_each(|j| {
788 a.at_mut(res_col, j)
789 .iter_mut()
790 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
791 })
792 }
793}
794
795unsafe impl<B: Backend + CPUAVX> VecZnxFillDistF64Impl<B> for B {
796 fn vec_znx_fill_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
797 _module: &Module<B>,
798 basek: usize,
799 res: &mut R,
800 res_col: usize,
801 k: usize,
802 source: &mut Source,
803 dist: D,
804 bound: f64,
805 ) where
806 R: VecZnxToMut,
807 {
808 let mut a: VecZnx<&mut [u8]> = res.to_mut();
809 assert!(
810 (bound.log2().ceil() as i64) < 64,
811 "invalid bound: ceil(log2(bound))={} > 63",
812 (bound.log2().ceil() as i64)
813 );
814
815 let limb: usize = k.div_ceil(basek) - 1;
816 let basek_rem: usize = (limb + 1) * basek - k;
817
818 if basek_rem != 0 {
819 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
820 let mut dist_f64: f64 = dist.sample(source);
821 while dist_f64.abs() > bound {
822 dist_f64 = dist.sample(source)
823 }
824 *a = (dist_f64.round() as i64) << basek_rem;
825 });
826 } else {
827 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
828 let mut dist_f64: f64 = dist.sample(source);
829 while dist_f64.abs() > bound {
830 dist_f64 = dist.sample(source)
831 }
832 *a = dist_f64.round() as i64
833 });
834 }
835 }
836}
837
838unsafe impl<B: Backend + CPUAVX> VecZnxAddDistF64Impl<B> for B {
839 fn vec_znx_add_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
840 _module: &Module<B>,
841 basek: usize,
842 res: &mut R,
843 res_col: usize,
844 k: usize,
845 source: &mut Source,
846 dist: D,
847 bound: f64,
848 ) where
849 R: VecZnxToMut,
850 {
851 let mut a: VecZnx<&mut [u8]> = res.to_mut();
852 assert!(
853 (bound.log2().ceil() as i64) < 64,
854 "invalid bound: ceil(log2(bound))={} > 63",
855 (bound.log2().ceil() as i64)
856 );
857
858 let limb: usize = k.div_ceil(basek) - 1;
859 let basek_rem: usize = (limb + 1) * basek - k;
860
861 if basek_rem != 0 {
862 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
863 let mut dist_f64: f64 = dist.sample(source);
864 while dist_f64.abs() > bound {
865 dist_f64 = dist.sample(source)
866 }
867 *a += (dist_f64.round() as i64) << basek_rem;
868 });
869 } else {
870 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
871 let mut dist_f64: f64 = dist.sample(source);
872 while dist_f64.abs() > bound {
873 dist_f64 = dist.sample(source)
874 }
875 *a += dist_f64.round() as i64
876 });
877 }
878 }
879}
880
881unsafe impl<B: Backend + CPUAVX> VecZnxFillNormalImpl<B> for B {
882 fn vec_znx_fill_normal_impl<R>(
883 module: &Module<B>,
884 basek: usize,
885 res: &mut R,
886 res_col: usize,
887 k: usize,
888 source: &mut Source,
889 sigma: f64,
890 bound: f64,
891 ) where
892 R: VecZnxToMut,
893 {
894 module.vec_znx_fill_dist_f64(
895 basek,
896 res,
897 res_col,
898 k,
899 source,
900 Normal::new(0.0, sigma).unwrap(),
901 bound,
902 );
903 }
904}
905
906unsafe impl<B: Backend + CPUAVX> VecZnxAddNormalImpl<B> for B {
907 fn vec_znx_add_normal_impl<R>(
908 module: &Module<B>,
909 basek: usize,
910 res: &mut R,
911 res_col: usize,
912 k: usize,
913 source: &mut Source,
914 sigma: f64,
915 bound: f64,
916 ) where
917 R: VecZnxToMut,
918 {
919 module.vec_znx_add_dist_f64(
920 basek,
921 res,
922 res_col,
923 k,
924 source,
925 Normal::new(0.0, sigma).unwrap(),
926 bound,
927 );
928 }
929}