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