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