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