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 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on the receiver.
257 ///
258 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
259 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
260 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
261 ///
262 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
263 /// `j` cols, the output is a [VecZnx] of `j` cols.
264 ///
265 /// If there is a mismatch between the dimensions the largest valid ones are used.
266 ///
267 /// ```text
268 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
269 /// |h i j|
270 /// |k l m|
271 /// ```
272 /// where each element is a [VecZnxDft].
273 ///
274 /// # Arguments
275 ///
276 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
277 /// * `a`: the left operand [VecZnx] of the vector matrix product.
278 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
279 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
280 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
281
282 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
283 ///
284 /// # Arguments
285 ///
286 /// * `c_cols`: number of cols of the output [VecZnxDft].
287 /// * `a_cols`: number of cols of the input [VecZnxDft].
288 /// * `rows`: number of rows of the input [VmpPMat].
289 /// * `cols`: number of cols of the input [VmpPMat].
290 fn vmp_apply_dft_to_dft_tmp_bytes(
291 &self,
292 c_cols: usize,
293 a_cols: usize,
294 rows: usize,
295 cols: usize,
296 ) -> usize;
297
298 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
299 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
300 ///
301 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
302 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
303 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
304 ///
305 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
306 /// `j` cols, the output is a [VecZnx] of `j` cols.
307 ///
308 /// If there is a mismatch between the dimensions the largest valid ones are used.
309 ///
310 /// ```text
311 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
312 /// |h i j|
313 /// |k l m|
314 /// ```
315 /// where each element is a [VecZnxDft].
316 ///
317 /// # Arguments
318 ///
319 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
320 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
321 /// * `b`: 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(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
324
325 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on top of the receiver instead of overwritting it.
326 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
327 ///
328 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
329 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
330 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
331 ///
332 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
333 /// `j` cols, the output is a [VecZnx] of `j` cols.
334 ///
335 /// If there is a mismatch between the dimensions the largest valid ones are used.
336 ///
337 /// ```text
338 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
339 /// |h i j|
340 /// |k l m|
341 /// ```
342 /// where each element is a [VecZnxDft].
343 ///
344 /// # Arguments
345 ///
346 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
347 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
348 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
349 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
350 fn vmp_apply_dft_to_dft_add(
351 &self,
352 c: &mut VecZnxDft,
353 a: &VecZnxDft,
354 b: &VmpPMat,
355 buf: &mut [u8],
356 );
357
358 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
359 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
360 ///
361 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
362 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
363 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
364 ///
365 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
366 /// `j` cols, the output is a [VecZnx] of `j` cols.
367 ///
368 /// If there is a mismatch between the dimensions the largest valid ones are used.
369 ///
370 /// ```text
371 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
372 /// |h i j|
373 /// |k l m|
374 /// ```
375 /// where each element is a [VecZnxDft].
376 ///
377 /// # Arguments
378 ///
379 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
380 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
381 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
382 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
383}
384
385impl VmpPMatOps for Module {
386 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize {
387 unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, cols as u64) as usize }
388 }
389
390 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
391 let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vmp_pmat(rows, cols));
392 let ptr: *mut u8 = data.as_mut_ptr();
393 VmpPMat {
394 data: data,
395 ptr: ptr,
396 n: self.n(),
397 cols: cols,
398 rows: rows,
399 backend: self.backend(),
400 }
401 }
402
403 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
404 unsafe { vmp::vmp_prepare_tmp_bytes(self.ptr, rows as u64, cols as u64) as usize }
405 }
406
407 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], tmp_bytes: &mut [u8]) {
408 debug_assert_eq!(a.len(), b.n * b.rows * b.cols);
409 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
410 #[cfg(debug_assertions)]
411 {
412 assert_alignement(tmp_bytes.as_ptr());
413 }
414 unsafe {
415 vmp::vmp_prepare_contiguous(
416 self.ptr,
417 b.as_mut_ptr() as *mut vmp_pmat_t,
418 a.as_ptr(),
419 b.rows() as u64,
420 b.cols() as u64,
421 tmp_bytes.as_mut_ptr(),
422 );
423 }
424 }
425
426 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], tmp_bytes: &mut [u8]) {
427 let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
428 #[cfg(debug_assertions)]
429 {
430 debug_assert_eq!(a.len(), b.rows);
431 a.iter().for_each(|ai| {
432 debug_assert_eq!(ai.len(), b.n * b.cols);
433 });
434 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
435 assert_alignement(tmp_bytes.as_ptr());
436 }
437 unsafe {
438 vmp::vmp_prepare_dblptr(
439 self.ptr,
440 b.as_mut_ptr() as *mut vmp_pmat_t,
441 ptrs.as_ptr(),
442 b.rows() as u64,
443 b.cols() as u64,
444 tmp_bytes.as_mut_ptr(),
445 );
446 }
447 }
448
449 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]) {
450 debug_assert_eq!(a.len(), b.cols() * self.n());
451 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
452 #[cfg(debug_assertions)]
453 {
454 assert_alignement(tmp_bytes.as_ptr());
455 }
456 unsafe {
457 vmp::vmp_prepare_row(
458 self.ptr,
459 b.as_mut_ptr() as *mut vmp_pmat_t,
460 a.as_ptr(),
461 row_i as u64,
462 b.rows() as u64,
463 b.cols() as u64,
464 tmp_bytes.as_mut_ptr(),
465 );
466 }
467 }
468
469 fn vmp_extract_row(&self, b: &mut VecZnxBig, a: &VmpPMat, row_i: usize) {
470 #[cfg(debug_assertions)]
471 {
472 assert_eq!(a.n(), b.n());
473 assert_eq!(a.cols(), b.cols());
474 }
475 unsafe {
476 vmp::vmp_extract_row(
477 self.ptr,
478 b.ptr as *mut vec_znx_big_t,
479 a.as_ptr() as *const vmp_pmat_t,
480 row_i as u64,
481 a.rows() as u64,
482 a.cols() as u64,
483 );
484 }
485 }
486
487 fn vmp_prepare_row_dft(&self, b: &mut VmpPMat, a: &VecZnxDft, row_i: usize) {
488 #[cfg(debug_assertions)]
489 {
490 assert_eq!(a.n(), b.n());
491 assert_eq!(a.cols(), b.cols());
492 }
493 unsafe {
494 vmp::vmp_prepare_row_dft(
495 self.ptr,
496 b.as_mut_ptr() as *mut vmp_pmat_t,
497 a.ptr as *const vec_znx_dft_t,
498 row_i as u64,
499 b.rows() as u64,
500 b.cols() as u64,
501 );
502 }
503 }
504
505 fn vmp_extract_row_dft(&self, b: &mut VecZnxDft, a: &VmpPMat, row_i: usize) {
506 #[cfg(debug_assertions)]
507 {
508 assert_eq!(a.n(), b.n());
509 assert_eq!(a.cols(), b.cols());
510 }
511 unsafe {
512 vmp::vmp_extract_row_dft(
513 self.ptr,
514 b.ptr as *mut vec_znx_dft_t,
515 a.as_ptr() as *const vmp_pmat_t,
516 row_i as u64,
517 a.rows() as u64,
518 a.cols() as u64,
519 );
520 }
521 }
522
523 fn vmp_apply_dft_tmp_bytes(
524 &self,
525 res_cols: usize,
526 a_cols: usize,
527 gct_rows: usize,
528 gct_cols: usize,
529 ) -> usize {
530 unsafe {
531 vmp::vmp_apply_dft_tmp_bytes(
532 self.ptr,
533 res_cols as u64,
534 a_cols as u64,
535 gct_rows as u64,
536 gct_cols as u64,
537 ) as usize
538 }
539 }
540
541 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
542 debug_assert!(
543 tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
544 );
545 #[cfg(debug_assertions)]
546 {
547 assert_alignement(tmp_bytes.as_ptr());
548 }
549 unsafe {
550 vmp::vmp_apply_dft(
551 self.ptr,
552 c.ptr as *mut vec_znx_dft_t,
553 c.cols() as u64,
554 a.as_ptr(),
555 a.cols() as u64,
556 a.n() as u64,
557 b.as_ptr() as *const vmp_pmat_t,
558 b.rows() as u64,
559 b.cols() as u64,
560 tmp_bytes.as_mut_ptr(),
561 )
562 }
563 }
564
565 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
566 debug_assert!(
567 tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
568 );
569 #[cfg(debug_assertions)]
570 {
571 assert_alignement(tmp_bytes.as_ptr());
572 }
573 unsafe {
574 vmp::vmp_apply_dft_add(
575 self.ptr,
576 c.ptr as *mut vec_znx_dft_t,
577 c.cols() as u64,
578 a.as_ptr(),
579 a.cols() as u64,
580 a.n() as u64,
581 b.as_ptr() as *const vmp_pmat_t,
582 b.rows() as u64,
583 b.cols() as u64,
584 tmp_bytes.as_mut_ptr(),
585 )
586 }
587 }
588
589 fn vmp_apply_dft_to_dft_tmp_bytes(
590 &self,
591 res_cols: usize,
592 a_cols: usize,
593 gct_rows: usize,
594 gct_cols: usize,
595 ) -> usize {
596 unsafe {
597 vmp::vmp_apply_dft_to_dft_tmp_bytes(
598 self.ptr,
599 res_cols as u64,
600 a_cols as u64,
601 gct_rows as u64,
602 gct_cols as u64,
603 ) as usize
604 }
605 }
606
607 fn vmp_apply_dft_to_dft(
608 &self,
609 c: &mut VecZnxDft,
610 a: &VecZnxDft,
611 b: &VmpPMat,
612 tmp_bytes: &mut [u8],
613 ) {
614 debug_assert!(
615 tmp_bytes.len()
616 >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
617 );
618 #[cfg(debug_assertions)]
619 {
620 assert_alignement(tmp_bytes.as_ptr());
621 }
622 unsafe {
623 vmp::vmp_apply_dft_to_dft(
624 self.ptr,
625 c.ptr as *mut vec_znx_dft_t,
626 c.cols() as u64,
627 a.ptr as *const vec_znx_dft_t,
628 a.cols() as u64,
629 b.as_ptr() as *const vmp_pmat_t,
630 b.rows() as u64,
631 b.cols() as u64,
632 tmp_bytes.as_mut_ptr(),
633 )
634 }
635 }
636
637 fn vmp_apply_dft_to_dft_add(
638 &self,
639 c: &mut VecZnxDft,
640 a: &VecZnxDft,
641 b: &VmpPMat,
642 tmp_bytes: &mut [u8],
643 ) {
644 debug_assert!(
645 tmp_bytes.len()
646 >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
647 );
648 #[cfg(debug_assertions)]
649 {
650 assert_alignement(tmp_bytes.as_ptr());
651 }
652 unsafe {
653 vmp::vmp_apply_dft_to_dft_add(
654 self.ptr,
655 c.ptr as *mut vec_znx_dft_t,
656 c.cols() as u64,
657 a.ptr as *const vec_znx_dft_t,
658 a.cols() as u64,
659 b.as_ptr() as *const vmp_pmat_t,
660 b.rows() as u64,
661 b.cols() as u64,
662 tmp_bytes.as_mut_ptr(),
663 )
664 }
665 }
666
667 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, tmp_bytes: &mut [u8]) {
668 debug_assert!(
669 tmp_bytes.len()
670 >= self.vmp_apply_dft_to_dft_tmp_bytes(b.cols(), b.cols(), a.rows(), a.cols())
671 );
672 #[cfg(debug_assertions)]
673 {
674 assert_alignement(tmp_bytes.as_ptr());
675 }
676 unsafe {
677 vmp::vmp_apply_dft_to_dft(
678 self.ptr,
679 b.ptr as *mut vec_znx_dft_t,
680 b.cols() as u64,
681 b.ptr as *mut vec_znx_dft_t,
682 b.cols() as u64,
683 a.as_ptr() as *const vmp_pmat_t,
684 a.rows() as u64,
685 a.cols() as u64,
686 tmp_bytes.as_mut_ptr(),
687 )
688 }
689 }
690}
691
692#[cfg(test)]
693mod tests {
694 use crate::{
695 alloc_aligned, Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps,
696 VecZnxOps, VmpPMat, VmpPMatOps,
697 };
698 use sampling::source::Source;
699
700 #[test]
701 fn vmp_prepare_row_dft() {
702 let module: Module = Module::new(32, crate::BACKEND::FFT64);
703 let vpmat_rows: usize = 4;
704 let vpmat_cols: usize = 5;
705 let log_base2k: usize = 8;
706 let mut a: VecZnx = module.new_vec_znx(vpmat_cols);
707 let mut a_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
708 let mut a_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
709 let mut b_big: VecZnxBig = module.new_vec_znx_big(vpmat_cols);
710 let mut b_dft: VecZnxDft = module.new_vec_znx_dft(vpmat_cols);
711 let mut vmpmat_0: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
712 let mut vmpmat_1: VmpPMat = module.new_vmp_pmat(vpmat_rows, vpmat_cols);
713
714 let mut tmp_bytes: Vec<u8> =
715 alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, vpmat_cols));
716
717 for row_i in 0..vpmat_rows {
718 let mut source: Source = Source::new([0u8; 32]);
719 module.fill_uniform(log_base2k, &mut a, vpmat_cols, &mut source);
720 module.vec_znx_dft(&mut a_dft, &a);
721 module.vmp_prepare_row(&mut vmpmat_0, &a.raw(), row_i, &mut tmp_bytes);
722
723 // Checks that prepare(vmp_pmat, a) = prepare_dft(vmp_pmat, a_dft)
724 module.vmp_prepare_row_dft(&mut vmpmat_1, &a_dft, row_i);
725 assert_eq!(vmpmat_0.raw::<u8>(), vmpmat_1.raw::<u8>());
726
727 // Checks that a_dft = extract_dft(prepare(vmp_pmat, a), b_dft)
728 module.vmp_extract_row_dft(&mut b_dft, &vmpmat_0, row_i);
729 assert_eq!(a_dft.raw::<u8>(&module), b_dft.raw::<u8>(&module));
730
731 // Checks that a_big = extract(prepare_dft(vmp_pmat, a_dft), b_big)
732 module.vmp_extract_row(&mut b_big, &vmpmat_0, row_i);
733 module.vec_znx_idft(&mut a_big, &a_dft, &mut tmp_bytes);
734 assert_eq!(a_big.raw::<i64>(&module), b_big.raw::<i64>(&module));
735 }
736
737 module.free();
738 }
739}