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