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