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