1use std::ptr::NonNull;
  2
  3use poulpy_hal::{
  4    layouts::{Backend, Module},
  5    oep::ModuleNewImpl,
  6    reference::{
  7        fft64::{
  8            reim::{
  9                ReimAdd, ReimAddInplace, ReimAddMul, ReimCopy, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimIFFTTable, ReimMul,
 10                ReimMulInplace, ReimNegate, ReimNegateInplace, ReimSub, ReimSubInplace, ReimSubNegateInplace, ReimToZnx,
 11                ReimToZnxInplace, ReimZero, reim_copy_ref, reim_zero_ref,
 12            },
 13            reim4::{
 14                Reim4Extract1Blk, Reim4Mat1ColProd, Reim4Mat2Cols2ndColProd, Reim4Mat2ColsProd, Reim4Save1Blk, Reim4Save2Blks,
 15            },
 16        },
 17        znx::{
 18            ZnxAdd, ZnxAddInplace, ZnxAutomorphism, ZnxCopy, ZnxExtractDigitAddMul, ZnxMulAddPowerOfTwo, ZnxMulPowerOfTwo,
 19            ZnxMulPowerOfTwoInplace, ZnxNegate, ZnxNegateInplace, ZnxNormalizeDigit, ZnxNormalizeFinalStep,
 20            ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly, ZnxNormalizeFirstStepInplace,
 21            ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxNormalizeMiddleStepInplace, ZnxRotate, ZnxSub,
 22            ZnxSubInplace, ZnxSubNegateInplace, ZnxSwitchRing, ZnxZero, znx_copy_ref, znx_rotate, znx_zero_ref,
 23        },
 24    },
 25};
 26
 27use crate::cpu_fft64_avx::{
 28    FFT64Avx,
 29    reim::{
 30        ReimFFTAvx, ReimIFFTAvx, reim_add_avx2_fma, reim_add_inplace_avx2_fma, reim_addmul_avx2_fma, reim_from_znx_i64_bnd50_fma,
 31        reim_mul_avx2_fma, reim_mul_inplace_avx2_fma, reim_negate_avx2_fma, reim_negate_inplace_avx2_fma, reim_sub_avx2_fma,
 32        reim_sub_inplace_avx2_fma, reim_sub_negate_inplace_avx2_fma, reim_to_znx_i64_inplace_bnd63_avx2_fma,
 33    },
 34    reim_to_znx_i64_bnd63_avx2_fma,
 35    reim4::{
 36        reim4_extract_1blk_from_reim_avx, reim4_save_1blk_to_reim_avx, reim4_save_2blk_to_reim_avx,
 37        reim4_vec_mat1col_product_avx, reim4_vec_mat2cols_2ndcol_product_avx, reim4_vec_mat2cols_product_avx,
 38    },
 39    znx_avx::{
 40        znx_add_avx, znx_add_inplace_avx, znx_automorphism_avx, znx_extract_digit_addmul_avx, znx_mul_add_power_of_two_avx,
 41        znx_mul_power_of_two_avx, znx_mul_power_of_two_inplace_avx, znx_negate_avx, znx_negate_inplace_avx,
 42        znx_normalize_digit_avx, znx_normalize_final_step_avx, znx_normalize_final_step_inplace_avx,
 43        znx_normalize_first_step_avx, znx_normalize_first_step_carry_only_avx, znx_normalize_first_step_inplace_avx,
 44        znx_normalize_middle_step_avx, znx_normalize_middle_step_carry_only_avx, znx_normalize_middle_step_inplace_avx,
 45        znx_sub_avx, znx_sub_inplace_avx, znx_sub_negate_inplace_avx, znx_switch_ring_avx,
 46    },
 47};
 48
 49#[repr(C)]
 50pub struct FFT64AvxHandle {
 51    table_fft: ReimFFTTable<f64>,
 52    table_ifft: ReimIFFTTable<f64>,
 53}
 54
 55impl Backend for FFT64Avx {
 56    type ScalarPrep = f64;
 57    type ScalarBig = i64;
 58    type Handle = FFT64AvxHandle;
 59    unsafe fn destroy(handle: NonNull<Self::Handle>) {
 60        unsafe {
 61            drop(Box::from_raw(handle.as_ptr()));
 62        }
 63    }
 64
 65    fn layout_big_word_count() -> usize {
 66        1
 67    }
 68
 69    fn layout_prep_word_count() -> usize {
 70        1
 71    }
 72}
 73
 74unsafe impl ModuleNewImpl<Self> for FFT64Avx {
 75    fn new_impl(n: u64) -> Module<Self> {
 76        if !std::arch::is_x86_feature_detected!("avx")
 77            || !std::arch::is_x86_feature_detected!("avx2")
 78            || !std::arch::is_x86_feature_detected!("fma")
 79        {
 80            panic!("arch must support avx2, avx and fma")
 81        }
 82
 83        let handle: FFT64AvxHandle = FFT64AvxHandle {
 84            table_fft: ReimFFTTable::new(n as usize >> 1),
 85            table_ifft: ReimIFFTTable::new(n as usize >> 1),
 86        };
 87        // Leak Box to get a stable NonNull pointer
 88        let ptr: NonNull<FFT64AvxHandle> = NonNull::from(Box::leak(Box::new(handle)));
 89        unsafe { Module::from_nonnull(ptr, n) }
 90    }
 91}
 92
 93pub trait FFT64ModuleHandle {
 94    fn get_fft_table(&self) -> &ReimFFTTable<f64>;
 95    fn get_ifft_table(&self) -> &ReimIFFTTable<f64>;
 96}
 97
 98impl FFT64ModuleHandle for Module<FFT64Avx> {
 99    fn get_fft_table(&self) -> &ReimFFTTable<f64> {
100        let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
101        &h.table_fft
102    }
103    fn get_ifft_table(&self) -> &ReimIFFTTable<f64> {
104        let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
105        &h.table_ifft
106    }
107}
108
109impl ZnxAdd for FFT64Avx {
110    #[inline(always)]
111    fn znx_add(res: &mut [i64], a: &[i64], b: &[i64]) {
112        unsafe {
113            znx_add_avx(res, a, b);
114        }
115    }
116}
117
118impl ZnxAddInplace for FFT64Avx {
119    #[inline(always)]
120    fn znx_add_inplace(res: &mut [i64], a: &[i64]) {
121        unsafe {
122            znx_add_inplace_avx(res, a);
123        }
124    }
125}
126
127impl ZnxSub for FFT64Avx {
128    #[inline(always)]
129    fn znx_sub(res: &mut [i64], a: &[i64], b: &[i64]) {
130        unsafe {
131            znx_sub_avx(res, a, b);
132        }
133    }
134}
135
136impl ZnxSubInplace for FFT64Avx {
137    #[inline(always)]
138    fn znx_sub_inplace(res: &mut [i64], a: &[i64]) {
139        unsafe {
140            znx_sub_inplace_avx(res, a);
141        }
142    }
143}
144
145impl ZnxSubNegateInplace for FFT64Avx {
146    #[inline(always)]
147    fn znx_sub_negate_inplace(res: &mut [i64], a: &[i64]) {
148        unsafe {
149            znx_sub_negate_inplace_avx(res, a);
150        }
151    }
152}
153
154impl ZnxAutomorphism for FFT64Avx {
155    #[inline(always)]
156    fn znx_automorphism(p: i64, res: &mut [i64], a: &[i64]) {
157        unsafe {
158            znx_automorphism_avx(p, res, a);
159        }
160    }
161}
162
163impl ZnxCopy for FFT64Avx {
164    #[inline(always)]
165    fn znx_copy(res: &mut [i64], a: &[i64]) {
166        znx_copy_ref(res, a);
167    }
168}
169
170impl ZnxNegate for FFT64Avx {
171    #[inline(always)]
172    fn znx_negate(res: &mut [i64], src: &[i64]) {
173        unsafe {
174            znx_negate_avx(res, src);
175        }
176    }
177}
178
179impl ZnxNegateInplace for FFT64Avx {
180    #[inline(always)]
181    fn znx_negate_inplace(res: &mut [i64]) {
182        unsafe {
183            znx_negate_inplace_avx(res);
184        }
185    }
186}
187
188impl ZnxMulAddPowerOfTwo for FFT64Avx {
189    #[inline(always)]
190    fn znx_muladd_power_of_two(k: i64, res: &mut [i64], a: &[i64]) {
191        unsafe {
192            znx_mul_add_power_of_two_avx(k, res, a);
193        }
194    }
195}
196
197impl ZnxMulPowerOfTwo for FFT64Avx {
198    #[inline(always)]
199    fn znx_mul_power_of_two(k: i64, res: &mut [i64], a: &[i64]) {
200        unsafe {
201            znx_mul_power_of_two_avx(k, res, a);
202        }
203    }
204}
205
206impl ZnxMulPowerOfTwoInplace for FFT64Avx {
207    #[inline(always)]
208    fn znx_mul_power_of_two_inplace(k: i64, res: &mut [i64]) {
209        unsafe {
210            znx_mul_power_of_two_inplace_avx(k, res);
211        }
212    }
213}
214
215impl ZnxRotate for FFT64Avx {
216    #[inline(always)]
217    fn znx_rotate(p: i64, res: &mut [i64], src: &[i64]) {
218        znx_rotate::<Self>(p, res, src);
219    }
220}
221
222impl ZnxZero for FFT64Avx {
223    #[inline(always)]
224    fn znx_zero(res: &mut [i64]) {
225        znx_zero_ref(res);
226    }
227}
228
229impl ZnxSwitchRing for FFT64Avx {
230    #[inline(always)]
231    fn znx_switch_ring(res: &mut [i64], a: &[i64]) {
232        unsafe {
233            znx_switch_ring_avx(res, a);
234        }
235    }
236}
237
238impl ZnxNormalizeFinalStep for FFT64Avx {
239    #[inline(always)]
240    fn znx_normalize_final_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
241        unsafe {
242            znx_normalize_final_step_avx(base2k, lsh, x, a, carry);
243        }
244    }
245}
246
247impl ZnxNormalizeFinalStepInplace for FFT64Avx {
248    #[inline(always)]
249    fn znx_normalize_final_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
250        unsafe {
251            znx_normalize_final_step_inplace_avx(base2k, lsh, x, carry);
252        }
253    }
254}
255
256impl ZnxNormalizeFirstStep for FFT64Avx {
257    #[inline(always)]
258    fn znx_normalize_first_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
259        unsafe {
260            znx_normalize_first_step_avx(base2k, lsh, x, a, carry);
261        }
262    }
263}
264
265impl ZnxNormalizeFirstStepCarryOnly for FFT64Avx {
266    #[inline(always)]
267    fn znx_normalize_first_step_carry_only(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
268        unsafe {
269            znx_normalize_first_step_carry_only_avx(base2k, lsh, x, carry);
270        }
271    }
272}
273
274impl ZnxNormalizeFirstStepInplace for FFT64Avx {
275    #[inline(always)]
276    fn znx_normalize_first_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
277        unsafe {
278            znx_normalize_first_step_inplace_avx(base2k, lsh, x, carry);
279        }
280    }
281}
282
283impl ZnxNormalizeMiddleStep for FFT64Avx {
284    #[inline(always)]
285    fn znx_normalize_middle_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
286        unsafe {
287            znx_normalize_middle_step_avx(base2k, lsh, x, a, carry);
288        }
289    }
290}
291
292impl ZnxNormalizeMiddleStepCarryOnly for FFT64Avx {
293    #[inline(always)]
294    fn znx_normalize_middle_step_carry_only(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
295        unsafe {
296            znx_normalize_middle_step_carry_only_avx(base2k, lsh, x, carry);
297        }
298    }
299}
300
301impl ZnxNormalizeMiddleStepInplace for FFT64Avx {
302    #[inline(always)]
303    fn znx_normalize_middle_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
304        unsafe {
305            znx_normalize_middle_step_inplace_avx(base2k, lsh, x, carry);
306        }
307    }
308}
309
310impl ZnxExtractDigitAddMul for FFT64Avx {
311    #[inline(always)]
312    fn znx_extract_digit_addmul(base2k: usize, lsh: usize, res: &mut [i64], src: &mut [i64]) {
313        unsafe {
314            znx_extract_digit_addmul_avx(base2k, lsh, res, src);
315        }
316    }
317}
318
319impl ZnxNormalizeDigit for FFT64Avx {
320    #[inline(always)]
321    fn znx_normalize_digit(base2k: usize, res: &mut [i64], src: &mut [i64]) {
322        unsafe {
323            znx_normalize_digit_avx(base2k, res, src);
324        }
325    }
326}
327
328impl ReimDFTExecute<ReimFFTTable<f64>, f64> for FFT64Avx {
329    #[inline(always)]
330    fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
331        ReimFFTAvx::reim_dft_execute(table, data);
332    }
333}
334
335impl ReimDFTExecute<ReimIFFTTable<f64>, f64> for FFT64Avx {
336    #[inline(always)]
337    fn reim_dft_execute(table: &ReimIFFTTable<f64>, data: &mut [f64]) {
338        ReimIFFTAvx::reim_dft_execute(table, data);
339    }
340}
341
342impl ReimFromZnx for FFT64Avx {
343    #[inline(always)]
344    fn reim_from_znx(res: &mut [f64], a: &[i64]) {
345        unsafe {
346            reim_from_znx_i64_bnd50_fma(res, a);
347        }
348    }
349}
350
351impl ReimToZnx for FFT64Avx {
352    #[inline(always)]
353    fn reim_to_znx(res: &mut [i64], divisor: f64, a: &[f64]) {
354        unsafe {
355            reim_to_znx_i64_bnd63_avx2_fma(res, divisor, a);
356        }
357    }
358}
359
360impl ReimToZnxInplace for FFT64Avx {
361    #[inline(always)]
362    fn reim_to_znx_inplace(res: &mut [f64], divisor: f64) {
363        unsafe {
364            reim_to_znx_i64_inplace_bnd63_avx2_fma(res, divisor);
365        }
366    }
367}
368
369impl ReimAdd for FFT64Avx {
370    #[inline(always)]
371    fn reim_add(res: &mut [f64], a: &[f64], b: &[f64]) {
372        unsafe {
373            reim_add_avx2_fma(res, a, b);
374        }
375    }
376}
377
378impl ReimAddInplace for FFT64Avx {
379    #[inline(always)]
380    fn reim_add_inplace(res: &mut [f64], a: &[f64]) {
381        unsafe {
382            reim_add_inplace_avx2_fma(res, a);
383        }
384    }
385}
386
387impl ReimSub for FFT64Avx {
388    #[inline(always)]
389    fn reim_sub(res: &mut [f64], a: &[f64], b: &[f64]) {
390        unsafe {
391            reim_sub_avx2_fma(res, a, b);
392        }
393    }
394}
395
396impl ReimSubInplace for FFT64Avx {
397    #[inline(always)]
398    fn reim_sub_inplace(res: &mut [f64], a: &[f64]) {
399        unsafe {
400            reim_sub_inplace_avx2_fma(res, a);
401        }
402    }
403}
404
405impl ReimSubNegateInplace for FFT64Avx {
406    #[inline(always)]
407    fn reim_sub_negate_inplace(res: &mut [f64], a: &[f64]) {
408        unsafe {
409            reim_sub_negate_inplace_avx2_fma(res, a);
410        }
411    }
412}
413
414impl ReimNegate for FFT64Avx {
415    #[inline(always)]
416    fn reim_negate(res: &mut [f64], a: &[f64]) {
417        unsafe {
418            reim_negate_avx2_fma(res, a);
419        }
420    }
421}
422
423impl ReimNegateInplace for FFT64Avx {
424    #[inline(always)]
425    fn reim_negate_inplace(res: &mut [f64]) {
426        unsafe {
427            reim_negate_inplace_avx2_fma(res);
428        }
429    }
430}
431
432impl ReimMul for FFT64Avx {
433    #[inline(always)]
434    fn reim_mul(res: &mut [f64], a: &[f64], b: &[f64]) {
435        unsafe {
436            reim_mul_avx2_fma(res, a, b);
437        }
438    }
439}
440
441impl ReimMulInplace for FFT64Avx {
442    #[inline(always)]
443    fn reim_mul_inplace(res: &mut [f64], a: &[f64]) {
444        unsafe {
445            reim_mul_inplace_avx2_fma(res, a);
446        }
447    }
448}
449
450impl ReimAddMul for FFT64Avx {
451    #[inline(always)]
452    fn reim_addmul(res: &mut [f64], a: &[f64], b: &[f64]) {
453        unsafe {
454            reim_addmul_avx2_fma(res, a, b);
455        }
456    }
457}
458
459impl ReimCopy for FFT64Avx {
460    #[inline(always)]
461    fn reim_copy(res: &mut [f64], a: &[f64]) {
462        reim_copy_ref(res, a);
463    }
464}
465
466impl ReimZero for FFT64Avx {
467    #[inline(always)]
468    fn reim_zero(res: &mut [f64]) {
469        reim_zero_ref(res);
470    }
471}
472
473impl Reim4Extract1Blk for FFT64Avx {
474    #[inline(always)]
475    fn reim4_extract_1blk(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
476        unsafe {
477            reim4_extract_1blk_from_reim_avx(m, rows, blk, dst, src);
478        }
479    }
480}
481
482impl Reim4Save1Blk for FFT64Avx {
483    #[inline(always)]
484    fn reim4_save_1blk<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
485        unsafe {
486            reim4_save_1blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
487        }
488    }
489}
490
491impl Reim4Save2Blks for FFT64Avx {
492    #[inline(always)]
493    fn reim4_save_2blks<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
494        unsafe {
495            reim4_save_2blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
496        }
497    }
498}
499
500impl Reim4Mat1ColProd for FFT64Avx {
501    #[inline(always)]
502    fn reim4_mat1col_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
503        unsafe {
504            reim4_vec_mat1col_product_avx(nrows, dst, u, v);
505        }
506    }
507}
508
509impl Reim4Mat2ColsProd for FFT64Avx {
510    #[inline(always)]
511    fn reim4_mat2cols_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
512        unsafe {
513            reim4_vec_mat2cols_product_avx(nrows, dst, u, v);
514        }
515    }
516}
517
518impl Reim4Mat2Cols2ndColProd for FFT64Avx {
519    #[inline(always)]
520    fn reim4_mat2cols_2ndcol_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
521        unsafe {
522            reim4_vec_mat2cols_2ndcol_product_avx(nrows, dst, u, v);
523        }
524    }
525}