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