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