poulpy/poulpy-hal/src/reference/fft64/vmp.rs
poulpy-hal/src/reference/fft64/vmp.rsRaw
1use crate::{
2 cast_mut,
3 layouts::{MatZnx, MatZnxToRef, VecZnx, VecZnxToRef, VmpPMatToMut, ZnxView, ZnxViewMut},
4 oep::VecZnxDftAllocBytesImpl,
5 reference::fft64::{
6 reim::{ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimZero},
7 reim4::{Reim4Extract1Blk, Reim4Mat1ColProd, Reim4Mat2Cols2ndColProd, Reim4Mat2ColsProd, Reim4Save1Blk, Reim4Save2Blks},
8 vec_znx_dft::vec_znx_dft_apply,
9 },
10};
11
12use crate::layouts::{Backend, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, VmpPMat, VmpPMatToRef, ZnxInfos};
13
14pub fn vmp_prepare_tmp_bytes(n: usize) -> usize {
15 n * size_of::<i64>()
16}
17
18pub fn vmp_prepare<R, A, BE>(table: &ReimFFTTable<f64>, pmat: &mut R, mat: &A, tmp: &mut [f64])
19where
20 BE: Backend<ScalarPrep = f64> + ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimFromZnx + Reim4Extract1Blk,
21 R: VmpPMatToMut<BE>,
22 A: MatZnxToRef,
23{
24 let mut res: crate::layouts::VmpPMat<&mut [u8], BE> = pmat.to_mut();
25 let a: MatZnx<&[u8]> = mat.to_ref();
26
27 #[cfg(debug_assertions)]
28 {
29 assert_eq!(a.n(), res.n());
30 assert_eq!(
31 res.cols_in(),
32 a.cols_in(),
33 "res.cols_in: {} != a.cols_in: {}",
34 res.cols_in(),
35 a.cols_in()
36 );
37 assert_eq!(
38 res.rows(),
39 a.rows(),
40 "res.rows: {} != a.rows: {}",
41 res.rows(),
42 a.rows()
43 );
44 assert_eq!(
45 res.cols_out(),
46 a.cols_out(),
47 "res.cols_out: {} != a.cols_out: {}",
48 res.cols_out(),
49 a.cols_out()
50 );
51 assert_eq!(
52 res.size(),
53 a.size(),
54 "res.size: {} != a.size: {}",
55 res.size(),
56 a.size()
57 );
58 }
59
60 let nrows: usize = a.cols_in() * a.rows();
61 let ncols: usize = a.cols_out() * a.size();
62 vmp_prepare_core::<BE>(table, res.raw_mut(), a.raw(), nrows, ncols, tmp);
63}
64
65pub(crate) fn vmp_prepare_core<REIM>(
66 table: &ReimFFTTable<f64>,
67 pmat: &mut [f64],
68 mat: &[i64],
69 nrows: usize,
70 ncols: usize,
71 tmp: &mut [f64],
72) where
73 REIM: ReimDFTExecute<ReimFFTTable<f64>, f64> + ReimFromZnx + Reim4Extract1Blk,
74{
75 let m: usize = table.m();
76 let n: usize = m << 1;
77
78 #[cfg(debug_assertions)]
79 {
80 assert!(n >= 8);
81 assert_eq!(mat.len(), n * nrows * ncols);
82 assert_eq!(pmat.len(), n * nrows * ncols);
83 assert_eq!(tmp.len(), vmp_prepare_tmp_bytes(n) / size_of::<i64>())
84 }
85
86 let offset: usize = nrows * ncols * 8;
87
88 for row_i in 0..nrows {
89 for col_i in 0..ncols {
90 let pos: usize = n * (row_i * ncols + col_i);
91
92 REIM::reim_from_znx(tmp, &mat[pos..pos + n]);
93 REIM::reim_dft_execute(table, tmp);
94
95 let dst: &mut [f64] = if col_i == (ncols - 1) && !ncols.is_multiple_of(2) {
96 &mut pmat[col_i * nrows * 8 + row_i * 8..]
97 } else {
98 &mut pmat[(col_i / 2) * (nrows * 16) + row_i * 16 + (col_i % 2) * 8..]
99 };
100
101 for blk_i in 0..m >> 2 {
102 REIM::reim4_extract_1blk(m, 1, blk_i, &mut dst[blk_i * offset..], tmp);
103 }
104 }
105 }
106}
107
108pub fn vmp_apply_dft_tmp_bytes(n: usize, a_size: usize, prows: usize, pcols_in: usize) -> usize {
109 let row_max: usize = (a_size).min(prows);
110 (16 + (n + 8) * row_max * pcols_in) * size_of::<f64>()
111}
112
113pub fn vmp_apply_dft<R, A, M, BE>(table: &ReimFFTTable<f64>, res: &mut R, a: &A, pmat: &M, tmp_bytes: &mut [f64])
114where
115 BE: Backend<ScalarPrep = f64>
116 + VecZnxDftAllocBytesImpl<BE>
117 + ReimDFTExecute<ReimFFTTable<f64>, f64>
118 + ReimZero
119 + Reim4Extract1Blk
120 + Reim4Mat1ColProd
121 + Reim4Mat2Cols2ndColProd
122 + Reim4Mat2ColsProd
123 + Reim4Save2Blks
124 + Reim4Save1Blk
125 + ReimFromZnx,
126 R: VecZnxDftToMut<BE>,
127 A: VecZnxToRef,
128 M: VmpPMatToRef<BE>,
129{
130 let a: VecZnx<&[u8]> = a.to_ref();
131 let pmat: VmpPMat<&[u8], BE> = pmat.to_ref();
132
133 let n: usize = a.n();
134 let cols: usize = pmat.cols_in();
135 let size: usize = a.size().min(pmat.rows());
136
137 #[cfg(debug_assertions)]
138 {
139 assert!(tmp_bytes.len() >= vmp_apply_dft_tmp_bytes(n, size, pmat.rows(), cols));
140 assert!(a.cols() <= cols);
141 }
142
143 let (data, tmp_bytes) = tmp_bytes.split_at_mut(BE::vec_znx_dft_alloc_bytes_impl(n, cols, size));
144
145 let mut a_dft: VecZnxDft<&mut [u8], BE> = VecZnxDft::from_data(cast_mut(data), n, cols, size);
146
147 let offset: usize = cols - a.cols();
148 for j in 0..cols {
149 vec_znx_dft_apply(table, 1, 0, &mut a_dft, j, &a, offset + j);
150 }
151
152 vmp_apply_dft_to_dft(res, &a_dft, &pmat, tmp_bytes);
153}
154
155pub fn vmp_apply_dft_to_dft_tmp_bytes(a_size: usize, prows: usize, pcols_in: usize) -> usize {
156 let row_max: usize = (a_size).min(prows);
157 (16 + 8 * row_max * pcols_in) * size_of::<f64>()
158}
159
160pub fn vmp_apply_dft_to_dft<R, A, M, BE>(res: &mut R, a: &A, pmat: &M, tmp_bytes: &mut [f64])
161where
162 BE: Backend<ScalarPrep = f64>
163 + ReimZero
164 + Reim4Extract1Blk
165 + Reim4Mat1ColProd
166 + Reim4Mat2Cols2ndColProd
167 + Reim4Mat2ColsProd
168 + Reim4Save2Blks
169 + Reim4Save1Blk,
170 R: VecZnxDftToMut<BE>,
171 A: VecZnxDftToRef<BE>,
172 M: VmpPMatToRef<BE>,
173{
174 use crate::layouts::{ZnxView, ZnxViewMut};
175
176 let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
177 let a: VecZnxDft<&[u8], BE> = a.to_ref();
178 let pmat: VmpPMat<&[u8], BE> = pmat.to_ref();
179
180 #[cfg(debug_assertions)]
181 {
182 assert_eq!(res.n(), pmat.n());
183 assert_eq!(a.n(), pmat.n());
184 assert_eq!(res.cols(), pmat.cols_out());
185 assert_eq!(a.cols(), pmat.cols_in());
186 }
187
188 let n: usize = res.n();
189 let nrows: usize = pmat.cols_in() * pmat.rows();
190 let ncols: usize = pmat.cols_out() * pmat.size();
191
192 let pmat_raw: &[f64] = pmat.raw();
193 let a_raw: &[f64] = a.raw();
194 let res_raw: &mut [f64] = res.raw_mut();
195
196 vmp_apply_dft_to_dft_core::<true, BE>(n, res_raw, a_raw, pmat_raw, 0, nrows, ncols, tmp_bytes)
197}
198
199pub fn vmp_apply_dft_to_dft_add<R, A, M, BE>(res: &mut R, a: &A, pmat: &M, limb_offset: usize, tmp_bytes: &mut [f64])
200where
201 BE: Backend<ScalarPrep = f64>
202 + ReimZero
203 + Reim4Extract1Blk
204 + Reim4Mat1ColProd
205 + Reim4Mat2Cols2ndColProd
206 + Reim4Mat2ColsProd
207 + Reim4Save2Blks
208 + Reim4Save1Blk,
209 R: VecZnxDftToMut<BE>,
210 A: VecZnxDftToRef<BE>,
211 M: VmpPMatToRef<BE>,
212{
213 use crate::layouts::{ZnxView, ZnxViewMut};
214
215 let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
216 let a: VecZnxDft<&[u8], BE> = a.to_ref();
217 let pmat: VmpPMat<&[u8], BE> = pmat.to_ref();
218
219 #[cfg(debug_assertions)]
220 {
221 assert_eq!(res.n(), pmat.n());
222 assert_eq!(a.n(), pmat.n());
223 assert_eq!(res.cols(), pmat.cols_out());
224 assert_eq!(a.cols(), pmat.cols_in());
225 }
226
227 let n: usize = res.n();
228 let nrows: usize = pmat.cols_in() * pmat.rows();
229 let ncols: usize = pmat.cols_out() * pmat.size();
230
231 let pmat_raw: &[f64] = pmat.raw();
232 let a_raw: &[f64] = a.raw();
233 let res_raw: &mut [f64] = res.raw_mut();
234
235 vmp_apply_dft_to_dft_core::<false, BE>(
236 n,
237 res_raw,
238 a_raw,
239 pmat_raw,
240 limb_offset,
241 nrows,
242 ncols,
243 tmp_bytes,
244 )
245}
246
247#[allow(clippy::too_many_arguments)]
248fn vmp_apply_dft_to_dft_core<const OVERWRITE: bool, REIM>(
249 n: usize,
250 res: &mut [f64],
251 a: &[f64],
252 pmat: &[f64],
253 limb_offset: usize,
254 nrows: usize,
255 ncols: usize,
256 tmp_bytes: &mut [f64],
257) where
258 REIM: ReimZero
259 + Reim4Extract1Blk
260 + Reim4Mat1ColProd
261 + Reim4Mat2Cols2ndColProd
262 + Reim4Mat2ColsProd
263 + Reim4Save2Blks
264 + Reim4Save1Blk,
265{
266 #[cfg(debug_assertions)]
267 {
268 assert!(n >= 8);
269 assert!(n.is_power_of_two());
270 assert_eq!(pmat.len(), n * nrows * ncols);
271 assert!(res.len() & (n - 1) == 0);
272 assert!(a.len() & (n - 1) == 0);
273 }
274
275 let a_size: usize = a.len() / n;
276 let res_size: usize = res.len() / n;
277
278 let m: usize = n >> 1;
279
280 let (mat2cols_output, extracted_blk) = tmp_bytes.split_at_mut(16);
281
282 let row_max: usize = nrows.min(a_size);
283 let col_max: usize = ncols.min(res_size);
284
285 if limb_offset >= col_max {
286 if OVERWRITE {
287 REIM::reim_zero(res);
288 }
289 return;
290 }
291
292 for blk_i in 0..(m >> 2) {
293 let mat_blk_start: &[f64] = &pmat[blk_i * (8 * nrows * ncols)..];
294
295 REIM::reim4_extract_1blk(m, row_max, blk_i, extracted_blk, a);
296
297 if limb_offset.is_multiple_of(2) {
298 for (col_res, col_pmat) in (0..).step_by(2).zip((limb_offset..col_max - 1).step_by(2)) {
299 let col_offset: usize = col_pmat * (8 * nrows);
300 REIM::reim4_mat2cols_prod(
301 row_max,
302 mat2cols_output,
303 extracted_blk,
304 &mat_blk_start[col_offset..],
305 );
306 REIM::reim4_save_2blks::<OVERWRITE>(m, blk_i, &mut res[col_res * n..], mat2cols_output);
307 }
308 } else {
309 let col_offset: usize = (limb_offset - 1) * (8 * nrows);
310 REIM::reim4_mat2cols_2ndcol_prod(
311 row_max,
312 mat2cols_output,
313 extracted_blk,
314 &mat_blk_start[col_offset..],
315 );
316
317 REIM::reim4_save_1blk::<OVERWRITE>(m, blk_i, res, mat2cols_output);
318
319 for (col_res, col_pmat) in (1..)
320 .step_by(2)
321 .zip((limb_offset + 1..col_max - 1).step_by(2))
322 {
323 let col_offset: usize = col_pmat * (8 * nrows);
324 REIM::reim4_mat2cols_prod(
325 row_max,
326 mat2cols_output,
327 extracted_blk,
328 &mat_blk_start[col_offset..],
329 );
330 REIM::reim4_save_2blks::<OVERWRITE>(m, blk_i, &mut res[col_res * n..], mat2cols_output);
331 }
332 }
333
334 if !col_max.is_multiple_of(2) {
335 let last_col: usize = col_max - 1;
336 let col_offset: usize = last_col * (8 * nrows);
337
338 if last_col >= limb_offset {
339 if ncols == col_max {
340 REIM::reim4_mat1col_prod(
341 row_max,
342 mat2cols_output,
343 extracted_blk,
344 &mat_blk_start[col_offset..],
345 );
346 } else {
347 REIM::reim4_mat2cols_prod(
348 row_max,
349 mat2cols_output,
350 extracted_blk,
351 &mat_blk_start[col_offset..],
352 );
353 }
354 REIM::reim4_save_1blk::<OVERWRITE>(
355 m,
356 blk_i,
357 &mut res[(last_col - limb_offset) * n..],
358 mat2cols_output,
359 );
360 }
361 }
362 }
363
364 REIM::reim_zero(&mut res[col_max * n..]);
365}