1use poulpy_hal::{
  2    layouts::{
  3        Backend, Data, Module, Scratch, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef,
  4        VecZnxToRef,
  5    },
  6    oep::{
  7        VecZnxDftAddImpl, VecZnxDftAddInplaceImpl, VecZnxDftAllocBytesImpl, VecZnxDftAllocImpl, VecZnxDftApplyImpl,
  8        VecZnxDftCopyImpl, VecZnxDftFromBytesImpl, VecZnxDftSubImpl, VecZnxDftSubInplaceImpl, VecZnxDftSubNegateInplaceImpl,
  9        VecZnxDftZeroImpl, VecZnxIdftApplyConsumeImpl, VecZnxIdftApplyImpl, VecZnxIdftApplyTmpAImpl, VecZnxIdftApplyTmpBytesImpl,
 10    },
 11    reference::fft64::vec_znx_dft::{
 12        vec_znx_dft_add, vec_znx_dft_add_inplace, vec_znx_dft_apply, vec_znx_dft_copy, vec_znx_dft_sub, vec_znx_dft_sub_inplace,
 13        vec_znx_dft_sub_negate_inplace, vec_znx_dft_zero, vec_znx_idft_apply, vec_znx_idft_apply_consume,
 14        vec_znx_idft_apply_tmpa,
 15    },
 16};
 17
 18use crate::cpu_fft64_avx::{FFT64Avx, module::FFT64ModuleHandle};
 19
 20unsafe impl VecZnxDftFromBytesImpl<Self> for FFT64Avx {
 21    fn vec_znx_dft_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<Self> {
 22        VecZnxDft::<Vec<u8>, Self>::from_bytes(n, cols, size, bytes)
 23    }
 24}
 25
 26unsafe impl VecZnxDftAllocBytesImpl<Self> for FFT64Avx {
 27    fn vec_znx_dft_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
 28        Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Avx as Backend>::ScalarPrep>()
 29    }
 30}
 31
 32unsafe impl VecZnxDftAllocImpl<Self> for FFT64Avx {
 33    fn vec_znx_dft_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxDftOwned<Self> {
 34        VecZnxDftOwned::alloc(n, cols, size)
 35    }
 36}
 37
 38unsafe impl VecZnxIdftApplyTmpBytesImpl<Self> for FFT64Avx {
 39    fn vec_znx_idft_apply_tmp_bytes_impl(_module: &Module<Self>) -> usize {
 40        0
 41    }
 42}
 43
 44unsafe impl VecZnxIdftApplyImpl<Self> for FFT64Avx {
 45    fn vec_znx_idft_apply_impl<R, A>(
 46        module: &Module<Self>,
 47        res: &mut R,
 48        res_col: usize,
 49        a: &A,
 50        a_col: usize,
 51        _scratch: &mut Scratch<Self>,
 52    ) where
 53        R: VecZnxBigToMut<Self>,
 54        A: VecZnxDftToRef<Self>,
 55    {
 56        vec_znx_idft_apply(module.get_ifft_table(), res, res_col, a, a_col);
 57    }
 58}
 59
 60unsafe impl VecZnxIdftApplyTmpAImpl<Self> for FFT64Avx {
 61    fn vec_znx_idft_apply_tmpa_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
 62    where
 63        R: VecZnxBigToMut<Self>,
 64        A: VecZnxDftToMut<Self>,
 65    {
 66        vec_znx_idft_apply_tmpa(module.get_ifft_table(), res, res_col, a, a_col);
 67    }
 68}
 69
 70unsafe impl VecZnxIdftApplyConsumeImpl<Self> for FFT64Avx {
 71    fn vec_znx_idft_apply_consume_impl<D: Data>(module: &Module<Self>, res: VecZnxDft<D, FFT64Avx>) -> VecZnxBig<D, FFT64Avx>
 72    where
 73        VecZnxDft<D, FFT64Avx>: VecZnxDftToMut<Self>,
 74    {
 75        vec_znx_idft_apply_consume(module.get_ifft_table(), res)
 76    }
 77}
 78
 79unsafe impl VecZnxDftApplyImpl<Self> for FFT64Avx {
 80    fn vec_znx_dft_apply_impl<R, A>(
 81        module: &Module<Self>,
 82        step: usize,
 83        offset: usize,
 84        res: &mut R,
 85        res_col: usize,
 86        a: &A,
 87        a_col: usize,
 88    ) where
 89        R: VecZnxDftToMut<Self>,
 90        A: VecZnxToRef,
 91    {
 92        vec_znx_dft_apply(module.get_fft_table(), step, offset, res, res_col, a, a_col);
 93    }
 94}
 95
 96unsafe impl VecZnxDftAddImpl<Self> for FFT64Avx {
 97    fn vec_znx_dft_add_impl<R, A, B>(
 98        _module: &Module<Self>,
 99        res: &mut R,
100        res_col: usize,
101        a: &A,
102        a_col: usize,
103        b: &B,
104        b_col: usize,
105    ) where
106        R: VecZnxDftToMut<Self>,
107        A: VecZnxDftToRef<Self>,
108        B: VecZnxDftToRef<Self>,
109    {
110        vec_znx_dft_add(res, res_col, a, a_col, b, b_col);
111    }
112}
113
114unsafe impl VecZnxDftAddInplaceImpl<Self> for FFT64Avx {
115    fn vec_znx_dft_add_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
116    where
117        R: VecZnxDftToMut<Self>,
118        A: VecZnxDftToRef<Self>,
119    {
120        vec_znx_dft_add_inplace(res, res_col, a, a_col);
121    }
122}
123
124unsafe impl VecZnxDftSubImpl<Self> for FFT64Avx {
125    fn vec_znx_dft_sub_impl<R, A, B>(
126        _module: &Module<Self>,
127        res: &mut R,
128        res_col: usize,
129        a: &A,
130        a_col: usize,
131        b: &B,
132        b_col: usize,
133    ) where
134        R: VecZnxDftToMut<Self>,
135        A: VecZnxDftToRef<Self>,
136        B: VecZnxDftToRef<Self>,
137    {
138        vec_znx_dft_sub(res, res_col, a, a_col, b, b_col);
139    }
140}
141
142unsafe impl VecZnxDftSubInplaceImpl<Self> for FFT64Avx {
143    fn vec_znx_dft_sub_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
144    where
145        R: VecZnxDftToMut<Self>,
146        A: VecZnxDftToRef<Self>,
147    {
148        vec_znx_dft_sub_inplace(res, res_col, a, a_col);
149    }
150}
151
152unsafe impl VecZnxDftSubNegateInplaceImpl<Self> for FFT64Avx {
153    fn vec_znx_dft_sub_negate_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
154    where
155        R: VecZnxDftToMut<Self>,
156        A: VecZnxDftToRef<Self>,
157    {
158        vec_znx_dft_sub_negate_inplace(res, res_col, a, a_col);
159    }
160}
161
162unsafe impl VecZnxDftCopyImpl<Self> for FFT64Avx {
163    fn vec_znx_dft_copy_impl<R, A>(
164        _module: &Module<Self>,
165        step: usize,
166        offset: usize,
167        res: &mut R,
168        res_col: usize,
169        a: &A,
170        a_col: usize,
171    ) where
172        R: VecZnxDftToMut<Self>,
173        A: VecZnxDftToRef<Self>,
174    {
175        vec_znx_dft_copy(step, offset, res, res_col, a, a_col);
176    }
177}
178
179unsafe impl VecZnxDftZeroImpl<Self> for FFT64Avx {
180    fn vec_znx_dft_zero_impl<R>(_module: &Module<Self>, res: &mut R)
181    where
182        R: VecZnxDftToMut<Self>,
183    {
184        vec_znx_dft_zero(res);
185    }
186}