poulpy/base2k/src/vmp.rs
base2k/src/vmp.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::Free;
7use crate::{Module, VecZnx, VecZnxDft};
8use std::cmp::min;
9
10/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
11/// stored as a 3D matrix in the DFT domain in a single contiguous array.
12/// Each row of the [VmpPMat] can be seen as a [VecZnxDft].
13///
14/// The backend array of [VmpPMat] is allocate in C,
15/// and thus must be manually freed.
16///
17/// [VmpPMat] is used to permform a vector matrix product between a [VecZnx] and a [VmpPMat].
18/// See the trait [VectorMatrixProduct] for additional information.
19pub struct VmpPMat {
20 /// The pointer to the C memory.
21 pub data: *mut vmp_pmat_t,
22 /// The number of [VecZnxDft].
23 pub rows: usize,
24 /// The number of limbs in each [VecZnxDft].
25 pub cols: usize,
26 /// The ring degree of each [VecZnxDft].
27 pub n: usize,
28}
29
30impl VmpPMat {
31 /// Returns the pointer to the [vmp_pmat_t].
32 pub fn data(&self) -> *mut vmp_pmat_t {
33 self.data
34 }
35
36 /// Returns the number of rows of the [VmpPMat].
37 /// The number of rows (i.e. of [VecZnx]) of the [VmpPMat].
38 pub fn rows(&self) -> usize {
39 self.rows
40 }
41
42 /// Returns the number of cols of the [VmpPMat].
43 /// The number of cols refers to the number of limbs
44 /// of the prepared [VecZnx].
45 pub fn cols(&self) -> usize {
46 self.cols
47 }
48
49 /// Returns the ring dimension of the [VmpPMat].
50 pub fn n(&self) -> usize {
51 self.n
52 }
53
54 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
55 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
56 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
57 pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
58 let mut res: Vec<T> = vec![T::default(); self.n];
59
60 if self.n < 8 {
61 res.copy_from_slice(
62 &self.get_backend_array::<T>()[(row + col * self.rows()) * self.n()
63 ..(row + col * self.rows()) * (self.n() + 1)],
64 );
65 } else {
66 (0..self.n >> 3).for_each(|blk| {
67 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.get_array(row, col, blk)[..8]);
68 });
69 }
70
71 res
72 }
73
74 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
75 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
76 fn get_array<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
77 let nrows: usize = self.rows();
78 let ncols: usize = self.cols();
79 if col == (ncols - 1) && (ncols & 1 == 1) {
80 &self.get_backend_array::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
81 } else {
82 &self.get_backend_array::<T>()[blk * nrows * ncols * 8
83 + (col / 2) * (2 * nrows) * 8
84 + row * 2 * 8
85 + (col % 2) * 8..]
86 }
87 }
88
89 /// Returns a non-mutable reference of [T] of the entire contiguous array of the [VmpPMat].
90 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
91 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
92 /// The length of the returned array is rows * cols * n.
93 pub fn get_backend_array<T>(&self) -> &[T] {
94 let ptr: *const T = self.data as *const T;
95 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
96 unsafe { &std::slice::from_raw_parts(ptr, len) }
97 }
98}
99
100impl Free for VmpPMat {
101 fn free(self) {
102 unsafe { delete_vmp_pmat(self.data) };
103 drop(self);
104 }
105}
106
107/// This trait implements methods for vector matrix product,
108/// that is, multiplying a [VecZnx] with a [VmpPMat].
109pub trait VectorMatrixProduct {
110 /// Allocates a new [VmpPMat] with the given number of rows and columns.
111 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
112
113 /// Returns the number of bytes needed as scratch space for [VectorMatrixProduct::vmp_prepare_contiguous].
114 fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize;
115
116 /// Prepares a [VmpPMat] from a contiguous array of [i64].
117 /// The helper struct [Matrix3D] can be used to contruct and populate
118 /// the appropriate contiguous array.
119 ///
120 /// The size of buf can be obtained with [VectorMatrixProduct::vmp_prepare_contiguous_tmp_bytes].
121 ///
122 /// # Example
123 /// ```
124 /// use base2k::{Module, Matrix3D, VmpPMat, FFT64, Free};
125 /// use base2k::vmp::VectorMatrixProduct;
126 /// use std::cmp::min;
127 ///
128 /// let n: usize = 1024;
129 /// let module = Module::new::<FFT64>(n);
130 /// let rows = 5;
131 /// let cols = 6;
132 ///
133 /// let mut b_mat: Matrix3D<i64> = Matrix3D::new(rows, cols, n);
134 ///
135 /// // Populates the i-th row of b_math with X^1 * 2^(i * log_w) (here log_w is undefined)
136 /// (0..min(rows, cols)).for_each(|i| {
137 /// b_mat.at_mut(i, i)[1] = 1 as i64;
138 /// });
139 ///
140 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_contiguous_tmp_bytes(rows, cols)];
141 ///
142 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
143 /// module.vmp_prepare_contiguous(&mut vmp_pmat, &b_mat.data, &mut buf);
144 ///
145 /// vmp_pmat.free() // don't forget to free the memory once vmp_pmat is not needed anymore.
146 /// ```
147 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]);
148
149 /// Prepares a [VmpPMat] from a vector of [VecZnx].
150 ///
151 /// The size of buf can be obtained with [VectorMatrixProduct::vmp_prepare_contiguous_tmp_bytes].
152 ///
153 /// # Example
154 /// ```
155 /// use base2k::{Module, FFT64, Matrix3D, VmpPMat, VecZnx, Free};
156 /// use base2k::vmp::VectorMatrixProduct;
157 /// use std::cmp::min;
158 ///
159 /// let n: usize = 1024;
160 /// let module: Module = Module::new::<FFT64>(n);
161 /// let rows: usize = 5;
162 /// let cols: usize = 6;
163 ///
164 /// let mut vecznx: Vec<VecZnx>= Vec::new();
165 /// (0..rows).for_each(|_|{
166 /// vecznx.push(module.new_vec_znx(cols));
167 /// });
168 ///
169 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_contiguous_tmp_bytes(rows, cols)];
170 ///
171 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
172 /// module.vmp_prepare_dblptr(&mut vmp_pmat, &vecznx, &mut buf);
173 ///
174 /// vmp_pmat.free();
175 /// module.free();
176 /// ```
177 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]);
178
179 /// Returns the size of the stratch space necessary for [VectorMatrixProduct::vmp_apply_dft].
180 fn vmp_apply_dft_tmp_bytes(
181 &self,
182 c_limbs: usize,
183 a_limbs: usize,
184 rows: usize,
185 cols: usize,
186 ) -> usize;
187
188 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
189 /// The size of `buf` is given by [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
190 ///
191 /// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::svp_apply_dft]
192 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
193 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
194 ///
195 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
196 /// `j` cols, the output is a [VecZnx] of `j` limbs.
197 ///
198 /// If there is a mismatch between the dimensions the largest valid ones are used.
199 ///
200 /// ```text
201 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
202 /// |h i j|
203 /// |k l m|
204 /// ```
205 /// where each element is a [VecZnxDft].
206 ///
207 /// # Example
208 /// ```
209 /// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
210 /// use base2k::vmp::VectorMatrixProduct;
211 ///
212 /// let n = 1024;
213 ///
214 /// let module: Module = Module::new::<FFT64>(n);
215 /// let limbs: usize = 5;
216 ///
217 /// let rows: usize = limbs;
218 /// let cols: usize = limbs + 1;
219 /// let c_limbs: usize = cols;
220 /// let a_limbs: usize = limbs;
221 /// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
222 ///
223 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
224 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
225 ///
226 /// let a: VecZnx = module.new_vec_znx(limbs);
227 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
228 /// module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
229 ///
230 /// c_dft.free();
231 /// vmp_pmat.free();
232 /// module.free();
233 /// ```
234 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
235
236 /// Returns the size of the stratch space necessary for [VectorMatrixProduct::vmp_apply_dft_to_dft].
237 fn vmp_apply_dft_to_dft_tmp_bytes(
238 &self,
239 c_limbs: usize,
240 a_limbs: usize,
241 rows: usize,
242 cols: usize,
243 ) -> usize;
244
245 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
246 /// The size of `buf` is given by [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
247 ///
248 /// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::svp_apply_dft]
249 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
250 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
251 ///
252 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
253 /// `j` cols, the output is a [VecZnx] of `j` limbs.
254 ///
255 /// If there is a mismatch between the dimensions the largest valid ones are used.
256 ///
257 /// ```text
258 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
259 /// |h i j|
260 /// |k l m|
261 /// ```
262 /// where each element is a [VecZnxDft].
263 ///
264 /// # Example
265 /// ```
266 /// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
267 /// use base2k::vmp::VectorMatrixProduct;
268 ///
269 /// let n = 1024;
270 ///
271 /// let module: Module = Module::new::<FFT64>(n);
272 /// let limbs: usize = 5;
273 ///
274 /// let rows: usize = limbs;
275 /// let cols: usize = limbs + 1;
276 /// let c_limbs: usize = cols;
277 /// let a_limbs: usize = limbs;
278 /// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
279 ///
280 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
281 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
282 ///
283 /// let a_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
284 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
285 /// module.vmp_apply_dft_to_dft(&mut c_dft, &a_dft, &vmp_pmat, &mut buf);
286 ///
287 /// a_dft.free();
288 /// c_dft.free();
289 /// vmp_pmat.free();
290 /// module.free();
291 /// ```
292 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
293
294 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
295 /// The size of `buf` is given by [VectorMatrixProduct::vmp_apply_dft_to_dft_tmp_bytes].
296 ///
297 /// A vector matrix product is equivalent to a sum of [ScalarVectorProduct::svp_apply_dft]
298 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
299 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
300 ///
301 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
302 /// `j` cols, the output is a [VecZnx] of `j` limbs.
303 ///
304 /// If there is a mismatch between the dimensions the largest valid ones are used.
305 ///
306 /// ```text
307 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
308 /// |h i j|
309 /// |k l m|
310 /// ```
311 /// where each element is a [VecZnxDft].
312 ///
313 /// # Example
314 /// ```
315 /// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, FFT64, Free};
316 /// use base2k::vmp::VectorMatrixProduct;
317 ///
318 /// let n = 1024;
319 ///
320 /// let module: Module = Module::new::<FFT64>(n);
321 /// let limbs: usize = 5;
322 ///
323 /// let rows: usize = limbs;
324 /// let cols: usize = limbs + 1;
325 /// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(limbs, limbs, rows, cols);
326 ///
327 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
328 /// let a: VecZnx = module.new_vec_znx(limbs);
329 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
330 ///
331 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
332 /// module.vmp_apply_dft_to_dft_inplace(&mut c_dft, &vmp_pmat, &mut buf);
333 ///
334 /// c_dft.free();
335 /// vmp_pmat.free();
336 /// module.free();
337 /// ```
338 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
339}
340
341impl VectorMatrixProduct for Module {
342 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
343 unsafe {
344 VmpPMat {
345 data: new_vmp_pmat(self.0, rows as u64, cols as u64),
346 rows,
347 cols,
348 n: self.n(),
349 }
350 }
351 }
352 fn vmp_prepare_contiguous_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
353 unsafe { vmp_prepare_contiguous_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
354 }
355
356 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
357 unsafe {
358 vmp_prepare_contiguous(
359 self.0,
360 b.data(),
361 a.as_ptr(),
362 b.rows() as u64,
363 b.cols() as u64,
364 buf.as_mut_ptr(),
365 );
366 }
367 }
368
369 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &Vec<VecZnx>, buf: &mut [u8]) {
370 let rows: usize = b.rows();
371 let cols: usize = b.cols();
372
373 let mut mat: Matrix3D<i64> = Matrix3D::<i64>::new(rows, cols, self.n());
374
375 (0..min(rows, a.len())).for_each(|i| {
376 mat.set_row(i, &a[i].data);
377 });
378
379 self.vmp_prepare_contiguous(b, &mat.data, buf);
380
381 /*
382 NOT IMPLEMENTED IN SPQLIOS
383 let mut ptrs: Vec<*const i64> = a.iter().map(|v| v.data.as_ptr()).collect();
384 unsafe {
385 vmp_prepare_dblptr(
386 self.0,
387 b.data(),
388 ptrs.as_mut_ptr(),
389 b.rows() as u64,
390 b.cols() as u64,
391 buf.as_mut_ptr(),
392 );
393 }
394 */
395 }
396
397 fn vmp_apply_dft_tmp_bytes(
398 &self,
399 c_limbs: usize,
400 a_limbs: usize,
401 rows: usize,
402 cols: usize,
403 ) -> usize {
404 unsafe {
405 vmp_apply_dft_tmp_bytes(
406 self.0,
407 c_limbs as u64,
408 a_limbs as u64,
409 rows as u64,
410 cols as u64,
411 ) as usize
412 }
413 }
414
415 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
416 unsafe {
417 vmp_apply_dft(
418 self.0,
419 c.0,
420 c.limbs() as u64,
421 a.as_ptr(),
422 a.limbs() as u64,
423 a.n() as u64,
424 b.data(),
425 b.rows() as u64,
426 b.cols() as u64,
427 buf.as_mut_ptr(),
428 )
429 }
430 }
431
432 fn vmp_apply_dft_to_dft_tmp_bytes(
433 &self,
434 c_limbs: usize,
435 a_limbs: usize,
436 rows: usize,
437 cols: usize,
438 ) -> usize {
439 unsafe {
440 vmp_apply_dft_to_dft_tmp_bytes(
441 self.0,
442 c_limbs as u64,
443 a_limbs as u64,
444 rows as u64,
445 cols as u64,
446 ) as usize
447 }
448 }
449
450 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]) {
451 unsafe {
452 vmp_apply_dft_to_dft(
453 self.0,
454 c.0,
455 c.limbs() as u64,
456 a.0,
457 a.limbs() as u64,
458 b.data(),
459 b.rows() as u64,
460 b.cols() as u64,
461 buf.as_mut_ptr(),
462 )
463 }
464 }
465
466 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
467 unsafe {
468 vmp_apply_dft_to_dft(
469 self.0,
470 b.0,
471 b.limbs() as u64,
472 b.0,
473 b.limbs() as u64,
474 a.data(),
475 a.rows() as u64,
476 a.cols() as u64,
477 buf.as_mut_ptr(),
478 )
479 }
480 }
481}
482
483/// A helper struture that stores a 3D matrix as a contiguous array.
484/// To be passed to [VectorMatrixProduct::vmp_prepare_contiguous].
485///
486/// rows: index of the i-th base2K power.
487/// cols: index of the j-th limb of the i-th row.
488/// n : polynomial degree.
489///
490/// A [Matrix3D] can be seen as a vector of [VecZnx].
491pub struct Matrix3D<T> {
492 pub data: Vec<T>,
493 pub rows: usize,
494 pub cols: usize,
495 pub n: usize,
496}
497
498impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
499 /// Allocates a new [Matrix3D] with the respective dimensions.
500 ///
501 /// # Example
502 /// ```
503 /// use base2k::Matrix3D;
504 ///
505 /// let rows = 5; // #decomp
506 /// let cols = 5; // #limbs
507 /// let n = 1024; // #coeffs
508 ///
509 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
510 /// ```
511 pub fn new(rows: usize, cols: usize, n: usize) -> Self {
512 let size = rows * cols * n;
513 Self {
514 data: vec![T::default(); size],
515 rows,
516 cols,
517 n,
518 }
519 }
520
521 /// Returns a non-mutable reference to the entry (row, col) of the [Matrix3D].
522 /// The returned array is of size n.
523 ///
524 /// # Example
525 /// ```
526 /// use base2k::Matrix3D;
527 ///
528 /// let rows = 5; // #decomp
529 /// let cols = 5; // #limbs
530 /// let n = 1024; // #coeffs
531 ///
532 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
533 ///
534 /// let elem: &[i64] = mat.at(4, 4); // size n
535 /// ```
536 pub fn at(&self, row: usize, col: usize) -> &[T] {
537 assert!(row < self.rows && col < self.cols);
538 let idx: usize = row * (self.n * self.cols) + col * self.n;
539 &self.data[idx..idx + self.n]
540 }
541
542 /// Returns a mutable reference of the array at the (row, col) entry of the [Matrix3D].
543 /// The returned array is of size n.
544 ///
545 /// # Example
546 /// ```
547 /// use base2k::Matrix3D;
548 ///
549 /// let rows = 5; // #decomp
550 /// let cols = 5; // #limbs
551 /// let n = 1024; // #coeffs
552 ///
553 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
554 ///
555 /// let elem: &mut [i64] = mat.at_mut(4, 4); // size n
556 /// ```
557 pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
558 assert!(row < self.rows && col < self.cols);
559 let idx: usize = row * (self.n * self.cols) + col * self.n;
560 &mut self.data[idx..idx + self.n]
561 }
562
563 /// Sets the entry \[row\] of the [Matrix3D].
564 /// Typicall this is used to assign a [VecZnx] to the i-th row
565 /// of the [Matrix3D].
566 ///
567 /// # Example
568 /// ```
569 /// use base2k::{Matrix3D, VecZnx};
570 ///
571 /// let rows = 5; // #decomp
572 /// let cols = 5; // #limbs
573 /// let n = 1024; // #coeffs
574 ///
575 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
576 ///
577 /// let a: VecZnx = VecZnx::new(n, cols);
578 ///
579 /// mat.set_row(1, &a.data);
580 /// ```
581 pub fn set_row(&mut self, row: usize, a: &[T]) {
582 assert!(
583 row < self.rows,
584 "invalid argument row: row={} > self.rows={}",
585 row,
586 self.rows
587 );
588 let idx: usize = row * (self.n * self.cols);
589 let size: usize = min(a.len(), self.cols * self.n);
590 self.data[idx..idx + size].copy_from_slice(&a[..size]);
591 }
592}