poulpy/base2k/src/vmp.rs
base2k/src/vmp.rsRaw
1use crate::ffi::vec_znx_big::vec_znx_big_t;
2use crate::ffi::vec_znx_dft::vec_znx_dft_t;
3use crate::ffi::vmp::{self, vmp_pmat_t};
4use crate::{Backend, FFT64, Infos, Module, VecZnx, VecZnxBig, VecZnxDft, alloc_aligned, assert_alignement};
5use std::marker::PhantomData;
6
7/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
8/// stored as a 3D matrix in the DFT domain in a single contiguous array.
9/// Each col of the [VmpPMat] can be seen as a collection of [VecZnxDft].
10///
11/// [VmpPMat] is used to permform a vector matrix product between a [VecZnx]/[VecZnxDft] and a [VmpPMat].
12/// See the trait [VmpPMatOps] for additional information.
13pub struct VmpPMat<B: Backend> {
14 /// Raw data, is empty if borrowing scratch space.
15 data: Vec<u8>,
16 /// Pointer to data. Can point to scratch space.
17 ptr: *mut u8,
18 /// The ring degree of each polynomial.
19 n: usize,
20 /// Number of rows
21 rows: usize,
22 /// Number of cols
23 cols: usize,
24 /// The number of small polynomials
25 limbs: usize,
26 _marker: PhantomData<B>,
27}
28
29impl<B: Backend> Infos for VmpPMat<B> {
30 fn n(&self) -> usize {
31 self.n
32 }
33
34 fn log_n(&self) -> usize {
35 (usize::BITS - (self.n() - 1).leading_zeros()) as _
36 }
37
38 fn rows(&self) -> usize {
39 self.rows
40 }
41
42 fn cols(&self) -> usize {
43 self.cols
44 }
45
46 fn limbs(&self) -> usize {
47 self.limbs
48 }
49
50 fn poly_count(&self) -> usize {
51 self.rows * self.cols * self.limbs
52 }
53}
54
55impl VmpPMat<FFT64> {
56
57 fn new(module: &Module<FFT64>, rows: usize, cols: usize, limbs: usize) -> VmpPMat<FFT64> {
58 let mut data: Vec<u8> = alloc_aligned::<u8>(module.bytes_of_vmp_pmat(rows, cols, limbs));
59 let ptr: *mut u8 = data.as_mut_ptr();
60 VmpPMat::<FFT64> {
61 data: data,
62 ptr: ptr,
63 n: module.n(),
64 rows: rows,
65 cols: cols,
66 limbs: limbs,
67 _marker: PhantomData,
68 }
69 }
70
71 pub fn as_ptr(&self) -> *const u8 {
72 self.ptr
73 }
74
75 pub fn as_mut_ptr(&self) -> *mut u8 {
76 self.ptr
77 }
78
79 pub fn borrowed(&self) -> bool {
80 self.data.len() == 0
81 }
82
83 /// Returns a non-mutable reference to the entire contiguous array of the [VmpPMat].
84 pub fn raw(&self) -> &[f64] {
85 let ptr: *const f64 = self.ptr as *const f64;
86 let size: usize = self.n() * self.poly_count();
87 unsafe { &std::slice::from_raw_parts(ptr, size) }
88 }
89
90 /// Returns a mutable reference of to the entire contiguous array of the [VmpPMat].
91 pub fn raw_mut(&self) -> &mut [f64] {
92 let ptr: *mut f64 = self.ptr as *mut f64;
93 let size: usize = self.n() * self.poly_count();
94 unsafe { std::slice::from_raw_parts_mut(ptr, size) }
95 }
96
97 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
98 ///
99 /// # Arguments
100 ///
101 /// * `row`: row index (i).
102 /// * `col`: col index (j).
103 pub fn at(&self, row: usize, col: usize) -> Vec<f64> {
104 let mut res: Vec<f64> = alloc_aligned(self.n);
105
106 if self.n < 8 {
107 res.copy_from_slice(&self.raw()[(row + col * self.rows()) * self.n()..(row + col * self.rows()) * (self.n() + 1)]);
108 } else {
109 (0..self.n >> 3).for_each(|blk| {
110 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
111 });
112 }
113
114 res
115 }
116
117 fn at_block(&self, row: usize, col: usize, blk: usize) -> &[f64] {
118 let nrows: usize = self.rows();
119 let nsize: usize = self.limbs();
120 if col == (nsize - 1) && (nsize & 1 == 1) {
121 &self.raw()[blk * nrows * nsize * 8 + col * nrows * 8 + row * 8..]
122 } else {
123 &self.raw()[blk * nrows * nsize * 8 + (col / 2) * (2 * nrows) * 8 + row * 2 * 8 + (col % 2) * 8..]
124 }
125 }
126}
127
128/// This trait implements methods for vector matrix product,
129/// that is, multiplying a [VecZnx] with a [VmpPMat].
130pub trait VmpPMatOps<B: Backend> {
131 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize, limbs: usize) -> usize;
132
133 /// Allocates a new [VmpPMat] with the given number of rows and columns.
134 ///
135 /// # Arguments
136 ///
137 /// * `rows`: number of rows (number of [VecZnxDft]).
138 /// * `size`: number of size (number of size of each [VecZnxDft]).
139 fn new_vmp_pmat(&self, rows: usize, cols: usize, limbs: usize) -> VmpPMat<B>;
140
141 /// Returns the number of bytes needed as scratch space for [VmpPMatOps::vmp_prepare_contiguous].
142 ///
143 /// # Arguments
144 ///
145 /// * `rows`: number of rows of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
146 /// * `size`: number of size of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
147 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize, size: usize) -> usize;
148
149 /// Prepares a [VmpPMat] from a contiguous array of [i64].
150 /// The helper struct [Matrix3D] can be used to contruct and populate
151 /// the appropriate contiguous array.
152 ///
153 /// # Arguments
154 ///
155 /// * `b`: [VmpPMat] on which the values are encoded.
156 /// * `a`: the contiguous array of [i64] of the 3D matrix to encode on the [VmpPMat].
157 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
158 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat<B>, a: &[i64], buf: &mut [u8]);
159
160 /// Prepares the ith-row of [VmpPMat] from a [VecZnx].
161 ///
162 /// # Arguments
163 ///
164 /// * `b`: [VmpPMat] on which the values are encoded.
165 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
166 /// * `row_i`: the index of the row to prepare.
167 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
168 ///
169 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
170 fn vmp_prepare_row(&self, b: &mut VmpPMat<B>, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]);
171
172 /// Extracts the ith-row of [VmpPMat] into a [VecZnxBig].
173 ///
174 /// # Arguments
175 ///
176 /// * `b`: the [VecZnxBig] to on which to extract the row of the [VmpPMat].
177 /// * `a`: [VmpPMat] on which the values are encoded.
178 /// * `row_i`: the index of the row to extract.
179 fn vmp_extract_row(&self, b: &mut VecZnxBig<B>, a: &VmpPMat<B>, row_i: usize);
180
181 /// Prepares the ith-row of [VmpPMat] from a [VecZnxDft].
182 ///
183 /// # Arguments
184 ///
185 /// * `b`: [VmpPMat] on which the values are encoded.
186 /// * `a`: the [VecZnxDft] to encode on the [VmpPMat].
187 /// * `row_i`: the index of the row to prepare.
188 ///
189 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
190 fn vmp_prepare_row_dft(&self, b: &mut VmpPMat<B>, a: &VecZnxDft<B>, row_i: usize);
191
192 /// Extracts the ith-row of [VmpPMat] into a [VecZnxDft].
193 ///
194 /// # Arguments
195 ///
196 /// * `b`: the [VecZnxDft] to on which to extract the row of the [VmpPMat].
197 /// * `a`: [VmpPMat] on which the values are encoded.
198 /// * `row_i`: the index of the row to extract.
199 fn vmp_extract_row_dft(&self, b: &mut VecZnxDft<B>, a: &VmpPMat<B>, row_i: usize);
200
201 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft].
202 ///
203 /// # Arguments
204 ///
205 /// * `c_size`: number of size of the output [VecZnxDft].
206 /// * `a_size`: number of size of the input [VecZnx].
207 /// * `rows`: number of rows of the input [VmpPMat].
208 /// * `size`: number of size of the input [VmpPMat].
209 fn vmp_apply_dft_tmp_bytes(&self, c_size: usize, a_size: usize, rows: usize, size: usize) -> usize;
210
211 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
212 ///
213 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
214 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
215 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
216 ///
217 /// As such, given an input [VecZnx] of `i` size and a [VmpPMat] of `i` rows and
218 /// `j` size, the output is a [VecZnx] of `j` size.
219 ///
220 /// If there is a mismatch between the dimensions the largest valid ones are used.
221 ///
222 /// ```text
223 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
224 /// |h i j|
225 /// |k l m|
226 /// ```
227 /// where each element is a [VecZnxDft].
228 ///
229 /// # Arguments
230 ///
231 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
232 /// * `a`: the left operand [VecZnx] of the vector matrix product.
233 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
234 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
235 fn vmp_apply_dft(&self, c: &mut VecZnxDft<B>, a: &VecZnx, b: &VmpPMat<B>, buf: &mut [u8]);
236
237 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on the receiver.
238 ///
239 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
240 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
241 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
242 ///
243 /// As such, given an input [VecZnx] of `i` size and a [VmpPMat] of `i` rows and
244 /// `j` size, the output is a [VecZnx] of `j` size.
245 ///
246 /// If there is a mismatch between the dimensions the largest valid ones are used.
247 ///
248 /// ```text
249 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
250 /// |h i j|
251 /// |k l m|
252 /// ```
253 /// where each element is a [VecZnxDft].
254 ///
255 /// # Arguments
256 ///
257 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
258 /// * `a`: the left operand [VecZnx] of the vector matrix product.
259 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
260 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
261 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft<B>, a: &VecZnx, b: &VmpPMat<B>, buf: &mut [u8]);
262
263 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
264 ///
265 /// # Arguments
266 ///
267 /// * `c_size`: number of size of the output [VecZnxDft].
268 /// * `a_size`: number of size of the input [VecZnxDft].
269 /// * `rows`: number of rows of the input [VmpPMat].
270 /// * `size`: number of size of the input [VmpPMat].
271 fn vmp_apply_dft_to_dft_tmp_bytes(&self, c_size: usize, a_size: usize, rows: usize, size: usize) -> usize;
272
273 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
274 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
275 ///
276 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
277 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
278 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
279 ///
280 /// As such, given an input [VecZnx] of `i` size and a [VmpPMat] of `i` rows and
281 /// `j` size, the output is a [VecZnx] of `j` size.
282 ///
283 /// If there is a mismatch between the dimensions the largest valid ones are used.
284 ///
285 /// ```text
286 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
287 /// |h i j|
288 /// |k l m|
289 /// ```
290 /// where each element is a [VecZnxDft].
291 ///
292 /// # Arguments
293 ///
294 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
295 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
296 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
297 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
298 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft<B>, a: &VecZnxDft<B>, b: &VmpPMat<B>, buf: &mut [u8]);
299
300 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] and adds on top of the receiver instead of overwritting it.
301 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
302 ///
303 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
304 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
305 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
306 ///
307 /// As such, given an input [VecZnx] of `i` size and a [VmpPMat] of `i` rows and
308 /// `j` size, the output is a [VecZnx] of `j` size.
309 ///
310 /// If there is a mismatch between the dimensions the largest valid ones are used.
311 ///
312 /// ```text
313 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
314 /// |h i j|
315 /// |k l m|
316 /// ```
317 /// where each element is a [VecZnxDft].
318 ///
319 /// # Arguments
320 ///
321 /// * `c`: the operand on which the output of the vector matrix product is added, as a [VecZnxDft].
322 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
323 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
324 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
325 fn vmp_apply_dft_to_dft_add(&self, c: &mut VecZnxDft<B>, a: &VecZnxDft<B>, b: &VmpPMat<B>, buf: &mut [u8]);
326
327 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
328 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
329 ///
330 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
331 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
332 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
333 ///
334 /// As such, given an input [VecZnx] of `i` size and a [VmpPMat] of `i` rows and
335 /// `j` size, the output is a [VecZnx] of `j` size.
336 ///
337 /// If there is a mismatch between the dimensions the largest valid ones are used.
338 ///
339 /// ```text
340 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
341 /// |h i j|
342 /// |k l m|
343 /// ```
344 /// where each element is a [VecZnxDft].
345 ///
346 /// # Arguments
347 ///
348 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
349 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
350 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
351 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft<B>, a: &VmpPMat<B>, buf: &mut [u8]);
352}
353
354impl VmpPMatOps<FFT64> for Module<FFT64> {
355
356 fn new_vmp_pmat(&self, rows: usize, cols: usize, limbs: usize) -> VmpPMat<FFT64> {
357 VmpPMat::<FFT64>::new(self, rows, cols, limbs)
358 }
359
360 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize, limbs: usize) -> usize {
361 unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, (limbs* cols) as u64) as usize }
362 }
363
364 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize, size: usize) -> usize {
365 unsafe { vmp::vmp_prepare_tmp_bytes(self.ptr, rows as u64, (size * cols) as u64) as usize }
366 }
367
368 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat<FFT64>, a: &[i64], tmp_bytes: &mut [u8]) {
369
370 #[cfg(debug_assertions)]
371 {
372 assert_eq!(a.len(), b.n() * b.poly_count());
373 assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols(), b.limbs()));
374 assert_alignement(tmp_bytes.as_ptr());
375 }
376 unsafe {
377 vmp::vmp_prepare_contiguous(
378 self.ptr,
379 b.as_mut_ptr() as *mut vmp_pmat_t,
380 a.as_ptr(),
381 b.rows() as u64,
382 (b.limbs()*b.cols()) as u64,
383 tmp_bytes.as_mut_ptr(),
384 );
385 }
386 }
387
388 fn vmp_prepare_row(&self, b: &mut VmpPMat<FFT64>, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]) {
389 #[cfg(debug_assertions)]
390 {
391 assert_eq!(a.len(), b.limbs() * self.n() * b.cols());
392 assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols(), b.limbs()));
393 assert_alignement(tmp_bytes.as_ptr());
394 }
395 unsafe {
396 vmp::vmp_prepare_row(
397 self.ptr,
398 b.as_mut_ptr() as *mut vmp_pmat_t,
399 a.as_ptr(),
400 row_i as u64,
401 b.rows() as u64,
402 (b.limbs()*b.cols()) as u64,
403 tmp_bytes.as_mut_ptr(),
404 );
405 }
406 }
407
408 fn vmp_extract_row(&self, b: &mut VecZnxBig<FFT64>, a: &VmpPMat<FFT64>, row_i: usize) {
409 #[cfg(debug_assertions)]
410 {
411 assert_eq!(a.n(), b.n());
412 assert_eq!(a.limbs(), b.limbs());
413 assert_eq!(a.cols(), b.cols());
414 }
415 unsafe {
416 vmp::vmp_extract_row(
417 self.ptr,
418 b.ptr as *mut vec_znx_big_t,
419 a.as_ptr() as *const vmp_pmat_t,
420 row_i as u64,
421 a.rows() as u64,
422 (a.limbs()*a.cols()) as u64,
423 );
424 }
425 }
426
427 fn vmp_prepare_row_dft(&self, b: &mut VmpPMat<FFT64>, a: &VecZnxDft<FFT64>, row_i: usize) {
428 #[cfg(debug_assertions)]
429 {
430 assert_eq!(a.n(), b.n());
431 assert_eq!(a.limbs(), b.limbs());
432 }
433 unsafe {
434 vmp::vmp_prepare_row_dft(
435 self.ptr,
436 b.as_mut_ptr() as *mut vmp_pmat_t,
437 a.ptr as *const vec_znx_dft_t,
438 row_i as u64,
439 b.rows() as u64,
440 b.limbs() as u64,
441 );
442 }
443 }
444
445 fn vmp_extract_row_dft(&self, b: &mut VecZnxDft<FFT64>, a: &VmpPMat<FFT64>, row_i: usize) {
446 #[cfg(debug_assertions)]
447 {
448 assert_eq!(a.n(), b.n());
449 assert_eq!(a.limbs(), b.limbs());
450 }
451 unsafe {
452 vmp::vmp_extract_row_dft(
453 self.ptr,
454 b.ptr as *mut vec_znx_dft_t,
455 a.as_ptr() as *const vmp_pmat_t,
456 row_i as u64,
457 a.rows() as u64,
458 a.limbs() as u64,
459 );
460 }
461 }
462
463 fn vmp_apply_dft_tmp_bytes(&self, res_size: usize, a_size: usize, gct_rows: usize, gct_size: usize) -> usize {
464 unsafe {
465 vmp::vmp_apply_dft_tmp_bytes(
466 self.ptr,
467 res_size as u64,
468 a_size as u64,
469 gct_rows as u64,
470 gct_size as u64,
471 ) as usize
472 }
473 }
474
475 fn vmp_apply_dft(&self, c: &mut VecZnxDft<FFT64>, a: &VecZnx, b: &VmpPMat<FFT64>, tmp_bytes: &mut [u8]) {
476 debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.limbs(), a.limbs(), b.rows(), b.limbs()));
477 #[cfg(debug_assertions)]
478 {
479 assert_alignement(tmp_bytes.as_ptr());
480 }
481 unsafe {
482 vmp::vmp_apply_dft(
483 self.ptr,
484 c.ptr as *mut vec_znx_dft_t,
485 c.limbs() as u64,
486 a.as_ptr(),
487 a.limbs() as u64,
488 (a.n() * a.cols()) as u64,
489 b.as_ptr() as *const vmp_pmat_t,
490 b.rows() as u64,
491 b.limbs() as u64,
492 tmp_bytes.as_mut_ptr(),
493 )
494 }
495 }
496
497 fn vmp_apply_dft_add(&self, c: &mut VecZnxDft<FFT64>, a: &VecZnx, b: &VmpPMat<FFT64>, tmp_bytes: &mut [u8]) {
498 debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.limbs(), a.limbs(), b.rows(), b.limbs()));
499 #[cfg(debug_assertions)]
500 {
501 assert_alignement(tmp_bytes.as_ptr());
502 }
503 unsafe {
504 vmp::vmp_apply_dft_add(
505 self.ptr,
506 c.ptr as *mut vec_znx_dft_t,
507 c.limbs() as u64,
508 a.as_ptr(),
509 a.limbs() as u64,
510 (a.n() * a.limbs()) as u64,
511 b.as_ptr() as *const vmp_pmat_t,
512 b.rows() as u64,
513 b.limbs() as u64,
514 tmp_bytes.as_mut_ptr(),
515 )
516 }
517 }
518
519 fn vmp_apply_dft_to_dft_tmp_bytes(&self, res_size: usize, a_size: usize, gct_rows: usize, gct_size: usize) -> usize {
520 unsafe {
521 vmp::vmp_apply_dft_to_dft_tmp_bytes(
522 self.ptr,
523 res_size as u64,
524 a_size as u64,
525 gct_rows as u64,
526 gct_size as u64,
527 ) as usize
528 }
529 }
530
531 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft<FFT64>, a: &VecZnxDft<FFT64>, b: &VmpPMat<FFT64>, tmp_bytes: &mut [u8]) {
532 debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(c.limbs(), a.limbs(), b.rows(), b.limbs()));
533 #[cfg(debug_assertions)]
534 {
535 assert_alignement(tmp_bytes.as_ptr());
536 }
537 unsafe {
538 vmp::vmp_apply_dft_to_dft(
539 self.ptr,
540 c.ptr as *mut vec_znx_dft_t,
541 c.limbs() as u64,
542 a.ptr as *const vec_znx_dft_t,
543 a.limbs() as u64,
544 b.as_ptr() as *const vmp_pmat_t,
545 b.rows() as u64,
546 b.limbs() as u64,
547 tmp_bytes.as_mut_ptr(),
548 )
549 }
550 }
551
552 fn vmp_apply_dft_to_dft_add(&self, c: &mut VecZnxDft<FFT64>, a: &VecZnxDft<FFT64>, b: &VmpPMat<FFT64>, tmp_bytes: &mut [u8]) {
553 debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(c.limbs(), a.limbs(), b.rows(), b.limbs()));
554 #[cfg(debug_assertions)]
555 {
556 assert_alignement(tmp_bytes.as_ptr());
557 }
558 unsafe {
559 vmp::vmp_apply_dft_to_dft_add(
560 self.ptr,
561 c.ptr as *mut vec_znx_dft_t,
562 c.limbs() as u64,
563 a.ptr as *const vec_znx_dft_t,
564 a.limbs() as u64,
565 b.as_ptr() as *const vmp_pmat_t,
566 b.rows() as u64,
567 b.limbs() as u64,
568 tmp_bytes.as_mut_ptr(),
569 )
570 }
571 }
572
573 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft<FFT64>, a: &VmpPMat<FFT64>, tmp_bytes: &mut [u8]) {
574 debug_assert!(tmp_bytes.len() >= self.vmp_apply_dft_to_dft_tmp_bytes(b.limbs(), b.limbs(), a.rows(), a.limbs()));
575 #[cfg(debug_assertions)]
576 {
577 assert_alignement(tmp_bytes.as_ptr());
578 }
579 unsafe {
580 vmp::vmp_apply_dft_to_dft(
581 self.ptr,
582 b.ptr as *mut vec_znx_dft_t,
583 b.limbs() as u64,
584 b.ptr as *mut vec_znx_dft_t,
585 b.limbs() as u64,
586 a.as_ptr() as *const vmp_pmat_t,
587 a.rows() as u64,
588 a.limbs() as u64,
589 tmp_bytes.as_mut_ptr(),
590 )
591 }
592 }
593}
594
595#[cfg(test)]
596mod tests {
597 use crate::{
598 FFT64, Module, Sampling, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat, VmpPMatOps,
599 alloc_aligned,
600 };
601 use sampling::source::Source;
602
603 #[test]
604 fn vmp_prepare_row_dft() {
605 let module: Module<FFT64> = Module::<FFT64>::new(32);
606 let vpmat_rows: usize = 4;
607 let vpmat_size: usize = 5;
608 let log_base2k: usize = 8;
609 let mut a: VecZnx = module.new_vec_znx(1, vpmat_size);
610 let mut a_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, vpmat_size);
611 let mut a_big: VecZnxBig<FFT64> = module.new_vec_znx_big(1, vpmat_size);
612 let mut b_big: VecZnxBig<FFT64> = module.new_vec_znx_big(1, vpmat_size);
613 let mut b_dft: VecZnxDft<FFT64> = module.new_vec_znx_dft(1, vpmat_size);
614 let mut vmpmat_0: VmpPMat<FFT64> = module.new_vmp_pmat(vpmat_rows, 1, vpmat_size);
615 let mut vmpmat_1: VmpPMat<FFT64> = module.new_vmp_pmat(vpmat_rows, 1, vpmat_size);
616
617 let mut tmp_bytes: Vec<u8> = alloc_aligned(module.vmp_prepare_tmp_bytes(vpmat_rows, 1, vpmat_size));
618
619 for row_i in 0..vpmat_rows {
620 let mut source: Source = Source::new([0u8; 32]);
621 module.fill_uniform(log_base2k, &mut a, 0, vpmat_size, &mut source);
622 module.vec_znx_dft(&mut a_dft, &a);
623 module.vmp_prepare_row(&mut vmpmat_0, &a.raw(), row_i, &mut tmp_bytes);
624
625 // Checks that prepare(vmp_pmat, a) = prepare_dft(vmp_pmat, a_dft)
626 module.vmp_prepare_row_dft(&mut vmpmat_1, &a_dft, row_i);
627 assert_eq!(vmpmat_0.raw(), vmpmat_1.raw());
628
629 // Checks that a_dft = extract_dft(prepare(vmp_pmat, a), b_dft)
630 module.vmp_extract_row_dft(&mut b_dft, &vmpmat_0, row_i);
631 assert_eq!(a_dft.raw(), b_dft.raw());
632
633 // Checks that a_big = extract(prepare_dft(vmp_pmat, a_dft), b_big)
634 module.vmp_extract_row(&mut b_big, &vmpmat_0, row_i);
635 module.vec_znx_idft(&mut a_big, &a_dft, &mut tmp_bytes);
636 assert_eq!(a_big.raw(), b_big.raw());
637 }
638
639 module.free();
640 }
641}