poulpy/base2k/src/vec_znx.rs
base2k/src/vec_znx.rsRaw
1use crate::cast_mut;
2use crate::ffi::vec_znx;
3use crate::ffi::znx;
4use crate::{alloc_aligned, assert_alignement};
5use crate::{Infos, Module};
6use itertools::izip;
7use std::cmp::min;
8
9/// [VecZnx] represents a vector of small norm polynomials of Zn\[X\] with [i64] coefficients.
10/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
11/// in the memory.
12#[derive(Clone)]
13pub struct VecZnx {
14 /// Polynomial degree.
15 n: usize,
16
17 /// Number of columns.
18 cols: usize,
19
20 /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
21 data: Vec<i64>,
22
23 /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
24 ptr: *mut i64,
25}
26
27pub trait VecZnxVec {
28 fn dblptr(&self) -> Vec<&[i64]>;
29 fn dblptr_mut(&mut self) -> Vec<&mut [i64]>;
30}
31
32impl VecZnxVec for Vec<VecZnx> {
33 fn dblptr(&self) -> Vec<&[i64]> {
34 self.iter().map(|v| v.raw()).collect()
35 }
36
37 fn dblptr_mut(&mut self) -> Vec<&mut [i64]> {
38 self.iter_mut().map(|v| v.raw_mut()).collect()
39 }
40}
41
42pub fn bytes_of_vec_znx(n: usize, cols: usize) -> usize {
43 n * cols * 8
44}
45
46impl VecZnx {
47 /// Returns a new struct implementing [VecZnx] with the provided data as backing array.
48 ///
49 /// The struct will take ownership of buf[..[VecZnx::bytes_of]]
50 ///
51 /// User must ensure that data is properly alligned and that
52 /// the size of data is at least equal to [VecZnx::bytes_of].
53 pub fn from_bytes(n: usize, cols: usize, bytes: &mut [u8]) -> Self {
54 #[cfg(debug_assertions)]
55 {
56 assert_eq!(bytes.len(), Self::bytes_of(n, cols));
57 assert_alignement(bytes.as_ptr());
58 }
59 unsafe {
60 let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
61 let ptr: *mut i64 = bytes_i64.as_mut_ptr();
62 VecZnx {
63 n: n,
64 cols: cols,
65 data: Vec::from_raw_parts(bytes_i64.as_mut_ptr(), bytes.len(), bytes.len()),
66 ptr: ptr,
67 }
68 }
69 }
70
71 pub fn from_bytes_borrow(n: usize, cols: usize, bytes: &mut [u8]) -> Self {
72 #[cfg(debug_assertions)]
73 {
74 assert!(bytes.len() >= Self::bytes_of(n, cols));
75 assert_alignement(bytes.as_ptr());
76 }
77 VecZnx {
78 n: n,
79 cols: cols,
80 data: Vec::new(),
81 ptr: bytes.as_mut_ptr() as *mut i64,
82 }
83 }
84
85 pub fn bytes_of(n: usize, cols: usize) -> usize {
86 bytes_of_vec_znx(n, cols)
87 }
88
89 pub fn copy_from(&mut self, a: &VecZnx) {
90 copy_vec_znx_from(self, a);
91 }
92
93 pub fn raw(&self) -> &[i64] {
94 unsafe { std::slice::from_raw_parts(self.ptr, self.n * self.cols) }
95 }
96
97 pub fn borrowing(&self) -> bool {
98 self.data.len() == 0
99 }
100
101 pub fn raw_mut(&mut self) -> &mut [i64] {
102 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.n * self.cols) }
103 }
104
105 pub fn as_ptr(&self) -> *const i64 {
106 self.ptr
107 }
108
109 pub fn as_mut_ptr(&mut self) -> *mut i64 {
110 self.ptr
111 }
112
113 pub fn at(&self, i: usize) -> &[i64] {
114 let n: usize = self.n();
115 &self.raw()[n * i..n * (i + 1)]
116 }
117
118 pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
119 let n: usize = self.n();
120 &mut self.raw_mut()[n * i..n * (i + 1)]
121 }
122
123 pub fn at_ptr(&self, i: usize) -> *const i64 {
124 self.ptr.wrapping_add(i * self.n)
125 }
126
127 pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
128 self.ptr.wrapping_add(i * self.n)
129 }
130
131 pub fn zero(&mut self) {
132 unsafe { znx::znx_zero_i64_ref((self.n * self.cols) as u64, self.ptr) }
133 }
134
135 pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
136 normalize(log_base2k, self, carry)
137 }
138
139 pub fn rsh(&mut self, log_base2k: usize, k: usize, carry: &mut [u8]) {
140 rsh(log_base2k, self, k, carry)
141 }
142
143 pub fn switch_degree(&self, a: &mut VecZnx) {
144 switch_degree(a, self)
145 }
146
147 pub fn print(&self, cols: usize, n: usize) {
148 (0..cols).for_each(|i| println!("{}: {:?}", i, &self.at(i)[..n]))
149 }
150}
151
152impl Infos for VecZnx {
153 /// Returns the base 2 logarithm of the [VecZnx] degree.
154 fn log_n(&self) -> usize {
155 (usize::BITS - (self.n - 1).leading_zeros()) as _
156 }
157
158 /// Returns the [VecZnx] degree.
159 fn n(&self) -> usize {
160 self.n
161 }
162
163 /// Returns the number of cols of the [VecZnx].
164 fn cols(&self) -> usize {
165 self.cols
166 }
167
168 /// Returns the number of rows of the [VecZnx].
169 fn rows(&self) -> usize {
170 1
171 }
172}
173
174/// Copies the coefficients of `a` on the receiver.
175/// Copy is done with the minimum size matching both backing arrays.
176pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
177 let data_a: &[i64] = a.raw();
178 let data_b: &mut [i64] = b.raw_mut();
179 let size = min(data_b.len(), data_a.len());
180 data_b[..size].copy_from_slice(&data_a[..size])
181}
182
183impl VecZnx {
184 /// Allocates a new [VecZnx] composed of #cols polynomials of Z\[X\].
185 pub fn new(n: usize, cols: usize) -> Self {
186 let mut data: Vec<i64> = alloc_aligned::<i64>(n * cols);
187 let ptr: *mut i64 = data.as_mut_ptr();
188 Self {
189 n: n,
190 cols: cols,
191 data: data,
192 ptr: ptr,
193 }
194 }
195
196 /// Truncates the precision of the [VecZnx] by k bits.
197 ///
198 /// # Arguments
199 ///
200 /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
201 /// * `k`: the number of bits of precision to drop.
202 pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
203 if k == 0 {
204 return;
205 }
206
207 if !self.borrowing() {
208 self.data
209 .truncate((self.cols() - k / log_base2k) * self.n());
210 }
211
212 self.cols -= k / log_base2k;
213
214 let k_rem: usize = k % log_base2k;
215
216 if k_rem != 0 {
217 let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
218 self.at_mut(self.cols() - 1)
219 .iter_mut()
220 .for_each(|x: &mut i64| *x &= mask)
221 }
222 }
223}
224
225pub fn switch_degree(b: &mut VecZnx, a: &VecZnx) {
226 let (n_in, n_out) = (a.n(), b.n());
227 let (gap_in, gap_out): (usize, usize);
228
229 if n_in > n_out {
230 (gap_in, gap_out) = (n_in / n_out, 1)
231 } else {
232 (gap_in, gap_out) = (1, n_out / n_in);
233 b.zero();
234 }
235
236 let cols = min(a.cols(), b.cols());
237
238 (0..cols).for_each(|i| {
239 izip!(
240 a.at(i).iter().step_by(gap_in),
241 b.at_mut(i).iter_mut().step_by(gap_out)
242 )
243 .for_each(|(x_in, x_out)| *x_out = *x_in);
244 });
245}
246
247fn normalize(log_base2k: usize, a: &mut VecZnx, tmp_bytes: &mut [u8]) {
248 let n: usize = a.n();
249
250 debug_assert!(
251 tmp_bytes.len() >= n * 8,
252 "invalid tmp_bytes: tmp_bytes.len()={} < self.n()={}",
253 tmp_bytes.len(),
254 n
255 );
256 #[cfg(debug_assertions)]
257 {
258 assert_alignement(tmp_bytes.as_ptr())
259 }
260
261 let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
262
263 unsafe {
264 znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
265 (0..a.cols()).rev().for_each(|i| {
266 znx::znx_normalize(
267 n as u64,
268 log_base2k as u64,
269 a.at_mut_ptr(i),
270 carry_i64.as_mut_ptr(),
271 a.at_mut_ptr(i),
272 carry_i64.as_mut_ptr(),
273 )
274 });
275 }
276}
277
278pub fn rsh(log_base2k: usize, a: &mut VecZnx, k: usize, tmp_bytes: &mut [u8]) {
279 let n: usize = a.n();
280
281 debug_assert!(
282 tmp_bytes.len() >> 3 >= n,
283 "invalid carry: carry.len()/8={} < self.n()={}",
284 tmp_bytes.len() >> 3,
285 n
286 );
287
288 #[cfg(debug_assertions)]
289 {
290 assert_alignement(tmp_bytes.as_ptr())
291 }
292
293 let cols: usize = a.cols();
294 let cols_steps: usize = k / log_base2k;
295
296 a.raw_mut().rotate_right(n * cols_steps);
297 unsafe {
298 znx::znx_zero_i64_ref((n * cols_steps) as u64, a.as_mut_ptr());
299 }
300
301 let k_rem = k % log_base2k;
302
303 if k_rem != 0 {
304 let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
305
306 unsafe {
307 znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
308 }
309
310 let log_base2k: usize = log_base2k;
311
312 (cols_steps..cols).for_each(|i| {
313 izip!(carry_i64.iter_mut(), a.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
314 *xi += *ci << log_base2k;
315 *ci = get_base_k_carry(*xi, k_rem);
316 *xi = (*xi - *ci) >> k_rem;
317 });
318 })
319 }
320}
321
322#[inline(always)]
323fn get_base_k_carry(x: i64, k: usize) -> i64 {
324 (x << 64 - k) >> (64 - k)
325}
326
327pub trait VecZnxOps {
328 /// Allocates a new [VecZnx].
329 ///
330 /// # Arguments
331 ///
332 /// * `cols`: the number of cols.
333 fn new_vec_znx(&self, cols: usize) -> VecZnx;
334
335 /// Returns the minimum number of bytes necessary to allocate
336 /// a new [VecZnx] through [VecZnx::from_bytes].
337 fn bytes_of_vec_znx(&self, cols: usize) -> usize;
338
339 fn vec_znx_normalize_tmp_bytes(&self) -> usize;
340
341 /// c <- a + b.
342 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
343
344 /// b <- b + a.
345 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
346
347 /// c <- a - b.
348 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
349
350 /// b <- b - a.
351 fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx);
352
353 /// b <- -a.
354 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
355
356 /// b <- -b.
357 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
358
359 /// b <- a * X^k (mod X^{n} + 1)
360 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
361
362 /// a <- a * X^k (mod X^{n} + 1)
363 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
364
365 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
366 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx, a_cols: usize);
367
368 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
369 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx, a_cols: usize);
370
371 /// Splits b into subrings and copies them them into a.
372 ///
373 /// # Panics
374 ///
375 /// This method requires that all [VecZnx] of b have the same ring degree
376 /// and that b.n() * b.len() <= a.n()
377 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
378
379 /// Merges the subrings a into b.
380 ///
381 /// # Panics
382 ///
383 /// This method requires that all [VecZnx] of a have the same ring degree
384 /// and that a.n() * a.len() <= b.n()
385 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
386}
387
388impl VecZnxOps for Module {
389 fn new_vec_znx(&self, cols: usize) -> VecZnx {
390 VecZnx::new(self.n(), cols)
391 }
392
393 fn bytes_of_vec_znx(&self, cols: usize) -> usize {
394 self.n() * cols * 8
395 }
396
397 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
398 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
399 }
400
401 // c <- a + b
402 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
403 unsafe {
404 vec_znx::vec_znx_add(
405 self.ptr,
406 c.as_mut_ptr(),
407 c.cols() as u64,
408 c.n() as u64,
409 a.as_ptr(),
410 a.cols() as u64,
411 a.n() as u64,
412 b.as_ptr(),
413 b.cols() as u64,
414 b.n() as u64,
415 )
416 }
417 }
418
419 // b <- a + b
420 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
421 unsafe {
422 vec_znx::vec_znx_add(
423 self.ptr,
424 b.as_mut_ptr(),
425 b.cols() as u64,
426 b.n() as u64,
427 a.as_ptr(),
428 a.cols() as u64,
429 a.n() as u64,
430 b.as_ptr(),
431 b.cols() as u64,
432 b.n() as u64,
433 )
434 }
435 }
436
437 // c <- a + b
438 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
439 unsafe {
440 vec_znx::vec_znx_sub(
441 self.ptr,
442 c.as_mut_ptr(),
443 c.cols() as u64,
444 c.n() as u64,
445 a.as_ptr(),
446 a.cols() as u64,
447 a.n() as u64,
448 b.as_ptr(),
449 b.cols() as u64,
450 b.n() as u64,
451 )
452 }
453 }
454
455 // b <- a + b
456 fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
457 unsafe {
458 vec_znx::vec_znx_sub(
459 self.ptr,
460 b.as_mut_ptr(),
461 b.cols() as u64,
462 b.n() as u64,
463 a.as_ptr(),
464 a.cols() as u64,
465 a.n() as u64,
466 b.as_ptr(),
467 b.cols() as u64,
468 b.n() as u64,
469 )
470 }
471 }
472
473 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
474 unsafe {
475 vec_znx::vec_znx_negate(
476 self.ptr,
477 b.as_mut_ptr(),
478 b.cols() as u64,
479 b.n() as u64,
480 a.as_ptr(),
481 a.cols() as u64,
482 a.n() as u64,
483 )
484 }
485 }
486
487 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
488 unsafe {
489 vec_znx::vec_znx_negate(
490 self.ptr,
491 a.as_mut_ptr(),
492 a.cols() as u64,
493 a.n() as u64,
494 a.as_ptr(),
495 a.cols() as u64,
496 a.n() as u64,
497 )
498 }
499 }
500
501 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
502 unsafe {
503 vec_znx::vec_znx_rotate(
504 self.ptr,
505 k,
506 b.as_mut_ptr(),
507 b.cols() as u64,
508 b.n() as u64,
509 a.as_ptr(),
510 a.cols() as u64,
511 a.n() as u64,
512 )
513 }
514 }
515
516 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
517 unsafe {
518 vec_znx::vec_znx_rotate(
519 self.ptr,
520 k,
521 a.as_mut_ptr(),
522 a.cols() as u64,
523 a.n() as u64,
524 a.as_ptr(),
525 a.cols() as u64,
526 a.n() as u64,
527 )
528 }
529 }
530
531 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
532 ///
533 /// # Arguments
534 ///
535 /// * `a`: input.
536 /// * `b`: output.
537 /// * `k`: the power to which to map each coefficients.
538 /// * `a_cols`: the number of a_cols on which to apply the mapping.
539 ///
540 /// # Panics
541 ///
542 /// The method will panic if the argument `a` is greater than `a.cols()`.
543 ///
544 /// # Example
545 /// ```
546 /// use base2k::{Module, MODULETYPE, VecZnx, Encoding, Infos, VecZnxOps};
547 /// use itertools::izip;
548 ///
549 /// let n: usize = 8; // polynomial degree
550 /// let module = Module::new(n, MODULETYPE::FFT64);
551 /// let mut a: VecZnx = module.new_vec_znx(2);
552 /// let mut b: VecZnx = module.new_vec_znx(2);
553 /// let mut c: VecZnx = module.new_vec_znx(2);
554 ///
555 /// (0..a.cols()).for_each(|i|{
556 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
557 /// *x = i as i64
558 /// })
559 /// });
560 ///
561 /// module.vec_znx_automorphism(-1, &mut b, &a, 1); // X^i -> X^(-i)
562 /// let col = c.at_mut(0);
563 /// (1..col.len()).for_each(|i|{
564 /// col[n-i] = -(i as i64)
565 /// });
566 /// izip!(b.raw().iter(), c.raw().iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
567 /// ```
568 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx, a_cols: usize) {
569 debug_assert_eq!(a.n(), self.n());
570 debug_assert_eq!(b.n(), self.n());
571 debug_assert!(a.cols() >= a_cols);
572 unsafe {
573 vec_znx::vec_znx_automorphism(
574 self.ptr,
575 k,
576 b.as_mut_ptr(),
577 b.cols() as u64,
578 b.n() as u64,
579 a.as_ptr(),
580 a_cols as u64,
581 a.n() as u64,
582 );
583 }
584 }
585
586 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
587 ///
588 /// # Arguments
589 ///
590 /// * `a`: input and output.
591 /// * `k`: the power to which to map each coefficients.
592 /// * `a_cols`: the number of cols on which to apply the mapping.
593 ///
594 /// # Panics
595 ///
596 /// The method will panic if the argument `cols` is greater than `self.cols()`.
597 ///
598 /// # Example
599 /// ```
600 /// use base2k::{Module, MODULETYPE, VecZnx, Encoding, Infos, VecZnxOps};
601 /// use itertools::izip;
602 ///
603 /// let n: usize = 8; // polynomial degree
604 /// let module = Module::new(n, MODULETYPE::FFT64);
605 /// let mut a: VecZnx = VecZnx::new(n, 2);
606 /// let mut b: VecZnx = VecZnx::new(n, 2);
607 ///
608 /// (0..a.cols()).for_each(|i|{
609 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
610 /// *x = i as i64
611 /// })
612 /// });
613 ///
614 /// module.vec_znx_automorphism_inplace(-1, &mut a, 1); // X^i -> X^(-i)
615 /// let col = b.at_mut(0);
616 /// (1..col.len()).for_each(|i|{
617 /// col[n-i] = -(i as i64)
618 /// });
619 /// izip!(a.raw().iter(), b.raw().iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
620 /// ```
621 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx, a_cols: usize) {
622 debug_assert_eq!(a.n(), self.n());
623 debug_assert!(a.cols() >= a_cols);
624 unsafe {
625 vec_znx::vec_znx_automorphism(
626 self.ptr,
627 k,
628 a.as_mut_ptr(),
629 a.cols() as u64,
630 a.n() as u64,
631 a.as_ptr(),
632 a_cols as u64,
633 a.n() as u64,
634 );
635 }
636 }
637
638 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
639 let (n_in, n_out) = (a.n(), b[0].n());
640
641 debug_assert!(
642 n_out < n_in,
643 "invalid a: output ring degree should be smaller"
644 );
645 b[1..].iter().for_each(|bi| {
646 debug_assert_eq!(
647 bi.n(),
648 n_out,
649 "invalid input a: all VecZnx must have the same degree"
650 )
651 });
652
653 b.iter_mut().enumerate().for_each(|(i, bi)| {
654 if i == 0 {
655 switch_degree(bi, a);
656 self.vec_znx_rotate(-1, buf, a);
657 } else {
658 switch_degree(bi, buf);
659 self.vec_znx_rotate_inplace(-1, buf);
660 }
661 })
662 }
663
664 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
665 let (n_in, n_out) = (b.n(), a[0].n());
666
667 debug_assert!(
668 n_out < n_in,
669 "invalid a: output ring degree should be smaller"
670 );
671 a[1..].iter().for_each(|ai| {
672 debug_assert_eq!(
673 ai.n(),
674 n_out,
675 "invalid input a: all VecZnx must have the same degree"
676 )
677 });
678
679 a.iter().enumerate().for_each(|(_, ai)| {
680 switch_degree(b, ai);
681 self.vec_znx_rotate_inplace(-1, b);
682 });
683
684 self.vec_znx_rotate_inplace(a.len() as i64, b);
685 }
686}