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