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::{Infos, Module};
5use crate::{alloc_aligned, assert_alignement};
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 pub n: usize,
16
17 /// Number of columns.
18 pub cols: usize,
19
20 /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
21 pub data: Vec<i64>,
22
23 /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
24 pub 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 <- a - b.
351 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx);
352
353 /// b <- b - a.
354 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx);
355
356 /// b <- -a.
357 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
358
359 /// b <- -b.
360 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
361
362 /// b <- a * X^k (mod X^{n} + 1)
363 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
364
365 /// a <- a * X^k (mod X^{n} + 1)
366 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
367
368 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
369 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
370
371 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
372 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
373
374 /// Splits b into subrings and copies them them into a.
375 ///
376 /// # Panics
377 ///
378 /// This method requires that all [VecZnx] of b have the same ring degree
379 /// and that b.n() * b.len() <= a.n()
380 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
381
382 /// Merges the subrings a into b.
383 ///
384 /// # Panics
385 ///
386 /// This method requires that all [VecZnx] of a have the same ring degree
387 /// and that a.n() * a.len() <= b.n()
388 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
389}
390
391impl VecZnxOps for Module {
392 fn new_vec_znx(&self, cols: usize) -> VecZnx {
393 VecZnx::new(self.n(), cols)
394 }
395
396 fn bytes_of_vec_znx(&self, cols: usize) -> usize {
397 self.n() * cols * 8
398 }
399
400 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
401 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
402 }
403
404 // c <- a + b
405 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
406 unsafe {
407 vec_znx::vec_znx_add(
408 self.ptr,
409 c.as_mut_ptr(),
410 c.cols() as u64,
411 c.n() as u64,
412 a.as_ptr(),
413 a.cols() as u64,
414 a.n() as u64,
415 b.as_ptr(),
416 b.cols() as u64,
417 b.n() as u64,
418 )
419 }
420 }
421
422 // b <- a + b
423 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
424 unsafe {
425 vec_znx::vec_znx_add(
426 self.ptr,
427 b.as_mut_ptr(),
428 b.cols() as u64,
429 b.n() as u64,
430 a.as_ptr(),
431 a.cols() as u64,
432 a.n() as u64,
433 b.as_ptr(),
434 b.cols() as u64,
435 b.n() as u64,
436 )
437 }
438 }
439
440 // c <- a + b
441 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
442 unsafe {
443 vec_znx::vec_znx_sub(
444 self.ptr,
445 c.as_mut_ptr(),
446 c.cols() as u64,
447 c.n() as u64,
448 a.as_ptr(),
449 a.cols() as u64,
450 a.n() as u64,
451 b.as_ptr(),
452 b.cols() as u64,
453 b.n() as u64,
454 )
455 }
456 }
457
458 // b <- a - b
459 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
460 unsafe {
461 vec_znx::vec_znx_sub(
462 self.ptr,
463 b.as_mut_ptr(),
464 b.cols() as u64,
465 b.n() as u64,
466 a.as_ptr(),
467 a.cols() as u64,
468 a.n() as u64,
469 b.as_ptr(),
470 b.cols() as u64,
471 b.n() as u64,
472 )
473 }
474 }
475
476 // b <- b - a
477 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
478 unsafe {
479 vec_znx::vec_znx_sub(
480 self.ptr,
481 b.as_mut_ptr(),
482 b.cols() as u64,
483 b.n() as u64,
484 b.as_ptr(),
485 b.cols() as u64,
486 b.n() as u64,
487 a.as_ptr(),
488 a.cols() as u64,
489 a.n() as u64,
490 )
491 }
492 }
493
494 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
495 unsafe {
496 vec_znx::vec_znx_negate(
497 self.ptr,
498 b.as_mut_ptr(),
499 b.cols() as u64,
500 b.n() as u64,
501 a.as_ptr(),
502 a.cols() as u64,
503 a.n() as u64,
504 )
505 }
506 }
507
508 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
509 unsafe {
510 vec_znx::vec_znx_negate(
511 self.ptr,
512 a.as_mut_ptr(),
513 a.cols() as u64,
514 a.n() as u64,
515 a.as_ptr(),
516 a.cols() as u64,
517 a.n() as u64,
518 )
519 }
520 }
521
522 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
523 unsafe {
524 vec_znx::vec_znx_rotate(
525 self.ptr,
526 k,
527 b.as_mut_ptr(),
528 b.cols() as u64,
529 b.n() as u64,
530 a.as_ptr(),
531 a.cols() as u64,
532 a.n() as u64,
533 )
534 }
535 }
536
537 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
538 unsafe {
539 vec_znx::vec_znx_rotate(
540 self.ptr,
541 k,
542 a.as_mut_ptr(),
543 a.cols() as u64,
544 a.n() as u64,
545 a.as_ptr(),
546 a.cols() as u64,
547 a.n() as u64,
548 )
549 }
550 }
551
552 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
553 ///
554 /// # Arguments
555 ///
556 /// * `a`: input.
557 /// * `b`: output.
558 /// * `k`: the power to which to map each coefficients.
559 /// * `a_cols`: the number of a_cols on which to apply the mapping.
560 ///
561 /// # Panics
562 ///
563 /// The method will panic if the argument `a` is greater than `a.cols()`.
564 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
565 debug_assert_eq!(a.n(), self.n());
566 debug_assert_eq!(b.n(), self.n());
567 unsafe {
568 vec_znx::vec_znx_automorphism(
569 self.ptr,
570 k,
571 b.as_mut_ptr(),
572 b.cols() as u64,
573 b.n() as u64,
574 a.as_ptr(),
575 a.cols() as u64,
576 a.n() as u64,
577 );
578 }
579 }
580
581 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
582 ///
583 /// # Arguments
584 ///
585 /// * `a`: input and output.
586 /// * `k`: the power to which to map each coefficients.
587 /// * `a_cols`: the number of cols on which to apply the mapping.
588 ///
589 /// # Panics
590 ///
591 /// The method will panic if the argument `cols` is greater than `self.cols()`.
592 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
593 debug_assert_eq!(a.n(), self.n());
594 unsafe {
595 vec_znx::vec_znx_automorphism(
596 self.ptr,
597 k,
598 a.as_mut_ptr(),
599 a.cols() as u64,
600 a.n() as u64,
601 a.as_ptr(),
602 a.cols() as u64,
603 a.n() as u64,
604 );
605 }
606 }
607
608 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
609 let (n_in, n_out) = (a.n(), b[0].n());
610
611 debug_assert!(
612 n_out < n_in,
613 "invalid a: output ring degree should be smaller"
614 );
615 b[1..].iter().for_each(|bi| {
616 debug_assert_eq!(
617 bi.n(),
618 n_out,
619 "invalid input a: all VecZnx must have the same degree"
620 )
621 });
622
623 b.iter_mut().enumerate().for_each(|(i, bi)| {
624 if i == 0 {
625 switch_degree(bi, a);
626 self.vec_znx_rotate(-1, buf, a);
627 } else {
628 switch_degree(bi, buf);
629 self.vec_znx_rotate_inplace(-1, buf);
630 }
631 })
632 }
633
634 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
635 let (n_in, n_out) = (b.n(), a[0].n());
636
637 debug_assert!(
638 n_out < n_in,
639 "invalid a: output ring degree should be smaller"
640 );
641 a[1..].iter().for_each(|ai| {
642 debug_assert_eq!(
643 ai.n(),
644 n_out,
645 "invalid input a: all VecZnx must have the same degree"
646 )
647 });
648
649 a.iter().enumerate().for_each(|(_, ai)| {
650 switch_degree(b, ai);
651 self.vec_znx_rotate_inplace(-1, b);
652 });
653
654 self.vec_znx_rotate_inplace(a.len() as i64, b);
655 }
656}