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