1use crate::cpu_fft64_avx::FFT64Avx;
  2use poulpy_hal::{
  3    api::{TakeSlice, VecZnxBigAutomorphismInplaceTmpBytes, VecZnxBigNormalizeTmpBytes},
  4    layouts::{
  5        Backend, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef,
  6        ZnxInfos, ZnxView, ZnxViewMut,
  7    },
  8    oep::{
  9        TakeSliceImpl, VecZnxBigAddImpl, VecZnxBigAddInplaceImpl, VecZnxBigAddNormalImpl, VecZnxBigAddSmallImpl,
 10        VecZnxBigAddSmallInplaceImpl, VecZnxBigAllocBytesImpl, VecZnxBigAllocImpl, VecZnxBigAutomorphismImpl,
 11        VecZnxBigAutomorphismInplaceImpl, VecZnxBigAutomorphismInplaceTmpBytesImpl, VecZnxBigFromBytesImpl,
 12        VecZnxBigFromSmallImpl, VecZnxBigNegateImpl, VecZnxBigNegateInplaceImpl, VecZnxBigNormalizeImpl,
 13        VecZnxBigNormalizeTmpBytesImpl, VecZnxBigSubABInplaceImpl, VecZnxBigSubBAInplaceImpl, VecZnxBigSubImpl,
 14        VecZnxBigSubSmallAImpl, VecZnxBigSubSmallAInplaceImpl, VecZnxBigSubSmallBImpl, VecZnxBigSubSmallBInplaceImpl,
 15    },
 16    reference::{
 17        fft64::vec_znx_big::{
 18            vec_znx_big_add, vec_znx_big_add_inplace, vec_znx_big_add_normal_ref, vec_znx_big_add_small,
 19            vec_znx_big_add_small_inplace, vec_znx_big_automorphism, vec_znx_big_automorphism_inplace,
 20            vec_znx_big_automorphism_inplace_tmp_bytes, vec_znx_big_negate, vec_znx_big_negate_inplace, vec_znx_big_normalize,
 21            vec_znx_big_normalize_tmp_bytes, vec_znx_big_sub, vec_znx_big_sub_ab_inplace, vec_znx_big_sub_ba_inplace,
 22            vec_znx_big_sub_small_a, vec_znx_big_sub_small_a_inplace, vec_znx_big_sub_small_b, vec_znx_big_sub_small_b_inplace,
 23        },
 24        znx::{znx_copy_ref, znx_zero_ref},
 25    },
 26    source::Source,
 27};
 28
 29unsafe impl VecZnxBigAllocBytesImpl<Self> for FFT64Avx {
 30    fn vec_znx_big_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
 31        Self::layout_big_word_count() * n * cols * size * size_of::<f64>()
 32    }
 33}
 34
 35unsafe impl VecZnxBigAllocImpl<Self> for FFT64Avx {
 36    fn vec_znx_big_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxBigOwned<Self> {
 37        VecZnxBig::alloc(n, cols, size)
 38    }
 39}
 40
 41unsafe impl VecZnxBigFromBytesImpl<Self> for FFT64Avx {
 42    fn vec_znx_big_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<Self> {
 43        VecZnxBig::from_bytes(n, cols, size, bytes)
 44    }
 45}
 46
 47unsafe impl VecZnxBigFromSmallImpl<Self> for FFT64Avx {
 48    fn vec_znx_big_from_small_impl<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 49    where
 50        R: VecZnxBigToMut<Self>,
 51        A: VecZnxToRef,
 52    {
 53        let mut res: VecZnxBig<&mut [u8], FFT64Avx> = res.to_mut();
 54        let a: VecZnx<&[u8]> = a.to_ref();
 55
 56        #[cfg(debug_assertions)]
 57        {
 58            assert_eq!(res.n(), a.n());
 59        }
 60
 61        let res_size: usize = res.size();
 62        let a_size: usize = a.size();
 63
 64        let min_size: usize = res_size.min(a_size);
 65
 66        for j in 0..min_size {
 67            znx_copy_ref(res.at_mut(res_col, j), a.at(a_col, j));
 68        }
 69
 70        for j in min_size..res_size {
 71            znx_zero_ref(res.at_mut(res_col, j));
 72        }
 73    }
 74}
 75
 76unsafe impl VecZnxBigAddNormalImpl<Self> for FFT64Avx {
 77    fn add_normal_impl<R: VecZnxBigToMut<Self>>(
 78        _module: &Module<Self>,
 79        basek: usize,
 80        res: &mut R,
 81        res_col: usize,
 82        k: usize,
 83        source: &mut Source,
 84        sigma: f64,
 85        bound: f64,
 86    ) {
 87        vec_znx_big_add_normal_ref(basek, res, res_col, k, sigma, bound, source);
 88    }
 89}
 90
 91unsafe impl VecZnxBigAddImpl<Self> for FFT64Avx {
 92    /// Adds `a` to `b` and stores the result on `c`.
 93    fn vec_znx_big_add_impl<R, A, B>(
 94        _module: &Module<Self>,
 95        res: &mut R,
 96        res_col: usize,
 97        a: &A,
 98        a_col: usize,
 99        b: &B,
100        b_col: usize,
101    ) where
102        R: VecZnxBigToMut<Self>,
103        A: VecZnxBigToRef<Self>,
104        B: VecZnxBigToRef<Self>,
105    {
106        vec_znx_big_add(res, res_col, a, a_col, b, b_col);
107    }
108}
109
110unsafe impl VecZnxBigAddInplaceImpl<Self> for FFT64Avx {
111    /// Adds `a` to `b` and stores the result on `b`.
112    fn vec_znx_big_add_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
113    where
114        R: VecZnxBigToMut<Self>,
115        A: VecZnxBigToRef<Self>,
116    {
117        vec_znx_big_add_inplace(res, res_col, a, a_col);
118    }
119}
120
121unsafe impl VecZnxBigAddSmallImpl<Self> for FFT64Avx {
122    /// Adds `a` to `b` and stores the result on `c`.
123    fn vec_znx_big_add_small_impl<R, A, B>(
124        _module: &Module<Self>,
125        res: &mut R,
126        res_col: usize,
127        a: &A,
128        a_col: usize,
129        b: &B,
130        b_col: usize,
131    ) where
132        R: VecZnxBigToMut<Self>,
133        A: VecZnxBigToRef<Self>,
134        B: VecZnxToRef,
135    {
136        vec_znx_big_add_small(res, res_col, a, a_col, b, b_col);
137    }
138}
139
140unsafe impl VecZnxBigAddSmallInplaceImpl<Self> for FFT64Avx {
141    /// Adds `a` to `b` and stores the result on `b`.
142    fn vec_znx_big_add_small_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
143    where
144        R: VecZnxBigToMut<Self>,
145        A: VecZnxToRef,
146    {
147        vec_znx_big_add_small_inplace(res, res_col, a, a_col);
148    }
149}
150
151unsafe impl VecZnxBigSubImpl<Self> for FFT64Avx {
152    /// Subtracts `a` to `b` and stores the result on `c`.
153    fn vec_znx_big_sub_impl<R, A, B>(
154        _module: &Module<Self>,
155        res: &mut R,
156        res_col: usize,
157        a: &A,
158        a_col: usize,
159        b: &B,
160        b_col: usize,
161    ) where
162        R: VecZnxBigToMut<Self>,
163        A: VecZnxBigToRef<Self>,
164        B: VecZnxBigToRef<Self>,
165    {
166        vec_znx_big_sub(res, res_col, a, a_col, b, b_col);
167    }
168}
169
170unsafe impl VecZnxBigSubABInplaceImpl<Self> for FFT64Avx {
171    /// Subtracts `a` from `b` and stores the result on `b`.
172    fn vec_znx_big_sub_ab_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
173    where
174        R: VecZnxBigToMut<Self>,
175        A: VecZnxBigToRef<Self>,
176    {
177        vec_znx_big_sub_ab_inplace(res, res_col, a, a_col);
178    }
179}
180
181unsafe impl VecZnxBigSubBAInplaceImpl<Self> for FFT64Avx {
182    /// Subtracts `b` from `a` and stores the result on `b`.
183    fn vec_znx_big_sub_ba_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
184    where
185        R: VecZnxBigToMut<Self>,
186        A: VecZnxBigToRef<Self>,
187    {
188        vec_znx_big_sub_ba_inplace(res, res_col, a, a_col);
189    }
190}
191
192unsafe impl VecZnxBigSubSmallAImpl<Self> for FFT64Avx {
193    /// Subtracts `b` from `a` and stores the result on `c`.
194    fn vec_znx_big_sub_small_a_impl<R, A, B>(
195        _module: &Module<Self>,
196        res: &mut R,
197        res_col: usize,
198        a: &A,
199        a_col: usize,
200        b: &B,
201        b_col: usize,
202    ) where
203        R: VecZnxBigToMut<Self>,
204        A: VecZnxToRef,
205        B: VecZnxBigToRef<Self>,
206    {
207        vec_znx_big_sub_small_a(res, res_col, a, a_col, b, b_col);
208    }
209}
210
211unsafe impl VecZnxBigSubSmallAInplaceImpl<Self> for FFT64Avx {
212    /// Subtracts `a` from `res` and stores the result on `res`.
213    fn vec_znx_big_sub_small_a_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
214    where
215        R: VecZnxBigToMut<Self>,
216        A: VecZnxToRef,
217    {
218        vec_znx_big_sub_small_a_inplace(res, res_col, a, a_col);
219    }
220}
221
222unsafe impl VecZnxBigSubSmallBImpl<Self> for FFT64Avx {
223    /// Subtracts `b` from `a` and stores the result on `c`.
224    fn vec_znx_big_sub_small_b_impl<R, A, B>(
225        _module: &Module<Self>,
226        res: &mut R,
227        res_col: usize,
228        a: &A,
229        a_col: usize,
230        b: &B,
231        b_col: usize,
232    ) where
233        R: VecZnxBigToMut<Self>,
234        A: VecZnxBigToRef<Self>,
235        B: VecZnxToRef,
236    {
237        vec_znx_big_sub_small_b(res, res_col, a, a_col, b, b_col);
238    }
239}
240
241unsafe impl VecZnxBigSubSmallBInplaceImpl<Self> for FFT64Avx {
242    /// Subtracts `res` from `a` and stores the result on `res`.
243    fn vec_znx_big_sub_small_b_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
244    where
245        R: VecZnxBigToMut<Self>,
246        A: VecZnxToRef,
247    {
248        vec_znx_big_sub_small_b_inplace(res, res_col, a, a_col);
249    }
250}
251
252unsafe impl VecZnxBigNegateImpl<Self> for FFT64Avx {
253    fn vec_znx_big_negate_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
254    where
255        R: VecZnxBigToMut<Self>,
256        A: VecZnxBigToRef<Self>,
257    {
258        vec_znx_big_negate(res, res_col, a, a_col);
259    }
260}
261
262unsafe impl VecZnxBigNegateInplaceImpl<Self> for FFT64Avx {
263    fn vec_znx_big_negate_inplace_impl<R>(_module: &Module<Self>, res: &mut R, res_col: usize)
264    where
265        R: VecZnxBigToMut<Self>,
266    {
267        vec_znx_big_negate_inplace(res, res_col);
268    }
269}
270
271unsafe impl VecZnxBigNormalizeTmpBytesImpl<Self> for FFT64Avx {
272    fn vec_znx_big_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
273        vec_znx_big_normalize_tmp_bytes(module.n())
274    }
275}
276
277unsafe impl VecZnxBigNormalizeImpl<Self> for FFT64Avx
278where
279    Self: TakeSliceImpl<Self>,
280{
281    fn vec_znx_big_normalize_impl<R, A>(
282        module: &Module<Self>,
283        basek: usize,
284        res: &mut R,
285        res_col: usize,
286        a: &A,
287        a_col: usize,
288        scratch: &mut Scratch<Self>,
289    ) where
290        R: VecZnxToMut,
291        A: VecZnxBigToRef<Self>,
292    {
293        let (carry, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes() / size_of::<i64>());
294        vec_znx_big_normalize(basek, res, res_col, a, a_col, carry);
295    }
296}
297
298unsafe impl VecZnxBigAutomorphismImpl<Self> for FFT64Avx {
299    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
300    fn vec_znx_big_automorphism_impl<R, A>(_module: &Module<Self>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
301    where
302        R: VecZnxBigToMut<Self>,
303        A: VecZnxBigToRef<Self>,
304    {
305        vec_znx_big_automorphism(p, res, res_col, a, a_col);
306    }
307}
308
309unsafe impl VecZnxBigAutomorphismInplaceTmpBytesImpl<Self> for FFT64Avx {
310    fn vec_znx_big_automorphism_inplace_tmp_bytes_impl(module: &Module<Self>) -> usize {
311        vec_znx_big_automorphism_inplace_tmp_bytes(module.n())
312    }
313}
314
315unsafe impl VecZnxBigAutomorphismInplaceImpl<Self> for FFT64Avx
316where
317    Module<Self>: VecZnxBigAutomorphismInplaceTmpBytes,
318{
319    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
320    fn vec_znx_big_automorphism_inplace_impl<R>(
321        module: &Module<Self>,
322        p: i64,
323        res: &mut R,
324        res_col: usize,
325        scratch: &mut Scratch<Self>,
326    ) where
327        R: VecZnxBigToMut<Self>,
328    {
329        let (tmp, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes() / size_of::<i64>());
330        vec_znx_big_automorphism_inplace(p, res, res_col, tmp);
331    }
332}