1use crate::{
  2    hal::{
  3        api::{ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut},
  4        layouts::{
  5            Data, DataRef, Module, ScalarZnxToRef, SvpPPol, SvpPPolBytesOf, SvpPPolOwned, SvpPPolToMut, SvpPPolToRef, VecZnxDft,
  6            VecZnxDftToMut, VecZnxDftToRef,
  7        },
  8        oep::{SvpApplyImpl, SvpApplyInplaceImpl, SvpPPolAllocBytesImpl, SvpPPolAllocImpl, SvpPPolFromBytesImpl, SvpPrepareImpl},
  9    },
 10    implementation::cpu_spqlios::{
 11        ffi::{svp, vec_znx_dft::vec_znx_dft_t},
 12        module_fft64::FFT64,
 13    },
 14};
 15
 16const SVP_PPOL_FFT64_WORD_SIZE: usize = 1;
 17
 18impl<D: Data> SvpPPolBytesOf for SvpPPol<D, FFT64> {
 19    fn bytes_of(n: usize, cols: usize) -> usize {
 20        SVP_PPOL_FFT64_WORD_SIZE * n * cols * size_of::<f64>()
 21    }
 22}
 23
 24impl<D: Data> ZnxSliceSize for SvpPPol<D, FFT64> {
 25    fn sl(&self) -> usize {
 26        SVP_PPOL_FFT64_WORD_SIZE * self.n()
 27    }
 28}
 29
 30impl<D: DataRef> ZnxView for SvpPPol<D, FFT64> {
 31    type Scalar = f64;
 32}
 33
 34unsafe impl SvpPPolFromBytesImpl<Self> for FFT64 {
 35    fn svp_ppol_from_bytes_impl(n: usize, cols: usize, bytes: Vec<u8>) -> SvpPPolOwned<Self> {
 36        SvpPPolOwned::from_bytes(n, cols, bytes)
 37    }
 38}
 39
 40unsafe impl SvpPPolAllocImpl<Self> for FFT64 {
 41    fn svp_ppol_alloc_impl(n: usize, cols: usize) -> SvpPPolOwned<Self> {
 42        SvpPPolOwned::alloc(n, cols)
 43    }
 44}
 45
 46unsafe impl SvpPPolAllocBytesImpl<Self> for FFT64 {
 47    fn svp_ppol_alloc_bytes_impl(n: usize, cols: usize) -> usize {
 48        SvpPPol::<Vec<u8>, Self>::bytes_of(n, cols)
 49    }
 50}
 51
 52unsafe impl SvpPrepareImpl<Self> for FFT64 {
 53    fn svp_prepare_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
 54    where
 55        R: SvpPPolToMut<Self>,
 56        A: ScalarZnxToRef,
 57    {
 58        unsafe {
 59            svp::svp_prepare(
 60                module.ptr(),
 61                res.to_mut().at_mut_ptr(res_col, 0) as *mut svp::svp_ppol_t,
 62                a.to_ref().at_ptr(a_col, 0),
 63            )
 64        }
 65    }
 66}
 67
 68unsafe impl SvpApplyImpl<Self> for FFT64 {
 69    fn svp_apply_impl<R, A, B>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 70    where
 71        R: VecZnxDftToMut<Self>,
 72        A: SvpPPolToRef<Self>,
 73        B: VecZnxDftToRef<Self>,
 74    {
 75        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
 76        let a: SvpPPol<&[u8], Self> = a.to_ref();
 77        let b: VecZnxDft<&[u8], Self> = b.to_ref();
 78        unsafe {
 79            svp::svp_apply_dft_to_dft(
 80                module.ptr(),
 81                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
 82                res.size() as u64,
 83                res.cols() as u64,
 84                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
 85                b.at_ptr(b_col, 0) as *const vec_znx_dft_t,
 86                b.size() as u64,
 87                b.cols() as u64,
 88            )
 89        }
 90    }
 91}
 92
 93unsafe impl SvpApplyInplaceImpl for FFT64 {
 94    fn svp_apply_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
 95    where
 96        R: VecZnxDftToMut<Self>,
 97        A: SvpPPolToRef<Self>,
 98    {
 99        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
100        let a: SvpPPol<&[u8], Self> = a.to_ref();
101        unsafe {
102            svp::svp_apply_dft_to_dft(
103                module.ptr(),
104                res.at_mut_ptr(res_col, 0) as *mut vec_znx_dft_t,
105                res.size() as u64,
106                res.cols() as u64,
107                a.at_ptr(a_col, 0) as *const svp::svp_ppol_t,
108                res.at_ptr(res_col, 0) as *const vec_znx_dft_t,
109                res.size() as u64,
110                res.cols() as u64,
111            )
112        }
113    }
114}