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