poulpy/poulpy-hal/src/reference/fft64/vmp.rs
poulpy-hal/src/reference/fft64/vmp.rsRaw
1use crate::{
2 cast_mut,
3 layouts::{DataViewMut, 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_bytes_of_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_zero<R, BE: Backend>(res: &mut R)
161where
162 R: VmpPMatToMut<BE>,
163{
164 res.to_mut().data_mut().fill(0);
165}
166
167pub fn vmp_apply_dft_to_dft<R, A, M, BE>(res: &mut R, a: &A, pmat: &M, tmp_bytes: &mut [f64])
168where
169 BE: Backend<ScalarPrep = f64>
170 + ReimZero
171 + Reim4Extract1Blk
172 + Reim4Mat1ColProd
173 + Reim4Mat2Cols2ndColProd
174 + Reim4Mat2ColsProd
175 + Reim4Save2Blks
176 + Reim4Save1Blk,
177 R: VecZnxDftToMut<BE>,
178 A: VecZnxDftToRef<BE>,
179 M: VmpPMatToRef<BE>,
180{
181 use crate::layouts::{ZnxView, ZnxViewMut};
182
183 let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
184 let a: VecZnxDft<&[u8], BE> = a.to_ref();
185 let pmat: VmpPMat<&[u8], BE> = pmat.to_ref();
186
187 #[cfg(debug_assertions)]
188 {
189 assert_eq!(res.n(), pmat.n());
190 assert_eq!(a.n(), pmat.n());
191 assert_eq!(res.cols(), pmat.cols_out());
192 assert_eq!(a.cols(), pmat.cols_in());
193 }
194
195 let n: usize = res.n();
196 let nrows: usize = pmat.cols_in() * pmat.rows();
197 let ncols: usize = pmat.cols_out() * pmat.size();
198
199 let pmat_raw: &[f64] = pmat.raw();
200 let a_raw: &[f64] = a.raw();
201 let res_raw: &mut [f64] = res.raw_mut();
202
203 vmp_apply_dft_to_dft_core::<true, BE>(n, res_raw, a_raw, pmat_raw, 0, nrows, ncols, tmp_bytes)
204}
205
206pub 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])
207where
208 BE: Backend<ScalarPrep = f64>
209 + ReimZero
210 + Reim4Extract1Blk
211 + Reim4Mat1ColProd
212 + Reim4Mat2Cols2ndColProd
213 + Reim4Mat2ColsProd
214 + Reim4Save2Blks
215 + Reim4Save1Blk,
216 R: VecZnxDftToMut<BE>,
217 A: VecZnxDftToRef<BE>,
218 M: VmpPMatToRef<BE>,
219{
220 use crate::layouts::{ZnxView, ZnxViewMut};
221
222 let mut res: VecZnxDft<&mut [u8], BE> = res.to_mut();
223 let a: VecZnxDft<&[u8], BE> = a.to_ref();
224 let pmat: VmpPMat<&[u8], BE> = pmat.to_ref();
225
226 #[cfg(debug_assertions)]
227 {
228 assert_eq!(res.n(), pmat.n());
229 assert_eq!(a.n(), pmat.n());
230 assert_eq!(res.cols(), pmat.cols_out());
231 assert_eq!(a.cols(), pmat.cols_in());
232 }
233
234 let n: usize = res.n();
235 let nrows: usize = pmat.cols_in() * pmat.rows();
236 let ncols: usize = pmat.cols_out() * pmat.size();
237
238 let pmat_raw: &[f64] = pmat.raw();
239 let a_raw: &[f64] = a.raw();
240 let res_raw: &mut [f64] = res.raw_mut();
241
242 vmp_apply_dft_to_dft_core::<false, BE>(
243 n,
244 res_raw,
245 a_raw,
246 pmat_raw,
247 limb_offset,
248 nrows,
249 ncols,
250 tmp_bytes,
251 )
252}
253
254#[allow(clippy::too_many_arguments)]
255fn vmp_apply_dft_to_dft_core<const OVERWRITE: bool, REIM>(
256 n: usize,
257 res: &mut [f64],
258 a: &[f64],
259 pmat: &[f64],
260 limb_offset: usize,
261 nrows: usize,
262 ncols: usize,
263 tmp_bytes: &mut [f64],
264) where
265 REIM: ReimZero
266 + Reim4Extract1Blk
267 + Reim4Mat1ColProd
268 + Reim4Mat2Cols2ndColProd
269 + Reim4Mat2ColsProd
270 + Reim4Save2Blks
271 + Reim4Save1Blk,
272{
273 #[cfg(debug_assertions)]
274 {
275 assert!(n >= 8);
276 assert!(n.is_power_of_two());
277 assert_eq!(pmat.len(), n * nrows * ncols);
278 assert!(res.len() & (n - 1) == 0);
279 assert!(a.len() & (n - 1) == 0);
280 }
281
282 let a_size: usize = a.len() / n;
283 let res_size: usize = res.len() / n;
284
285 let m: usize = n >> 1;
286
287 let (mat2cols_output, extracted_blk) = tmp_bytes.split_at_mut(16);
288
289 let row_max: usize = nrows.min(a_size);
290 let col_max: usize = ncols.min(res_size);
291
292 if limb_offset >= col_max {
293 if OVERWRITE {
294 REIM::reim_zero(res);
295 }
296 return;
297 }
298
299 for blk_i in 0..(m >> 2) {
300 let mat_blk_start: &[f64] = &pmat[blk_i * (8 * nrows * ncols)..];
301
302 REIM::reim4_extract_1blk(m, row_max, blk_i, extracted_blk, a);
303
304 if limb_offset.is_multiple_of(2) {
305 for (col_res, col_pmat) in (0..).step_by(2).zip((limb_offset..col_max - 1).step_by(2)) {
306 let col_offset: usize = col_pmat * (8 * nrows);
307 REIM::reim4_mat2cols_prod(
308 row_max,
309 mat2cols_output,
310 extracted_blk,
311 &mat_blk_start[col_offset..],
312 );
313 REIM::reim4_save_2blks::<OVERWRITE>(m, blk_i, &mut res[col_res * n..], mat2cols_output);
314 }
315 } else {
316 let col_offset: usize = (limb_offset - 1) * (8 * nrows);
317 REIM::reim4_mat2cols_2ndcol_prod(
318 row_max,
319 mat2cols_output,
320 extracted_blk,
321 &mat_blk_start[col_offset..],
322 );
323
324 REIM::reim4_save_1blk::<OVERWRITE>(m, blk_i, res, mat2cols_output);
325
326 for (col_res, col_pmat) in (1..)
327 .step_by(2)
328 .zip((limb_offset + 1..col_max - 1).step_by(2))
329 {
330 let col_offset: usize = col_pmat * (8 * nrows);
331 REIM::reim4_mat2cols_prod(
332 row_max,
333 mat2cols_output,
334 extracted_blk,
335 &mat_blk_start[col_offset..],
336 );
337 REIM::reim4_save_2blks::<OVERWRITE>(m, blk_i, &mut res[col_res * n..], mat2cols_output);
338 }
339 }
340
341 if !col_max.is_multiple_of(2) {
342 let last_col: usize = col_max - 1;
343 let col_offset: usize = last_col * (8 * nrows);
344
345 if last_col >= limb_offset {
346 if ncols == col_max {
347 REIM::reim4_mat1col_prod(
348 row_max,
349 mat2cols_output,
350 extracted_blk,
351 &mat_blk_start[col_offset..],
352 );
353 } else {
354 REIM::reim4_mat2cols_prod(
355 row_max,
356 mat2cols_output,
357 extracted_blk,
358 &mat_blk_start[col_offset..],
359 );
360 }
361 REIM::reim4_save_1blk::<OVERWRITE>(
362 m,
363 blk_i,
364 &mut res[(last_col - limb_offset) * n..],
365 mat2cols_output,
366 );
367 }
368 }
369 }
370
371 REIM::reim_zero(&mut res[col_max * n..]);
372}