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