1use poulpy_hal::{
  2    api::{TakeSlice, VmpPrepareTmpBytes},
  3    layouts::{
  4        Backend, MatZnx, MatZnxToRef, Module, Scratch, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, VmpPMat, VmpPMatOwned,
  5        VmpPMatToMut, VmpPMatToRef, ZnxInfos,
  6    },
  7    oep::{
  8        VmpApplyDftToDftAddImpl, VmpApplyDftToDftAddTmpBytesImpl, VmpApplyDftToDftImpl, VmpApplyDftToDftTmpBytesImpl,
  9        VmpPMatAllocBytesImpl, VmpPMatAllocImpl, VmpPrepareImpl, VmpPrepareTmpBytesImpl, VmpZeroImpl,
 10    },
 11    reference::fft64::vmp::{
 12        vmp_apply_dft_to_dft, vmp_apply_dft_to_dft_add, vmp_apply_dft_to_dft_tmp_bytes, vmp_prepare, vmp_prepare_tmp_bytes,
 13        vmp_zero,
 14    },
 15};
 16
 17use crate::cpu_fft64_ref::{FFT64Ref, module::FFT64ModuleHandle};
 18
 19unsafe impl VmpPMatAllocBytesImpl<Self> for FFT64Ref {
 20    fn vmp_pmat_bytes_of_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 21        Self::layout_prep_word_count() * n * rows * cols_in * cols_out * size * size_of::<f64>()
 22    }
 23}
 24
 25unsafe impl VmpPMatAllocImpl<Self> for FFT64Ref {
 26    fn vmp_pmat_alloc_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<Self> {
 27        VmpPMatOwned::alloc(n, rows, cols_in, cols_out, size)
 28    }
 29}
 30
 31unsafe impl VmpApplyDftToDftImpl<Self> for FFT64Ref
 32where
 33    Scratch<Self>: TakeSlice,
 34    FFT64Ref: VmpApplyDftToDftTmpBytesImpl<Self>,
 35{
 36    fn vmp_apply_dft_to_dft_impl<R, A, C>(module: &Module<Self>, res: &mut R, a: &A, pmat: &C, scratch: &mut Scratch<Self>)
 37    where
 38        R: VecZnxDftToMut<Self>,
 39        A: VecZnxDftToRef<Self>,
 40        C: VmpPMatToRef<Self>,
 41    {
 42        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
 43        let a: VecZnxDft<&[u8], Self> = a.to_ref();
 44        let pmat: VmpPMat<&[u8], Self> = pmat.to_ref();
 45
 46        let (tmp, _) = scratch.take_slice(
 47            Self::vmp_apply_dft_to_dft_tmp_bytes_impl(
 48                module,
 49                res.size(),
 50                a.size(),
 51                pmat.rows(),
 52                pmat.cols_in(),
 53                pmat.cols_out(),
 54                pmat.size(),
 55            ) / size_of::<f64>(),
 56        );
 57        vmp_apply_dft_to_dft(&mut res, &a, &pmat, tmp);
 58    }
 59}
 60
 61unsafe impl VmpApplyDftToDftAddImpl<Self> for FFT64Ref
 62where
 63    Scratch<Self>: TakeSlice,
 64    FFT64Ref: VmpApplyDftToDftTmpBytesImpl<Self>,
 65{
 66    fn vmp_apply_dft_to_dft_add_impl<R, A, C>(
 67        module: &Module<Self>,
 68        res: &mut R,
 69        a: &A,
 70        pmat: &C,
 71        limb_offset: usize,
 72        scratch: &mut Scratch<Self>,
 73    ) where
 74        R: VecZnxDftToMut<Self>,
 75        A: VecZnxDftToRef<Self>,
 76        C: VmpPMatToRef<Self>,
 77    {
 78        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
 79        let a: VecZnxDft<&[u8], Self> = a.to_ref();
 80        let pmat: VmpPMat<&[u8], Self> = pmat.to_ref();
 81
 82        let (tmp, _) = scratch.take_slice(
 83            Self::vmp_apply_dft_to_dft_tmp_bytes_impl(
 84                module,
 85                res.size(),
 86                a.size(),
 87                pmat.rows(),
 88                pmat.cols_in(),
 89                pmat.cols_out(),
 90                pmat.size(),
 91            ) / size_of::<f64>(),
 92        );
 93        vmp_apply_dft_to_dft_add(&mut res, &a, &pmat, limb_offset * pmat.cols_out(), tmp);
 94    }
 95}
 96
 97unsafe impl VmpPrepareTmpBytesImpl<Self> for FFT64Ref {
 98    fn vmp_prepare_tmp_bytes_impl(module: &Module<Self>, _rows: usize, _cols_in: usize, _cols_out: usize, _size: usize) -> usize {
 99        vmp_prepare_tmp_bytes(module.n())
100    }
101}
102
103unsafe impl VmpPrepareImpl<Self> for FFT64Ref {
104    fn vmp_prepare_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
105    where
106        R: VmpPMatToMut<Self>,
107        A: MatZnxToRef,
108    {
109        {}
110        let mut res: VmpPMat<&mut [u8], Self> = res.to_mut();
111        let a: MatZnx<&[u8]> = a.to_ref();
112        let (tmp, _) =
113            scratch.take_slice(module.vmp_prepare_tmp_bytes(a.rows(), a.cols_in(), a.cols_out(), a.size()) / size_of::<f64>());
114        vmp_prepare(module.get_fft_table(), &mut res, &a, tmp);
115    }
116}
117
118unsafe impl VmpApplyDftToDftTmpBytesImpl<Self> for FFT64Ref {
119    fn vmp_apply_dft_to_dft_tmp_bytes_impl(
120        _module: &Module<Self>,
121        _res_size: usize,
122        a_size: usize,
123        b_rows: usize,
124        b_cols_in: usize,
125        _b_cols_out: usize,
126        _b_size: usize,
127    ) -> usize {
128        vmp_apply_dft_to_dft_tmp_bytes(a_size, b_rows, b_cols_in)
129    }
130}
131
132unsafe impl VmpApplyDftToDftAddTmpBytesImpl<Self> for FFT64Ref {
133    fn vmp_apply_dft_to_dft_add_tmp_bytes_impl(
134        _module: &Module<Self>,
135        _res_size: usize,
136        a_size: usize,
137        b_rows: usize,
138        b_cols_in: usize,
139        _b_cols_out: usize,
140        _b_size: usize,
141    ) -> usize {
142        vmp_apply_dft_to_dft_tmp_bytes(a_size, b_rows, b_cols_in)
143    }
144}
145
146unsafe impl VmpZeroImpl<Self> for FFT64Ref {
147    fn vmp_zero_impl<R>(_module: &Module<Self>, res: &mut R)
148    where
149        R: VmpPMatToMut<Self>,
150    {
151        vmp_zero(res);
152    }
153}