poulpy/base2k/src/znx_base.rs
base2k/src/znx_base.rsRaw
1use crate::{Backend, Module, alloc_aligned, assert_alignement, cast_mut};
2use itertools::izip;
3use std::cmp::min;
4
5pub struct ZnxBase {
6 /// The ring degree
7 pub n: usize,
8
9 /// The number of rows (in the third dimension)
10 pub rows: usize,
11
12 /// The number of polynomials
13 pub cols: usize,
14
15 /// The number of size per polynomial (a.k.a small polynomials).
16 pub size: usize,
17
18 /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
19 pub data: Vec<u8>,
20
21 /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
22 pub ptr: *mut u8,
23}
24
25impl ZnxBase {
26 pub fn from_bytes(n: usize, rows: usize, cols: usize, size: usize, mut bytes: Vec<u8>) -> Self {
27 let mut res: Self = Self::from_bytes_borrow(n, rows, cols, size, &mut bytes);
28 res.data = bytes;
29 res
30 }
31
32 pub fn from_bytes_borrow(n: usize, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
33 #[cfg(debug_assertions)]
34 {
35 assert_eq!(n & (n - 1), 0, "n must be a power of two");
36 assert!(n > 0, "n must be greater than 0");
37 assert!(rows > 0, "rows must be greater than 0");
38 assert!(cols > 0, "cols must be greater than 0");
39 assert!(size > 0, "size must be greater than 0");
40 }
41 Self {
42 n: n,
43 rows: rows,
44 cols: cols,
45 size: size,
46 data: Vec::new(),
47 ptr: bytes.as_mut_ptr(),
48 }
49 }
50}
51
52pub trait GetZnxBase {
53 fn znx(&self) -> &ZnxBase;
54 fn znx_mut(&mut self) -> &mut ZnxBase;
55}
56
57pub trait ZnxInfos {
58 /// Returns the ring degree of the polynomials.
59 fn n(&self) -> usize;
60
61 /// Returns the base two logarithm of the ring dimension of the polynomials.
62 fn log_n(&self) -> usize {
63 (usize::BITS - (self.n() - 1).leading_zeros()) as _
64 }
65
66 /// Returns the number of rows.
67 fn rows(&self) -> usize;
68
69 /// Returns the number of polynomials in each row.
70 fn cols(&self) -> usize;
71
72 /// Returns the number of size per polynomial.
73 fn size(&self) -> usize;
74
75 /// Returns the total number of small polynomials.
76 fn poly_count(&self) -> usize {
77 self.rows() * self.cols() * self.size()
78 }
79
80 /// Returns the slice size, which is the offset between
81 /// two size of the same column.
82 fn sl(&self) -> usize;
83}
84
85// pub trait ZnxSliceSize {}
86
87//(Jay) TODO: Remove ZnxAlloc
88// pub trait ZnxAlloc<B: Backend>
89// where
90// Self: Sized + ZnxInfos,
91// {
92// type Scalar;
93// fn new(module: &Module<B>, rows: usize, cols: usize, size: usize) -> Self {
94// let bytes: Vec<u8> = alloc_aligned::<u8>(Self::bytes_of(module, rows, cols, size));
95// Self::from_bytes(module, rows, cols, size, bytes)
96// }
97
98// fn from_bytes(module: &Module<B>, rows: usize, cols: usize, size: usize, mut bytes: Vec<u8>) -> Self {
99// let mut res: Self = Self::from_bytes_borrow(module, rows, cols, size, &mut bytes);
100// res.znx_mut().data = bytes;
101// res
102// }
103
104// fn from_bytes_borrow(module: &Module<B>, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self;
105
106// fn bytes_of(module: &Module<B>, rows: usize, cols: usize, size: usize) -> usize;
107// }
108
109pub trait DataView {
110 type D;
111 fn data(&self) -> &Self::D;
112}
113
114pub trait DataViewMut: DataView {
115 fn data_mut(&mut self) -> &mut Self::D;
116}
117
118pub trait ZnxView: ZnxInfos + DataView<D: AsRef<[u8]>> {
119 type Scalar: Copy;
120
121 /// Returns a non-mutable pointer to the underlying coefficients array.
122 fn as_ptr(&self) -> *const Self::Scalar {
123 self.data().as_ref().as_ptr() as *const Self::Scalar
124 }
125
126 /// Returns a non-mutable reference to the entire underlying coefficient array.
127 fn raw(&self) -> &[Self::Scalar] {
128 unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
129 }
130
131 /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
132 fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
133 #[cfg(debug_assertions)]
134 {
135 assert!(i < self.cols());
136 assert!(j < self.size());
137 }
138 let offset: usize = self.n() * (j * self.cols() + i);
139 unsafe { self.as_ptr().add(offset) }
140 }
141
142 /// Returns non-mutable reference to the (i, j)-th small polynomial.
143 fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
144 unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
145 }
146}
147
148pub trait ZnxViewMut: ZnxView + DataViewMut<D: AsMut<[u8]>> {
149 /// Returns a mutable pointer to the underlying coefficients array.
150 fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
151 self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
152 }
153
154 /// Returns a mutable reference to the entire underlying coefficient array.
155 fn raw_mut(&mut self) -> &mut [Self::Scalar] {
156 unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
157 }
158
159 /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
160 fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
161 #[cfg(debug_assertions)]
162 {
163 assert!(i < self.cols());
164 assert!(j < self.size());
165 }
166 let offset: usize = self.n() * (j * self.cols() + i);
167 unsafe { self.as_mut_ptr().add(offset) }
168 }
169
170 /// Returns mutable reference to the (i, j)-th small polynomial.
171 fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
172 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
173 }
174}
175
176//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
177impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: AsMut<[u8]>> {}
178
179use std::convert::TryFrom;
180use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
181pub trait Num:
182 Copy
183 + Default
184 + PartialEq
185 + PartialOrd
186 + Add<Output = Self>
187 + Sub<Output = Self>
188 + Mul<Output = Self>
189 + Div<Output = Self>
190 + Neg<Output = Self>
191 + AddAssign
192{
193 const BITS: u32;
194}
195
196impl Num for i64 {
197 const BITS: u32 = 64;
198}
199
200impl Num for i128 {
201 const BITS: u32 = 128;
202}
203
204impl Num for f64 {
205 const BITS: u32 = 64;
206}
207
208pub trait ZnxZero: ZnxViewMut
209where
210 Self: Sized,
211{
212 fn zero(&mut self) {
213 unsafe {
214 std::ptr::write_bytes(
215 self.as_mut_ptr(),
216 0,
217 self.n() * size_of::<Self::Scalar>() * self.poly_count(),
218 );
219 }
220 }
221
222 fn zero_at(&mut self, i: usize, j: usize) {
223 unsafe {
224 std::ptr::write_bytes(
225 self.at_mut_ptr(i, j),
226 0,
227 self.n() * size_of::<Self::Scalar>(),
228 );
229 }
230 }
231}
232
233// Blanket implementations
234impl<T> ZnxZero for T where T: ZnxViewMut {}
235// impl<T> ZnxRsh for T where T: ZnxZero {}
236
237pub fn switch_degree<S: Copy, DMut: ZnxViewMut<Scalar = S> + ZnxZero, D: ZnxView<Scalar = S>>(
238 b: &mut DMut,
239 col_b: usize,
240 a: &D,
241 col_a: usize,
242) {
243 let (n_in, n_out) = (a.n(), b.n());
244 let (gap_in, gap_out): (usize, usize);
245
246 if n_in > n_out {
247 (gap_in, gap_out) = (n_in / n_out, 1)
248 } else {
249 (gap_in, gap_out) = (1, n_out / n_in);
250 b.zero();
251 }
252
253 let size: usize = min(a.size(), b.size());
254
255 (0..size).for_each(|i| {
256 izip!(
257 a.at(col_a, i).iter().step_by(gap_in),
258 b.at_mut(col_b, i).iter_mut().step_by(gap_out)
259 )
260 .for_each(|(x_in, x_out)| *x_out = *x_in);
261 });
262}
263
264// (Jay)TODO: implement rsh for VecZnx, VecZnxBig
265// pub trait ZnxRsh: ZnxZero {
266// fn rsh(&mut self, k: usize, log_base2k: usize, col: usize, carry: &mut [u8]) {
267// rsh(k, log_base2k, self, col, carry)
268// }
269// }
270// pub fn rsh<V: ZnxRsh + ZnxZero>(k: usize, log_base2k: usize, a: &mut V, a_col: usize, tmp_bytes: &mut [u8]) {
271// let n: usize = a.n();
272// let size: usize = a.size();
273// let cols: usize = a.cols();
274
275// #[cfg(debug_assertions)]
276// {
277// assert!(
278// tmp_bytes.len() >= rsh_tmp_bytes::<V::Scalar>(n),
279// "invalid carry: carry.len()/size_ofSelf::Scalar={} < rsh_tmp_bytes({}, {})",
280// tmp_bytes.len() / size_of::<V::Scalar>(),
281// n,
282// size,
283// );
284// assert_alignement(tmp_bytes.as_ptr());
285// }
286
287// let size: usize = a.size();
288// let steps: usize = k / log_base2k;
289
290// a.raw_mut().rotate_right(n * steps * cols);
291// (0..cols).for_each(|i| {
292// (0..steps).for_each(|j| {
293// a.zero_at(i, j);
294// })
295// });
296
297// let k_rem: usize = k % log_base2k;
298
299// if k_rem != 0 {
300// let carry: &mut [V::Scalar] = cast_mut(tmp_bytes);
301
302// unsafe {
303// std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
304// }
305
306// let log_base2k_t: V::Scalar = V::Scalar::try_from(log_base2k).unwrap();
307// let shift: V::Scalar = V::Scalar::try_from(V::Scalar::BITS as usize - k_rem).unwrap();
308// let k_rem_t: V::Scalar = V::Scalar::try_from(k_rem).unwrap();
309
310// (steps..size).for_each(|i| {
311// izip!(carry.iter_mut(), a.at_mut(a_col, i).iter_mut()).for_each(|(ci, xi)| {
312// *xi += *ci << log_base2k_t;
313// *ci = get_base_k_carry(*xi, shift);
314// *xi = (*xi - *ci) >> k_rem_t;
315// });
316// })
317// }
318// }
319
320// #[inline(always)]
321// fn get_base_k_carry<T: Num>(x: T, shift: T) -> T {
322// (x << shift) >> shift
323// }
324
325// pub fn rsh_tmp_bytes<T: Num>(n: usize) -> usize {
326// n * std::mem::size_of::<T>()
327// }
328
329// pub trait ZnxLayout: ZnxInfos {
330// type Scalar;
331
332// /// Returns true if the receiver is only borrowing the data.
333// fn borrowing(&self) -> bool {
334// self.znx().data.len() == 0
335// }
336
337// /// Returns a non-mutable pointer to the underlying coefficients array.
338// fn as_ptr(&self) -> *const Self::Scalar {
339// self.znx().ptr as *const Self::Scalar
340// }
341
342// /// Returns a mutable pointer to the underlying coefficients array.
343// fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
344// self.znx_mut().ptr as *mut Self::Scalar
345// }
346
347// /// Returns a non-mutable reference to the entire underlying coefficient array.
348// fn raw(&self) -> &[Self::Scalar] {
349// unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
350// }
351
352// /// Returns a mutable reference to the entire underlying coefficient array.
353// fn raw_mut(&mut self) -> &mut [Self::Scalar] {
354// unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
355// }
356
357// /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
358// fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
359// #[cfg(debug_assertions)]
360// {
361// assert!(i < self.cols());
362// assert!(j < self.size());
363// }
364// let offset: usize = self.n() * (j * self.cols() + i);
365// unsafe { self.as_ptr().add(offset) }
366// }
367
368// /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
369// fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
370// #[cfg(debug_assertions)]
371// {
372// assert!(i < self.cols());
373// assert!(j < self.size());
374// }
375// let offset: usize = self.n() * (j * self.cols() + i);
376// unsafe { self.as_mut_ptr().add(offset) }
377// }
378
379// /// Returns non-mutable reference to the (i, j)-th small polynomial.
380// fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
381// unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
382// }
383
384// /// Returns mutable reference to the (i, j)-th small polynomial.
385// fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
386// unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
387// }
388// }