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