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