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