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