1use crate::{
  2    layouts::{
  3        Backend, CnvPVecL, CnvPVecLToMut, CnvPVecLToRef, CnvPVecR, CnvPVecRToMut, CnvPVecRToRef, VecZnx, VecZnxBig,
  4        VecZnxBigToMut, VecZnxDft, VecZnxDftToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut, ZnxZero,
  5    },
  6    reference::fft64::{
  7        reim::{ReimAdd, ReimCopy, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimZero},
  8        reim4::{
  9            Reim4Convolution, Reim4Convolution1Coeff, Reim4Convolution2Coeffs, Reim4Extract1BlkContiguous,
 10            Reim4Save1BlkContiguous,
 11        },
 12        vec_znx_dft::vec_znx_dft_apply,
 13    },
 14};
 15
 16pub fn convolution_prepare_left<R, A, T, BE>(table: &ReimFFTTable<f64>, res: &mut R, a: &A, tmp: &mut T)
 17where
 18    BE: Backend<ScalarPrep = f64>
 19        + ReimZero
 20        + Reim4Extract1BlkContiguous
 21        + ReimDFTExecute<ReimFFTTable<f64>, f64>
 22        + ReimFromZnx
 23        + ReimZero,
 24    R: CnvPVecLToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
 25    A: VecZnxToRef,
 26    T: VecZnxDftToMut<BE>,
 27{
 28    convolution_prepare(table, res, a, tmp)
 29}
 30
 31pub fn convolution_prepare_right<R, A, T, BE>(table: &ReimFFTTable<f64>, res: &mut R, a: &A, tmp: &mut T)
 32where
 33    BE: Backend<ScalarPrep = f64>
 34        + ReimZero
 35        + Reim4Extract1BlkContiguous
 36        + ReimDFTExecute<ReimFFTTable<f64>, f64>
 37        + ReimFromZnx
 38        + ReimZero,
 39    R: CnvPVecRToMut<BE> + ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
 40    A: VecZnxToRef,
 41    T: VecZnxDftToMut<BE>,
 42{
 43    convolution_prepare(table, res, a, tmp)
 44}
 45
 46fn convolution_prepare<R, A, T, BE>(table: &ReimFFTTable<f64>, res: &mut R, a: &A, tmp: &mut T)
 47where
 48    BE: Backend<ScalarPrep = f64>
 49        + ReimZero
 50        + Reim4Extract1BlkContiguous
 51        + ReimDFTExecute<ReimFFTTable<f64>, f64>
 52        + ReimFromZnx
 53        + ReimZero,
 54    R: ZnxInfos + ZnxViewMut<Scalar = BE::ScalarPrep>,
 55    A: VecZnxToRef,
 56    T: VecZnxDftToMut<BE>,
 57{
 58    let a: &VecZnx<&[u8]> = &a.to_ref();
 59    let tmp: &mut VecZnxDft<&mut [u8], BE> = &mut tmp.to_mut();
 60
 61    let cols: usize = res.cols();
 62    assert_eq!(a.cols(), cols, "a.cols():{} != res.cols():{cols}", a.cols());
 63
 64    let res_size: usize = res.size();
 65    let min_size: usize = res_size.min(a.size());
 66
 67    let m: usize = a.n() >> 1;
 68
 69    let n: usize = table.m() << 1;
 70
 71    let res_raw: &mut [f64] = res.raw_mut();
 72
 73    for i in 0..cols {
 74        vec_znx_dft_apply(table, 1, 0, tmp, 0, a, i);
 75
 76        let tmp_raw: &[f64] = tmp.raw();
 77        let res_col: &mut [f64] = &mut res_raw[i * n * res_size..];
 78
 79        for blk_i in 0..m / 4 {
 80            BE::reim4_extract_1blk_contiguous(m, min_size, blk_i, &mut res_col[blk_i * res_size * 8..], tmp_raw);
 81            BE::reim_zero(&mut res_col[blk_i * res_size * 8 + min_size * 8..(blk_i + 1) * res_size * 8]);
 82        }
 83    }
 84}
 85
 86pub fn convolution_by_const_apply_tmp_bytes(res_size: usize, a_size: usize, b_size: usize) -> usize {
 87    let min_size: usize = res_size.min(a_size + b_size - 1);
 88    size_of::<i64>() * (min_size + a_size) * 8
 89}
 90
 91pub fn convolution_by_const_apply<R, A, BE>(
 92    res: &mut R,
 93    res_offset: usize,
 94    res_col: usize,
 95    a: &A,
 96    a_col: usize,
 97    b: &[i64],
 98    tmp: &mut [i64],
 99) where
