poulpy/base2k/src/vmp.rs
base2k/src/vmp.rsRaw
1use crate::ffi::vmp;
2use crate::{Infos, Module, VecZnx, 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 limbs 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 limbs 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 mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)];
157 ///
158 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
159 /// module.vmp_prepare_dblptr(&mut vmp_pmat, &vecznx, &mut buf);
160 ///
161 /// vmp_pmat.free();
162 /// module.free();
163 /// ```
164 fn vmp_prepare_dblptr<T: VecZnxApi + Infos>(&self, b: &mut VmpPMat, a: &Vec<T>, buf: &mut [u8]);
165
166 /// Prepares the ith-row of [VmpPMat] from a vector of [VecZnx].
167 ///
168 /// # Arguments
169 ///
170 /// * `b`: [VmpPMat] on which the values are encoded.
171 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
172 /// * `row_i`: the index of the row to prepare.
173 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
174 ///
175 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
176 /// /// # Example
177 /// ```
178 /// use base2k::{Module, FFT64, VmpPMat, VmpPMatOps, VecZnx, VecZnxOps, Free};
179 /// use std::cmp::min;
180 ///
181 /// let n: usize = 1024;
182 /// let module: Module = Module::new::<FFT64>(n);
183 /// let rows: usize = 5;
184 /// let cols: usize = 6;
185 ///
186 /// let vecznx = module.new_vec_znx(cols);
187 ///
188 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)];
189 ///
190 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
191 /// module.vmp_prepare_row(&mut vmp_pmat, &vecznx, 0, &mut buf);
192 ///
193 /// vmp_pmat.free();
194 /// module.free();
195 /// ```
196 fn vmp_prepare_row<T: VecZnxApi + Infos>(
197 &self,
198 b: &mut VmpPMat,
199 a: &T,
200 row_i: usize,
201 tmp_bytes: &mut [u8],
202 );
203
204 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft].
205 ///
206 /// # Arguments
207 ///
208 /// * `c_limbs`: number of limbs of the output [VecZnxDft].
209 /// * `a_limbs`: number of limbs of the input [VecZnx].
210 /// * `rows`: number of rows of the input [VmpPMat].
211 /// * `cols`: number of cols of the input [VmpPMat].
212 fn vmp_apply_dft_tmp_bytes(
213 &self,
214 c_limbs: usize,
215 a_limbs: usize,
216 rows: usize,
217 cols: usize,
218 ) -> usize;
219
220 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
221 ///
222 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
223 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
224 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
225 ///
226 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
227 /// `j` cols, the output is a [VecZnx] of `j` limbs.
228 ///
229 /// If there is a mismatch between the dimensions the largest valid ones are used.
230 ///
231 /// ```text
232 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
233 /// |h i j|
234 /// |k l m|
235 /// ```
236 /// where each element is a [VecZnxDft].
237 ///
238 /// # Arguments
239 ///
240 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
241 /// * `a`: the left operand [VecZnx] of the vector matrix product.
242 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
243 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
244 ///
245 /// # Example
246 /// ```
247 /// use base2k::{Module, VecZnx, VecZnxOps, VecZnxDft, VecZnxDftOps, VmpPMat, VmpPMatOps, FFT64, Free, VecZnxApi};
248 ///
249 /// let n = 1024;
250 ///
251 /// let module: Module = Module::new::<FFT64>(n);
252 /// let limbs: usize = 5;
253 ///
254 /// let rows: usize = limbs;
255 /// let cols: usize = limbs + 1;
256 /// let c_limbs: usize = cols;
257 /// let a_limbs: usize = limbs;
258 /// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
259 ///
260 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
261 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
262 ///
263 /// let a: VecZnx = module.new_vec_znx(limbs);
264 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
265 /// module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
266 ///
267 /// c_dft.free();
268 /// vmp_pmat.free();
269 /// module.free();
270 /// ```
271 fn vmp_apply_dft<T: VecZnxApi + Infos>(
272 &self,
273 c: &mut VecZnxDft,
274 a: &T,
275 b: &VmpPMat,
276 buf: &mut [u8],
277 );
278
279 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
280 ///
281 /// # Arguments
282 ///
283 /// * `c_limbs`: number of limbs of the output [VecZnxDft].
284 /// * `a_limbs`: number of limbs of the input [VecZnxDft].
285 /// * `rows`: number of rows of the input [VmpPMat].
286 /// * `cols`: number of cols of the input [VmpPMat].
287 fn vmp_apply_dft_to_dft_tmp_bytes(
288 &self,
289 c_limbs: usize,
290 a_limbs: usize,
291 rows: usize,
292 cols: usize,
293 ) -> usize;
294
295 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
296 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
297 ///
298 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
299 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
300 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
301 ///
302 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
303 /// `j` cols, the output is a [VecZnx] of `j` limbs.
304 ///
305 /// If there is a mismatch between the dimensions the largest valid ones are used.
306 ///
307 /// ```text
308 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
309 /// |h i j|
310 /// |k l m|
311 /// ```
312 /// where each element is a [VecZnxDft].
313 ///
314 /// # Arguments
315 ///
316 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
317 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
318 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
319 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
320 ///
321 /// # Example
322 /// ```
323 /// use base2k::{Module, VecZnx, VecZnxDft, VecZnxDftOps, VmpPMat, VmpPMatOps, FFT64, Free};
324 ///
325 /// let n = 1024;
326 ///
327 /// let module: Module = Module::new::<FFT64>(n);
328 /// let limbs: usize = 5;
329 ///
330 /// let rows: usize = limbs;
331 /// let cols: usize = limbs + 1;
332 /// let c_limbs: usize = cols;
333 /// let a_limbs: usize = limbs;
334 /// let tmp_bytes: usize = module.vmp_apply_dft_to_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
335 ///
336 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
337 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
338 ///
339 /// let a_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
340 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
341 /// module.vmp_apply_dft_to_dft(&mut c_dft, &a_dft, &vmp_pmat, &mut buf);
342 ///
343 /// a_dft.free();
344 /// c_dft.free();
345 /// vmp_pmat.free();
346 /// module.free();
347 /// ```
348 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
349
350 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
351 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
352 ///
353 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
354 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
355 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
356 ///
357 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
358 /// `j` cols, the output is a [VecZnx] of `j` limbs.
359 ///
360 /// If there is a mismatch between the dimensions the largest valid ones are used.
361 ///
362 /// ```text
363 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
364 /// |h i j|
365 /// |k l m|
366 /// ```
367 /// where each element is a [VecZnxDft].
368 ///
369 /// # Arguments
370 ///
371 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
372 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
373 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
374 ///
375 /// # Example
376 /// ```
377 /// use base2k::{Module, VecZnx, VecZnxOps, VecZnxDft, VmpPMat, VmpPMatOps, FFT64, Free, VecZnxApi, VecZnxDftOps};
378 ///
379 /// let n = 1024;
380 ///
381 /// let module: Module = Module::new::<FFT64>(n);
382 /// let limbs: usize = 5;
383 ///
384 /// let rows: usize = limbs;
385 /// let cols: usize = limbs + 1;
386 /// let tmp_bytes: usize = module.vmp_apply_dft_to_dft_tmp_bytes(limbs, limbs, rows, cols);
387 ///
388 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
389 /// let a: VecZnx = module.new_vec_znx(limbs);
390 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
391 ///
392 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
393 /// module.vmp_apply_dft_to_dft_inplace(&mut c_dft, &vmp_pmat, &mut buf);
394 ///
395 /// c_dft.free();
396 /// vmp_pmat.free();
397 /// module.free();
398 /// ```
399 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
400}
401
402impl VmpPMatOps for Module {
403 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
404 unsafe {
405 VmpPMat {
406 data: vmp::new_vmp_pmat(self.0, rows as u64, cols as u64),
407 rows,
408 cols,
409 n: self.n(),
410 }
411 }
412 }
413
414 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
415 unsafe { vmp::vmp_prepare_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
416 }
417
418 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
419 unsafe {
420 vmp::vmp_prepare_contiguous(
421 self.0,
422 b.data(),
423 a.as_ptr(),
424 b.rows() as u64,
425 b.cols() as u64,
426 buf.as_mut_ptr(),
427 );
428 }
429 }
430
431 fn vmp_prepare_dblptr<T: VecZnxApi + Infos>(
432 &self,
433 b: &mut VmpPMat,
434 a: &Vec<T>,
435 buf: &mut [u8],
436 ) {
437 let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
438 unsafe {
439 vmp::vmp_prepare_dblptr(
440 self.0,
441 b.data(),
442 ptrs.as_ptr(),
443 b.rows() as u64,
444 b.cols() as u64,
445 buf.as_mut_ptr(),
446 );
447 }
448 }
449
450 fn vmp_prepare_row<T: VecZnxApi + Infos>(
451 &self,
452 b: &mut VmpPMat,
453 a: &T,
454 row_i: usize,
455 buf: &mut [u8],
456 ) {
457 unsafe {
458 vmp::vmp_prepare_row(
459 self.0,
460 b.data(),
461 a.as_ptr(),
462 row_i as u64,
463 b.rows() as u64,
464 b.cols() as u64,
465 buf.as_mut_ptr(),
466 );
467 }
468 }
469
470 fn vmp_apply_dft_tmp_bytes(
471 &self,
472 c_limbs: usize,
473 a_limbs: usize,
474 rows: usize,
475 cols: usize,
476 ) -> usize {
477 unsafe {
478 vmp::vmp_apply_dft_tmp_bytes(
479 self.0,
480 c_limbs as u64,
481 a_limbs as u64,
482 rows as u64,
483 cols as u64,
484 ) as usize
485 }
486 }
487
488 fn vmp_apply_dft<T: VecZnxApi + Infos>(
489 &self,
490 c: &mut VecZnxDft,
491 a: &T,
492 b: &VmpPMat,
493 buf: &mut [u8],
494 ) {
495 unsafe {
496 vmp::vmp_apply_dft(
497 self.0,
498 c.0,
499 c.limbs() as u64,
500 a.as_ptr(),
501 a.limbs() as u64,
502 a.n() as u64,
503 b.data(),
504 b.rows() as u64,
505 b.cols() as u64,
506 buf.as_mut_ptr(),
507 )
508 }
509 }
510
511 fn vmp_apply_dft_to_dft_tmp_bytes(
512 &self,
513 c_limbs: usize,
514 a_limbs: usize,
515 rows: usize,
516 cols: usize,
517 ) -> usize {
518 unsafe {
519 vmp::vmp_apply_dft_to_dft_tmp_bytes(
520 self.0,
521 c_limbs as u64,
522 a_limbs as u64,
523 rows as u64,
524 cols as u64,
525 ) as usize
526 }
527 }
528
529 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]) {
530 unsafe {
531 vmp::vmp_apply_dft_to_dft(
532 self.0,
533 c.0,
534 c.limbs() as u64,
535 a.0,
536 a.limbs() as u64,
537 b.data(),
538 b.rows() as u64,
539 b.cols() as u64,
540 buf.as_mut_ptr(),
541 )
542 }
543 }
544
545 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
546 unsafe {
547 vmp::vmp_apply_dft_to_dft(
548 self.0,
549 b.0,
550 b.limbs() as u64,
551 b.0,
552 b.limbs() as u64,
553 a.data(),
554 a.rows() as u64,
555 a.cols() as u64,
556 buf.as_mut_ptr(),
557 )
558 }
559 }
560}