poulpy/base2k/src/vmp.rs
base2k/src/vmp.rsRaw
1use crate::ffi::vec_znx_dft::vec_znx_dft_t;
2use crate::ffi::vmp::{self, vmp_pmat_t};
3use crate::{alloc_aligned, assert_alignement, Infos, Module, VecZnx, VecZnxDft, MODULETYPE};
4
5/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
6/// stored as a 3D matrix in the DFT domain in a single contiguous array.
7/// Each row of the [VmpPMat] can be seen as a [VecZnxDft].
8///
9/// The backend array of [VmpPMat] is allocate in C,
10/// and thus must be manually freed.
11///
12/// [VmpPMat] is used to permform a vector matrix product between a [VecZnx] and a [VmpPMat].
13/// See the trait [VmpPMatOps] for additional information.
14pub struct VmpPMat {
15 /// Raw data, is empty if borrowing scratch space.
16 data: Vec<u8>,
17 /// Pointer to data. Can point to scratch space.
18 ptr: *mut u8,
19 /// The number of [VecZnxDft].
20 rows: usize,
21 /// The number of cols in each [VecZnxDft].
22 cols: usize,
23 /// The ring degree of each [VecZnxDft].
24 n: usize,
25
26 backend: MODULETYPE,
27}
28
29impl Infos for VmpPMat {
30 /// Returns the ring dimension of the [VmpPMat].
31 fn n(&self) -> usize {
32 self.n
33 }
34
35 fn log_n(&self) -> usize {
36 (usize::BITS - (self.n() - 1).leading_zeros()) as _
37 }
38
39 /// Returns the number of rows (i.e. of [VecZnxDft]) of the [VmpPMat]
40 fn rows(&self) -> usize {
41 self.rows
42 }
43
44 /// Returns the number of cols of the [VmpPMat].
45 /// The number of cols refers to the number of cols
46 /// of each [VecZnxDft].
47 /// This method is equivalent to [Self::cols].
48 fn cols(&self) -> usize {
49 self.cols
50 }
51}
52
53impl VmpPMat {
54 pub fn as_ptr(&self) -> *const u8 {
55 self.ptr
56 }
57
58 pub fn as_mut_ptr(&self) -> *mut u8 {
59 self.ptr
60 }
61
62 pub fn borrowed(&self) -> bool{
63 self.data.len() == 0
64 }
65
66 /// Returns a non-mutable reference of `T` of the entire contiguous array of the [VmpPMat].
67 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
68 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
69 /// The length of the returned array is rows * cols * n.
70 pub fn raw<T>(&self) -> &[T] {
71 let ptr: *const T = self.ptr as *const T;
72 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
73 unsafe { &std::slice::from_raw_parts(ptr, len) }
74 }
75
76 /// Returns a non-mutable reference of `T` of the entire contiguous array of the [VmpPMat].
77 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
78 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
79 /// The length of the returned array is rows * cols * n.
80 pub fn raw_mut<T>(&self) -> &mut [T] {
81 let ptr: *mut T = self.ptr as *mut T;
82 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
83 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
84 }
85
86 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
87 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
88 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
89 ///
90 /// # Arguments
91 ///
92 /// * `row`: row index (i).
93 /// * `col`: col index (j).
94 pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
95 let mut res: Vec<T> = alloc_aligned(self.n);
96
97 if self.n < 8 {
98 res.copy_from_slice(
99 &self.raw::<T>()[(row + col * self.rows()) * self.n()
100 ..(row + col * self.rows()) * (self.n() + 1)],
101 );
102 } else {
103 (0..self.n >> 3).for_each(|blk| {
104 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.at_block(row, col, blk)[..8]);
105 });
106 }
107
108 res
109 }
110
111 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
112 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
113 fn at_block<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
114 let nrows: usize = self.rows();
115 let ncols: usize = self.cols();
116 if col == (ncols - 1) && (ncols & 1 == 1) {
117 &self.raw::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
118 } else {
119 &self.raw::<T>()[blk * nrows * ncols * 8
120 + (col / 2) * (2 * nrows) * 8
121 + row * 2 * 8
122 + (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 {
130 fn bytes_of_vmp_pmat(&self, rows: usize, cols: 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 /// * `cols`: number of cols (number of cols of each [VecZnxDft]).
138 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
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 /// * `cols`: number of cols of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
146 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: 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, a: &[i64], buf: &mut [u8]);
158
159 /// Prepares a [VmpPMat] from a vector of [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 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
166 ///
167 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
168 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]);
169
170 /// Prepares the ith-row of [VmpPMat] from a vector of [VecZnx].
171 ///
172 /// # Arguments
173 ///
174 /// * `b`: [VmpPMat] on which the values are encoded.
175 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
176 /// * `row_i`: the index of the row to prepare.
177 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
178 ///
179 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
180 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]);
181
182 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft].
183 ///
184 /// # Arguments
185 ///
186 /// * `c_cols`: number of cols of the output [VecZnxDft].
187 /// * `a_cols`: number of cols of the input [VecZnx].
188 /// * `rows`: number of rows of the input [VmpPMat].
189 /// * `cols`: number of cols of the input [VmpPMat].
190 fn vmp_apply_dft_tmp_bytes(
191 &self,
192 c_cols: usize,
193 a_cols: usize,
194 rows: usize,
195 cols: usize,
196 ) -> usize;
197
198 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
199 ///
200 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
201 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
202 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
203 ///
204 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
205 /// `j` cols, the output is a [VecZnx] of `j` cols.
206 ///
207 /// If there is a mismatch between the dimensions the largest valid ones are used.
208 ///
209 /// ```text
210 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
211 /// |h i j|
212 /// |k l m|
213 /// ```
214 /// where each element is a [VecZnxDft].
215 ///
216 /// # Arguments
217 ///
218 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
219 /// * `a`: the left operand [VecZnx] of the vector matrix product.
220 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
221 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
222 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
223
224 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
225 ///
226 /// # Arguments
227 ///
228 /// * `c_cols`: number of cols of the output [VecZnxDft].
229 /// * `a_cols`: number of cols of the input [VecZnxDft].
230 /// * `rows`: number of rows of the input [VmpPMat].
231 /// * `cols`: number of cols of the input [VmpPMat].
232 fn vmp_apply_dft_to_dft_tmp_bytes(
233 &self,
234 c_cols: usize,
235 a_cols: usize,
236 rows: usize,
237 cols: usize,
238 ) -> usize;
239
240 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
241 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
242 ///
243 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
244 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
245 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
246 ///
247 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
248 /// `j` cols, the output is a [VecZnx] of `j` cols.
249 ///
250 /// If there is a mismatch between the dimensions the largest valid ones are used.
251 ///
252 /// ```text
253 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
254 /// |h i j|
255 /// |k l m|
256 /// ```
257 /// where each element is a [VecZnxDft].
258 ///
259 /// # Arguments
260 ///
261 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
262 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
263 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
264 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
265 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
266
267 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
268 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
269 ///
270 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
271 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
272 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
273 ///
274 /// As such, given an input [VecZnx] of `i` cols and a [VmpPMat] of `i` rows and
275 /// `j` cols, the output is a [VecZnx] of `j` cols.
276 ///
277 /// If there is a mismatch between the dimensions the largest valid ones are used.
278 ///
279 /// ```text
280 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
281 /// |h i j|
282 /// |k l m|
283 /// ```
284 /// where each element is a [VecZnxDft].
285 ///
286 /// # Arguments
287 ///
288 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
289 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
290 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
291 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
292}
293
294impl VmpPMatOps for Module {
295 fn bytes_of_vmp_pmat(&self, rows: usize, cols: usize) -> usize {
296 unsafe { vmp::bytes_of_vmp_pmat(self.ptr, rows as u64, cols as u64) as usize }
297 }
298
299 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
300 let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vmp_pmat(rows, cols));
301 let ptr: *mut u8 = data.as_mut_ptr();
302 VmpPMat {
303 data: data,
304 ptr: ptr,
305 n: self.n(),
306 cols: cols,
307 rows: rows,
308 backend: self.backend(),
309 }
310 }
311
312 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
313 unsafe { vmp::vmp_prepare_tmp_bytes(self.ptr, rows as u64, cols as u64) as usize }
314 }
315
316 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], tmp_bytes: &mut [u8]) {
317 debug_assert_eq!(a.len(), b.n * b.rows * b.cols);
318 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
319 #[cfg(debug_assertions)]
320 {
321 assert_alignement(tmp_bytes.as_ptr());
322 }
323 unsafe {
324 vmp::vmp_prepare_contiguous(
325 self.ptr,
326 b.as_mut_ptr() as *mut vmp_pmat_t,
327 a.as_ptr(),
328 b.rows() as u64,
329 b.cols() as u64,
330 tmp_bytes.as_mut_ptr(),
331 );
332 }
333 }
334
335 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], tmp_bytes: &mut [u8]) {
336 let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
337 #[cfg(debug_assertions)]
338 {
339 debug_assert_eq!(a.len(), b.rows);
340 a.iter().for_each(|ai| {
341 debug_assert_eq!(ai.len(), b.n * b.cols);
342 });
343 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
344 assert_alignement(tmp_bytes.as_ptr());
345 }
346 unsafe {
347 vmp::vmp_prepare_dblptr(
348 self.ptr,
349 b.as_mut_ptr() as *mut vmp_pmat_t,
350 ptrs.as_ptr(),
351 b.rows() as u64,
352 b.cols() as u64,
353 tmp_bytes.as_mut_ptr(),
354 );
355 }
356 }
357
358 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]) {
359 debug_assert_eq!(a.len(), b.cols() * self.n());
360 debug_assert!(tmp_bytes.len() >= self.vmp_prepare_tmp_bytes(b.rows(), b.cols()));
361 #[cfg(debug_assertions)]
362 {
363 assert_alignement(tmp_bytes.as_ptr());
364 }
365 unsafe {
366 vmp::vmp_prepare_row(
367 self.ptr,
368 b.as_mut_ptr() as *mut vmp_pmat_t,
369 a.as_ptr(),
370 row_i as u64,
371 b.rows() as u64,
372 b.cols() as u64,
373 tmp_bytes.as_mut_ptr(),
374 );
375 }
376 }
377
378 fn vmp_apply_dft_tmp_bytes(
379 &self,
380 res_cols: usize,
381 a_cols: usize,
382 gct_rows: usize,
383 gct_cols: usize,
384 ) -> usize {
385 unsafe {
386 vmp::vmp_apply_dft_tmp_bytes(
387 self.ptr,
388 res_cols as u64,
389 a_cols as u64,
390 gct_rows as u64,
391 gct_cols as u64,
392 ) as usize
393 }
394 }
395
396 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, tmp_bytes: &mut [u8]) {
397 debug_assert!(
398 tmp_bytes.len() >= self.vmp_apply_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
399 );
400 #[cfg(debug_assertions)]
401 {
402 assert_alignement(tmp_bytes.as_ptr());
403 }
404 unsafe {
405 vmp::vmp_apply_dft(
406 self.ptr,
407 c.ptr as *mut vec_znx_dft_t,
408 c.cols() as u64,
409 a.as_ptr(),
410 a.cols() as u64,
411 a.n() as u64,
412 b.as_ptr() as *const vmp_pmat_t,
413 b.rows() as u64,
414 b.cols() as u64,
415 tmp_bytes.as_mut_ptr(),
416 )
417 }
418 }
419
420 fn vmp_apply_dft_to_dft_tmp_bytes(
421 &self,
422 res_cols: usize,
423 a_cols: usize,
424 gct_rows: usize,
425 gct_cols: usize,
426 ) -> usize {
427 unsafe {
428 vmp::vmp_apply_dft_to_dft_tmp_bytes(
429 self.ptr,
430 res_cols as u64,
431 a_cols as u64,
432 gct_rows as u64,
433 gct_cols as u64,
434 ) as usize
435 }
436 }
437
438 fn vmp_apply_dft_to_dft(
439 &self,
440 c: &mut VecZnxDft,
441 a: &VecZnxDft,
442 b: &VmpPMat,
443 tmp_bytes: &mut [u8],
444 ) {
445 debug_assert!(
446 tmp_bytes.len()
447 >= self.vmp_apply_dft_to_dft_tmp_bytes(c.cols(), a.cols(), b.rows(), b.cols())
448 );
449 #[cfg(debug_assertions)]
450 {
451 assert_alignement(tmp_bytes.as_ptr());
452 }
453 unsafe {
454 vmp::vmp_apply_dft_to_dft(
455 self.ptr,
456 c.ptr as *mut vec_znx_dft_t,
457 c.cols() as u64,
458 a.ptr as *const vec_znx_dft_t,
459 a.cols() as u64,
460 b.as_ptr() as *const vmp_pmat_t,
461 b.rows() as u64,
462 b.cols() as u64,
463 tmp_bytes.as_mut_ptr(),
464 )
465 }
466 }
467
468 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, tmp_bytes: &mut [u8]) {
469 debug_assert!(
470 tmp_bytes.len()
471 >= self.vmp_apply_dft_to_dft_tmp_bytes(b.cols(), b.cols(), a.rows(), a.cols())
472 );
473 #[cfg(debug_assertions)]
474 {
475 assert_alignement(tmp_bytes.as_ptr());
476 }
477 unsafe {
478 vmp::vmp_apply_dft_to_dft(
479 self.ptr,
480 b.ptr as *mut vec_znx_dft_t,
481 b.cols() as u64,
482 b.ptr as *mut vec_znx_dft_t,
483 b.cols() as u64,
484 a.as_ptr() as *const vmp_pmat_t,
485 a.rows() as u64,
486 a.cols() as u64,
487 tmp_bytes.as_mut_ptr(),
488 )
489 }
490 }
491}