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 rug::{
4 Assign, Float,
5 float::Round,
6 ops::{AddAssignRound, DivAssignRound, SubAssignRound},
7};
8use sampling::source::Source;
9
10use crate::{
11 hal::{
12 api::{
13 TakeSlice, TakeVecZnx, VecZnxAddDistF64, VecZnxCopy, VecZnxDecodeVecFloat, VecZnxFillDistF64,
14 VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxSwithcDegree, ZnxInfos, ZnxSliceSize, ZnxView,
15 ZnxViewMut, ZnxZero,
16 },
17 layouts::{Backend, Module, ScalarZnx, ScalarZnxToRef, Scratch, VecZnx, VecZnxOwned, VecZnxToMut, VecZnxToRef},
18 oep::{
19 VecZnxAddDistF64Impl, VecZnxAddImpl, VecZnxAddInplaceImpl, VecZnxAddNormalImpl, VecZnxAddScalarInplaceImpl,
20 VecZnxAllocBytesImpl, VecZnxAllocImpl, VecZnxAutomorphismImpl, VecZnxAutomorphismInplaceImpl, VecZnxCopyImpl,
21 VecZnxDecodeCoeffsi64Impl, VecZnxDecodeVecFloatImpl, VecZnxDecodeVeci64Impl, VecZnxEncodeCoeffsi64Impl,
22 VecZnxEncodeVeci64Impl, VecZnxFillDistF64Impl, VecZnxFillNormalImpl, VecZnxFillUniformImpl, VecZnxFromBytesImpl,
23 VecZnxLshInplaceImpl, VecZnxMergeImpl, VecZnxMulXpMinusOneImpl, VecZnxMulXpMinusOneInplaceImpl, VecZnxNegateImpl,
24 VecZnxNegateInplaceImpl, VecZnxNormalizeImpl, VecZnxNormalizeInplaceImpl, VecZnxNormalizeTmpBytesImpl,
25 VecZnxRotateImpl, VecZnxRotateInplaceImpl, VecZnxRshInplaceImpl, VecZnxSplitImpl, VecZnxStdImpl,
26 VecZnxSubABInplaceImpl, VecZnxSubBAInplaceImpl, VecZnxSubImpl, VecZnxSubScalarInplaceImpl, VecZnxSwithcDegreeImpl,
27 },
28 },
29 implementation::cpu_spqlios::{
30 CPUAVX,
31 ffi::{module::module_info_t, vec_znx, znx},
32 },
33};
34
35unsafe impl<B: Backend> VecZnxAllocImpl<B> for B
36where
37 B: CPUAVX,
38{
39 fn vec_znx_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxOwned {
40 VecZnxOwned::alloc::<i64>(n, cols, size)
41 }
42}
43
44unsafe impl<B: Backend> VecZnxFromBytesImpl<B> for B
45where
46 B: CPUAVX,
47{
48 fn vec_znx_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned {
49 VecZnxOwned::from_bytes::<i64>(n, cols, size, bytes)
50 }
51}
52
53unsafe impl<B: Backend> VecZnxAllocBytesImpl<B> for B
54where
55 B: CPUAVX,
56{
57 fn vec_znx_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
58 VecZnxOwned::alloc_bytes::<i64>(n, cols, size)
59 }
60}
61
62unsafe impl<B: Backend> VecZnxNormalizeTmpBytesImpl<B> for B
63where
64 B: CPUAVX,
65{
66 fn vec_znx_normalize_tmp_bytes_impl(module: &Module<B>, n: usize) -> usize {
67 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(module.ptr() as *const module_info_t, n as u64) as usize }
68 }
69}
70
71unsafe impl<B: Backend> VecZnxNormalizeImpl<B> for B
72where
73 B: CPUAVX,
74{
75 fn vec_znx_normalize_impl<R, A>(
76 module: &Module<B>,
77 basek: usize,
78 res: &mut R,
79 res_col: usize,
80 a: &A,
81 a_col: usize,
82 scratch: &mut Scratch<B>,
83 ) where
84 R: VecZnxToMut,
85 A: VecZnxToRef,
86 {
87 let a: VecZnx<&[u8]> = a.to_ref();
88 let mut res: VecZnx<&mut [u8]> = res.to_mut();
89
90 #[cfg(debug_assertions)]
91 {
92 assert_eq!(res.n(), a.n());
93 }
94
95 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes(a.n()));
96
97 unsafe {
98 vec_znx::vec_znx_normalize_base2k(
99 module.ptr() as *const module_info_t,
100 a.n() as u64,
101 basek as u64,
102 res.at_mut_ptr(res_col, 0),
103 res.size() as u64,
104 res.sl() as u64,
105 a.at_ptr(a_col, 0),
106 a.size() as u64,
107 a.sl() as u64,
108 tmp_bytes.as_mut_ptr(),
109 );
110 }
111 }
112}
113
114unsafe impl<B: Backend> VecZnxNormalizeInplaceImpl<B> for B
115where
116 B: CPUAVX,
117{
118 fn vec_znx_normalize_inplace_impl<A>(module: &Module<B>, basek: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<B>)
119 where
120 A: VecZnxToMut,
121 {
122 let mut a: VecZnx<&mut [u8]> = a.to_mut();
123
124 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes(a.n()));
125
126 unsafe {
127 vec_znx::vec_znx_normalize_base2k(
128 module.ptr() as *const module_info_t,
129 a.n() as u64,
130 basek as u64,
131 a.at_mut_ptr(a_col, 0),
132 a.size() as u64,
133 a.sl() as u64,
134 a.at_ptr(a_col, 0),
135 a.size() as u64,
136 a.sl() as u64,
137 tmp_bytes.as_mut_ptr(),
138 );
139 }
140 }
141}
142
143unsafe impl<B: Backend> VecZnxAddImpl<B> for B
144where
145 B: CPUAVX,
146{
147 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)
148 where
149 R: VecZnxToMut,
150 A: VecZnxToRef,
151 C: VecZnxToRef,
152 {
153 let a: VecZnx<&[u8]> = a.to_ref();
154 let b: VecZnx<&[u8]> = b.to_ref();
155 let mut res: VecZnx<&mut [u8]> = res.to_mut();
156
157 #[cfg(debug_assertions)]
158 {
159 assert_eq!(a.n(), module.n());
160 assert_eq!(b.n(), module.n());
161 assert_eq!(res.n(), module.n());
162 assert_ne!(a.as_ptr(), b.as_ptr());
163 }
164 unsafe {
165 vec_znx::vec_znx_add(
166 module.ptr() as *const module_info_t,
167 res.at_mut_ptr(res_col, 0),
168 res.size() as u64,
169 res.sl() as u64,
170 a.at_ptr(a_col, 0),
171 a.size() as u64,
172 a.sl() as u64,
173 b.at_ptr(b_col, 0),
174 b.size() as u64,
175 b.sl() as u64,
176 )
177 }
178 }
179}
180
181unsafe impl<B: Backend> VecZnxAddInplaceImpl<B> for B
182where
183 B: CPUAVX,
184{
185 fn vec_znx_add_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
186 where
187 R: VecZnxToMut,
188 A: VecZnxToRef,
189 {
190 let a: VecZnx<&[u8]> = a.to_ref();
191 let mut res: VecZnx<&mut [u8]> = res.to_mut();
192
193 #[cfg(debug_assertions)]
194 {
195 assert_eq!(a.n(), module.n());
196 assert_eq!(res.n(), module.n());
197 }
198 unsafe {
199 vec_znx::vec_znx_add(
200 module.ptr() as *const module_info_t,
201 res.at_mut_ptr(res_col, 0),
202 res.size() as u64,
203 res.sl() as u64,
204 a.at_ptr(a_col, 0),
205 a.size() as u64,
206 a.sl() as u64,
207 res.at_ptr(res_col, 0),
208 res.size() as u64,
209 res.sl() as u64,
210 )
211 }
212 }
213}
214
215unsafe impl<B: Backend> VecZnxAddScalarInplaceImpl<B> for B
216where
217 B: CPUAVX,
218{
219 fn vec_znx_add_scalar_inplace_impl<R, A>(
220 module: &Module<B>,
221 res: &mut R,
222 res_col: usize,
223 res_limb: usize,
224 a: &A,
225 a_col: usize,
226 ) where
227 R: VecZnxToMut,
228 A: ScalarZnxToRef,
229 {
230 let mut res: VecZnx<&mut [u8]> = res.to_mut();
231 let a: ScalarZnx<&[u8]> = a.to_ref();
232
233 #[cfg(debug_assertions)]
234 {
235 assert_eq!(a.n(), module.n());
236 assert_eq!(res.n(), module.n());
237 }
238
239 unsafe {
240 vec_znx::vec_znx_add(
241 module.ptr() as *const module_info_t,
242 res.at_mut_ptr(res_col, res_limb),
243 1 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 res.at_ptr(res_col, res_limb),
249 1 as u64,
250 res.sl() as u64,
251 )
252 }
253 }
254}
255
256unsafe impl<B: Backend> VecZnxSubImpl<B> for B
257where
258 B: CPUAVX,
259{
260 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)
261 where
262 R: VecZnxToMut,
263 A: VecZnxToRef,
264 C: VecZnxToRef,
265 {
266 let a: VecZnx<&[u8]> = a.to_ref();
267 let b: VecZnx<&[u8]> = b.to_ref();
268 let mut res: VecZnx<&mut [u8]> = res.to_mut();
269
270 #[cfg(debug_assertions)]
271 {
272 assert_eq!(a.n(), module.n());
273 assert_eq!(b.n(), module.n());
274 assert_eq!(res.n(), module.n());
275 assert_ne!(a.as_ptr(), b.as_ptr());
276 }
277 unsafe {
278 vec_znx::vec_znx_sub(
279 module.ptr() as *const module_info_t,
280 res.at_mut_ptr(res_col, 0),
281 res.size() as u64,
282 res.sl() as u64,
283 a.at_ptr(a_col, 0),
284 a.size() as u64,
285 a.sl() as u64,
286 b.at_ptr(b_col, 0),
287 b.size() as u64,
288 b.sl() as u64,
289 )
290 }
291 }
292}
293
294unsafe impl<B: Backend> VecZnxSubABInplaceImpl<B> for B
295where
296 B: CPUAVX,
297{
298 fn vec_znx_sub_ab_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
299 where
300 R: VecZnxToMut,
301 A: VecZnxToRef,
302 {
303 let a: VecZnx<&[u8]> = a.to_ref();
304 let mut res: VecZnx<&mut [u8]> = res.to_mut();
305 #[cfg(debug_assertions)]
306 {
307 assert_eq!(a.n(), module.n());
308 assert_eq!(res.n(), module.n());
309 }
310 unsafe {
311 vec_znx::vec_znx_sub(
312 module.ptr() as *const module_info_t,
313 res.at_mut_ptr(res_col, 0),
314 res.size() as u64,
315 res.sl() as u64,
316 res.at_ptr(res_col, 0),
317 res.size() as u64,
318 res.sl() as u64,
319 a.at_ptr(a_col, 0),
320 a.size() as u64,
321 a.sl() as u64,
322 )
323 }
324 }
325}
326
327unsafe impl<B: Backend> VecZnxSubBAInplaceImpl<B> for B
328where
329 B: CPUAVX,
330{
331 fn vec_znx_sub_ba_inplace_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
332 where
333 R: VecZnxToMut,
334 A: VecZnxToRef,
335 {
336 let a: VecZnx<&[u8]> = a.to_ref();
337 let mut res: VecZnx<&mut [u8]> = res.to_mut();
338 #[cfg(debug_assertions)]
339 {
340 assert_eq!(a.n(), module.n());
341 assert_eq!(res.n(), module.n());
342 }
343 unsafe {
344 vec_znx::vec_znx_sub(
345 module.ptr() as *const module_info_t,
346 res.at_mut_ptr(res_col, 0),
347 res.size() 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, 0),
353 res.size() as u64,
354 res.sl() as u64,
355 )
356 }
357 }
358}
359
360unsafe impl<B: Backend> VecZnxSubScalarInplaceImpl<B> for B
361where
362 B: CPUAVX,
363{
364 fn vec_znx_sub_scalar_inplace_impl<R, A>(
365 module: &Module<B>,
366 res: &mut R,
367 res_col: usize,
368 res_limb: usize,
369 a: &A,
370 a_col: usize,
371 ) where
372 R: VecZnxToMut,
373 A: ScalarZnxToRef,
374 {
375 let mut res: VecZnx<&mut [u8]> = res.to_mut();
376 let a: ScalarZnx<&[u8]> = a.to_ref();
377
378 #[cfg(debug_assertions)]
379 {
380 assert_eq!(a.n(), module.n());
381 assert_eq!(res.n(), module.n());
382 }
383
384 unsafe {
385 vec_znx::vec_znx_sub(
386 module.ptr() as *const module_info_t,
387 res.at_mut_ptr(res_col, res_limb),
388 1 as u64,
389 res.sl() as u64,
390 a.at_ptr(a_col, 0),
391 a.size() as u64,
392 a.sl() as u64,
393 res.at_ptr(res_col, res_limb),
394 1 as u64,
395 res.sl() as u64,
396 )
397 }
398 }
399}
400
401unsafe impl<B: Backend> VecZnxNegateImpl<B> for B
402where
403 B: CPUAVX,
404{
405 fn vec_znx_negate_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
406 where
407 R: VecZnxToMut,
408 A: VecZnxToRef,
409 {
410 let a: VecZnx<&[u8]> = a.to_ref();
411 let mut res: VecZnx<&mut [u8]> = res.to_mut();
412 #[cfg(debug_assertions)]
413 {
414 assert_eq!(a.n(), module.n());
415 assert_eq!(res.n(), module.n());
416 }
417 unsafe {
418 vec_znx::vec_znx_negate(
419 module.ptr() as *const module_info_t,
420 res.at_mut_ptr(res_col, 0),
421 res.size() as u64,
422 res.sl() as u64,
423 a.at_ptr(a_col, 0),
424 a.size() as u64,
425 a.sl() as u64,
426 )
427 }
428 }
429}
430
431unsafe impl<B: Backend> VecZnxNegateInplaceImpl<B> for B
432where
433 B: CPUAVX,
434{
435 fn vec_znx_negate_inplace_impl<A>(module: &Module<B>, a: &mut A, a_col: usize)
436 where
437 A: VecZnxToMut,
438 {
439 let mut a: VecZnx<&mut [u8]> = a.to_mut();
440 #[cfg(debug_assertions)]
441 {
442 assert_eq!(a.n(), module.n());
443 }
444 unsafe {
445 vec_znx::vec_znx_negate(
446 module.ptr() as *const module_info_t,
447 a.at_mut_ptr(a_col, 0),
448 a.size() as u64,
449 a.sl() as u64,
450 a.at_ptr(a_col, 0),
451 a.size() as u64,
452 a.sl() as u64,
453 )
454 }
455 }
456}
457
458unsafe impl<B: Backend> VecZnxLshInplaceImpl<B> for B
459where
460 B: CPUAVX,
461{
462 fn vec_znx_lsh_inplace_impl<A>(_module: &Module<B>, basek: usize, k: usize, a: &mut A)
463 where
464 A: VecZnxToMut,
465 {
466 vec_znx_lsh_inplace_ref(basek, k, a)
467 }
468}
469
470pub fn vec_znx_lsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
471where
472 A: VecZnxToMut,
473{
474 let mut a: VecZnx<&mut [u8]> = a.to_mut();
475
476 let n: usize = a.n();
477 let cols: usize = a.cols();
478 let size: usize = a.size();
479 let steps: usize = k / basek;
480
481 a.raw_mut().rotate_left(n * steps * cols);
482 (0..cols).for_each(|i| {
483 (size - steps..size).for_each(|j| {
484 a.zero_at(i, j);
485 })
486 });
487
488 let k_rem: usize = k % basek;
489
490 if k_rem != 0 {
491 let shift: usize = i64::BITS as usize - k_rem;
492 (0..cols).for_each(|i| {
493 (0..steps).for_each(|j| {
494 a.at_mut(i, j).iter_mut().for_each(|xi| {
495 *xi <<= shift;
496 });
497 });
498 });
499 }
500}
501
502unsafe impl<B: Backend> VecZnxRshInplaceImpl<B> for B
503where
504 B: CPUAVX,
505{
506 fn vec_znx_rsh_inplace_impl<A>(_module: &Module<B>, basek: usize, k: usize, a: &mut A)
507 where
508 A: VecZnxToMut,
509 {
510 vec_znx_rsh_inplace_ref(basek, k, a)
511 }
512}
513
514pub fn vec_znx_rsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
515where
516 A: VecZnxToMut,
517{
518 let mut a: VecZnx<&mut [u8]> = a.to_mut();
519 let n: usize = a.n();
520 let cols: usize = a.cols();
521 let size: usize = a.size();
522 let steps: usize = k / basek;
523
524 a.raw_mut().rotate_right(n * steps * cols);
525 (0..cols).for_each(|i| {
526 (0..steps).for_each(|j| {
527 a.zero_at(i, j);
528 })
529 });
530
531 let k_rem: usize = k % basek;
532
533 if k_rem != 0 {
534 let mut carry: Vec<i64> = vec![0i64; n]; // ALLOC (but small so OK)
535 let shift: usize = i64::BITS as usize - k_rem;
536 (0..cols).for_each(|i| {
537 carry.fill(0);
538 (steps..size).for_each(|j| {
539 izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
540 *xi += *ci << basek;
541 *ci = (*xi << shift) >> shift;
542 *xi = (*xi - *ci) >> k_rem;
543 });
544 });
545 })
546 }
547}
548
549unsafe impl<B: Backend> VecZnxRotateImpl<B> for B
550where
551 B: CPUAVX,
552{
553 fn vec_znx_rotate_impl<R, A>(_module: &Module<B>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
554 where
555 R: VecZnxToMut,
556 A: VecZnxToRef,
557 {
558 let a: VecZnx<&[u8]> = a.to_ref();
559 let mut res: VecZnx<&mut [u8]> = res.to_mut();
560 #[cfg(debug_assertions)]
561 {
562 assert_eq!(res.n(), a.n());
563 }
564 unsafe {
565 (0..a.size()).for_each(|j| {
566 znx::znx_rotate_i64(
567 a.n() as u64,
568 k,
569 res.at_mut_ptr(res_col, j),
570 a.at_ptr(a_col, j),
571 );
572 });
573 }
574 }
575}
576
577unsafe impl<B: Backend> VecZnxRotateInplaceImpl<B> for B
578where
579 B: CPUAVX,
580{
581 fn vec_znx_rotate_inplace_impl<A>(_module: &Module<B>, k: i64, a: &mut A, a_col: usize)
582 where
583 A: VecZnxToMut,
584 {
585 let mut a: VecZnx<&mut [u8]> = a.to_mut();
586 unsafe {
587 (0..a.size()).for_each(|j| {
588 znx::znx_rotate_inplace_i64(a.n() as u64, k, a.at_mut_ptr(a_col, j));
589 });
590 }
591 }
592}
593
594unsafe impl<B: Backend> VecZnxAutomorphismImpl<B> for B
595where
596 B: CPUAVX,
597{
598 fn vec_znx_automorphism_impl<R, A>(module: &Module<B>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
599 where
600 R: VecZnxToMut,
601 A: VecZnxToRef,
602 {
603 let a: VecZnx<&[u8]> = a.to_ref();
604 let mut res: VecZnx<&mut [u8]> = res.to_mut();
605 #[cfg(debug_assertions)]
606 {
607 assert_eq!(a.n(), module.n());
608 assert_eq!(res.n(), module.n());
609 }
610 unsafe {
611 vec_znx::vec_znx_automorphism(
612 module.ptr() as *const module_info_t,
613 k,
614 res.at_mut_ptr(res_col, 0),
615 res.size() as u64,
616 res.sl() as u64,
617 a.at_ptr(a_col, 0),
618 a.size() as u64,
619 a.sl() as u64,
620 )
621 }
622 }
623}
624
625unsafe impl<B: Backend> VecZnxAutomorphismInplaceImpl<B> for B
626where
627 B: CPUAVX,
628{
629 fn vec_znx_automorphism_inplace_impl<A>(module: &Module<B>, k: i64, a: &mut A, a_col: usize)
630 where
631 A: VecZnxToMut,
632 {
633 let mut a: VecZnx<&mut [u8]> = a.to_mut();
634 #[cfg(debug_assertions)]
635 {
636 assert_eq!(a.n(), module.n());
637 assert!(
638 k & 1 != 0,
639 "invalid galois element: must be odd but is {}",
640 k
641 );
642 }
643 unsafe {
644 vec_znx::vec_znx_automorphism(
645 module.ptr() as *const module_info_t,
646 k,
647 a.at_mut_ptr(a_col, 0),
648 a.size() as u64,
649 a.sl() as u64,
650 a.at_ptr(a_col, 0),
651 a.size() as u64,
652 a.sl() as u64,
653 )
654 }
655 }
656}
657
658unsafe impl<B: Backend> VecZnxMulXpMinusOneImpl<B> for B
659where
660 B: CPUAVX,
661{
662 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)
663 where
664 R: VecZnxToMut,
665 A: VecZnxToRef,
666 {
667 let a: VecZnx<&[u8]> = a.to_ref();
668 let mut res: VecZnx<&mut [u8]> = res.to_mut();
669 #[cfg(debug_assertions)]
670 {
671 assert_eq!(a.n(), module.n());
672 assert_eq!(res.n(), module.n());
673 }
674 unsafe {
675 vec_znx::vec_znx_mul_xp_minus_one(
676 module.ptr() as *const module_info_t,
677 p,
678 res.at_mut_ptr(res_col, 0),
679 res.size() as u64,
680 res.sl() as u64,
681 a.at_ptr(a_col, 0),
682 a.size() as u64,
683 a.sl() as u64,
684 )
685 }
686 }
687}
688
689unsafe impl<B: Backend> VecZnxMulXpMinusOneInplaceImpl<B> for B
690where
691 B: CPUAVX,
692{
693 fn vec_znx_mul_xp_minus_one_inplace_impl<R>(module: &Module<B>, p: i64, res: &mut R, res_col: usize)
694 where
695 R: VecZnxToMut,
696 {
697 let mut res: VecZnx<&mut [u8]> = res.to_mut();
698 #[cfg(debug_assertions)]
699 {
700 assert_eq!(res.n(), module.n());
701 }
702 unsafe {
703 vec_znx::vec_znx_mul_xp_minus_one(
704 module.ptr() as *const module_info_t,
705 p,
706 res.at_mut_ptr(res_col, 0),
707 res.size() as u64,
708 res.sl() as u64,
709 res.at_ptr(res_col, 0),
710 res.size() as u64,
711 res.sl() as u64,
712 )
713 }
714 }
715}
716
717unsafe impl<B: Backend> VecZnxSplitImpl<B> for B
718where
719 B: CPUAVX,
720{
721 fn vec_znx_split_impl<R, A>(
722 module: &Module<B>,
723 res: &mut Vec<R>,
724 res_col: usize,
725 a: &A,
726 a_col: usize,
727 scratch: &mut Scratch<B>,
728 ) where
729 R: VecZnxToMut,
730 A: VecZnxToRef,
731 {
732 vec_znx_split_ref(module, res, res_col, a, a_col, scratch)
733 }
734}
735
736pub fn vec_znx_split_ref<R, A, B: Backend>(
737 module: &Module<B>,
738 res: &mut Vec<R>,
739 res_col: usize,
740 a: &A,
741 a_col: usize,
742 scratch: &mut Scratch<B>,
743) where
744 B: CPUAVX,
745 R: VecZnxToMut,
746 A: VecZnxToRef,
747{
748 let a: VecZnx<&[u8]> = a.to_ref();
749
750 let (n_in, n_out) = (a.n(), res[0].to_mut().n());
751
752 let (mut buf, _) = scratch.take_vec_znx(module, 1, a.size());
753
754 debug_assert!(
755 n_out < n_in,
756 "invalid a: output ring degree should be smaller"
757 );
758 res[1..].iter_mut().for_each(|bi| {
759 debug_assert_eq!(
760 bi.to_mut().n(),
761 n_out,
762 "invalid input a: all VecZnx must have the same degree"
763 )
764 });
765
766 res.iter_mut().enumerate().for_each(|(i, bi)| {
767 if i == 0 {
768 module.vec_znx_switch_degree(bi, res_col, &a, a_col);
769 module.vec_znx_rotate(-1, &mut buf, 0, &a, a_col);
770 } else {
771 module.vec_znx_switch_degree(bi, res_col, &mut buf, a_col);
772 module.vec_znx_rotate_inplace(-1, &mut buf, a_col);
773 }
774 })
775}
776
777unsafe impl<B: Backend> VecZnxMergeImpl<B> for B
778where
779 B: CPUAVX,
780{
781 fn vec_znx_merge_impl<R, A>(module: &Module<B>, res: &mut R, res_col: usize, a: Vec<A>, a_col: usize)
782 where
783 R: VecZnxToMut,
784 A: VecZnxToRef,
785 {
786 vec_znx_merge_ref(module, res, res_col, a, a_col)
787 }
788}
789
790pub fn vec_znx_merge_ref<R, A, B: Backend>(module: &Module<B>, res: &mut R, res_col: usize, a: Vec<A>, a_col: usize)
791where
792 B: CPUAVX,
793 R: VecZnxToMut,
794 A: VecZnxToRef,
795{
796 let mut res: VecZnx<&mut [u8]> = res.to_mut();
797
798 let (n_in, n_out) = (res.n(), a[0].to_ref().n());
799
800 debug_assert!(
801 n_out < n_in,
802 "invalid a: output ring degree should be smaller"
803 );
804 a[1..].iter().for_each(|ai| {
805 debug_assert_eq!(
806 ai.to_ref().n(),
807 n_out,
808 "invalid input a: all VecZnx must have the same degree"
809 )
810 });
811
812 a.iter().enumerate().for_each(|(_, ai)| {
813 module.vec_znx_switch_degree(&mut res, res_col, ai, a_col);
814 module.vec_znx_rotate_inplace(-1, &mut res, res_col);
815 });
816
817 module.vec_znx_rotate_inplace(a.len() as i64, &mut res, res_col);
818}
819
820unsafe impl<B: Backend> VecZnxSwithcDegreeImpl<B> for B
821where
822 B: CPUAVX,
823{
824 fn vec_znx_switch_degree_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_switch_degree_ref(module, res, res_col, a, a_col)
830 }
831}
832
833pub fn vec_znx_switch_degree_ref<R, A, B: Backend>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
834where
835 B: CPUAVX,
836 R: VecZnxToMut,
837 A: VecZnxToRef,
838{
839 let a: VecZnx<&[u8]> = a.to_ref();
840 let mut res: VecZnx<&mut [u8]> = res.to_mut();
841
842 let (n_in, n_out) = (a.n(), res.n());
843
844 if n_in == n_out {
845 module.vec_znx_copy(&mut res, res_col, &a, a_col);
846 return;
847 }
848
849 let (gap_in, gap_out): (usize, usize);
850 if n_in > n_out {
851 (gap_in, gap_out) = (n_in / n_out, 1)
852 } else {
853 (gap_in, gap_out) = (1, n_out / n_in);
854 res.zero();
855 }
856
857 let size: usize = a.size().min(res.size());
858
859 (0..size).for_each(|i| {
860 izip!(
861 a.at(a_col, i).iter().step_by(gap_in),
862 res.at_mut(res_col, i).iter_mut().step_by(gap_out)
863 )
864 .for_each(|(x_in, x_out)| *x_out = *x_in);
865 });
866}
867
868unsafe impl<B: Backend> VecZnxCopyImpl<B> for B
869where
870 B: CPUAVX,
871{
872 fn vec_znx_copy_impl<R, A>(_module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
873 where
874 R: VecZnxToMut,
875 A: VecZnxToRef,
876 {
877 vec_znx_copy_ref(res, res_col, a, a_col)
878 }
879}
880
881pub fn vec_znx_copy_ref<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
882where
883 R: VecZnxToMut,
884 A: VecZnxToRef,
885{
886 let mut res_mut: VecZnx<&mut [u8]> = res.to_mut();
887 let a_ref: VecZnx<&[u8]> = a.to_ref();
888
889 let min_size: usize = res_mut.size().min(a_ref.size());
890
891 (0..min_size).for_each(|j| {
892 res_mut
893 .at_mut(res_col, j)
894 .copy_from_slice(a_ref.at(a_col, j));
895 });
896 (min_size..res_mut.size()).for_each(|j| {
897 res_mut.zero_at(res_col, j);
898 })
899}
900
901unsafe impl<B: Backend> VecZnxStdImpl<B> for B
902where
903 B: CPUAVX,
904{
905 fn vec_znx_std_impl<A>(module: &Module<B>, basek: usize, a: &A, a_col: usize) -> f64
906 where
907 A: VecZnxToRef,
908 {
909 let a: VecZnx<&[u8]> = a.to_ref();
910 let prec: u32 = (a.size() * basek) as u32;
911 let mut data: Vec<Float> = (0..a.n()).map(|_| Float::with_val(prec, 0)).collect();
912 module.decode_vec_float(basek, &a, a_col, &mut data);
913 // std = sqrt(sum((xi - avg)^2) / n)
914 let mut avg: Float = Float::with_val(prec, 0);
915 data.iter().for_each(|x| {
916 avg.add_assign_round(x, Round::Nearest);
917 });
918 avg.div_assign_round(Float::with_val(prec, data.len()), Round::Nearest);
919 data.iter_mut().for_each(|x| {
920 x.sub_assign_round(&avg, Round::Nearest);
921 });
922 let mut std: Float = Float::with_val(prec, 0);
923 data.iter().for_each(|x| std += x * x);
924 std.div_assign_round(Float::with_val(prec, data.len()), Round::Nearest);
925 std = std.sqrt();
926 std.to_f64()
927 }
928}
929
930unsafe impl<B: Backend> VecZnxFillUniformImpl<B> for B
931where
932 B: CPUAVX,
933{
934 fn vec_znx_fill_uniform_impl<R>(_module: &Module<B>, basek: usize, res: &mut R, res_col: usize, k: usize, source: &mut Source)
935 where
936 R: VecZnxToMut,
937 {
938 let mut a: VecZnx<&mut [u8]> = res.to_mut();
939 let base2k: u64 = 1 << basek;
940 let mask: u64 = base2k - 1;
941 let base2k_half: i64 = (base2k >> 1) as i64;
942 (0..k.div_ceil(basek)).for_each(|j| {
943 a.at_mut(res_col, j)
944 .iter_mut()
945 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
946 })
947 }
948}
949
950unsafe impl<B: Backend> VecZnxFillDistF64Impl<B> for B
951where
952 B: CPUAVX,
953{
954 fn vec_znx_fill_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
955 _module: &Module<B>,
956 basek: usize,
957 res: &mut R,
958 res_col: usize,
959 k: usize,
960 source: &mut Source,
961 dist: D,
962 bound: f64,
963 ) where
964 R: VecZnxToMut,
965 {
966 let mut a: VecZnx<&mut [u8]> = res.to_mut();
967 assert!(
968 (bound.log2().ceil() as i64) < 64,
969 "invalid bound: ceil(log2(bound))={} > 63",
970 (bound.log2().ceil() as i64)
971 );
972
973 let limb: usize = k.div_ceil(basek) - 1;
974 let basek_rem: usize = (limb + 1) * basek - k;
975
976 if basek_rem != 0 {
977 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
978 let mut dist_f64: f64 = dist.sample(source);
979 while dist_f64.abs() > bound {
980 dist_f64 = dist.sample(source)
981 }
982 *a = (dist_f64.round() as i64) << basek_rem;
983 });
984 } else {
985 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
986 let mut dist_f64: f64 = dist.sample(source);
987 while dist_f64.abs() > bound {
988 dist_f64 = dist.sample(source)
989 }
990 *a = dist_f64.round() as i64
991 });
992 }
993 }
994}
995
996unsafe impl<B: Backend> VecZnxAddDistF64Impl<B> for B
997where
998 B: CPUAVX,
999{
1000 fn vec_znx_add_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
1001 _module: &Module<B>,
1002 basek: usize,
1003 res: &mut R,
1004 res_col: usize,
1005 k: usize,
1006 source: &mut Source,
1007 dist: D,
1008 bound: f64,
1009 ) where
1010 R: VecZnxToMut,
1011 {
1012 let mut a: VecZnx<&mut [u8]> = res.to_mut();
1013 assert!(
1014 (bound.log2().ceil() as i64) < 64,
1015 "invalid bound: ceil(log2(bound))={} > 63",
1016 (bound.log2().ceil() as i64)
1017 );
1018
1019 let limb: usize = k.div_ceil(basek) - 1;
1020 let basek_rem: usize = (limb + 1) * basek - k;
1021
1022 if basek_rem != 0 {
1023 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
1024 let mut dist_f64: f64 = dist.sample(source);
1025 while dist_f64.abs() > bound {
1026 dist_f64 = dist.sample(source)
1027 }
1028 *a += (dist_f64.round() as i64) << basek_rem;
1029 });
1030 } else {
1031 a.at_mut(res_col, limb).iter_mut().for_each(|a| {
1032 let mut dist_f64: f64 = dist.sample(source);
1033 while dist_f64.abs() > bound {
1034 dist_f64 = dist.sample(source)
1035 }
1036 *a += dist_f64.round() as i64
1037 });
1038 }
1039 }
1040}
1041
1042unsafe impl<B: Backend> VecZnxFillNormalImpl<B> for B
1043where
1044 B: CPUAVX,
1045{
1046 fn vec_znx_fill_normal_impl<R>(
1047 module: &Module<B>,
1048 basek: usize,
1049 res: &mut R,
1050 res_col: usize,
1051 k: usize,
1052 source: &mut Source,
1053 sigma: f64,
1054 bound: f64,
1055 ) where
1056 R: VecZnxToMut,
1057 {
1058 module.vec_znx_fill_dist_f64(
1059 basek,
1060 res,
1061 res_col,
1062 k,
1063 source,
1064 Normal::new(0.0, sigma).unwrap(),
1065 bound,
1066 );
1067 }
1068}
1069
1070unsafe impl<B: Backend> VecZnxAddNormalImpl<B> for B
1071where
1072 B: CPUAVX,
1073{
1074 fn vec_znx_add_normal_impl<R>(
1075 module: &Module<B>,
1076 basek: usize,
1077 res: &mut R,
1078 res_col: usize,
1079 k: usize,
1080 source: &mut Source,
1081 sigma: f64,
1082 bound: f64,
1083 ) where
1084 R: VecZnxToMut,
1085 {
1086 module.vec_znx_add_dist_f64(
1087 basek,
1088 res,
1089 res_col,
1090 k,
1091 source,
1092 Normal::new(0.0, sigma).unwrap(),
1093 bound,
1094 );
1095 }
1096}
1097
1098unsafe impl<B: Backend> VecZnxEncodeVeci64Impl<B> for B
1099where
1100 B: CPUAVX,
1101{
1102 fn encode_vec_i64_impl<R>(
1103 _module: &Module<B>,
1104 basek: usize,
1105 res: &mut R,
1106 res_col: usize,
1107 k: usize,
1108 data: &[i64],
1109 log_max: usize,
1110 ) where
1111 R: VecZnxToMut,
1112 {
1113 let size: usize = k.div_ceil(basek);
1114
1115 #[cfg(debug_assertions)]
1116 {
1117 let a: VecZnx<&mut [u8]> = res.to_mut();
1118 assert!(
1119 size <= a.size(),
1120 "invalid argument k: k.div_ceil(basek)={} > a.size()={}",
1121 size,
1122 a.size()
1123 );
1124 assert!(res_col < a.cols());
1125 assert!(data.len() <= a.n())
1126 }
1127
1128 let data_len: usize = data.len();
1129 let mut a: VecZnx<&mut [u8]> = res.to_mut();
1130 let k_rem: usize = basek - (k % basek);
1131
1132 // Zeroes coefficients of the i-th column
1133 (0..a.size()).for_each(|i| {
1134 a.zero_at(res_col, i);
1135 });
1136
1137 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
1138 // values on the last limb.
1139 // Else we decompose values base2k.
1140 if log_max + k_rem < 63 || k_rem == basek {
1141 a.at_mut(res_col, size - 1)[..data_len].copy_from_slice(&data[..data_len]);
1142 } else {
1143 let mask: i64 = (1 << basek) - 1;
1144 let steps: usize = size.min(log_max.div_ceil(basek));
1145 (size - steps..size)
1146 .rev()
1147 .enumerate()
1148 .for_each(|(i, i_rev)| {
1149 let shift: usize = i * basek;
1150 izip!(a.at_mut(res_col, i_rev).iter_mut(), data.iter()).for_each(|(y, x)| *y = (x >> shift) & mask);
1151 })
1152 }
1153
1154 // Case where self.prec % self.k != 0.
1155 if k_rem != basek {
1156 let steps: usize = size.min(log_max.div_ceil(basek));
1157 (size - steps..size).rev().for_each(|i| {
1158 a.at_mut(res_col, i)[..data_len]
1159 .iter_mut()
1160 .for_each(|x| *x <<= k_rem);
1161 })
1162 }
1163 }
1164}
1165
1166unsafe impl<B: Backend> VecZnxEncodeCoeffsi64Impl<B> for B
1167where
1168 B: CPUAVX,
1169{
1170 fn encode_coeff_i64_impl<R>(
1171 _module: &Module<B>,
1172 basek: usize,
1173 res: &mut R,
1174 res_col: usize,
1175 k: usize,
1176 i: usize,
1177 data: i64,
1178 log_max: usize,
1179 ) where
1180 R: VecZnxToMut,
1181 {
1182 let size: usize = k.div_ceil(basek);
1183
1184 #[cfg(debug_assertions)]
1185 {
1186 let a: VecZnx<&mut [u8]> = res.to_mut();
1187 assert!(i < a.n());
1188 assert!(
1189 size <= a.size(),
1190 "invalid argument k: k.div_ceil(basek)={} > a.size()={}",
1191 size,
1192 a.size()
1193 );
1194 assert!(res_col < a.cols());
1195 }
1196
1197 let k_rem: usize = basek - (k % basek);
1198 let mut a: VecZnx<&mut [u8]> = res.to_mut();
1199 (0..a.size()).for_each(|j| a.at_mut(res_col, j)[i] = 0);
1200
1201 // If 2^{basek} * 2^{k_rem} < 2^{63}-1, then we can simply copy
1202 // values on the last limb.
1203 // Else we decompose values base2k.
1204 if log_max + k_rem < 63 || k_rem == basek {
1205 a.at_mut(res_col, size - 1)[i] = data;
1206 } else {
1207 let mask: i64 = (1 << basek) - 1;
1208 let steps: usize = size.min(log_max.div_ceil(basek));
1209 (size - steps..size)
1210 .rev()
1211 .enumerate()
1212 .for_each(|(j, j_rev)| {
1213 a.at_mut(res_col, j_rev)[i] = (data >> (j * basek)) & mask;
1214 })
1215 }
1216
1217 // Case where prec % k != 0.
1218 if k_rem != basek {
1219 let steps: usize = size.min(log_max.div_ceil(basek));
1220 (size - steps..size).rev().for_each(|j| {
1221 a.at_mut(res_col, j)[i] <<= k_rem;
1222 })
1223 }
1224 }
1225}
1226
1227unsafe impl<B: Backend> VecZnxDecodeVeci64Impl<B> for B
1228where
1229 B: CPUAVX,
1230{
1231 fn decode_vec_i64_impl<R>(_module: &Module<B>, basek: usize, res: &R, res_col: usize, k: usize, data: &mut [i64])
1232 where
1233 R: VecZnxToRef,
1234 {
1235 let size: usize = k.div_ceil(basek);
1236 #[cfg(debug_assertions)]
1237 {
1238 let a: VecZnx<&[u8]> = res.to_ref();
1239 assert!(
1240 data.len() >= a.n(),
1241 "invalid data: data.len()={} < a.n()={}",
1242 data.len(),
1243 a.n()
1244 );
1245 assert!(res_col < a.cols());
1246 }
1247
1248 let a: VecZnx<&[u8]> = res.to_ref();
1249 data.copy_from_slice(a.at(res_col, 0));
1250 let rem: usize = basek - (k % basek);
1251 if k < basek {
1252 data.iter_mut().for_each(|x| *x >>= rem);
1253 } else {
1254 (1..size).for_each(|i| {
1255 if i == size - 1 && rem != basek {
1256 let k_rem: usize = basek - rem;
1257 izip!(a.at(res_col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
1258 *y = (*y << k_rem) + (x >> rem);
1259 });
1260 } else {
1261 izip!(a.at(res_col, i).iter(), data.iter_mut()).for_each(|(x, y)| {
1262 *y = (*y << basek) + x;
1263 });
1264 }
1265 })
1266 }
1267 }
1268}
1269
1270unsafe impl<B: Backend> VecZnxDecodeCoeffsi64Impl<B> for B
1271where
1272 B: CPUAVX,
1273{
1274 fn decode_coeff_i64_impl<R>(_module: &Module<B>, basek: usize, res: &R, res_col: usize, k: usize, i: usize) -> i64
1275 where
1276 R: VecZnxToRef,
1277 {
1278 #[cfg(debug_assertions)]
1279 {
1280 let a: VecZnx<&[u8]> = res.to_ref();
1281 assert!(i < a.n());
1282 assert!(res_col < a.cols())
1283 }
1284
1285 let a: VecZnx<&[u8]> = res.to_ref();
1286 let size: usize = k.div_ceil(basek);
1287 let mut res: i64 = 0;
1288 let rem: usize = basek - (k % basek);
1289 (0..size).for_each(|j| {
1290 let x: i64 = a.at(res_col, j)[i];
1291 if j == size - 1 && rem != basek {
1292 let k_rem: usize = basek - rem;
1293 res = (res << k_rem) + (x >> rem);
1294 } else {
1295 res = (res << basek) + x;
1296 }
1297 });
1298 res
1299 }
1300}
1301
1302unsafe impl<B: Backend> VecZnxDecodeVecFloatImpl<B> for B
1303where
1304 B: CPUAVX,
1305{
1306 fn decode_vec_float_impl<R>(_module: &Module<B>, basek: usize, res: &R, res_col: usize, data: &mut [Float])
1307 where
1308 R: VecZnxToRef,
1309 {
1310 #[cfg(debug_assertions)]
1311 {
1312 let a: VecZnx<&[u8]> = res.to_ref();
1313 assert!(
1314 data.len() >= a.n(),
1315 "invalid data: data.len()={} < a.n()={}",
1316 data.len(),
1317 a.n()
1318 );
1319 assert!(res_col < a.cols());
1320 }
1321
1322 let a: VecZnx<&[u8]> = res.to_ref();
1323 let size: usize = a.size();
1324 let prec: u32 = (basek * size) as u32;
1325
1326 // 2^{basek}
1327 let base = Float::with_val(prec, (1 << basek) as f64);
1328
1329 // y[i] = sum x[j][i] * 2^{-basek*j}
1330 (0..size).for_each(|i| {
1331 if i == 0 {
1332 izip!(a.at(res_col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
1333 y.assign(*x);
1334 *y /= &base;
1335 });
1336 } else {
1337 izip!(a.at(res_col, size - i - 1).iter(), data.iter_mut()).for_each(|(x, y)| {
1338 *y += Float::with_val(prec, *x);
1339 *y /= &base;
1340 });
1341 }
1342 });
1343 }
1344}