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