poulpy/base2k/src/vector_matrix_product.rs
base2k/src/vector_matrix_product.rsRaw
1use crate::ffi::vmp::{
2 delete_vmp_pmat, new_vmp_pmat, vmp_apply_dft, vmp_apply_dft_tmp_bytes, vmp_apply_dft_to_dft,
3 vmp_apply_dft_to_dft_tmp_bytes, vmp_pmat_t, vmp_prepare_contiguous,
4 vmp_prepare_contiguous_tmp_bytes,
5};
6use crate::{Module, VecZnx, VecZnxDft};
7use std::cmp::min;
8
9/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
10/// stored as a 3D matrix in the DFT domain in a single contiguous array.
11pub struct VmpPMat {
12 /// The pointer to the C memory.
13 pub data: *mut vmp_pmat_t,
14 /// The number of [VecZnx].
15 pub rows: usize,
16 /// The number of limbs in each [VecZnx].
17 pub cols: usize,
18 /// The ring degree of each [VecZnx].
19 pub n: usize,
20}
21
22impl VmpPMat {
23
24 /// Returns the pointer to the [vmp_pmat_t].
25 pub fn data(&self) -> *mut vmp_pmat_t {
26 self.data
27 }
28
29 /// Returns the number of rows of the [VmpPMat].
30 /// The number of rows (i.e. of [VecZnx]) of the [VmpPMat].
31 pub fn rows(&self) -> usize {
32 self.rows
33 }
34
35 /// Returns the number of cols of the [VmpPMat].
36 /// The number of cols refers to the number of limbs
37 /// of the prepared [VecZnx].
38 pub fn cols(&self) -> usize {
39 self.cols
40 }
41
42 /// Returns the ring dimension of the [VmpPMat].
43 pub fn n(&self) -> usize {
44 self.n
45 }
46
47 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
48 /// When using FFT64 as backend, T should be f64.
49 /// When using NTT120 as backend, T should be i64.
50 pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
51 let mut res: Vec<T> = vec![T::default(); self.n];
52
53 if self.n < 8 {
54 res.copy_from_slice(
55 &self.get_backend_array::<T>()[(row + col * self.rows()) * self.n()
56 ..(row + col * self.rows()) * (self.n() + 1)],
57 );
58 } else {
59 (0..self.n >> 3).for_each(|blk| {
60 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.get_array(row, col, blk)[..8]);
61 });
62 }
63
64 res
65 }
66
67 /// When using FFT64 as backend, T should be f64.
68 /// When using NTT120 as backend, T should be i64.
69 fn get_array<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
70 let nrows: usize = self.rows();
71 let ncols: usize = self.cols();
72 if col == (ncols - 1) && (ncols & 1 == 1) {
73 &self.get_backend_array::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
74 } else {
75 &self.get_backend_array::<T>()[blk * nrows * ncols * 8
76 + (col / 2) * (2 * nrows) * 8
77 + row * 2 * 8
78 + (col % 2) * 8..]
79 }
80 }
81
82 /// Returns a non-mutable reference of T to the entire contiguous array of the [VmpPMat].
83 /// When using FFT64 as backend, T should be f64.
84 /// When using NTT120 as backend, T should be i64.
85 /// The length of the returned array is rows * cols * n.
86 pub fn get_backend_array<T>(&self) -> &[T] {
87 let ptr: *const T = self.data as *const T;
88 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
89 unsafe { &std::slice::from_raw_parts(ptr, len) }
90 }
91
92 /// frees the memory and self destructs.
93 pub fn delete(self) {
94 unsafe { delete_vmp_pmat(self.data) };
95 drop(self);
96 }
97}
98
99impl Module {
100
101 /// Allocates a new [VmpPMat] with the given number of rows and columns.
102 pub fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
103 unsafe {
104 VmpPMat {
105 data: new_vmp_pmat(self.0, rows as u64, cols as u64),
106 rows,
107 cols,
108 n: self.n(),
109 }
110 }
111 }
112
113 /// Returns the number of bytes needed as scratch space for [Self::vmp_prepare_contiguous].
114 pub fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
115 unsafe { vmp_prepare_contiguous_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
116 }
117
118 /// Prepares a [VmpPMat] given a contiguous array of [i64].
119 /// The helper struct [Matrix3D] can be used to contruct the
120 /// appropriate contiguous array.
121 ///
122 /// # Example
123 /// ```
124 /// let mut b_mat: Matrix3D<i64> = Matrix3D::new(rows, cols, n);
125 ///
126 /// (0..min(rows, cols)).for_each(|i| {
127 /// b_mat.at_mut(i, i)[1] = 1 as i64;
128 /// });
129 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
130 /// module.vmp_prepare_contiguous(&mut vmp_pmat, &b_mat.data, &mut buf);
131 /// ```
132 pub fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
133 unsafe {
134 vmp_prepare_contiguous(
135 self.0,
136 b.data(),
137 a.as_ptr(),
138 b.rows() as u64,
139 b.cols() as u64,
140 buf.as_mut_ptr(),
141 );
142 }
143 }
144
145 pub fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]) {
146 let rows: usize = b.rows();
147 let cols: usize = b.cols();
148
149 let mut mat: Matrix3D<i64> = Matrix3D::<i64>::new(rows, cols, self.n());
150
151 (0..min(rows, a.len())).for_each(|i| {
152 mat.set_row(i, &a[i].data);
153 });
154
155 self.vmp_prepare_contiguous(b, &mat.data, buf);
156
157 /*
158 NOT IMPLEMENTED IN SPQLIOS
159 let mut ptrs: Vec<*const i64> = a.iter().map(|v| v.data.as_ptr()).collect();
160 unsafe {
161 vmp_prepare_dblptr(
162 self.0,
163 b.data(),
164 ptrs.as_mut_ptr(),
165 b.rows() as u64,
166 b.cols() as u64,
167 buf.as_mut_ptr(),
168 );
169 }
170 */
171 }
172
173 pub fn vmp_apply_dft_tmp_bytes(
174 &self,
175 c_limbs: usize,
176 a_limbs: usize,
177 rows: usize,
178 cols: usize,
179 ) -> usize {
180 unsafe {
181 vmp_apply_dft_tmp_bytes(
182 self.0,
183 c_limbs as u64,
184 a_limbs as u64,
185 rows as u64,
186 cols as u64,
187 ) as usize
188 }
189 }
190
191 pub fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
192 unsafe {
193 vmp_apply_dft(
194 self.0,
195 c.0,
196 c.limbs() as u64,
197 a.as_ptr(),
198 a.limbs() as u64,
199 a.n() as u64,
200 b.data(),
201 b.rows() as u64,
202 b.cols() as u64,
203 buf.as_mut_ptr(),
204 )
205 }
206 }
207
208 pub fn vmp_apply_dft_to_dft_tmp_bytes(
209 &self,
210 c_limbs: usize,
211 a_limbs: usize,
212 rows: usize,
213 cols: usize,
214 ) -> usize {
215 unsafe {
216 vmp_apply_dft_to_dft_tmp_bytes(
217 self.0,
218 c_limbs as u64,
219 a_limbs as u64,
220 rows as u64,
221 cols as u64,
222 ) as usize
223 }
224 }
225
226 pub fn vmp_apply_dft_to_dft(
227 &self,
228 c: &mut VecZnxDft,
229 a: &VecZnxDft,
230 b: &VmpPMat,
231 buf: &mut [u8],
232 ) {
233 unsafe {
234 vmp_apply_dft_to_dft(
235 self.0,
236 c.0,
237 c.limbs() as u64,
238 a.0,
239 a.limbs() as u64,
240 b.data(),
241 b.rows() as u64,
242 b.cols() as u64,
243 buf.as_mut_ptr(),
244 )
245 }
246 }
247
248 pub fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
249 unsafe {
250 vmp_apply_dft_to_dft(
251 self.0,
252 b.0,
253 b.limbs() as u64,
254 b.0,
255 b.limbs() as u64,
256 a.data(),
257 a.rows() as u64,
258 a.cols() as u64,
259 buf.as_mut_ptr(),
260 )
261 }
262 }
263}
264
265/// A helper struture that stores a 3D matrix as a contiguous array.
266/// To be passed to [Module::vmp_prepare_contiguous].
267///
268/// rows: index of the i-th base2K power.
269/// cols: index of the j-th limb of the i-th row.
270/// n : polynomial degree.
271///
272/// A [Matrix3D] can be seen as a vector of [VecZnx].
273pub struct Matrix3D<T> {
274 pub data: Vec<T>,
275 pub rows: usize,
276 pub cols: usize,
277 pub n: usize,
278}
279
280impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
281 /// Allocates a new [Matrix3D] with the respective dimensions.
282 ///
283 /// # Example
284 /// ```
285 /// let rows = 5; // #decomp
286 /// let cols = 5; // #limbs
287 /// let n = 1024; // #coeffs
288 ///
289 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
290 /// ```
291 pub fn new(rows: usize, cols: usize, n: usize) -> Self {
292 let size = rows * cols * n;
293 Self {
294 data: vec![T::default(); size],
295 rows,
296 cols,
297 n,
298 }
299 }
300
301 /// Returns a non-mutable reference to the entry (row, col) of the [Matrix3D].
302 /// The returned array is of size n.
303 ///
304 /// # Example
305 /// ```
306 /// let rows = 5; // #decomp
307 /// let cols = 5; // #limbs
308 /// let n = 1024; // #coeffs
309 ///
310 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
311 ///
312 /// let elem: &[i64] = mat.at(5, 5); // size n
313 /// ```
314 pub fn at(&self, row: usize, col: usize) -> &[T] {
315 assert!(row <= self.rows && col <= self.cols);
316 let idx: usize = row * (self.n * self.cols) + col * self.n;
317 &self.data[idx..idx + self.n]
318 }
319
320 /// Returns a mutable reference of the array at the (row, col) entry of the [Matrix3D].
321 /// The returned array is of size n.
322 ///
323 /// # Example
324 /// ```
325 /// let rows = 5; // #decomp
326 /// let cols = 5; // #limbs
327 /// let n = 1024; // #coeffs
328 ///
329 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
330 ///
331 /// let elem: &mut [i64] = mat.at_mut(5, 5); // size n
332 /// ```
333 pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
334 assert!(row <= self.rows && col <= self.cols);
335 let idx: usize = row * (self.n * self.cols) + col * self.n;
336 &mut self.data[idx..idx + self.n]
337 }
338
339 /// Sets the entry \[row\] of the [Matrix3D].
340 /// Typicall this is used to assign a [VecZnx] to the i-th row
341 /// of the [Matrix3D].
342 ///
343 /// # Example
344 /// ```
345 /// let rows = 5; // #decomp
346 /// let cols = 5; // #limbs
347 /// let n = 1024; // #coeffs
348 ///
349 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
350 ///
351 /// let a: Vec<i64> = VecZnx::new(n, cols);
352 ///
353 /// mat.set_row(1, &a.data);
354 /// ```
355 pub fn set_row(&mut self, row: usize, a: &[T]) {
356 assert!(
357 row < self.rows,
358 "invalid argument row: row={} > self.rows={}",
359 row,
360 self.rows
361 );
362 let idx: usize = row * (self.n * self.cols);
363 let size: usize = min(a.len(), self.cols * self.n);
364 self.data[idx..idx + size].copy_from_slice(&a[..size]);
365 }
366}