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