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(module: &Module<FFT64>, rows: usize, cols_in: usize, cols_out: usize, size: usize) -> usize {
 45        unsafe {
 46            vmp::vmp_prepare_tmp_bytes(
 47                module.ptr(),
 48                (rows * cols_in) as u64,
 49                (cols_out * size) as u64,
 50            ) as usize
 51        }
 52    }
 53}
 54
 55unsafe impl VmpPMatPrepareImpl<FFT64> for FFT64 {
 56    fn vmp_prepare_impl<R, A>(module: &Module<FFT64>, res: &mut R, a: &A, scratch: &mut Scratch<FFT64>)
 57    where
 58        R: VmpPMatToMut<FFT64>,
 59        A: MatZnxToRef,
 60    {
 61        let mut res: VmpPMat<&mut [u8], FFT64> = res.to_mut();
 62        let a: MatZnx<&[u8]> = a.to_ref();
 63
 64        #[cfg(debug_assertions)]
 65        {
 66            assert_eq!(a.n(), res.n());
 67            assert_eq!(
 68                res.cols_in(),
 69                a.cols_in(),
 70                "res.cols_in: {} != a.cols_in: {}",
 71                res.cols_in(),
 72                a.cols_in()
 73            );
 74            assert_eq!(
 75                res.rows(),
 76                a.rows(),
 77                "res.rows: {} != a.rows: {}",
 78                res.rows(),
 79                a.rows()
 80            );
 81            assert_eq!(
 82                res.cols_out(),
 83                a.cols_out(),
 84                "res.cols_out: {} != a.cols_out: {}",
 85                res.cols_out(),
 86                a.cols_out()
 87            );
 88            assert_eq!(
 89                res.size(),
 90                a.size(),
 91                "res.size: {} != a.size: {}",
 92                res.size(),
 93                a.size()
 94            );
 95        }
 96
 97        let (tmp_bytes, _) = scratch.take_slice(module.vmp_prepare_tmp_bytes(a.rows(), a.cols_in(), a.cols_out(), a.size()));
 98
 99        unsafe {
100            vmp::vmp_prepare_contiguous(
101                module.ptr(),
102                res.as_mut_ptr() as *mut vmp::vmp_pmat_t,
103                a.as_ptr(),
104                (a.rows() * a.cols_in()) as u64,
105                (a.size() * a.cols_out()) as u64,
106                tmp_bytes.as_mut_ptr(),
107            );
108        }
109    }
110}
111
112unsafe impl VmpApplyTmpBytesImpl<FFT64> for FFT64 {
113    fn vmp_apply_tmp_bytes_impl(
114        module: &Module<FFT64>,
115        res_size: usize,
116        a_size: usize,
117        b_rows: usize,
118        b_cols_in: usize,
119        b_cols_out: usize,
120        b_size: usize,
121    ) -> usize {
122        unsafe {
123            vmp::vmp_apply_dft_to_dft_tmp_bytes(
124                module.ptr(),
125                (res_size * b_cols_out) as u64,
126                (a_size * b_cols_in) as u64,
127                (b_rows * b_cols_in) as u64,
128                (b_size * b_cols_out) as u64,
129            ) as usize
130        }
131    }
132}
133
134unsafe impl VmpApplyImpl<FFT64> for FFT64 {
135    fn vmp_apply_impl<R, A, C>(module: &Module<FFT64>, res: &mut R, a: &A, b: &C, scratch: &mut Scratch<FFT64>)
136    where
137        R: VecZnxDftToMut<FFT64>,
138        A: VecZnxDftToRef<FFT64>,
139        C: VmpPMatToRef<FFT64>,
140    {
141        let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
142        let a: VecZnxDft<&[u8], _> = a.to_ref();
143        let b: VmpPMat<&[u8], _> = b.to_ref();
144
145        #[cfg(debug_assertions)]
146        {
147            assert_eq!(b.n(), res.n());
148            assert_eq!(a.n(), res.n());
149            assert_eq!(
150                res.cols(),
151                b.cols_out(),
152                "res.cols(): {} != b.cols_out: {}",
153                res.cols(),
154                b.cols_out()
155            );
156            assert_eq!(
157                a.cols(),
158                b.cols_in(),
159                "a.cols(): {} != b.cols_in: {}",
160                a.cols(),
161                b.cols_in()
162            );
163        }
164
165        let (tmp_bytes, _) = scratch.take_slice(module.vmp_apply_tmp_bytes(
166            res.size(),
167            a.size(),
168            b.rows(),
169            b.cols_in(),
170            b.cols_out(),
171            b.size(),
172        ));
173        unsafe {
174            vmp::vmp_apply_dft_to_dft(
175                module.ptr(),
176                res.as_mut_ptr() as *mut vec_znx_dft_t,
177                (res.size() * res.cols()) as u64,
178                a.as_ptr() as *const vec_znx_dft_t,
179                (a.size() * a.cols()) as u64,
180                b.as_ptr() as *const vmp::vmp_pmat_t,
181                (b.rows() * b.cols_in()) as u64,
182                (b.size() * b.cols_out()) as u64,
183                tmp_bytes.as_mut_ptr(),
184            )
185        }
186    }
187}
188
189unsafe impl VmpApplyAddTmpBytesImpl<FFT64> for FFT64 {
190    fn vmp_apply_add_tmp_bytes_impl(
191        module: &Module<FFT64>,
192        res_size: usize,
193        a_size: usize,
194        b_rows: usize,
195        b_cols_in: usize,
196        b_cols_out: usize,
197        b_size: usize,
198    ) -> usize {
199        unsafe {
200            vmp::vmp_apply_dft_to_dft_tmp_bytes(
201                module.ptr(),
202                (res_size * b_cols_out) as u64,
203                (a_size * b_cols_in) as u64,
204                (b_rows * b_cols_in) as u64,
205                (b_size * b_cols_out) as u64,
206            ) as usize
207        }
208    }
209}
210
211unsafe impl VmpApplyAddImpl<FFT64> for FFT64 {
212    fn vmp_apply_add_impl<R, A, C>(module: &Module<FFT64>, res: &mut R, a: &A, b: &C, scale: usize, scratch: &mut Scratch<FFT64>)
213    where
214        R: VecZnxDftToMut<FFT64>,
215        A: VecZnxDftToRef<FFT64>,
216        C: VmpPMatToRef<FFT64>,
217    {
218        let mut res: VecZnxDft<&mut [u8], _> = res.to_mut();
219        let a: VecZnxDft<&[u8], _> = a.to_ref();
220        let b: VmpPMat<&[u8], _> = b.to_ref();
221
222        #[cfg(debug_assertions)]
223        {
224            assert_eq!(b.n(), res.n());
225            assert_eq!(a.n(), res.n());
226            assert_eq!(
227                res.cols(),
228                b.cols_out(),
229                "res.cols(): {} != b.cols_out: {}",
230                res.cols(),
231                b.cols_out()
232            );
233            assert_eq!(
234                a.cols(),
235                b.cols_in(),
236                "a.cols(): {} != b.cols_in: {}",
237                a.cols(),
238                b.cols_in()
239            );
240        }
241
242        let (tmp_bytes, _) = scratch.take_slice(module.vmp_apply_tmp_bytes(
243            res.size(),
244            a.size(),
245            b.rows(),
246            b.cols_in(),
247            b.cols_out(),
248            b.size(),
249        ));
250        unsafe {
251            vmp::vmp_apply_dft_to_dft_add(
252                module.ptr(),
253                res.as_mut_ptr() as *mut vec_znx_dft_t,
254                (res.size() * res.cols()) as u64,
255                a.as_ptr() as *const vec_znx_dft_t,
256                (a.size() * a.cols()) as u64,
257                b.as_ptr() as *const vmp::vmp_pmat_t,
258                (b.rows() * b.cols_in()) as u64,
259                (b.size() * b.cols_out()) as u64,
260                (scale * b.cols_out()) as u64,
261                tmp_bytes.as_mut_ptr(),
262            )
263        }
264    }
265}