100    BE: Backend<ScalarBig = i64>
101        + I64ConvolutionByConst1Coeff
102        + I64ConvolutionByConst2Coeffs
103        + I64Extract1BlkContiguous
104        + I64Save1BlkContiguous,
105    R: VecZnxBigToMut<BE>,
106    A: VecZnxToRef,
107{
108    let res: &mut VecZnxBig<&mut [u8], BE> = &mut res.to_mut();
109    let a: &VecZnx<&[u8]> = &a.to_ref();
110
111    let n: usize = res.n();
112    assert_eq!(a.n(), n);
113
114    let res_size: usize = res.size();
115    let a_size: usize = a.size();
116    let b_size: usize = b.len();
117
118    let bound: usize = a_size + b_size - 1;
119    let min_size: usize = res_size.min(bound);
120    let offset: usize = res_offset.min(bound);
121
122    let a_sl: usize = n * a.cols();
123    let res_sl: usize = n * res.cols();
124
125    let res_raw: &mut [i64] = res.raw_mut();
126    let a_raw: &[i64] = a.raw();
127
128    let a_idx: usize = n * a_col;
129    let res_idx: usize = n * res_col;
130
131    let (res_blk, a_blk) = tmp[..(min_size + a_size) * 8].split_at_mut(min_size * 8);
132
133    for blk_i in 0..n / 8 {
134        BE::i64_extract_1blk_contiguous(a_sl, a_idx, a_size, blk_i, a_blk, a_raw);
135        BE::i64_convolution_by_const(res_blk, min_size, offset, a_blk, a_size, b);
136        BE::i64_save_1blk_contiguous(res_sl, res_idx, min_size, blk_i, res_raw, res_blk);
137    }
138
139    for j in min_size..res_size {
140        res.zero_at(res_col, j);
141    }
142}
143
144pub fn convolution_apply_dft_tmp_bytes(res_size: usize, a_size: usize, b_size: usize) -> usize {
145    let min_size: usize = res_size.min(a_size + b_size - 1);
146    size_of::<f64>() * 8 * min_size
147}
148
149#[allow(clippy::too_many_arguments)]
150pub fn convolution_apply_dft<R, A, B, BE>(
151    res: &mut R,
152    res_offset: usize,
153    res_col: usize,
154    a: &A,
155    a_col: usize,
156    b: &B,
157    b_col: usize,
158    tmp: &mut [f64],
159) where
160    BE: Backend<ScalarPrep = f64> + Reim4Save1BlkContiguous + Reim4Convolution1Coeff + Reim4Convolution2Coeffs,
161    R: VecZnxDftToMut<BE>,
162    A: CnvPVecLToRef<BE>,
163    B: CnvPVecRToRef<BE>,
164{
165    let res: &mut VecZnxDft<&mut [u8], BE> = &mut res.to_mut();
166    let a: &CnvPVecL<&[u8], BE> = &a.to_ref();
167    let b: &CnvPVecR<&[u8], BE> = &b.to_ref();
168
169    let n: usize = res.n();
170    assert_eq!(a.n(), n);
171    assert_eq!(b.n(), n);
172    let m: usize = n >> 1;
173
174    let res_size: usize = res.size();
175    let a_size: usize = a.size();
176    let b_size: usize = b.size();
177
178    let bound: usize = a_size + b_size - 1;
179    let min_size: usize = res_size.min(bound);
180    let offset: usize = res_offset.min(bound);
181
182    let dst: &mut [f64] = res.raw_mut();
183    let a_raw: &[f64] = a.raw();
184    let b_raw: &[f64] = b.raw();
185
186    let mut a_idx: usize = a_col * n * a_size;
187    let mut b_idx: usize = b_col * n * b_size;
188    let a_offset: usize = a_size * 8;
189    let b_offset: usize = b_size * 8;
190    for blk_i in 0..m / 4 {
191        BE::reim4_convolution(tmp, min_size, offset, &a_raw[a_idx..], a_size, &b_raw[b_idx..], b_size);
192        BE::reim4_save_1blk_contiguous(m, min_size, blk_i, dst, tmp);
193        a_idx += a_offset;
194        b_idx += b_offset;
195    }
196
197    for j in min_size..res_size {
198        res.zero_at(res_col, j);
199    }
200}
201
202pub fn convolution_pairwise_apply_dft_tmp_bytes(res_size: usize, a_size: usize, b_size: usize) -> usize {
203    convolution_apply_dft_tmp_bytes(res_size, a_size, b_size) + (a_size + b_size) * size_of::<f64>() * 8
204}
205
206#[allow(clippy::too_many_arguments)]
207pub fn convolution_pairwise_apply_dft<R, A, B, BE>(
208    res: &mut R,
209    res_offset: usize,
210    res_col: usize,
211    a: &A,
212    b: &B,
213    col_i: usize,
214    col_j: usize,
215    tmp: &mut [f64],
216) where
217    BE: Backend<ScalarPrep = f64>
218        + ReimZero
219        + ReimAdd
220        + ReimCopy
221        + Reim4Save1BlkContiguous
222        + Reim4Convolution1Coeff
223        + Reim4Convolution2Coeffs,
224    R: VecZnxDftToMut<BE>,
225    A: CnvPVecLToRef<BE>,
226    B: CnvPVecRToRef<BE>,
227{
228    if col_i == col_j {
229        convolution_apply_dft(res, res_offset, res_col, a, col_i, b, col_j, tmp);
230        return;
231    }
232
233    let res: &mut VecZnxDft<&mut [u8], BE> = &mut res.to_mut();
234    let a: &CnvPVecL<&[u8], BE> = &a.to_ref();
235    let b: &CnvPVecR<&[u8], BE> = &b.to_ref();
236
237    let n: usize = res.n();
238    let m: usize = n >> 1;
239
240    assert_eq!(a.n(), n);
241    assert_eq!(b.n(), n);
242
243    let res_size: usize = res.size();
244    let a_size: usize = a.size();
245    let b_size: usize = b.size();
246
247    assert_eq!(
248        tmp.len(),
249        convolution_pairwise_apply_dft_tmp_bytes(res_size, a_size, b_size) / size_of::<f64>()
250    );
251
252    let bound: usize = a_size + b_size - 1;
253    let min_size: usize = res_size.min(bound);
254    let offset: usize = res_offset.min(bound);
255
256    let res_raw: &mut [f64] = res.raw_mut();
257    let a_raw: &[f64] = a.raw();
258    let b_raw: &[f64] = b.raw();
259
260    let a_row_size: usize = a_size * 8;
261    let b_row_size: usize = b_size * 8;
262
263    let mut a0_idx: usize = col_i * n * a_size;
264    let mut a1_idx: usize = col_j * n * a_size;
265    let mut b0_idx: usize = col_i * n * b_size;
266    let mut b1_idx: usize = col_j * n * b_size;
267
268    let (tmp_a, tmp) = tmp.split_at_mut(a_row_size);
269    let (tmp_b, tmp_res) = tmp.split_at_mut(b_row_size);
270
271    for blk_i in 0..m / 4 {
272        let a0: &[f64] = &a_raw[a0_idx..];
273        let a1: &[f64] = &a_raw[a1_idx..];
274        let b0: &[f64] = &b_raw[b0_idx..];
275        let b1: &[f64] = &b_raw[b1_idx..];
276
277        BE::reim_add(tmp_a, &a0[..a_row_size], &a1[..a_row_size]);
278        BE::reim_add(tmp_b, &b0[..b_row_size], &b1[..b_row_size]);
279
280        BE::reim4_convolution(tmp_res, min_size, offset, tmp_a, a_size, tmp_b, b_size);
281        BE::reim4_save_1blk_contiguous(m, min_size, blk_i, res_raw, tmp_res);
282
283        a0_idx += a_row_size;
284        a1_idx += a_row_size;
285        b0_idx += b_row_size;
286        b1_idx += b_row_size;
287    }
288
289    for j in min_size..res_size {
290        res.zero_at(res_col, j);
291    }
292}
293
294pub trait I64Extract1BlkContiguous {
295    fn i64_extract_1blk_contiguous(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]);
296}
297
298pub trait I64Save1BlkContiguous {
299    fn i64_save_1blk_contiguous(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]);
300}
301
302#[inline(always)]
303pub fn i64_extract_1blk_contiguous_ref(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
304    debug_assert!(blk < (n >> 3));
305    debug_assert!(dst.len() >= rows * 8, "dst.len(): {} < rows*8: {}", dst.len(), 8 * rows);
306
307    let offset: usize = offset + (blk << 3);
308
309    // src = 8-values chunks spaced by n, dst = sequential 8-values chunks
310    let src_rows = src.chunks_exact(n).take(rows);
311    let dst_chunks = dst.chunks_exact_mut(8).take(rows);
312
313    for (dst_chunk, src_row) in dst_chunks.zip(src_rows) {
314        dst_chunk.copy_from_slice(&src_row[offset..offset + 8]);
315    }
316}
317
318#[inline(always)]
319pub fn i64_save_1blk_contiguous_ref(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
320    debug_assert!(blk < (n >> 3));
321    debug_assert!(src.len() >= rows * 8);
322
323    let offset: usize = offset + (blk << 3);
324
325    // dst = 4-values chunks spaced by m, src = sequential 4-values chunks
326    let dst_rows = dst.chunks_exact_mut(n).take(rows);
327    let src_chunks = src.chunks_exact(8).take(rows);
328
329    for (dst_row, src_chunk) in dst_rows.zip(src_chunks) {
330        dst_row[offset..offset + 8].copy_from_slice(src_chunk);
331    }
332}
333
334pub trait I64ConvolutionByConst1Coeff {
335    fn i64_convolution_by_const_1coeff(k: usize, dst: &mut [i64; 8], a: &[i64], a_size: usize, b: &[i64]);
336}
337
338#[inline(always)]
339pub fn i64_convolution_by_const_1coeff_ref(k: usize, dst: &mut [i64; 8], a: &[i64], a_size: usize, b: &[i64]) {
340    dst.fill(0);
341
342    let b_size: usize = b.len();
343
344    if k >= a_size + b_size {
345        return;
346    }
347    let j_min: usize = k.saturating_sub(a_size - 1);
348    let j_max: usize = (k + 1).min(b_size);
349
350    for j in j_min..j_max {
351        let ai: &[i64] = &a[8 * (k - j)..];
352        let bi: i64 = b[j];
353
354        dst[0] = dst[0].wrapping_add(ai[0].wrapping_mul(bi));
355        dst[1] = dst[1].wrapping_add(ai[1].wrapping_mul(bi));
356        dst[2] = dst[2].wrapping_add(ai[2].wrapping_mul(bi));
357        dst[3] = dst[3].wrapping_add(ai[3].wrapping_mul(bi));
358        dst[4] = dst[4].wrapping_add(ai[4].wrapping_mul(bi));
359        dst[5] = dst[5].wrapping_add(ai[5].wrapping_mul(bi));
360        dst[6] = dst[6].wrapping_add(ai[6].wrapping_mul(bi));
361        dst[7] = dst[7].wrapping_add(ai[7].wrapping_mul(bi));
362    }
363}
364
365#[inline(always)]
366pub(crate) fn as_arr_i64<const size: usize>(x: &[i64]) -> &[i64; size] {
367    debug_assert!(x.len() >= size, "x.len():{} < size:{}", x.len(), size);
368    unsafe { &*(x.as_ptr() as *const [i64; size]) }
369}
370
371#[inline(always)]
372pub(crate) fn as_arr_i64_mut<const size: usize>(x: &mut [i64]) -> &mut [i64; size] {
373    debug_assert!(x.len() >= size, "x.len():{} < size:{}", x.len(), size);
374    unsafe { &mut *(x.as_mut_ptr() as *mut [i64; size]) }
375}
376
377pub trait I64ConvolutionByConst2Coeffs {
378    fn i64_convolution_by_const_2coeffs(k: usize, dst: &mut [i64; 16], a: &[i64], a_size: usize, b: &[i64]);
379}
380
381#[inline(always)]
382pub fn i64_convolution_by_const_2coeffs_ref(k: usize, dst: &mut [i64; 16], a: &[i64], a_size: usize, b: &[i64]) {
383    i64_convolution_by_const_1coeff_ref(k, as_arr_i64_mut(&mut dst[..8]), a, a_size, b);
384    i64_convolution_by_const_1coeff_ref(k + 1, as_arr_i64_mut(&mut dst[8..]), a, a_size, b);
385}
386
387impl<BE: Backend> I64ConvolutionByConst<BE> for BE where Self: I64ConvolutionByConst1Coeff + I64ConvolutionByConst2Coeffs {}
388
389pub trait I64ConvolutionByConst<BE: Backend>
390where
391    BE: I64ConvolutionByConst1Coeff + I64ConvolutionByConst2Coeffs,
392{
393    fn i64_convolution_by_const(dst: &mut [i64], dst_size: usize, offset: usize, a: &[i64], a_size: usize, b: &[i64]) {
394        assert!(a_size > 0);
395
396        for k in (0..dst_size - 1).step_by(2) {
397            BE::i64_convolution_by_const_2coeffs(k + offset, as_arr_i64_mut(&mut dst[8 * k..]), a, a_size, b);
398        }
399
400        if !dst_size.is_multiple_of(2) {
401            let k: usize = dst_size - 1;
402            BE::i64_convolution_by_const_1coeff(k + offset, as_arr_i64_mut(&mut dst[8 * k..]), a, a_size, b);
403        }
404    }
405}