1use poulpy_hal::{
  2    api::{TakeSlice, VmpApplyDftToDftTmpBytes, VmpPrepareTmpBytes},
  3    layouts::{
  4        Backend, MatZnx, MatZnxToRef, Module, Scratch, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef, VmpPMat, VmpPMatOwned,
  5        VmpPMatToMut, VmpPMatToRef, ZnxInfos, ZnxView, ZnxViewMut,
  6    },
  7    oep::{
  8        VmpApplyDftToDftAddImpl, VmpApplyDftToDftAddTmpBytesImpl, VmpApplyDftToDftImpl, VmpApplyDftToDftTmpBytesImpl,
  9        VmpPMatAllocBytesImpl, VmpPMatAllocImpl, VmpPMatFromBytesImpl, VmpPrepareImpl, VmpPrepareTmpBytesImpl, VmpZeroImpl,
 10    },
 11    reference::fft64::vmp::vmp_zero,
 12};
 13
 14use crate::cpu_spqlios::{
 15    FFT64Spqlios,
 16    ffi::{vec_znx_dft::vec_znx_dft_t, vmp},
 17};
 18
 19unsafe impl VmpPMatAllocBytesImpl<Self> for FFT64Spqlios {
 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 VmpPMatFromBytesImpl<Self> for FFT64Spqlios {
 26    fn vmp_pmat_from_bytes_impl(
 27        n: usize,
 28        rows: usize,
 29        cols_in: usize,
 30        cols_out: usize,
 31        size: usize,
 32        bytes: Vec<u8>,
 33    ) -> VmpPMatOwned<Self> {
 34        VmpPMatOwned::from_bytes(n, rows, cols_in, cols_out, size, bytes)
 35    }
 36}
 37
 38unsafe impl VmpPMatAllocImpl<Self> for FFT64Spqlios {
 39    fn vmp_pmat_alloc_impl(n: usize, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> VmpPMatOwned<Self> {
 40        VmpPMatOwned::alloc(n, rows, cols_in, cols_out, size)
 41    }
 42}
 43
 44unsafe impl VmpPrepareTmpBytesImpl<Self> for FFT64Spqlios {
 45    fn vmp_prepare_tmp_bytes_impl(module: &Module<Self>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 46        unsafe {
 47            vmp::vmp_prepare_tmp_bytes(
 48                module.ptr(),
 49                (rows * cols_in) as u64,
 50                (cols_out * size) as u64,
 51            ) as usize
 52        }
 53    }
 54}
 55
 56unsafe impl VmpPrepareImpl<Self> for FFT64Spqlios {
 57    fn vmp_prepare_impl<R, A>(module: &Module<Self>, res: &mut R, a: &A, scratch: &mut Scratch<Self>)
 58    where
 59        R: VmpPMatToMut<Self>,
 60        A: MatZnxToRef,
 61    {
 62        let mut res: VmpPMat<&mut [u8], Self> = res.to_mut();
 63        let a: MatZnx<&[u8]> = a.to_ref();
 64
 65        #[cfg(debug_assertions)]
 66        {
 67            assert_eq!(a.n(), res.n());
 68            assert_eq!(
 69                res.cols_in(),
 70                a.cols_in(),
 71                "res.cols_in: {} != a.cols_in: {}",
 72                res.cols_in(),
 73                a.cols_in()
 74            );
 75            assert_eq!(
 76                res.rows(),
 77                a.rows(),
 78                "res.rows: {} != a.rows: {}",
 79                res.rows(),
 80                a.rows()
 81            );
 82            assert_eq!(
 83                res.cols_out(),
 84                a.cols_out(),
 85                "res.cols_out: {} != a.cols_out: {}",
 86                res.cols_out(),
 87                a.cols_out()
 88            );
 89            assert_eq!(
 90                res.size(),
 91                a.size(),
 92                "res.size: {} != a.size: {}",
 93                res.size(),
 94                a.size()
 95            );
 96        }
 97
 98        let (tmp_bytes, _) = scratch.take_slice(module.vmp_prepare_tmp_bytes(a.rows(), a.cols_in(), a.cols_out(), a.size()));
 99
100        unsafe {
101            vmp::vmp_prepare_contiguous(
102                module.ptr(),
103                res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
104                a.as_ptr(),
105                (a.rows() * a.cols_in()) as u64,
106                (a.size() * a.cols_out()) as u64,
107                tmp_bytes.as_mut_ptr(),
108            );
109        }
110    }
111}
112
113unsafe impl VmpApplyDftToDftTmpBytesImpl<Self> for FFT64Spqlios {
114    fn vmp_apply_dft_to_dft_tmp_bytes_impl(
115        module: &Module<Self>,
116        res_size: usize,
117        a_size: usize,
118        b_rows: usize,
119        b_cols_in: usize,
120        b_cols_out: usize,
121        b_size: usize,
122    ) -> usize {
123        unsafe {
124            vmp::vmp_apply_dft_to_dft_tmp_bytes(
125                module.ptr(),
126                (res_size * b_cols_out) as u64,
127                (a_size * b_cols_in) as u64,
128                (b_rows * b_cols_in) as u64,
129                (b_size * b_cols_out) as u64,
130            ) as usize
131        }
132    }
133}
134
135unsafe impl VmpApplyDftToDftImpl<Self> for FFT64Spqlios {
136    fn vmp_apply_dft_to_dft_impl<R, A, C>(module: &Module<Self>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<Self>)
137    where
138        R: VecZnxDftToMut<Self>,
139        A: VecZnxDftToRef<Self>,
140        C: VmpPMatToRef<Self>,
141    {
142        let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
143        let a: VecZnxDft<&[u8], _> = a.to_ref();
144        let b: VmpPMat<&[u8], _> = b.to_ref();
145
146        #[cfg(debug_assertions)]
147        {
148            assert_eq!(b.n(), res.n());
149            assert_eq!(a.n(), res.n());
150            assert_eq!(
151                res.cols(),
152                b.cols_out(),
153                "res.cols(): {} != b.cols_out: {}",
154                res.cols(),
155                b.cols_out()
156            );
157            assert_eq!(
158                a.cols(),
159                b.cols_in(),
160                "a.cols(): {} != b.cols_in: {}",
161                a.cols(),
162                b.cols_in()
163            );
164        }
165
166        let (tmp_bytes, _) = scratch.take_slice(module.vmp_apply_dft_to_dft_tmp_bytes(
167            res.size(),
168            a.size(),
169            b.rows(),
170            b.cols_in(),
171            b.cols_out(),
172            b.size(),
173        ));
174        unsafe {
175            vmp::vmp_apply_dft_to_dft(
176                module.ptr(),
177                res.as_mut_ptr() as *mut vec_znx_dft_t,
178                (res.size() * res.cols()) as u64,
179                a.as_ptr() as *const vec_znx_dft_t,
180                (a.size() * a.cols()) as u64,
181                b.as_ptr() as *const vmp::vmp_pmat_t,
182                (b.rows() * b.cols_in()) as u64,
183                (b.size() * b.cols_out()) as u64,
184                tmp_bytes.as_mut_ptr(),
185            )
186        }
187    }
188}
189
190unsafe impl VmpApplyDftToDftAddTmpBytesImpl<Self> for FFT64Spqlios {
191    fn vmp_apply_dft_to_dft_add_tmp_bytes_impl(
192        module: &Module<Self>,
193        res_size: usize,
194        a_size: usize,
195        b_rows: usize,
196        b_cols_in: usize,
197        b_cols_out: usize,
198        b_size: usize,
199    ) -> usize {
200        unsafe {
201            vmp::vmp_apply_dft_to_dft_tmp_bytes(
202                module.ptr(),
203                (res_size * b_cols_out) as u64,
204                (a_size * b_cols_in) as u64,
205                (b_rows * b_cols_in) as u64,
206                (b_size * b_cols_out) as u64,
207            ) as usize
208        }
209    }
210}
211
212unsafe impl VmpApplyDftToDftAddImpl<Self> for FFT64Spqlios {
213    fn vmp_apply_dft_to_dft_add_impl<R, A, C>(
214        module: &Module<Self>,
215        res: &mut R,
216        a: &A,
217        b: &C,
218        scale: usize,
219        scratch: &mut Scratch<Self>,
220    ) where
221        R: VecZnxDftToMut<Self>,
222        A: VecZnxDftToRef<Self>,
223        C: VmpPMatToRef<Self>,
224    {
225        let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
226        let a: VecZnxDft<&[u8], _> = a.to_ref();
227        let b: VmpPMat<&[u8], _> = b.to_ref();
228
229        #[cfg(debug_assertions)]
230        {
231            assert_eq!(b.n(), res.n());
232            assert_eq!(a.n(), res.n());
233            assert_eq!(
234                res.cols(),
235                b.cols_out(),
236                "res.cols(): {} != b.cols_out: {}",
237                res.cols(),
238                b.cols_out()
239            );
240            assert_eq!(
241                a.cols(),
242                b.cols_in(),
243                "a.cols(): {} != b.cols_in: {}",
244                a.cols(),
245                b.cols_in()
246            );
247        }
248
249        let (tmp_bytes, _) = scratch.take_slice(module.vmp_apply_dft_to_dft_tmp_bytes(
250            res.size(),
251            a.size(),
252            b.rows(),
253            b.cols_in(),
254            b.cols_out(),
255            b.size(),
256        ));
257        unsafe {
258            vmp::vmp_apply_dft_to_dft_add(
259                module.ptr(),
260                res.as_mut_ptr() as *mut vec_znx_dft_t,
261                (res.size() * res.cols()) as u64,
262                a.as_ptr() as *const vec_znx_dft_t,
263                (a.size() * a.cols()) as u64,
264                b.as_ptr() as *const vmp::vmp_pmat_t,
265                (b.rows() * b.cols_in()) as u64,
266                (b.size() * b.cols_out()) as u64,
267                (scale * b.cols_out()) as u64,
268                tmp_bytes.as_mut_ptr(),
269            )
270        }
271    }
272}
273
274unsafe impl VmpZeroImpl<Self> for FFT64Spqlios {
275    fn vmp_zero_impl<R>(_module: &Module<Self>, res: &mut R)
276    where
277        R: VmpPMatToMut<Self>,
278    {
279        vmp_zero(res);
280    }
281}