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