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, MatZnxDftOwned, MatZnxDftToMut, MatZnxDftToRef, 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 /// * `b`: [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: MatZnxDftToMut<BACKEND>,
58 A: VecZnxDftToRef<BACKEND>;
59
60 /// Extracts the ith-row of [MatZnxDft] into a [VecZnxDft].
61 ///
62 /// # Arguments
63 ///
64 /// * `b`: 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<BACKEND>,
70 A: MatZnxDftToRef<BACKEND>;
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<BACKEND>,
100 A: VecZnxDftToRef<BACKEND>,
101 B: MatZnxDftToRef<BACKEND>;
102}
103
104impl<B: Backend> MatZnxDftAlloc<B> for Module<B> {
105 fn bytes_of_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
106 MatZnxDftOwned::bytes_of(self, rows, cols_in, cols_out, size)
107 }
108
109 fn new_mat_znx_dft(&self, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> MatZnxDftOwned<B> {
110 MatZnxDftOwned::new(self, rows, cols_in, cols_out, size)
111 }
112
113 fn new_mat_znx_dft_from_bytes(
114 &self,
115 rows: usize,
116 cols_in: usize,
117 cols_out: usize,
118 size: usize,
119 bytes: Vec<u8>,
120 ) -> MatZnxDftOwned<B> {
121 MatZnxDftOwned::new_from_bytes(self, rows, cols_in, cols_out, size, bytes)
122 }
123}
124
125impl<BACKEND: Backend> MatZnxDftScratch for Module<BACKEND> {
126 fn vmp_apply_tmp_bytes(
127 &self,
128 res_size: usize,
129 a_size: usize,
130 b_rows: usize,
131 b_cols_in: usize,
132 b_cols_out: usize,
133 b_size: usize,
134 ) -> usize {
135 unsafe {
136 vmp::vmp_apply_dft_to_dft_tmp_bytes(
137 self.ptr,
138 (res_size * b_cols_out) as u64,
139 (a_size * b_cols_in) as u64,
140 (b_rows * b_cols_in) as u64,
141 (b_size * b_cols_out) as u64,
142 ) as usize
143 }
144 }
145}
146
147impl MatZnxDftOps<FFT64> for Module<FFT64> {
148 fn vmp_prepare_row<R, A>(&self, res: &mut R, res_row: usize, res_col_in: usize, a: &A)
149 where
150 R: MatZnxDftToMut<FFT64>,
151 A: VecZnxDftToRef<FFT64>,
152 {
153 let mut res: MatZnxDft<&mut [u8], _> = res.to_mut();
154 let a: VecZnxDft<&[u8], _> = a.to_ref();
155
156 #[cfg(debug_assertions)]
157 {
158 assert_eq!(res.n(), self.n());
159 assert_eq!(a.n(), self.n());
160 assert_eq!(
161 a.cols(),
162 res.cols_out(),
163 "a.cols(): {} != res.cols_out(): {}",
164 a.cols(),
165 res.cols_out()
166 );
167 assert!(
168 res_row < res.rows(),
169 "res_row: {} >= res.rows(): {}",
170 res_row,
171 res.rows()
172 );
173 assert!(
174 res_col_in < res.cols_in(),
175 "res_col_in: {} >= res.cols_in(): {}",
176 res_col_in,
177 res.cols_in()
178 );
179 assert_eq!(
180 res.size(),
181 a.size(),
182 "res.size(): {} != a.size(): {}",
183 res.size(),
184 a.size()
185 );
186 }
187
188 unsafe {
189 vmp::vmp_prepare_row_dft(
190 self.ptr,
191 res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
192 a.as_ptr() as *const vec_znx_dft_t,
193 (res_row * res.cols_in() + res_col_in) as u64,
194 (res.rows() * res.cols_in()) as u64,
195 (res.size() * res.cols_out()) as u64,
196 );
197 }
198 }
199
200 fn vmp_extract_row<R, A>(&self, res: &mut R, a: &A, a_row: usize, a_col_in: usize)
201 where
202 R: VecZnxDftToMut<FFT64>,
203 A: MatZnxDftToRef<FFT64>,
204 {
205 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
206 let a: MatZnxDft<&[u8], _> = a.to_ref();
207
208 #[cfg(debug_assertions)]
209 {
210 assert_eq!(res.n(), self.n());
211 assert_eq!(a.n(), self.n());
212 assert_eq!(
213 res.cols(),
214 a.cols_out(),
215 "res.cols(): {} != a.cols_out(): {}",
216 res.cols(),
217 a.cols_out()
218 );
219 assert!(
220 a_row < a.rows(),
221 "a_row: {} >= a.rows(): {}",
222 a_row,
223 a.rows()
224 );
225 assert!(
226 a_col_in < a.cols_in(),
227 "a_col_in: {} >= a.cols_in(): {}",
228 a_col_in,
229 a.cols_in()
230 );
231 assert_eq!(
232 res.size(),
233 a.size(),
234 "res.size(): {} != a.size(): {}",
235 res.size(),
236 a.size()
237 );
238 }
239 unsafe {
240 vmp::vmp_extract_row_dft(
241 self.ptr,
242 res.as_mut_ptr() as *mut vec_znx_dft_t,
243 a.as_ptr() as *const vmp::vmp_pmat_t,
244 (a_row * a.cols_in() + a_col_in) as u64,
245 (a.rows() * a.cols_in()) as u64,
246 (a.size() * a.cols_out()) as u64,
247 );
248 }
249 }
250
251 fn vmp_apply<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch)
252 where
253 R: VecZnxDftToMut<FFT64>,
254 A: VecZnxDftToRef<FFT64>,
255 B: MatZnxDftToRef<FFT64>,
256 {
257 let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
258 let a: VecZnxDft<&[u8], _> = a.to_ref();
259 let b: MatZnxDft<&[u8], _> = b.to_ref();
260
261 #[cfg(debug_assertions)]
262 {
263 assert_eq!(res.n(), self.n());
264 assert_eq!(b.n(), self.n());
265 assert_eq!(a.n(), self.n());
266 assert_eq!(
267 res.cols(),
268 b.cols_out(),
269 "res.cols(): {} != b.cols_out: {}",
270 res.cols(),
271 b.cols_out()
272 );
273 assert_eq!(
274 a.cols(),
275 b.cols_in(),
276 "a.cols(): {} != b.cols_in: {}",
277 a.cols(),
278 b.cols_in()
279 );
280 }
281
282 let (tmp_bytes, _) = scratch.tmp_slice(self.vmp_apply_tmp_bytes(
283 res.size(),
284 a.size(),
285 b.rows(),
286 b.cols_in(),
287 b.cols_out(),
288 b.size(),
289 ));
290 unsafe {
291 vmp::vmp_apply_dft_to_dft(
292 self.ptr,
293 res.as_mut_ptr() as *mut vec_znx_dft_t,
294 (res.size() * res.cols()) as u64,
295 a.as_ptr() as *const vec_znx_dft_t,
296 (a.size() * a.cols()) as u64,
297 b.as_ptr() as *const vmp::vmp_pmat_t,
298 (b.rows() * b.cols_in()) as u64,
299 (b.size() * b.cols_out()) as u64,
300 tmp_bytes.as_mut_ptr(),
301 )
302 }
303 }
304}
305#[cfg(test)]
306mod tests {
307 use crate::{
308 Encoding, FFT64, FillUniform, MatZnxDft, MatZnxDftOps, Module, ScratchOwned, VecZnx, VecZnxAlloc, VecZnxBig,
309 VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, ZnxInfos, ZnxView, ZnxViewMut,
310 };
311 use sampling::source::Source;
312
313 use super::{MatZnxDftAlloc, MatZnxDftScratch};
314
315 #[test]
316 fn vmp_prepare_row() {
317 let module: Module<FFT64> = Module::<FFT64>::new(16);
318 let log_base2k: usize = 8;
319 let mat_rows: usize = 4;
320 let mat_cols_in: usize = 2;
321 let mat_cols_out: usize = 2;
322 let mat_size: usize = 5;
323 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
324 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
325 let mut b_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
326 let mut mat: MatZnxDft<Vec<u8>, FFT64> = module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
327
328 for col_in in 0..mat_cols_in {
329 for row_i in 0..mat_rows {
330 let mut source: Source = Source::new([0u8; 32]);
331 (0..mat_cols_out).for_each(|col_out| {
332 a.fill_uniform(log_base2k, col_out, mat_size, &mut source);
333 module.vec_znx_dft(&mut a_dft, col_out, &a, col_out);
334 });
335 module.vmp_prepare_row(&mut mat, row_i, col_in, &a_dft);
336 module.vmp_extract_row(&mut b_dft, &mat, row_i, col_in);
337 assert_eq!(a_dft.raw(), b_dft.raw());
338 }
339 }
340
341 module.free();
342 }
343
344 #[test]
345 fn vmp_apply() {
346 let log_n: i32 = 5;
347 let n: usize = 1 << log_n;
348
349 let module: Module<FFT64> = Module::<FFT64>::new(n);
350 let log_base2k: usize = 15;
351 let a_size: usize = 5;
352 let mat_size: usize = 6;
353 let res_size: usize = 5;
354
355 [1, 2].iter().for_each(|in_cols| {
356 [1, 2].iter().for_each(|out_cols| {
357 let a_cols: usize = *in_cols;
358 let res_cols: usize = *out_cols;
359
360 let mat_rows: usize = a_size;
361 let mat_cols_in: usize = a_cols;
362 let mat_cols_out: usize = res_cols;
363 let res_cols: usize = mat_cols_out;
364
365 let mut scratch: ScratchOwned = ScratchOwned::new(
366 module.vmp_apply_tmp_bytes(
367 res_size,
368 a_size,
369 mat_rows,
370 mat_cols_in,
371 mat_cols_out,
372 mat_size,
373 ) | module.vec_znx_big_normalize_tmp_bytes(),
374 );
375
376 let mut a: VecZnx<Vec<u8>> = module.new_vec_znx(a_cols, a_size);
377
378 (0..a_cols).for_each(|i| {
379 a.at_mut(i, 2)[i + 1] = 1;
380 });
381
382 let mut mat_znx_dft: MatZnxDft<Vec<u8>, FFT64> =
383 module.new_mat_znx_dft(mat_rows, mat_cols_in, mat_cols_out, mat_size);
384
385 let mut c_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(mat_cols_out, mat_size);
386 let mut c_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(mat_cols_out, mat_size);
387
388 let mut tmp: VecZnx<Vec<u8>> = module.new_vec_znx(mat_cols_out, mat_size);
389
390 // Construts a [VecZnxMatDft] that performs cyclic rotations on each submatrix.
391 (0..a.size()).for_each(|row_i| {
392 (0..mat_cols_in).for_each(|col_in_i| {
393 (0..mat_cols_out).for_each(|col_out_i| {
394 let idx = 1 + col_in_i * mat_cols_out + col_out_i;
395 tmp.at_mut(col_out_i, row_i)[idx] = 1 as i64; // X^{idx}
396 module.vec_znx_dft(&mut c_dft, col_out_i, &tmp, col_out_i);
397 tmp.at_mut(col_out_i, row_i)[idx] = 0 as i64;
398 });
399 module.vmp_prepare_row(&mut mat_znx_dft, row_i, col_in_i, &c_dft);
400 });
401 });
402
403 let mut a_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(a_cols, a_size);
404 (0..a_cols).for_each(|i| {
405 module.vec_znx_dft(&mut a_dft, i, &a, i);
406 });
407
408 module.vmp_apply(&mut c_dft, &a_dft, &mat_znx_dft, scratch.borrow());
409
410 let mut res_have_vi64: Vec<i64> = vec![i64::default(); n];
411
412 let mut res_have: VecZnx<Vec<u8>> = module.new_vec_znx(res_cols, res_size);
413 (0..mat_cols_out).for_each(|i| {
414 module.vec_znx_idft_tmp_a(&mut c_big, i, &mut c_dft, i);
415 module.vec_znx_big_normalize(log_base2k, &mut res_have, i, &c_big, i, scratch.borrow());
416 });
417
418 (0..mat_cols_out).for_each(|col_i| {
419 let mut res_want_vi64: Vec<i64> = vec![i64::default(); n];
420 (0..a_cols).for_each(|i| {
421 res_want_vi64[(i + 1) + (1 + i * mat_cols_out + col_i)] = 1;
422 });
423 res_have.decode_vec_i64(col_i, log_base2k, log_base2k * 3, &mut res_have_vi64);
424 assert_eq!(res_have_vi64, res_want_vi64);
425 });
426 });
427 });
428
429 module.free();
430 }
431}