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 #[warn(dead_code)]
30 backend: BACKEND,
31}
32
33impl Infos for VmpPMat {
34 /// Returns the ring dimension of the [VmpPMat].
35 fn n(&self) -> usize {
36 self.n
37 }
38
39 fn log_n(&self) -> usize {
40 (usize::BITS - (self.n() - 1).leading_zeros()) as _
41 }
42
43 /// Returns the number of rows (i.e. of [VecZnxDft]) of the [VmpPMat]
44 fn rows(&self) -> usize {
45 self.rows
46 }
47
48 /// Returns the number of cols of the [VmpPMat].
49 /// The number of cols refers to the number of cols
50 /// of each [VecZnxDft].
51 /// This method is equivalent to [Self::cols].
52 fn cols(&self) -> usize {
53 self.cols
54 }
55}
56
57impl VmpPMat {
58 pub fn as_ptr(&self) -> *const u8 {
59 self.ptr
60 }
61
62 pub fn as_mut_ptr(&self) -> *mut u8 {
63 self.ptr
64 }
65
66 pub fn borrowed(&self) -> bool {
67 self.data.len() == 0
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 raw<T>(&self) -> &[T] {
75 let ptr: *const T = self.ptr 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 /// Returns a non-mutable reference of `T` of the entire contiguous array of the [VmpPMat].
81 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
82 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
83 /// The length of the returned array is rows * cols * n.
84 pub fn raw_mut<T>(&self) -> &mut [T] {
85 let ptr: *mut T = self.ptr as *mut T;
86 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
87 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
88 }
89
90 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
91 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
92 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
93 ///
94 /// # Arguments
95 ///
96 /// * `row`: row index (i).
97 /// * `col`: col index (j).
98 pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
99 let mut res: Vec<T> = alloc_aligned(self.n);
100
101 if self.n < 8 {
102 res.copy_from_slice(
103 &self.raw::<T>()[(row + col * self.rows()) * self.n()
104 ..(row + col * self.rows()) * (self.n() + 1)],
105 );
106 } else {
107 (0..self.n >> 3).for_each(|blk| {
108 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
109 });
110 }
111
112 res
113 }
114
115 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
116 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
117 fn at_block<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
118 let nrows: usize = self.rows();
119 let ncols: usize = self.cols();
120 if col == (ncols - 1) && (ncols & 1 == 1) {
121 &self.raw::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
122 } else {
123 &self.raw::<T>()[blk * nrows * ncols * 8
124 + (col / 2) * (2 * nrows) * 8
125 + row * 2 * 8
126 + (col % 2) * 8..]
127 }
128 }
129}
130
131/// This trait implements methods for vector matrix product,
132/// that is, multiplying a [VecZnx] with a [VmpPMat].
133pub trait VmpPMatOps {
134 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize;
135
136 /// Allocates a new [VmpPMat] with the given number of rows and columns.
137 ///
138 /// # Arguments
139 ///
140 /// * `rows`: number of rows (number of [VecZnxDft]).
141 /// * `cols`: number of cols (number of cols of each [VecZnxDft]).
142 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
143
144 /// Returns the number of bytes needed as scratch space for [VmpPMatOps::vmp_prepare_contiguous].
145 ///
146 /// # Arguments
147 ///
148 /// * `rows`: number of rows of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
149 /// * `cols`: number of cols of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
150 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize;
151
152 /// Prepares a [VmpPMat] from a contiguous array of [i64].
153 /// The helper struct [Matrix3D] can be used to contruct and populate
154 /// the appropriate contiguous array.
155 ///
156 /// # Arguments
157 ///
158 /// * `b`: [VmpPMat] on which the values are encoded.
159 /// * `a`: the contiguous array of [i64] of the 3D matrix to encode on the [VmpPMat].
160 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
161 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]);
162
163 /// Prepares a [VmpPMat] from a vector of [VecZnx].
164 ///
165 /// # Arguments
166 ///
167 /// * `b`: [VmpPMat] on which the values are encoded.
168 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
169 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
170 ///
171 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
172 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]);
173
174 /// Prepares the ith-row of [VmpPMat] from a [VecZnx].
175 ///
176 /// # Arguments
177 ///
178 /// * `b`: [VmpPMat] on which the values are encoded.
179 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
180 /// * `row_i`: the index of the row to prepare.
181 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
182 ///
183 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
184 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]);
185
186 /// Extracts the ith-row of [VmpPMat] into a [VecZnxBig].
187 ///
188 /// # Arguments
189 ///
190 /// * `b`: the [VecZnxBig] to on which to extract the row of the [VmpPMat].
191 /// * `a`: [VmpPMat] on which the values are encoded.
192 /// * `row_i`: the index of the row to extract.
193 fn vmp_extract_row(&self, b: &mut VecZnxBig, a: &VmpPMat, row_i: usize);
194
195 /// Prepares the ith-row of [VmpPMat] from a [VecZnxDft].
196 ///
197 /// # Arguments
198 ///
199 /// * `b`: [VmpPMat] on which the values are encoded.
200 /// * `a`: the [VecZnxDft] to encode on the [VmpPMat].
201 /// * `row_i`: the index of the row to prepare.
202 ///
203 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
204 fn vmp_prepare_row_dft(&self, b: &mut VmpPMat, a: &VecZnxDft, row_i: usize);
205
206 /// Extracts the ith-row of [VmpPMat] into a [VecZnxDft].
207 ///
208 /// # Arguments
209 ///
210 /// * `b`: the [VecZnxDft] to on which to extract the row of the [VmpPMat].
211 /// * `a`: [VmpPMat] on which the values are encoded.
212 /// * `row_i`: the index of the row to extract.
213 fn vmp_extract_row_dft(&self, b: &mut VecZnxDft, a: &VmpPMat, row_i: usize);
214
215 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft].
216 ///
217 /// # Arguments
218 ///
219 /// * `c_cols`: number of cols of the output [VecZnxDft].
220 /// * `a_cols`: number of cols of the input [VecZnx].
221 /// * `rows`: number of rows of the input [VmpPMat].
222 /// * `cols`: number of cols of the input [VmpPMat].
223 fn vmp_apply_dft_tmp_bytes(
224 &self,
225 c_cols: usize,
226 a_cols: usize,
227 rows: usize,
228 cols: usize,
229 ) -> usize;
230
231 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
232 ///
233 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
234 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
235 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
236 ///
237 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
238 /// `j` cols, the output is a [VecZnx] of `j` cols.
239 ///
240 /// If there is a mismatch between the dimensions the largest valid ones are used.
241 ///
242 /// ```text
243 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
244 /// |h i j|
245 /// |k l m|
246 /// ```
247 /// where each element is a [VecZnxDft].
248 ///
249 /// # Arguments
250 ///
251 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
252 /// * `a`: the left operand [VecZnx] of the vector matrix product.
253 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
254 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
255 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
256
257 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on the receiver.
258 ///
259 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
260 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
261 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
262 ///
263 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
264 /// `j` cols, the output is a [VecZnx] of `j` cols.
265 ///
266 /// If there is a mismatch between the dimensions the largest valid ones are used.
267 ///
268 /// ```text
269 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
270 /// |h i j|
271 /// |k l m|
272 /// ```
273 /// where each element is a [VecZnxDft].
274 ///
275 /// # Arguments
276 ///
277 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
278 /// * `a`: the left operand [VecZnx] of the vector matrix product.
279 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
280 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
281 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
282
283 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
284 ///
285 /// # Arguments
286 ///
287 /// * `c_cols`: number of cols of the output [VecZnxDft].
288 /// * `a_cols`: number of cols of the input [VecZnxDft].
289 /// * `rows`: number of rows of the input [VmpPMat].
290 /// * `cols`: number of cols of the input [VmpPMat].
291 fn vmp_apply_dft_to_dft_tmp_bytes(
292 &self,
293 c_cols: usize,
294 a_cols: usize,
295 rows: usize,
296 cols: usize,
297 ) -> usize;
298
299 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
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 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
321 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
322 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
323 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
324 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
325
326 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on top of the receiver instead of overwritting it.
327 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
328 ///
329 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
330 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
331 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
332 ///
333 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
334 /// `j` cols, the output is a [VecZnx] of `j` cols.
335 ///
336 /// If there is a mismatch between the dimensions the largest valid ones are used.
337 ///
338 /// ```text
339 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
340 /// |h i j|
341 /// |k l m|
342 /// ```
343 /// where each element is a [VecZnxDft].
344 ///
345 /// # Arguments
346 ///
347 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
348 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
349 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
350 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
351 fn vmp_apply_dft_to_dft_add(
352 &self,
353 c: &mut VecZnxDft,
354 a: &VecZnxDft,
355 b: &VmpPMat,
356 buf: &mut [u8],
357 );
358
359 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
360 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
361 ///
362 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
363 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
364 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
365 ///
366 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
367 /// `j` cols, the output is a [VecZnx] of `j` cols.
368 ///
369 /// If there is a mismatch between the dimensions the largest valid ones are used.
370 ///
371 /// ```text
372 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
373 /// |h i j|
374 /// |k l m|
375 /// ```
376 /// where each element is a [VecZnxDft].
377 ///
378 /// # Arguments
379 ///
380 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
381 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
382 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
383 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
384}
385
386impl VmpPMatOps for Module {
387 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize {
388 unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, cols as u64) as usize }
389 }
390
391 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
392 let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vmp_pmat(rows, cols));
393 let ptr: *mut u8 = data.as_mut_ptr();
394 VmpPMat {
395 data: data,
396 ptr: ptr,
397 n: self.n(),
398 cols: cols,
399 rows: rows,
400 backend: self.backend(),
401 }
402 }
403
404 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
405 unsafe { vmp::vmp_prepare_tmp_bytes(self.ptr, rows as u64, cols as u64) as usize }
406 }
407
408 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], tmp_bytes: &mut [u8]) {
409 debug_assert_eq!(a.len(), b.n * b.rows * b.cols);
410 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
411 #[cfg(debug_assertions)]
412 {
413 assert_alignement(tmp_bytes.as_ptr());
414 }
415 unsafe {
416 vmp::vmp_prepare_contiguous(
417 self.ptr,
418 b.as_mut_ptr() as *mut vmp_pmat_t,
419 a.as_ptr(),
420 b.rows() as u64,
421 b.cols() as u64,
422 tmp_bytes.as_mut_ptr(),
423 );
424 }
425 }
426
427 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], tmp_bytes: &mut [u8]) {
428 let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
429 #[cfg(debug_assertions)]
430 {
431 debug_assert_eq!(a.len(), b.rows);
432 a.iter().for_each(|ai| {
433 debug_assert_eq!(ai.len(), b.n * b.cols);
434 });
435 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
436 assert_alignement(tmp_bytes.as_ptr());
437 }
438 unsafe {
439 vmp::vmp_prepare_dblptr(
440 self.ptr,
441 b.as_mut_ptr() as *mut vmp_pmat_t,
442 ptrs.as_ptr(),
443 b.rows() as u64,
444 b.cols() as u64,
445 tmp_bytes.as_mut_ptr(),
446 );
447 }
448 }
449
450 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]) {
451 debug_assert_eq!(a.len(), b.cols() * self.n());
452 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
453 #[cfg(debug_assertions)]
454 {
455 assert_alignement(tmp_bytes.as_ptr());
456 }
457 unsafe {
458 vmp::vmp_prepare_row(
459 self.ptr,
460 b.as_mut_ptr() as *mut vmp_pmat_t,
461 a.as_ptr(),
462 row_i as u64,
463 b.rows() as u64,
464 b.cols() as u64,
465 tmp_bytes.as_mut_ptr(),
466 );
467 }
468 }
469
470 fn vmp_extract_row(&self, b: &mut VecZnxBig, a: &VmpPMat, row_i: usize) {
471 #[cfg(debug_assertions)]
472 {
473 assert_eq!(a.n(), b.n());
474 assert_eq!(a.cols(), b.cols());
475 }
476 unsafe {
477 vmp::vmp_extract_row(
478 self.ptr,
479 b.ptr as *mut vec_znx_big_t,
480 a.as_ptr() as *const vmp_pmat_t,
481 row_i as u64,
482 a.rows() as u64,
483 a.cols() as u64,
484 );
485 }
486 }
487
488 fn vmp_prepare_row_dft(&self, b: &mut VmpPMat, a: &VecZnxDft, row_i: usize) {
489 #[cfg(debug_assertions)]
490 {
491 assert_eq!(a.n(), b.n());
492 assert_eq!(a.cols(), b.cols());
493 }
494 unsafe {
495 vmp::vmp_prepare_row_dft(
496 self.ptr,
497 b.as_mut_ptr() as *mut vmp_pmat_t,
498 a.ptr as *const vec_znx_dft_t,
499 row_i as u64,
500 b.rows() as u64,
501 b.cols() as u64,
502 );
503 }
504 }
505
506 fn vmp_extract_row_dft(&self, b: &mut VecZnxDft, a: &VmpPMat, row_i: usize) {
507 #[cfg(debug_assertions)]
508 {
509 assert_eq!(a.n(), b.n());
510 assert_eq!(a.cols(), b.cols());
511 }
512 unsafe {
513 vmp::vmp_extract_row_dft(
514 self.ptr,
515 b.ptr as *mut vec_znx_dft_t,
516 a.as_ptr() as *const vmp_pmat_t,
517 row_i as u64,
518 a.rows() as u64,
519 a.cols() as u64,
520 );
521 }
522 }
523
524 fn vmp_apply_dft_tmp_bytes(
525 &self,
526 res_cols: usize,
527 a_cols: usize,
528 gct_rows: usize,
529 gct_cols: usize,
530 ) -> usize {
531 unsafe {
532 vmp::vmp_apply_dft_tmp_bytes(
533 self.ptr,
534 res_cols as u64,
535 a_cols as u64,
536 gct_rows as u64,
537 gct_cols as u64,
538 ) as usize
539 }
540 }
541
542 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
543 debug_assert!(
544 tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
545 );
546 #[cfg(debug_assertions)]
547 {
548 assert_alignement(tmp_bytes.as_ptr());
549 }
550 unsafe {
551 vmp::vmp_apply_dft(
552 self.ptr,
553 c.ptr as *mut vec_znx_dft_t,
554 c.cols() as u64,
555 a.as_ptr(),
556 a.cols() as u64,
557 a.n() as u64,
558 b.as_ptr() as *const vmp_pmat_t,
559 b.rows() as u64,
560 b.cols() as u64,
561 tmp_bytes.as_mut_ptr(),
562 )
563 }
564 }
565
566 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
567 debug_assert!(
568 tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
569 );
570 #[cfg(debug_assertions)]
571 {
572 assert_alignement(tmp_bytes.as_ptr());
573 }
574 unsafe {
575 vmp::vmp_apply_dft_add(
576 self.ptr,
577 c.ptr as *mut vec_znx_dft_t,
578 c.cols() as u64,
579 a.as_ptr(),
580 a.cols() as u64,
581 a.n() as u64,
582 b.as_ptr() as *const vmp_pmat_t,
583 b.rows() as u64,
584 b.cols() as u64,
585 tmp_bytes.as_mut_ptr(),
586 )
587 }
588 }
589
590 fn vmp_apply_dft_to_dft_tmp_bytes(
591 &self,
592 res_cols: usize,
593 a_cols: usize,
594 gct_rows: usize,
595 gct_cols: usize,
596 ) -> usize {
597 unsafe {
598 vmp::vmp_apply_dft_to_dft_tmp_bytes(
599 self.ptr,
600 res_cols as u64,
601 a_cols as u64,
602 gct_rows as u64,
603 gct_cols as u64,
604 ) as usize
605 }
606 }
607
608 fn vmp_apply_dft_to_dft(
609 &self,
610 c: &mut VecZnxDft,
611 a: &VecZnxDft,
612 b: &VmpPMat,
613 tmp_bytes: &mut [u8],
614 ) {
615 debug_assert!(
616 tmp_bytes.len()
617 >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
618 );
619 #[cfg(debug_assertions)]
620 {
621 assert_alignement(tmp_bytes.as_ptr());
622 }
623 unsafe {
624 vmp::vmp_apply_dft_to_dft(
625 self.ptr,
626 c.ptr as *mut vec_znx_dft_t,
627 c.cols() as u64,
628 a.ptr as *const vec_znx_dft_t,
629 a.cols() as u64,
630 b.as_ptr() as *const vmp_pmat_t,
631 b.rows() as u64,
632 b.cols() as u64,
633 tmp_bytes.as_mut_ptr(),
634 )
635 }
636 }
637
638 fn vmp_apply_dft_to_dft_add(
639 &self,
640 c: &mut VecZnxDft,
641 a: &VecZnxDft,
642 b: &VmpPMat,
643 tmp_bytes: &mut [u8],
644 ) {
645 debug_assert!(
646 tmp_bytes.len()
647 >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
648 );
649 #[cfg(debug_assertions)]
650 {
651 assert_alignement(tmp_bytes.as_ptr());
652 }
653 unsafe {
654 vmp::vmp_apply_dft_to_dft_add(
655 self.ptr,
656 c.ptr as *mut vec_znx_dft_t,
657 c.cols() as u64,
658 a.ptr as *const vec_znx_dft_t,
659 a.cols() as u64,
660 b.as_ptr() as *const vmp_pmat_t,
661 b.rows() as u64,
662 b.cols() as u64,
663 tmp_bytes.as_mut_ptr(),
664 )
665 }
666 }
667
668 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, tmp_bytes: &mut [u8]) {
669 debug_assert!(
670 tmp_bytes.len()
671 >= self.vmp_apply_dft_to_dft_tmp_bytes(b.cols(), b.cols(), a.rows(), a.cols())
672 );
673 #[cfg(debug_assertions)]
674 {
675 assert_alignement(tmp_bytes.as_ptr());
676 }
677 unsafe {
678 vmp::vmp_apply_dft_to_dft(
679 self.ptr,
680 b.ptr as *mut vec_znx_dft_t,
681 b.cols() as u64,
682 b.ptr as *mut vec_znx_dft_t,
683 b.cols() as u64,
684 a.as_ptr() as *const vmp_pmat_t,
685 a.rows() as u64,
686 a.cols() as u64,
687 tmp_bytes.as_mut_ptr(),
688 )
689 }
690 }
691}
692
693#[cfg(test)]
694mod tests {
695 use crate::{
696 alloc_aligned, Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps,
697 VecZnxOps, VmpPMat, VmpPMatOps,
698 };
699 use sampling::source::Source;
700
701 #[test]
702 fn vmp_prepare_row_dft() {
703 let module: Module = Module::new(32, crate::BACKEND::FFT64);
704 let vpmat_rows: usize = 4;
705 let vpmat_cols: usize = 5;
706 let log_base2k: usize = 8;
707 let mut a: VecZnx = module.new_vec_znx(vpmat_cols);
708 let mut a_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
709 let mut a_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
710 let mut b_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
711 let mut b_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
712 let mut vmpmat_0: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
713 let mut vmpmat_1: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
714
715 let mut tmp_bytes: Vec<u8> =
716 alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, vpmat_cols));
717
718 for row_i in 0..vpmat_rows {
719 let mut source: Source = Source::new([0u8; 32]);
720 module.fill_uniform(log_base2k, &mut a, vpmat_cols, &mut source);
721 module.vec_znx_dft(&mut a_dft, &a);
722 module.vmp_prepare_row(&mut vmpmat_0, &a.raw(), row_i, &mut tmp_bytes);
723
724 // Checks that prepare(vmp_pmat, a) = prepare_dft(vmp_pmat, a_dft)
725 module.vmp_prepare_row_dft(&mut vmpmat_1, &a_dft, row_i);
726 assert_eq!(vmpmat_0.raw::<u8>(), vmpmat_1.raw::<u8>());
727
728 // Checks that a_dft = extract_dft(prepare(vmp_pmat, a), b_dft)
729 module.vmp_extract_row_dft(&mut b_dft, &vmpmat_0, row_i);
730 assert_eq!(a_dft.raw::<u8>(&module), b_dft.raw::<u8>(&module));
731
732 // Checks that a_big = extract(prepare_dft(vmp_pmat, a_dft), b_big)
733 module.vmp_extract_row(&mut b_big, &vmpmat_0, row_i);
734 module.vec_znx_idft(&mut a_big, &a_dft, &mut tmp_bytes);
735 assert_eq!(a_big.raw::<i64>(&module), b_big.raw::<i64>(&module));
736 }
737
738 module.free();
739 }
740}