1use poulpy_hal::{
  2    api::{Convolution, ModuleN, ScratchTakeBasic, TakeSlice, VecZnxDftApply, VecZnxDftBytesOf},
  3    layouts::{
  4        Backend, CnvPVecL, CnvPVecLToMut, CnvPVecLToRef, CnvPVecR, CnvPVecRToMut, CnvPVecRToRef, Module, Scratch, VecZnx,
  5        VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftToMut, VecZnxToRef, ZnxInfos,
  6    },
  7    oep::{CnvPVecBytesOfImpl, CnvPVecLAllocImpl, ConvolutionImpl},
  8    reference::fft64::convolution::{
  9        convolution_apply_dft, convolution_apply_dft_tmp_bytes, convolution_by_const_apply, convolution_by_const_apply_tmp_bytes,
 10        convolution_pairwise_apply_dft, convolution_pairwise_apply_dft_tmp_bytes, convolution_prepare_left,
 11        convolution_prepare_right,
 12    },
 13};
 14
 15use crate::{FFT64Ref, module::FFT64ModuleHandle};
 16
 17unsafe impl CnvPVecLAllocImpl<Self> for FFT64Ref {
 18    fn cnv_pvec_left_alloc_impl(n: usize, cols: usize, size: usize) -> CnvPVecL<Vec<u8>, Self> {
 19        CnvPVecL::alloc(n, cols, size)
 20    }
 21
 22    fn cnv_pvec_right_alloc_impl(n: usize, cols: usize, size: usize) -> CnvPVecR<Vec<u8>, Self> {
 23        CnvPVecR::alloc(n, cols, size)
 24    }
 25}
 26
 27unsafe impl CnvPVecBytesOfImpl for FFT64Ref {
 28    fn bytes_of_cnv_pvec_left_impl(n: usize, cols: usize, size: usize) -> usize {
 29        Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Ref as Backend>::ScalarPrep>()
 30    }
 31
 32    fn bytes_of_cnv_pvec_right_impl(n: usize, cols: usize, size: usize) -> usize {
 33        Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Ref as Backend>::ScalarPrep>()
 34    }
 35}
 36
 37unsafe impl ConvolutionImpl<Self> for FFT64Ref
 38where
 39    Module<Self>: ModuleN + VecZnxDftBytesOf + VecZnxDftApply<Self>,
 40{
 41    fn cnv_prepare_left_tmp_bytes_impl(module: &Module<Self>, res_size: usize, a_size: usize) -> usize {
 42        module.bytes_of_vec_znx_dft(1, res_size.min(a_size))
 43    }
 44
 45    fn cnv_prepare_left_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
 46    where
 47        R: CnvPVecLToMut<Self>,
 48        A: VecZnxToRef,
 49    {
 50        let res: &mut CnvPVecL<&mut [u8], FFT64Ref> = &mut res.to_mut();
 51        let a: &VecZnx<&[u8]> = &a.to_ref();
 52        let (mut tmp, _) = scratch.take_vec_znx_dft(module, 1, res.size().min(a.size()));
 53        convolution_prepare_left(module.get_fft_table(), res, a, &mut tmp);
 54    }
 55
 56    fn cnv_prepare_right_tmp_bytes_impl(module: &Module<Self>, res_size: usize, a_size: usize) -> usize {
 57        module.bytes_of_vec_znx_dft(1, res_size.min(a_size))
 58    }
 59
 60    fn cnv_prepare_right_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
 61    where
 62        R: CnvPVecRToMut<Self>,
 63        A: VecZnxToRef,
 64    {
 65        let res: &mut CnvPVecR<&mut [u8], FFT64Ref> = &mut res.to_mut();
 66        let a: &VecZnx<&[u8]> = &a.to_ref();
 67        let (mut tmp, _) = scratch.take_vec_znx_dft(module, 1, res.size().min(a.size()));
 68        convolution_prepare_right(module.get_fft_table(), res, a, &mut tmp);
 69    }
 70
 71    fn cnv_apply_dft_tmp_bytes_impl(
 72        _module: &Module<Self>,
 73        res_size: usize,
 74        _res_offset: usize,
 75        a_size: usize,
 76        b_size: usize,
 77    ) -> usize {
 78        convolution_apply_dft_tmp_bytes(res_size, a_size, b_size)
 79    }
 80
 81    fn cnv_by_const_apply_tmp_bytes_impl(
 82        _module: &Module<Self>,
 83        res_size: usize,
 84        _res_offset: usize,
 85        a_size: usize,
 86        b_size: usize,
 87    ) -> usize {
 88        convolution_by_const_apply_tmp_bytes(res_size, a_size, b_size)
 89    }
 90
 91    fn cnv_by_const_apply_impl<R, A>(
 92        module: &Module<Self>,
 93        res: &mut R,
 94        res_offset: usize,
 95        res_col: usize,
 96        a: &A,
 97        a_col: usize,
 98        b: &[i64],
 99        scratch: &mut Scratch<Self>,
100    ) where
101        R: VecZnxBigToMut<Self>,
102        A: VecZnxToRef,
103    {
104        let res: &mut VecZnxBig<&mut [u8], Self> = &mut res.to_mut();
105        let a: &VecZnx<&[u8]> = &a.to_ref();
106        let (tmp, _) =
107            scratch.take_slice(module.cnv_by_const_apply_tmp_bytes(res.size(), res_offset, a.size(), b.len()) / size_of::<i64>());
108        convolution_by_const_apply(res, res_offset, res_col, a, a_col, b, tmp);
109    }
110
111    fn cnv_apply_dft_impl<R, A, B>(
112        module: &Module<Self>,
113        res: &mut R,
114        res_offset: usize,
115        res_col: usize,
116        a: &A,
117        a_col: usize,
118        b: &B,
119        b_col: usize,
120        scratch: &mut Scratch<Self>,
121    ) where
122        R: VecZnxDftToMut<Self>,
123        A: CnvPVecLToRef<Self>,
124        B: CnvPVecRToRef<Self>,
125    {
126        let res: &mut VecZnxDft<&mut [u8], FFT64Ref> = &mut res.to_mut();
127        let a: &CnvPVecL<&[u8], FFT64Ref> = &a.to_ref();
128        let b: &CnvPVecR<&[u8], FFT64Ref> = &b.to_ref();
129        let (tmp, _) =
130            scratch.take_slice(module.cnv_apply_dft_tmp_bytes(res.size(), res_offset, a.size(), b.size()) / size_of::<f64>());
131        convolution_apply_dft(res, res_offset, res_col, a, a_col, b, b_col, tmp);
132    }
133
134    fn cnv_pairwise_apply_dft_tmp_bytes(
135        _module: &Module<Self>,
136        res_size: usize,
137        _res_offset: usize,
138        a_size: usize,
139        b_size: usize,
140    ) -> usize {
141        convolution_pairwise_apply_dft_tmp_bytes(res_size, a_size, b_size)
142    }
143
144    fn cnv_pairwise_apply_dft_impl<R, A, B>(
145        module: &Module<Self>,
146        res: &mut R,
147        res_offset: usize,
148        res_col: usize,
149        a: &A,
150        b: &B,
151        col_0: usize,
152        col_1: usize,
153        scratch: &mut Scratch<Self>,
154    ) where
155        R: VecZnxDftToMut<Self>,
156        A: CnvPVecLToRef<Self>,
157        B: CnvPVecRToRef<Self>,
158    {
159        let res: &mut VecZnxDft<&mut [u8], FFT64Ref> = &mut res.to_mut();
160        let a: &CnvPVecL<&[u8], FFT64Ref> = &a.to_ref();
161        let b: &CnvPVecR<&[u8], FFT64Ref> = &b.to_ref();
162        let (tmp, _) = scratch
163            .take_slice(module.cnv_pairwise_apply_dft_tmp_bytes(res.size(), res_offset, a.size(), b.size()) / size_of::<f64>());
164        convolution_pairwise_apply_dft(res, res_offset, res_col, a, b, col_0, col_1, tmp);
165    }
166}