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