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, Scratch, VecZnxDft, VecZnxDftToMut,
6 VecZnxDftToRef,
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
43/// This trait implements methods for vector matrix product,
44/// that is, multiplying a [VecZnx] with a [MatZnxDft].
45pub trait MatZnxDftOps<BACKEND: Backend> {
46 /// Prepares the ith-row of [MatZnxDft] from a [VecZnxDft].
47 ///
48 /// # Arguments
49 ///
50 /// * `res`: [MatZnxDft] on which the values are encoded.
51 /// * `a`: the [VecZnxDft] to encode on the [MatZnxDft].
52 /// * `row_i`: the index of the row to prepare.
53 ///
54 /// The size of buf can be obtained with [MatZnxDftOps::vmp_prepare_tmp_bytes].
55 fn vmp_prepare_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
56 where
57 R: MatZnxToMut<FFT64>,
58 A: VecZnxDftToRef<FFT64>;
59
60 /// Extracts the ith-row of [MatZnxDft] into a [VecZnxDft].
61 ///
62 /// # Arguments
63 ///
64 /// * `res`: the [VecZnxDft] to on which to extract the row of the [MatZnxDft].
65 /// * `a`: [MatZnxDft] on which the values are encoded.
66 /// * `row_i`: the index of the row to extract.
67 fn vmp_extract_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
68 where
69 R: VecZnxDftToMut<FFT64>,
70 A: MatZnxToRef<FFT64>;
71
72 /// Applies the vector matrix product [VecZnxDft] x [MatZnxDft].
73 /// The size of `buf` is given by [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
74 ///
75 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
76 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
77 /// and each vector a [VecZnxDft] (row) of the [MatZnxDft].
78 ///
79 /// As such, given an input [VecZnx] of `i` size and a [MatZnxDft] of `i` rows and
80 /// `j` size, the output is a [VecZnx] of `j` size.
81 ///
82 /// If there is a mismatch between the dimensions the largest valid ones are used.
83 ///
84 /// ```text
85 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
86 /// |h i j|
87 /// |k l m|
88 /// ```
89 /// where each element is a [VecZnxDft].
90 ///
91 /// # Arguments
92 ///
93 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
94 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
95 /// * `b`: the right operand [MatZnxDft] of the vector matrix product.
96 /// * `buf`: scratch space, the size can be obtained with [MatZnxDftOps::vmp_apply_dft_to_dft_tmp_bytes].
97 fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
98 where
99 R: VecZnxDftToMut<FFT64>,
100 A: VecZnxDftToRef<FFT64>,
101 B: MatZnxToRef<FFT64>;
102
103 // Same as [MatZnxDftOps::vmp_apply] except result is added on R instead of overwritting R.
104 fn vmp_apply_add<R, A, B>(&self, res: &mut R, a: &A, b: &B, scale: usize, scratch: &mut Scratch)
105 where
106 R: VecZnxDftToMut<FFT64>,
107 A: VecZnxDftToRef<FFT64>,
108 B: MatZnxToRef<FFT64>;
109}
110
111impl<B: Backend> MatZnxDftAlloc<B> for Module<B> {
112 fn bytes_of_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
113 MatZnxDftOwned::bytes_of(self, rows, cols_in, cols_out, size)
114 }
115
116 fn new_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> MatZnxDftOwned<B> {
117 MatZnxDftOwned::new(self, rows, cols_in, cols_out, size)
118 }
119
120 fn new_mat_znx_dft_from_bytes(
121 &self,
122 rows: usize,
123 cols_in: usize,
124 cols_out: usize,
125 size: usize,
126 bytes: Vec<u8>,
127 ) -> MatZnxDftOwned<B> {
128 MatZnxDftOwned::new_from_bytes(self, rows, cols_in, cols_out, size, bytes)
129 }
130}
131
132impl<BACKEND: Backend> MatZnxDftScratch for Module<BACKEND> {
133 fn vmp_apply_tmp_bytes(
134 &self,
135 res_size: usize,
136 a_size: usize,
137 b_rows: usize,
138 b_cols_in: usize,
139 b_cols_out: usize,
140 b_size: usize,
141 ) -> usize {
142 unsafe {
143 vmp::vmp_apply_dft_to_dft_tmp_bytes(
144 self.ptr,
145 (res_size * b_cols_out) as u64,
146 (a_size * b_cols_in) as u64,
147 (b_rows * b_cols_in) as u64,
148 (b_size * b_cols_out) as u64,
149 ) as usize
150 }
151 }
152}
153
154impl MatZnxDftOps<FFT64> for Module<FFT64> {
155 fn vmp_prepare_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
156 where
157 R: MatZnxToMut<FFT64>,
158 A: VecZnxDftToRef<FFT64>,
159 {
160 let mut res: MatZnxDft<&mut [u8], FFT64> = res.to_mut();
161 let a: VecZnxDft<&[u8], _> = a.to_ref();
162
163 #[cfg(debug_assertions)]
164 {
165 assert_eq!(res.n(), self.n());
166 assert_eq!(a.n(), self.n());
167 assert_eq!(
168 a.cols(),
169 res.cols_out(),
170 "a.cols(): {} != res.cols_out(): {}",
171 a.cols(),
172 res.cols_out()
173 );
174 assert!(
175 res_row < res.rows(),
176 "res_row: {} >= res.rows(): {}",
177 res_row,
178 res.rows()
179 );
180 assert!(
181 res_col_in < res.cols_in(),
182 "res_col_in: {} >= res.cols_in(): {}",
183 res_col_in,
184 res.cols_in()
185 );
186 assert_eq!(
187 res.size(),
188 a.size(),
189 "res.size(): {} != a.size(): {}",
190 res.size(),
191 a.size()
192 );
193 }
194
195 unsafe {
196 vmp::vmp_prepare_row_dft(
197 self.ptr,
198 res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
199 a.as_ptr() as *const vec_znx_dft_t,
200 (res_row * res.cols_in() + res_col_in) as u64,
201 (res.rows() * res.cols_in()) as u64,
202 (res.size() * res.cols_out()) as u64,
203 );
204 }
205 }
206
207 fn vmp_extract_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
208 where
209 R: VecZnxDftToMut<FFT64>,
210 A: MatZnxToRef<FFT64>,
211 {
212 let mut res: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
213 let a: MatZnxDft<&[u8], _> = a.to_ref();
214
215 #[cfg(debug_assertions)]
216 {
217 assert_eq!(res.n(), self.n());
218 assert_eq!(a.n(), self.n());
219 assert_eq!(
220 res.cols(),
221 a.cols_out(),
222 "res.cols(): {} != a.cols_out(): {}",
223 res.cols(),
224 a.cols_out()
225 );
226 assert!(
227 a_row < a.rows(),
228 "a_row: {} >= a.rows(): {}",
229 a_row,
230 a.rows()
231 );
232 assert!(
233 a_col_in < a.cols_in(),
234 "a_col_in: {} >= a.cols_in(): {}",
235 a_col_in,
236 a.cols_in()
237 );
238 assert_eq!(
239 res.size(),
240 a.size(),
241 "res.size(): {} != a.size(): {}",
242 res.size(),
243 a.size()
244 );
245 }
246 unsafe {
247 vmp::vmp_extract_row_dft(
248 self.ptr,
249 res.as_mut_ptr() as *mut vec_znx_dft_t,
250 a.as_ptr() as *const vmp::vmp_pmat_t,
251 (a_row * a.cols_in() + a_col_in) as u64,
252 (a.rows() * a.cols_in()) as u64,
253 (a.size() * a.cols_out()) as u64,
254 );
255 }
256 }
257
258 fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
259 where
260 R: VecZnxDftToMut<FFT64>,
261 A: VecZnxDftToRef<FFT64>,
262 B: MatZnxToRef<FFT64>,
263 {
264 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
265 let a: VecZnxDft<&[u8], _> = a.to_ref();
266 let b: MatZnxDft<&[u8], _> = b.to_ref();
267
268 #[cfg(debug_assertions)]
269 {
270 assert_eq!(res.n(), self.n());
271 assert_eq!(b.n(), self.n());
272 assert_eq!(a.n(), self.n());
273 assert_eq!(
274 res.cols(),
275 b.cols_out(),
276 "res.cols(): {} != b.cols_out: {}",
277 res.cols(),
278 b.cols_out()
279 );
280 assert_eq!(
281 a.cols(),
282 b.cols_in(),
283 "a.cols(): {} != b.cols_in: {}",
284 a.cols(),
285 b.cols_in()
286 );
287 }
288
289 let (tmp_bytes, _) = scratch.tmp_slice(self.vmp_apply_tmp_bytes(
290 res.size(),
291 a.size(),
292 b.rows(),
293 b.cols_in(),
294 b.cols_out(),
295 b.size(),
296 ));
297 unsafe {
298 vmp::vmp_apply_dft_to_dft(
299 self.ptr,
300 res.as_mut_ptr() as *mut vec_znx_dft_t,
301 (res.size() * res.cols()) as u64,
302 a.as_ptr() as *const vec_znx_dft_t,
303 (a.size() * a.cols()) as u64,
304 b.as_ptr() as *const vmp::vmp_pmat_t,
305 (b.rows() * b.cols_in()) as u64,
306 (b.size() * b.cols_out()) as u64,
307 tmp_bytes.as_mut_ptr(),
308 )
309 }
310 }
311
312 fn vmp_apply_add<R, A, B>(&self, res: &mut R, a: &A, b: &B, scale: usize, scratch: &mut Scratch)
313 where
314 R: VecZnxDftToMut<FFT64>,
315 A: VecZnxDftToRef<FFT64>,
316 B: MatZnxToRef<FFT64>,
317 {
318 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
319 let a: VecZnxDft<&[u8], _> = a.to_ref();
320 let b: MatZnxDft<&[u8], _> = b.to_ref();
321
322 #[cfg(debug_assertions)]
323 {
324 assert_eq!(res.n(), self.n());
325 assert_eq!(b.n(), self.n());
326 assert_eq!(a.n(), self.n());
327 assert_eq!(
328 res.cols(),
329 b.cols_out(),
330 "res.cols(): {} != b.cols_out: {}",
331 res.cols(),
332 b.cols_out()
333 );
334 assert_eq!(
335 a.cols(),
336 b.cols_in(),
337 "a.cols(): {} != b.cols_in: {}",
338 a.cols(),
339 b.cols_in()
340 );
341 }
342
343 let (tmp_bytes, _) = scratch.tmp_slice(self.vmp_apply_tmp_bytes(
344 res.size(),
345 a.size(),
346 b.rows(),
347 b.cols_in(),
348 b.cols_out(),
349 b.size(),
350 ));
351 unsafe {
352 vmp::vmp_apply_dft_to_dft_add(
353 self.ptr,
354 res.as_mut_ptr() as *mut vec_znx_dft_t,
355 (res.size() * res.cols()) as u64,
356 a.as_ptr() as *const vec_znx_dft_t,
357 (a.size() * a.cols()) as u64,
358 b.as_ptr() as *const vmp::vmp_pmat_t,
359 (b.rows() * b.cols_in()) as u64,
360 (b.size() * b.cols_out()) as u64,
361 scale as u64,
362 tmp_bytes.as_mut_ptr(),
363 )
364 }
365 }
366}
367#[cfg(test)]
368mod tests {
369 use crate::{
370 Decoding, FFT64, FillUniform, MatZnxDft, MatZnxDftOps, Module, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxBig,
371 VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, ZnxInfos, ZnxView, ZnxViewMut,
372 ZnxZero,
373 };
374 use sampling::source::Source;
375
376 use super::{MatZnxDftAlloc, MatZnxDftScratch};
377
378 #[test]
379 fn vmp_prepare_row() {
380 let module: Module<FFT64> = Module::<FFT64>::new(16);
381 let basek: usize = 8;
382 let mat_rows: usize = 4;
383 let mat_cols_in: usize = 2;
384 let mat_cols_out: usize = 2;
385 let mat_size: usize = 5;
386 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
387 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
388 let mut b_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
389 let mut mat: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
390
391 for col_in in 0..mat_cols_in {
392 for row_i in 0..mat_rows {
393 let mut source: Source = Source::new([0u8; 32]);
394 (0..mat_cols_out).for_each(|col_out| {
395 a.fill_uniform(basek, col_out, mat_size, &mut source);
396 module.vec_znx_dft(&mut a_dft, col_out, &a, col_out);
397 });
398 module.vmp_prepare_row(&mut mat, row_i, col_in, &a_dft);
399 module.vmp_extract_row(&mut b_dft, &mat, row_i, col_in);
400 assert_eq!(a_dft.raw(), b_dft.raw());
401 }
402 }
403 }
404
405 #[test]
406 fn vmp_apply() {
407 let log_n: i32 = 5;
408 let n: usize = 1 << log_n;
409
410 let module: Module<FFT64> = Module::<FFT64>::new(n);
411 let basek: usize = 15;
412 let a_size: usize = 5;
413 let mat_size: usize = 6;
414 let res_size: usize = a_size;
415
416 [1, 2].iter().for_each(|in_cols| {
417 [1, 2].iter().for_each(|out_cols| {
418 let a_cols: usize = *in_cols;
419 let res_cols: usize = *out_cols;
420
421 let mat_rows: usize = a_size;
422 let mat_cols_in: usize = a_cols;
423 let mat_cols_out: usize = res_cols;
424
425 let mut scratch: ScratchOwned = ScratchOwned::new(
426 module.vmp_apply_tmp_bytes(
427 res_size,
428 a_size,
429 mat_rows,
430 mat_cols_in,
431 mat_cols_out,
432 mat_size,
433 ) | module.vec_znx_big_normalize_tmp_bytes(),
434 );
435
436 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
437
438 (0..a_cols).for_each(|i| {
439 a.at_mut(i, a_size - 1)[i + 1] = 1;
440 });
441
442 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
443 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
444
445 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
446 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
447
448 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
449
450 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
451 (0..a.size()).for_each(|row_i| {
452 (0..mat_cols_in).for_each(|col_in_i| {
453 (0..mat_cols_out).for_each(|col_out_i| {
454 let idx = 1 + col_in_i * mat_cols_out + col_out_i;
455 tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
456 module.vec_znx_dft(&mut c_dft, col_out_i, &tmp, col_out_i);
457 tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
458 });
459 module.vmp_prepare_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
460 });
461 });
462
463 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
464 (0..a_cols).for_each(|i| {
465 module.vec_znx_dft(&mut a_dft, i, &a, i);
466 });
467
468 module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
469
470 let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
471
472 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, res_size);
473 (0..mat_cols_out).for_each(|i| {
474 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
475 module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
476 });
477
478 (0..mat_cols_out).for_each(|col_i| {
479 let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
480 (0..a_cols).for_each(|i| {
481 res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
482 });
483 res_have.decode_vec_i64(col_i, basek, basek * a_size, &mut res_have_vi64);
484 assert_eq!(res_have_vi64, res_want_vi64);
485 });
486 });
487 });
488 }
489
490 #[test]
491 fn vmp_apply_add() {
492 let log_n: i32 = 5;
493 let n: usize = 1 << log_n;
494
495 let module: Module<FFT64> = Module::<FFT64>::new(n);
496 let basek: usize = 15;
497 let a_size: usize = 5;
498 let mat_size: usize = 6;
499 let res_size: usize = a_size;
500
501 [1, 2].iter().for_each(|in_cols| {
502 [1, 2].iter().for_each(|out_cols| {
503 let a_cols: usize = *in_cols;
504 let res_cols: usize = *out_cols;
505
506 let mat_rows: usize = a_size;
507 let mat_cols_in: usize = a_cols;
508 let mat_cols_out: usize = res_cols;
509
510 let mut scratch: ScratchOwned = ScratchOwned::new(
511 module.vmp_apply_tmp_bytes(
512 res_size,
513 a_size,
514 mat_rows,
515 mat_cols_in,
516 mat_cols_out,
517 mat_size,
518 ) | module.vec_znx_big_normalize_tmp_bytes(),
519 );
520
521 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
522
523 (0..a_cols).for_each(|i| {
524 a.at_mut(i, a_size - 1)[i + 1] = 1;
525 });
526
527 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
528 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
529
530 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
531 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
532
533 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
534
535 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
536 (0..a.size()).for_each(|row_i| {
537 (0..mat_cols_in).for_each(|col_in_i| {
538 (0..mat_cols_out).for_each(|col_out_i| {
539 let idx = 1 + col_in_i * mat_cols_out + col_out_i;
540 tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
541 module.vec_znx_dft(&mut c_dft, col_out_i, &tmp, col_out_i);
542 tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
543 });
544 module.vmp_prepare_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
545 });
546 });
547
548 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
549 (0..a_cols).for_each(|i| {
550 module.vec_znx_dft(&mut a_dft, i, &a, i);
551 });
552
553 c_dft.zero();
554 (0..c_dft.cols()).for_each(|i| {
555 module.vec_znx_dft(&mut c_dft, i, &a, 0);
556 });
557
558 module.vmp_apply_add(&mut c_dft, &a_dft, &mat_znx_dft, 0, scratch.borrow());
559
560 let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
561
562 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, res_size);
563 (0..mat_cols_out).for_each(|i| {
564 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
565 module.vec_znx_big_normalize(basek, &mut res_have, i, &c_big, i, scratch.borrow());
566 });
567
568 (0..mat_cols_out).for_each(|col_i| {
569 let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
570 (0..a_cols).for_each(|i| {
571 res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
572 });
573
574 res_want_vi64[1] += 1;
575
576 res_have.decode_vec_i64(col_i, basek, basek * a_size, &mut res_have_vi64);
577 assert_eq!(res_have_vi64, res_want_vi64);
578 });
579 });
580 });
581 }
582}