1use poulpy_hal::{
  2    api::{TakeSlice, VecZnxIdftApplyTmpBytes},
  3    layouts::{
  4        Backend, Data, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut,
  5        VecZnxDftToRef, VecZnxToRef, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut,
  6    },
  7    oep::{
  8        VecZnxDftAddImpl, VecZnxDftAddInplaceImpl, VecZnxDftAllocBytesImpl, VecZnxDftAllocImpl, VecZnxDftApplyImpl,
  9        VecZnxDftCopyImpl, VecZnxDftFromBytesImpl, VecZnxDftSubImpl, VecZnxDftSubInplaceImpl, VecZnxDftSubNegateInplaceImpl,
 10        VecZnxDftZeroImpl, VecZnxIdftApplyConsumeImpl, VecZnxIdftApplyImpl, VecZnxIdftApplyTmpAImpl, VecZnxIdftApplyTmpBytesImpl,
 11    },
 12    reference::{
 13        fft64::{
 14            reim::{ReimCopy, ReimZero, reim_copy_ref, reim_negate_inplace_ref, reim_negate_ref, reim_zero_ref},
 15            vec_znx_dft::vec_znx_dft_copy,
 16        },
 17        znx::znx_zero_ref,
 18    },
 19};
 20
 21use crate::cpu_spqlios::{
 22    FFT64Spqlios,
 23    ffi::{vec_znx_big, vec_znx_dft},
 24};
 25
 26unsafe impl VecZnxDftFromBytesImpl<Self> for FFT64Spqlios {
 27    fn vec_znx_dft_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<Self> {
 28        VecZnxDft::<Vec<u8>, Self>::from_bytes(n, cols, size, bytes)
 29    }
 30}
 31
 32unsafe impl VecZnxDftAllocBytesImpl<Self> for FFT64Spqlios {
 33    fn vec_znx_dft_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
 34        Self::layout_prep_word_count() * n * cols * size * size_of::<<FFT64Spqlios as Backend>::ScalarPrep>()
 35    }
 36}
 37
 38unsafe impl VecZnxDftAllocImpl<Self> for FFT64Spqlios {
 39    fn vec_znx_dft_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxDftOwned<Self> {
 40        VecZnxDftOwned::alloc(n, cols, size)
 41    }
 42}
 43
 44unsafe impl VecZnxIdftApplyTmpBytesImpl<Self> for FFT64Spqlios {
 45    fn vec_znx_idft_apply_tmp_bytes_impl(module: &Module<Self>) -> usize {
 46        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(module.ptr()) as usize }
 47    }
 48}
 49
 50unsafe impl VecZnxIdftApplyImpl<Self> for FFT64Spqlios {
 51    fn vec_znx_idft_apply_impl<R, A>(
 52        module: &Module<Self>,
 53        res: &mut R,
 54        res_col: usize,
 55        a: &A,
 56        a_col: usize,
 57        scratch: &mut Scratch<Self>,
 58    ) where
 59        R: VecZnxBigToMut<Self>,
 60        A: VecZnxDftToRef<Self>,
 61    {
 62        let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
 63        let a: VecZnxDft<&[u8], Self> = a.to_ref();
 64
 65        #[cfg(debug_assertions)]
 66        {
 67            assert_eq!(res.n(), a.n())
 68        }
 69
 70        let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_idft_apply_tmp_bytes());
 71
 72        let min_size: usize = res.size().min(a.size());
 73
 74        unsafe {
 75            (0..min_size).for_each(|j| {
 76                vec_znx_dft::vec_znx_idft(
 77                    module.ptr(),
 78                    res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
 79                    1_u64,
 80                    a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
 81                    1_u64,
 82                    tmp_bytes.as_mut_ptr(),
 83                )
 84            });
 85            (min_size..res.size()).for_each(|j| znx_zero_ref(res.at_mut(res_col, j)));
 86        }
 87    }
 88}
 89
 90unsafe impl VecZnxIdftApplyTmpAImpl<Self> for FFT64Spqlios {
 91    fn vec_znx_idft_apply_tmpa_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
 92    where
 93        R: VecZnxBigToMut<Self>,
 94        A: VecZnxDftToMut<Self>,
 95    {
 96        let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
 97        let mut a_mut: VecZnxDft<&mut [u8], Self> = a.to_mut();
 98
 99        let min_size: usize = res.size().min(a_mut.size());
100
101        unsafe {
102            (0..min_size).for_each(|j| {
103                vec_znx_dft::vec_znx_idft_tmp_a(
104                    module.ptr(),
105                    res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
106                    1_u64,
107                    a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
108                    1_u64,
109                )
110            });
111            (min_size..res.size()).for_each(|j| znx_zero_ref(res.at_mut(res_col, j)))
112        }
113    }
114}
115
116unsafe impl VecZnxIdftApplyConsumeImpl<Self> for FFT64Spqlios {
117    fn vec_znx_idft_apply_consume_impl<D: Data>(module: &Module<Self>, mut a: VecZnxDft<D, Self>) -> VecZnxBig<D, Self>
118    where
119        VecZnxDft<D, Self>: VecZnxDftToMut<Self>,
120    {
121        let mut a_mut: VecZnxDft<&mut [u8], Self> = a.to_mut();
122
123        unsafe {
124            // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
125            (0..a_mut.size()).for_each(|j| {
126                (0..a_mut.cols()).for_each(|i| {
127                    vec_znx_dft::vec_znx_idft_tmp_a(
128                        module.ptr(),
129                        a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
130                        1_u64,
131                        a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
132                        1_u64,
133                    )
134                });
135            });
136        }
137
138        a.into_big()
139    }
140}
141
142unsafe impl VecZnxDftApplyImpl<Self> for FFT64Spqlios {
143    fn vec_znx_dft_apply_impl<R, A>(
144        module: &Module<Self>,
145        step: usize,
146        offset: usize,
147        res: &mut R,
148        res_col: usize,
149        a: &A,
150        a_col: usize,
151    ) where
152        R: VecZnxDftToMut<Self>,
153        A: VecZnxToRef,
154    {
155        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
156        let a: VecZnx<&[u8]> = a.to_ref();
157        let steps: usize = a.size().div_ceil(step);
158        let min_steps: usize = res.size().min(steps);
159        unsafe {
160            (0..min_steps).for_each(|j| {
161                let limb: usize = offset + j * step;
162                if limb < a.size() {
163                    vec_znx_dft::vec_znx_dft(
164                        module.ptr(),
165                        res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
166                        1_u64,
167                        a.at_ptr(a_col, limb),
168                        1_u64,
169                        a.sl() as u64,
170                    )
171                }
172            });
173            (min_steps..res.size()).for_each(|j| reim_zero_ref(res.at_mut(res_col, j)));
174        }
175    }
176}
177
178unsafe impl VecZnxDftAddImpl<Self> for FFT64Spqlios {
179    fn vec_znx_dft_add_impl<R, A, D>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
180    where
181        R: VecZnxDftToMut<Self>,
182        A: VecZnxDftToRef<Self>,
183        D: VecZnxDftToRef<Self>,
184    {
185        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
186        let a: VecZnxDft<&[u8], Self> = a.to_ref();
187        let b: VecZnxDft<&[u8], Self> = b.to_ref();
188
189        let res_size: usize = res.size();
190        let a_size: usize = a.size();
191        let b_size: usize = b.size();
192
193        unsafe {
194            if a_size <= b_size {
195                let sum_size: usize = a_size.min(res_size);
196                let cpy_size: usize = b_size.min(res_size);
197
198                (0..sum_size).for_each(|j| {
199                    vec_znx_dft::vec_dft_add(
200                        module.ptr(),
201                        res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
202                        1,
203                        a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
204                        1,
205                        b.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
206                        1,
207                    );
208                });
209
210                for j in sum_size..cpy_size {
211                    reim_copy_ref(res.at_mut(res_col, j), b.at(b_col, j));
212                }
213
214                for j in cpy_size..res_size {
215                    reim_zero_ref(res.at_mut(res_col, j));
216                }
217            } else {
218                let sum_size: usize = b_size.min(res_size);
219                let cpy_size: usize = a_size.min(res_size);
220
221                (0..sum_size).for_each(|j| {
222                    vec_znx_dft::vec_dft_add(
223                        module.ptr(),
224                        res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
225                        1,
226                        a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
227                        1,
228                        b.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
229                        1,
230                    );
231                });
232
233                for j in sum_size..cpy_size {
234                    reim_copy_ref(res.at_mut(res_col, j), a.at(b_col, j));
235                }
236
237                for j in cpy_size..res_size {
238                    reim_zero_ref(res.at_mut(res_col, j));
239                }
240            }
241        }
242    }
243}
244
245unsafe impl VecZnxDftAddInplaceImpl<Self> for FFT64Spqlios {
246    fn vec_znx_dft_add_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
247    where
248        R: VecZnxDftToMut<Self>,
249        A: VecZnxDftToRef<Self>,
250    {
251        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
252        let a: VecZnxDft<&[u8], Self> = a.to_ref();
253
254        let min_size: usize = res.size().min(a.size());
255
256        unsafe {
257            (0..min_size).for_each(|j| {
258                vec_znx_dft::vec_dft_add(
259                    module.ptr(),
260                    res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
261                    1,
262                    res.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
263                    1,
264                    a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
265                    1,
266                );
267            });
268        }
269    }
270}
271
272unsafe impl VecZnxDftSubImpl<Self> for FFT64Spqlios {
273    fn vec_znx_dft_sub_impl<R, A, D>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
274    where
275        R: VecZnxDftToMut<Self>,
276        A: VecZnxDftToRef<Self>,
277        D: VecZnxDftToRef<Self>,
278    {
279        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
280        let a: VecZnxDft<&[u8], Self> = a.to_ref();
281        let b: VecZnxDft<&[u8], Self> = b.to_ref();
282
283        unsafe {
284            let res_size: usize = res.size();
285            let a_size: usize = a.size();
286            let b_size: usize = b.size();
287
288            if a_size <= b_size {
289                let sum_size: usize = a_size.min(res_size);
290                let cpy_size: usize = b_size.min(res_size);
291
292                (0..sum_size).for_each(|j| {
293                    vec_znx_dft::vec_dft_sub(
294                        module.ptr(),
295                        res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
296                        1,
297                        a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
298                        1,
299                        b.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
300                        1,
301                    );
302                });
303
304                for j in sum_size..cpy_size {
305                    reim_negate_ref(res.at_mut(res_col, j), b.at(b_col, j));
306                }
307
308                for j in cpy_size..res_size {
309                    reim_zero_ref(res.at_mut(res_col, j));
310                }
311            } else {
312                let sum_size: usize = b_size.min(res_size);
313                let cpy_size: usize = a_size.min(res_size);
314
315                (0..sum_size).for_each(|j| {
316                    vec_znx_dft::vec_dft_sub(
317                        module.ptr(),
318                        res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
319                        1,
320                        a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
321                        1,
322                        b.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
323                        1,
324                    );
325                });
326
327                for j in sum_size..cpy_size {
328                    reim_copy_ref(res.at_mut(res_col, j), a.at(a_col, j));
329                }
330
331                for j in cpy_size..res_size {
332                    reim_zero_ref(res.at_mut(res_col, j));
333                }
334            }
335        }
336    }
337}
338
339unsafe impl VecZnxDftSubInplaceImpl<Self> for FFT64Spqlios {
340    fn vec_znx_dft_sub_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
341    where
342        R: VecZnxDftToMut<Self>,
343        A: VecZnxDftToRef<Self>,
344    {
345        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
346        let a: VecZnxDft<&[u8], Self> = a.to_ref();
347
348        let min_size: usize = res.size().min(a.size());
349
350        unsafe {
351            (0..min_size).for_each(|j| {
352                vec_znx_dft::vec_dft_sub(
353                    module.ptr(),
354                    res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
355                    1,
356                    res.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
357                    1,
358                    a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
359                    1,
360                );
361            });
362        }
363    }
364}
365
366unsafe impl VecZnxDftSubNegateInplaceImpl<Self> for FFT64Spqlios {
367    fn vec_znx_dft_sub_negate_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
368    where
369        R: VecZnxDftToMut<Self>,
370        A: VecZnxDftToRef<Self>,
371    {
372        let mut res: VecZnxDft<&mut [u8], Self> = res.to_mut();
373        let a: VecZnxDft<&[u8], Self> = a.to_ref();
374
375        let min_size: usize = res.size().min(a.size());
376
377        unsafe {
378            (0..min_size).for_each(|j| {
379                vec_znx_dft::vec_dft_sub(
380                    module.ptr(),
381                    res.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
382                    1,
383                    a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
384                    1,
385                    res.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
386                    1,
387                );
388            });
389
390            for j in min_size..res.size() {
391                reim_negate_inplace_ref(res.at_mut(res_col, j));
392            }
393        }
394    }
395}
396
397unsafe impl VecZnxDftCopyImpl<Self> for FFT64Spqlios {
398    fn vec_znx_dft_copy_impl<R, A>(
399        _module: &Module<Self>,
400        step: usize,
401        offset: usize,
402        res: &mut R,
403        res_col: usize,
404        a: &A,
405        a_col: usize,
406    ) where
407        R: VecZnxDftToMut<Self>,
408        A: VecZnxDftToRef<Self>,
409    {
410        vec_znx_dft_copy(step, offset, res, res_col, a, a_col);
411    }
412}
413
414impl ReimCopy for FFT64Spqlios {
415    #[inline(always)]
416    fn reim_copy(res: &mut [f64], a: &[f64]) {
417        reim_copy_ref(res, a);
418    }
419}
420
421impl ReimZero for FFT64Spqlios {
422    #[inline(always)]
423    fn reim_zero(res: &mut [f64]) {
424        reim_zero_ref(res);
425    }
426}
427
428unsafe impl VecZnxDftZeroImpl<Self> for FFT64Spqlios {
429    fn vec_znx_dft_zero_impl<R>(_module: &Module<Self>, res: &mut R)
430    where
431        R: VecZnxDftToMut<Self>,
432    {
433        res.to_mut().data.fill(0);
434    }
435}