poulpy/backend/src/hal/api/znx_base.rs
backend/src/hal/api/znx_base.rsRaw
1use crate::hal::layouts::{Data, DataMut, DataRef};
2use rand_distr::num_traits::Zero;
3
4pub trait ZnxInfos {
5 /// Returns the ring degree of the polynomials.
6 fn n(&self) -> usize;
7
8 /// Returns the base two logarithm of the ring dimension of the polynomials.
9 fn log_n(&self) -> usize {
10 (usize::BITS - (self.n() - 1).leading_zeros()) as _
11 }
12
13 /// Returns the number of rows.
14 fn rows(&self) -> usize;
15
16 /// Returns the number of polynomials in each row.
17 fn cols(&self) -> usize;
18
19 /// Returns the number of size per polynomial.
20 fn size(&self) -> usize;
21
22 /// Returns the total number of small polynomials.
23 fn poly_count(&self) -> usize {
24 self.rows() * self.cols() * self.size()
25 }
26}
27
28pub trait ZnxSliceSize {
29 /// Returns the slice size, which is the offset between
30 /// two size of the same column.
31 fn sl(&self) -> usize;
32}
33
34pub trait DataView {
35 type D: Data;
36 fn data(&self) -> &Self::D;
37}
38
39pub trait DataViewMut: DataView {
40 fn data_mut(&mut self) -> &mut Self::D;
41}
42
43pub trait ZnxView: ZnxInfos + DataView<D: DataRef> {
44 type Scalar: Copy + Zero;
45
46 /// Returns a non-mutable pointer to the underlying coefficients array.
47 fn as_ptr(&self) -> *const Self::Scalar {
48 self.data().as_ref().as_ptr() as *const Self::Scalar
49 }
50
51 /// Returns a non-mutable reference to the entire underlying coefficient array.
52 fn raw(&self) -> &[Self::Scalar] {
53 unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
54 }
55
56 /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
57 fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
58 #[cfg(debug_assertions)]
59 {
60 assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
61 assert!(j < self.size(), "size: {} >= {}", j, self.size());
62 }
63 let offset: usize = self.n() * (j * self.cols() + i);
64 unsafe { self.as_ptr().add(offset) }
65 }
66
67 /// Returns non-mutable reference to the (i, j)-th small polynomial.
68 fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
69 unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
70 }
71}
72
73pub trait ZnxViewMut: ZnxView + DataViewMut<D: DataMut> {
74 /// Returns a mutable pointer to the underlying coefficients array.
75 fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
76 self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
77 }
78
79 /// Returns a mutable reference to the entire underlying coefficient array.
80 fn raw_mut(&mut self) -> &mut [Self::Scalar] {
81 unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
82 }
83
84 /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
85 fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
86 #[cfg(debug_assertions)]
87 {
88 assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
89 assert!(j < self.size(), "size: {} >= {}", j, self.size());
90 }
91 let offset: usize = self.n() * (j * self.cols() + i);
92 unsafe { self.as_mut_ptr().add(offset) }
93 }
94
95 /// Returns mutable reference to the (i, j)-th small polynomial.
96 fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
97 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
98 }
99}
100
101//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
102impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: DataMut> {}
103
104pub trait ZnxZero
105where
106 Self: Sized,
107{
108 fn zero(&mut self);
109 fn zero_at(&mut self, i: usize, j: usize);
110}