poulpy/backend/src/mat_znx_dft_ops.rs
backend/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, MatZnxDftOwned, MatZnxToMut, MatZnxToRef, Module, ScalarZnxAlloc, ScalarZnxDftAlloc,
6 ScalarZnxDftOps, Scratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps,
7};
8
9pub trait MatZnxDftAlloc<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_in: usize, cols_out: usize, size: usize) -> MatZnxDftOwned<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 ) -> MatZnxDftOwned<B>;
28}
29
30pub trait MatZnxDftScratch {
31 /// Returns the size of the stratch space necessary for [MatZnxDftOps::vmp_apply_dft_to_dft].
32 fn vmp_apply_tmp_bytes(
33 &self,
34 res_size: usize,
35 a_size: usize,
36 b_rows: usize,
37 b_cols_in: usize,
38 b_cols_out: usize,
39 b_size: usize,
40 ) -> usize;
41
42 fn mat_znx_dft_mul_x_pow_minus_one_scratch_space(&self, size: usize, cols_out: usize) -> usize;
43}
44
45/// This trait implements methods for vector matrix product,
46/// that is, multiplying a [VecZnx] with a [MatZnxDft].
47pub trait MatZnxDftOps<BACKEND: Backend> {
48 /// Prepares the ith-row of [MatZnxDft] from a [VecZnxDft].
49 ///
50 /// # Arguments
51 ///
52 /// * `res`: [MatZnxDft] on which the values are encoded.
53 /// * `a`: the [VecZnxDft] to encode on the [MatZnxDft].
54 /// * `row_i`: the index of the row to prepare.
55 ///
56 /// The size of buf can be obtained with [MatZnxDftOps::vmp_prepare_tmp_bytes].
57 fn mat_znx_dft_set_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
58 where
59 R: MatZnxToMut<FFT64>,
60 A: VecZnxDftToRef<FFT64>;
61
62 /// Extracts the ith-row of [MatZnxDft] into a [VecZnxDft].
63 ///
64 /// # Arguments
65 ///
66 /// * `res`: the [VecZnxDft] to on which to extract the row of the [MatZnxDft].
67 /// * `a`: [MatZnxDft] on which the values are encoded.
68 /// * `row_i`: the index of the row to extract.
69 fn mat_znx_dft_get_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
70 where
71 R: VecZnxDftToMut<FFT64>,
72 A: MatZnxToRef<FFT64>;
73
74 /// Multiplies A by (X^{k} - 1) and stores the result on R.
75 fn mat_znx_dft_mul_x_pow_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A, scratch: &mut Scratch)
76 where
77 R: MatZnxToMut<FFT64>,
78 A: MatZnxToRef<FFT64>;
79
80 /// Multiplies A by (X^{k} - 1).
81 fn mat_znx_dft_mul_x_pow_minus_one_inplace<A>(&self, k: i64, a: &mut A, scratch: &mut Scratch)
82 where
83 A: MatZnxToMut<FFT64>;
84
85 /// Applies the vector matrix product [VecZnxDft] x [MatZnxDft].
86 /// The size of `buf` is given by [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
87 ///
88 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
89 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
90 /// and each vector a [VecZnxDft] (row) of the [MatZnxDft].
91 ///
92 /// As such, given an input [VecZnx] of `i` size and a [MatZnxDft] of `i` rows and
93 /// `j` size, the output is a [VecZnx] of `j` size.
94 ///
95 /// If there is a mismatch between the dimensions the largest valid ones are used.
96 ///
97 /// ```text
98 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
99 /// |h i j|
100 /// |k l m|
101 /// ```
102 /// where each element is a [VecZnxDft].
103 ///
104 /// # Arguments
105 ///
106 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
107 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
108 /// * `b`: the right operand [MatZnxDft] of the vector matrix product.
109 /// * `buf`: scratch space, the size can be obtained with [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
110 fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
111 where
112 R: VecZnxDftToMut<FFT64>,
113 A: VecZnxDftToRef<FFT64>,
114 B: MatZnxToRef<FFT64>;
115
116 // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
117 fn vmp_apply_add<R, A, B>(&self, res: &mut R, a: &A, b: &B, scale: usize, scratch: &mut Scratch)
118 where
119 R: VecZnxDftToMut<FFT64>,
120 A: VecZnxDftToRef<FFT64>,
121 B: MatZnxToRef<FFT64>;
122}
123
124impl<B: Backend> MatZnxDftAlloc<B> for Module<B> {
125 fn bytes_of_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
126 MatZnxDftOwned::bytes_of(self, rows, cols_in, cols_out, size)
127 }
128
129 fn new_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> MatZnxDftOwned<B> {
130 MatZnxDftOwned::new(self, rows, cols_in, cols_out, size)
131 }
132
133 fn new_mat_znx_dft_from_bytes(
134 &self,
135 rows: usize,
136 cols_in: usize,
137 cols_out: usize,
138 size: usize,
139 bytes: Vec<u8>,
140 ) -> MatZnxDftOwned<B> {
141 MatZnxDftOwned::new_from_bytes(self, rows, cols_in, cols_out, size, bytes)
142 }
143}
144
145impl<BACKEND: Backend> MatZnxDftScratch for Module<BACKEND> {
146 fn vmp_apply_tmp_bytes(
147 &self,
148 res_size: usize,
149 a_size: usize,
150 b_rows: usize,
151 b_cols_in: usize,
152 b_cols_out: usize,
153 b_size: usize,
154 ) -> usize {
155 unsafe {
156 vmp::vmp_apply_dft_to_dft_tmp_bytes(
157 self.ptr,
158 (res_size * b_cols_out) as u64,
159 (a_size * b_cols_in) as u64,
160 (b_rows * b_cols_in) as u64,
161 (b_size * b_cols_out) as u64,
162 ) as usize
163 }
164 }
165
166 fn mat_znx_dft_mul_x_pow_minus_one_scratch_space(&self, size: usize, cols_out: usize) -> usize {
167 let xpm1_dft: usize = self.bytes_of_scalar_znx(1);
168 let xpm1: usize = self.bytes_of_scalar_znx_dft(1);
169 let tmp: usize = self.bytes_of_vec_znx_dft(cols_out, size);
170 xpm1_dft + (xpm1 | 2 * tmp)
171 }
172}
173
174impl MatZnxDftOps<FFT64> for Module<FFT64> {
175 fn mat_znx_dft_mul_x_pow_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A, scratch: &mut Scratch)
176 where
177 R: MatZnxToMut<FFT64>,
178 A: MatZnxToRef<FFT64>,
179 {
180 let mut res: MatZnxDft<&mut [u8], FFT64> = res.to_mut();
181 let a: MatZnxDft<&[u8], FFT64> = a.to_ref();
182
183 #[cfg(debug_assertions)]
184 {
185 assert_eq!(res.n(), self.n());
186 assert_eq!(a.n(), self.n());
187 assert_eq!(res.rows(), a.rows());
188 assert_eq!(res.cols_in(), a.cols_in());
189 assert_eq!(res.cols_out(), a.cols_out());
190 }
191
192 let (mut xpm1_dft, scratch1) = scratch.tmp_scalar_znx_dft(self, 1);
193
194 {
195 let (mut xpm1, _) = scratch1.tmp_scalar_znx(self, 1);
196 xpm1.data[0] = 1;
197 self.vec_znx_rotate_inplace(k, &mut xpm1, 0);
198 self.svp_prepare(&mut xpm1_dft, 0, &xpm1, 0);
199 }
200
201 let (mut tmp_0, scratch2) = scratch1.tmp_vec_znx_dft(self, res.cols_out(), res.size());
202 let (mut tmp_1, _) = scratch2.tmp_vec_znx_dft(self, res.cols_out(), res.size());
203
204 (0..res.rows()).for_each(|row_i| {
205 (0..res.cols_in()).for_each(|col_j| {
206 self.mat_znx_dft_get_row(&mut tmp_0, &a, row_i, col_j);
207
208 (0..tmp_0.cols()).for_each(|i| {
209 self.svp_apply(&mut tmp_1, i, &xpm1_dft, 0, &tmp_0, i);
210 self.vec_znx_dft_sub_ab_inplace(&mut tmp_1, i, &tmp_0, i);
211 });
212
213 self.mat_znx_dft_set_row(&mut res, row_i, col_j, &tmp_1);
214 });
215 })
216 }
217
218 fn mat_znx_dft_mul_x_pow_minus_one_inplace<A>(&self, k: i64, a: &mut A, scratch: &mut Scratch)
219 where
220 A: MatZnxToMut<FFT64>,
221 {
222 let mut a: MatZnxDft<&mut [u8], FFT64> = a.to_mut();
223
224 #[cfg(debug_assertions)]
225 {
226 assert_eq!(a.n(), self.n());
227 }
228
229 let (mut xpm1_dft, scratch1) = scratch.tmp_scalar_znx_dft(self, 1);
230
231 {
232 let (mut xpm1, _) = scratch1.tmp_scalar_znx(self, 1);
233 xpm1.data[0] = 1;
234 self.vec_znx_rotate_inplace(k, &mut xpm1, 0);
235 self.svp_prepare(&mut xpm1_dft, 0, &xpm1, 0);
236 }
237
238 let (mut tmp_0, scratch2) = scratch1.tmp_vec_znx_dft(self, a.cols_out(), a.size());
239 let (mut tmp_1, _) = scratch2.tmp_vec_znx_dft(self, a.cols_out(), a.size());
240
241 (0..a.rows()).for_each(|row_i| {
242 (0..a.cols_in()).for_each(|col_j| {
243 self.mat_znx_dft_get_row(&mut tmp_0, &a, row_i, col_j);
244
245 (0..tmp_0.cols()).for_each(|i| {
246 self.svp_apply(&mut tmp_1, i, &xpm1_dft, 0, &tmp_0, i);
247 self.vec_znx_dft_sub_ab_inplace(&mut tmp_1, i, &tmp_0, i);
248 });
249
250 self.mat_znx_dft_set_row(&mut a, row_i, col_j, &tmp_1);
251 });
252 })
253 }
254
255 fn mat_znx_dft_set_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
256 where
257 R: MatZnxToMut<FFT64>,
258 A: VecZnxDftToRef<FFT64>,
259 {
260 let mut res: MatZnxDft<&mut [u8], FFT64> = res.to_mut();
261 let a: VecZnxDft<&[u8], _> = a.to_ref();
262
263 #[cfg(debug_assertions)]
264 {
265 assert_eq!(res.n(), self.n());
266 assert_eq!(a.n(), self.n());
267 assert_eq!(
268 a.cols(),
269 res.cols_out(),
270 "a.cols(): {} != res.cols_out(): {}",
271 a.cols(),
272 res.cols_out()
273 );
274 assert!(
275 res_row < res.rows(),
276 "res_row: {} >= res.rows(): {}",
277 res_row,
278 res.rows()
279 );
280 assert!(
281 res_col_in < res.cols_in(),
282 "res_col_in: {} >= res.cols_in(): {}",
283 res_col_in,
284 res.cols_in()
285 );
286 assert_eq!(
287 res.size(),
288 a.size(),
289 "res.size(): {} != a.size(): {}",
290 res.size(),
291 a.size()
292 );
293 }
294
295 unsafe {
296 vmp::vmp_prepare_row_dft(
297 self.ptr,
298 res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
299 a.as_ptr() as *const vec_znx_dft_t,
300 (res_row * res.cols_in() + res_col_in) as u64,
301 (res.rows() * res.cols_in()) as u64,
302 (res.size() * res.cols_out()) as u64,
303 );
304 }
305 }
306
307 fn mat_znx_dft_get_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
308 where
309 R: VecZnxDftToMut<FFT64>,
310 A: MatZnxToRef<FFT64>,
311 {
312 let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
313 let a: MatZnxDft<&[u8], _> = a.to_ref();
314
315 #[cfg(debug_assertions)]
316 {
317 assert_eq!(res.n(), self.n());
318 assert_eq!(a.n(), self.n());
319 assert_eq!(
320 res.cols(),
321 a.cols_out(),
322 "res.cols(): {} != a.cols_out(): {}",
323 res.cols(),
324 a.cols_out()
325 );
326 assert!(
327 a_row < a.rows(),
328 "a_row: {} >= a.rows(): {}",
329 a_row,
330 a.rows()
331 );
332 assert!(
333 a_col_in < a.cols_in(),
334 "a_col_in: {} >= a.cols_in(): {}",
335 a_col_in,
336 a.cols_in()
337 );
338 assert_eq!(
339 res.size(),
340 a.size(),
341 "res.size(): {} != a.size(): {}",
342 res.size(),
343 a.size()
344 );
345 }
346 unsafe {
347 vmp::vmp_extract_row_dft(
348 self.ptr,
349 res.as_mut_ptr() as *mut vec_znx_dft_t,
350 a.as_ptr() as *const vmp::vmp_pmat_t,
351 (a_row * a.cols_in() + a_col_in) as u64,
352 (a.rows() * a.cols_in()) as u64,
353 (a.size() * a.cols_out()) as u64,
354 );
355 }
356 }
357
358 fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
359 where
360 R: VecZnxDftToMut<FFT64>,
361 A: VecZnxDftToRef<FFT64>,
362 B: MatZnxToRef<FFT64>,
363 {
364 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
365 let a: VecZnxDft<&[u8], _> = a.to_ref();
366 let b: MatZnxDft<&[u8], _> = b.to_ref();
367
368 #[cfg(debug_assertions)]
369 {
370 assert_eq!(res.n(), self.n());
371 assert_eq!(b.n(), self.n());
372 assert_eq!(a.n(), self.n());
373 assert_eq!(
374 res.cols(),
375 b.cols_out(),
376 "res.cols(): {} != b.cols_out: {}",
377 res.cols(),
378 b.cols_out()
379 );
380 assert_eq!(
381 a.cols(),
382 b.cols_in(),
383 "a.cols(): {} != b.cols_in: {}",
384 a.cols(),
385 b.cols_in()
386 );
387 }
388
389 let (tmp_bytes, _) = scratch.tmp_slice(self.vmp_apply_tmp_bytes(
390 res.size(),
391 a.size(),
392 b.rows(),
393 b.cols_in(),
394 b.cols_out(),
395 b.size(),
396 ));
397 unsafe {
398 vmp::vmp_apply_dft_to_dft(
399 self.ptr,
400 res.as_mut_ptr() as *mut vec_znx_dft_t,
401 (res.size() * res.cols()) as u64,
402 a.as_ptr() as *const vec_znx_dft_t,
403 (a.size() * a.cols()) as u64,
404 b.as_ptr() as *const vmp::vmp_pmat_t,
405 (b.rows() * b.cols_in()) as u64,
406 (b.size() * b.cols_out()) as u64,
407 tmp_bytes.as_mut_ptr(),
408 )
409 }
410 }
411
412 fn vmp_apply_add<R, A, B>(&self, res: &mut R, a: &A, b: &B, scale: usize, scratch: &mut Scratch)
413 where
414 R: VecZnxDftToMut<FFT64>,
415 A: VecZnxDftToRef<FFT64>,
416 B: MatZnxToRef<FFT64>,
417 {
418 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
419 let a: VecZnxDft<&[u8], _> = a.to_ref();
420 let b: MatZnxDft<&[u8], _> = b.to_ref();
421
422 #[cfg(debug_assertions)]
423 {
424 assert_eq!(res.n(), self.n());
425 assert_eq!(b.n(), self.n());
426 assert_eq!(a.n(), self.n());
427 assert_eq!(
428 res.cols(),
429 b.cols_out(),
430 "res.cols(): {} != b.cols_out: {}",
431 res.cols(),
432 b.cols_out()
433 );
434 assert_eq!(
435 a.cols(),
436 b.cols_in(),
437 "a.cols(): {} != b.cols_in: {}",
438 a.cols(),
439 b.cols_in()
440 );
441 }
442
443 let (tmp_bytes, _) = scratch.tmp_slice(self.vmp_apply_tmp_bytes(
444 res.size(),
445 a.size(),
446 b.rows(),
447 b.cols_in(),
448 b.cols_out(),
449 b.size(),
450 ));
451 unsafe {
452 vmp::vmp_apply_dft_to_dft_add(
453 self.ptr,
454 res.as_mut_ptr() as *mut vec_znx_dft_t,
455 (res.size() * res.cols()) as u64,
456 a.as_ptr() as *const vec_znx_dft_t,
457 (a.size() * a.cols()) as u64,
458 b.as_ptr() as *const vmp::vmp_pmat_t,
459 (b.rows() * b.cols_in()) as u64,
460 (b.size() * b.cols_out()) as u64,
461 (scale * b.cols_out()) as u64,
462 tmp_bytes.as_mut_ptr(),
463 )
464 }
465 }
466}
467#[cfg(test)]
468mod tests {
469 use crate::{
470 Decoding, FFT64, FillUniform, MatZnxDft, MatZnxDftOps, Module, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxBig,
471 VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, ZnxInfos, ZnxView,
472 ZnxViewMut, ZnxZero,
473 };
474 use sampling::source::Source;
475
476 use super::{MatZnxDftAlloc, MatZnxDftScratch};
477
478 #[test]
479 fn vmp_set_row() {
480 let module: Module<FFT64> = Module::<FFT64>::new(16);
481 let basek: usize = 8;
482 let mat_rows: usize = 4;
483 let mat_cols_in: usize = 2;
484 let mat_cols_out: usize = 2;
485 let mat_size: usize = 5;
486 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
487 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
488 let mut b_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
489 let mut mat: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
490
491 for col_in in 0..mat_cols_in {
492 for row_i in 0..mat_rows {
493 let mut source: Source = Source::new([0u8; 32]);
494 (0..mat_cols_out).for_each(|col_out| {
495 a.fill_uniform(basek, col_out, mat_size, &mut source);
496 module.vec_znx_dft(1, 0, &mut a_dft, col_out, &a, col_out);
497 });
498 module.mat_znx_dft_set_row(&mut mat, row_i, col_in, &a_dft);
499 module.mat_znx_dft_get_row(&mut b_dft, &mat, row_i, col_in);
500 assert_eq!(a_dft.raw(), b_dft.raw());
501 }
502 }
503 }
504
505 #[test]
506 fn vmp_apply() {
507 let log_n: i32 = 5;
508 let n: usize = 1 << log_n;
509
510 let module: Module<FFT64> = Module::<FFT64>::new(n);
511 let basek: usize = 15;
512 let a_size: usize = 5;
513 let mat_size: usize = 6;
514 let res_size: usize = a_size;
515
516 [1, 2].iter().for_each(|cols_in| {
517 [1, 2].iter().for_each(|cols_out| {
518 let a_cols: usize = *cols_in;
519 let res_cols: usize = *cols_out;
520
521 let mat_rows: usize = a_size;
522 let mat_cols_in: usize = a_cols;
523 let mat_cols_out: usize = res_cols;
524
525 let mut scratch: ScratchOwned = ScratchOwned::new(
526 module.vmp_apply_tmp_bytes(
527 res_size,
528 a_size,
529 mat_rows,
530 mat_cols_in,
531 mat_cols_out,
532 mat_size,
533 ) | module.vec_znx_big_normalize_tmp_bytes(),
534 );
535
536 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
537
538 (0..a_cols).for_each(|i| {
539 a.at_mut(i, a_size - 1)[i + 1] = 1;
540 });
541
542 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
543 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
544
545 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
546 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
547
548 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
549
550 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
551 (0..a.size()).for_each(|row_i| {
552 (0..mat_cols_in).for_each(|col_in_i| {
553 (0..mat_cols_out).for_each(|col_out_i| {
554 let idx = 1 + col_in_i * mat_cols_out + col_out_i;
555 tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
556 module.vec_znx_dft(1, 0, &mut c_dft, col_out_i, &tmp, col_out_i);
557 tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
558 });
559 module.mat_znx_dft_set_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
560 });
561 });
562
563 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
564 (0..a_cols).for_each(|i| {
565 module.vec_znx_dft(1, 0, &mut a_dft, i, &a, i);
566 });
567
568 module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
569
570 let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
571
572 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, res_size);
573 (0..mat_cols_out).for_each(|i| {
574 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
575 module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
576 });
577
578 (0..mat_cols_out).for_each(|col_i| {
579 let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
580 (0..a_cols).for_each(|i| {
581 res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
582 });
583 res_have.decode_vec_i64(col_i, basek, basek * a_size, &mut res_have_vi64);
584 assert_eq!(res_have_vi64, res_want_vi64);
585 });
586 });
587 });
588 }
589
590 #[test]
591 fn vmp_apply_add() {
592 let log_n: i32 = 4;
593 let n: usize = 1 << log_n;
594
595 let module: Module<FFT64> = Module::<FFT64>::new(n);
596 let basek: usize = 8;
597 let a_size: usize = 5;
598 let mat_size: usize = 5;
599 let res_size: usize = a_size;
600 let mut source: Source = Source::new([0u8; 32]);
601
602 [1, 2].iter().for_each(|cols_in| {
603 [1, 2].iter().for_each(|cols_out| {
604 (0..res_size).for_each(|shift| {
605 let a_cols: usize = *cols_in;
606 let res_cols: usize = *cols_out;
607
608 let mat_rows: usize = a_size;
609 let mat_cols_in: usize = a_cols;
610 let mat_cols_out: usize = res_cols;
611
612 let mut scratch: ScratchOwned = ScratchOwned::new(
613 module.vmp_apply_tmp_bytes(
614 res_size,
615 a_size,
616 mat_rows,
617 mat_cols_in,
618 mat_cols_out,
619 mat_size,
620 ) | module.vec_znx_big_normalize_tmp_bytes(),
621 );
622
623 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
624
625 (0..a_cols).for_each(|col_i| {
626 a.fill_uniform(basek, col_i, a.size(), &mut source);
627 });
628
629 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
630 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
631
632 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
633 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
634
635 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
636
637 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
638 (0..a.size()).for_each(|row_i| {
639 (0..mat_cols_in).for_each(|col_in_i| {
640 (0..mat_cols_out).for_each(|col_out_i| {
641 let idx: usize = 1 + col_in_i * mat_cols_out + col_out_i;
642 tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
643 module.vec_znx_dft(1, 0, &mut c_dft, col_out_i, &tmp, col_out_i);
644 tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
645 });
646 module.mat_znx_dft_set_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
647 });
648 });
649
650 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
651 (0..a_cols).for_each(|i| {
652 module.vec_znx_dft(1, 0, &mut a_dft, i, &a, i);
653 });
654
655 c_dft.zero();
656 (0..c_dft.cols()).for_each(|i| {
657 module.vec_znx_dft(1, 0, &mut c_dft, i, &a, 0);
658 });
659
660 module.vmp_apply_add(&mut c_dft, &a_dft, &mat_znx_dft, shift, scratch.borrow());
661
662 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, mat_size);
663 (0..mat_cols_out).for_each(|i| {
664 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
665 module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
666 });
667
668 let mut res_want: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, mat_size);
669
670 // Equivalent to vmp_add & scale
671 module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
672 (0..mat_cols_out).for_each(|i| {
673 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
674 module.vec_znx_big_normalize(basek, &mut res_want, i, &c_big, i, scratch.borrow());
675 });
676 module.vec_znx_shift_inplace(
677 basek,
678 (shift * basek) as i64,
679 &mut res_want,
680 scratch.borrow(),
681 );
682 (0..res_cols).for_each(|i| {
683 module.vec_znx_add_inplace(&mut res_want, i, &a, 0);
684 module.vec_znx_normalize_inplace(basek, &mut res_want, i, scratch.borrow());
685 });
686
687 assert_eq!(res_want, res_have);
688 });
689 });
690 });
691 }
692
693 #[test]
694 fn vmp_apply_digits() {
695 let log_n: i32 = 4;
696 let n: usize = 1 << log_n;
697
698 let module: Module<FFT64> = Module::<FFT64>::new(n);
699 let basek: usize = 8;
700 let a_size: usize = 6;
701 let mat_size: usize = 6;
702 let res_size: usize = a_size;
703
704 [1, 2].iter().for_each(|cols_in| {
705 [1, 2].iter().for_each(|cols_out| {
706 [1, 3, 6].iter().for_each(|digits| {
707 let mut source: Source = Source::new([0u8; 32]);
708
709 let a_cols: usize = *cols_in;
710 let res_cols: usize = *cols_out;
711
712 let mat_rows: usize = a_size;
713 let mat_cols_in: usize = a_cols;
714 let mat_cols_out: usize = res_cols;
715
716 let mut scratch: ScratchOwned = ScratchOwned::new(
717 module.vmp_apply_tmp_bytes(
718 res_size,
719 a_size,
720 mat_rows,
721 mat_cols_in,
722 mat_cols_out,
723 mat_size,
724 ) | module.vec_znx_big_normalize_tmp_bytes(),
725 );
726
727 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
728
729 (0..a_cols).for_each(|col_i| {
730 a.fill_uniform(basek, col_i, a.size(), &mut source);
731 });
732
733 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
734 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
735
736 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
737 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
738
739 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
740
741 let rows: usize = a.size() / digits;
742
743 let shift: usize = 1;
744
745 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
746 (0..rows).for_each(|row_i| {
747 (0..mat_cols_in).for_each(|col_in_i| {
748 (0..mat_cols_out).for_each(|col_out_i| {
749 let idx: usize = shift + col_in_i * mat_cols_out + col_out_i;
750 let limb: usize = (digits - 1) + row_i * digits;
751 tmp.at_mut(col_out_i, limb)[idx] = 1 as i64; // X^{idx}
752 module.vec_znx_dft(1, 0, &mut c_dft, col_out_i, &tmp, col_out_i);
753 tmp.at_mut(col_out_i, limb)[idx] = 0 as i64;
754 });
755 module.mat_znx_dft_set_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
756 });
757 });
758
759 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, (a_size + digits - 1) / digits);
760
761 (0..*digits).for_each(|di| {
762 (0..a_cols).for_each(|col_i| {
763 module.vec_znx_dft(*digits, digits - 1 - di, &mut a_dft, col_i, &a, col_i);
764 });
765
766 if di == 0 {
767 module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
768 } else {
769 module.vmp_apply_add(&mut c_dft, &a_dft, &mat_znx_dft, di, scratch.borrow());
770 }
771 });
772
773 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, mat_size);
774 (0..mat_cols_out).for_each(|i| {
775 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
776 module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
777 });
778
779 let mut res_want: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, mat_size);
780 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, mat_size);
781 (0..res_cols).for_each(|col_i| {
782 (0..a_cols).for_each(|j| {
783 module.vec_znx_rotate(
784 (col_i + j * mat_cols_out + shift) as i64,
785 &mut tmp,
786 0,
787 &a,
788 j,
789 );
790 module.vec_znx_add_inplace(&mut res_want, col_i, &tmp, 0);
791 });
792 module.vec_znx_normalize_inplace(basek, &mut res_want, col_i, scratch.borrow());
793 });
794
795 assert_eq!(res_have, res_want)
796 });
797 });
798 });
799 }
800
801 #[test]
802 fn mat_znx_dft_mul_x_pow_minus_one() {
803 let log_n: i32 = 5;
804 let n: usize = 1 << log_n;
805
806 let module: Module<FFT64> = Module::<FFT64>::new(n);
807 let basek: usize = 8;
808 let rows: usize = 2;
809 let cols_in: usize = 2;
810 let cols_out: usize = 2;
811 let size: usize = 4;
812
813 let mut scratch: ScratchOwned = ScratchOwned::new(module.mat_znx_dft_mul_x_pow_minus_one_scratch_space(size, cols_out));
814
815 let mut mat_want: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(rows, cols_in, cols_out, size);
816 let mut mat_have: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(rows, cols_in, cols_out, size);
817
818 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(1, size);
819 let mut tmp_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(cols_out, size);
820
821 let mut source: Source = Source::new([0u8; 32]);
822
823 (0..mat_want.rows()).for_each(|row_i| {
824 (0..mat_want.cols_in()).for_each(|col_i| {
825 (0..cols_out).for_each(|j| {
826 tmp.fill_uniform(basek, 0, size, &mut source);
827 module.vec_znx_dft(1, 0, &mut tmp_dft, j, &tmp, 0);
828 });
829
830 module.mat_znx_dft_set_row(&mut mat_want, row_i, col_i, &tmp_dft);
831 });
832 });
833
834 let k: i64 = 1;
835
836 module.mat_znx_dft_mul_x_pow_minus_one(k, &mut mat_have, &mat_want, scratch.borrow());
837
838 let mut have: VecZnx<Vec<u8>> = module.new_vec_znx(cols_out, size);
839 let mut want: VecZnx<Vec<u8>> = module.new_vec_znx(cols_out, size);
840 let mut tmp_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, size);
841
842 (0..mat_want.rows()).for_each(|row_i| {
843 (0..mat_want.cols_in()).for_each(|col_i| {
844 module.mat_znx_dft_get_row(&mut tmp_dft, &mat_want, row_i, col_i);
845
846 (0..cols_out).for_each(|j| {
847 module.vec_znx_idft(&mut tmp_big, 0, &tmp_dft, j, scratch.borrow());
848 // module.vec_znx_big_normalize(basek, &mut want, j, &tmp_big, 0, scratch.borrow());
849 module.vec_znx_big_normalize(basek, &mut tmp, 0, &tmp_big, 0, scratch.borrow());
850 module.vec_znx_rotate(k, &mut want, j, &tmp, 0);
851 module.vec_znx_sub_ab_inplace(&mut want, j, &tmp, 0);
852 module.vec_znx_normalize_inplace(basek, &mut want, j, scratch.borrow());
853 });
854
855 module.mat_znx_dft_get_row(&mut tmp_dft, &mat_have, row_i, col_i);
856
857 (0..cols_out).for_each(|j| {
858 module.vec_znx_idft(&mut tmp_big, 0, &tmp_dft, j, scratch.borrow());
859 module.vec_znx_big_normalize(basek, &mut have, j, &tmp_big, 0, scratch.borrow());
860 });
861
862 assert_eq!(have, want)
863 });
864 });
865 }
866}