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