1// ----------------------------------------------------------------------
 2// DISCLAIMER
 3//
 4// This module contains code that has been directly ported from the
 5// spqlios-arithmetic library
 6// (https://github.com/tfhe/spqlios-arithmetic), which is licensed
 7// under the Apache License, Version 2.0.
 8//
 9// The porting process from C to Rust was done with minimal changes
10// in order to preserve the semantics and performance characteristics
11// of the original implementation.
12//
13// Both Poulpy and spqlios-arithmetic are distributed under the terms
14// of the Apache License, Version 2.0. See the LICENSE file for details.
15//
16// ----------------------------------------------------------------------
17
18#![allow(bad_asm_style)]
19
20mod conversion;
21mod fft_avx2_fma;
22mod fft_vec_avx2_fma;
23mod ifft_avx2_fma;
24
25use std::arch::global_asm;
26
27pub(crate) use conversion::*;
28pub(crate) use fft_vec_avx2_fma::*;
29
30use poulpy_hal::reference::fft64::reim::{ReimDFTExecute, ReimFFTTable, ReimIFFTTable};
31use rand_distr::num_traits::{Float, FloatConst};
32
33use crate::reim::{fft_avx2_fma::fft_avx2_fma, ifft_avx2_fma::ifft_avx2_fma};
34
35global_asm!(
36    include_str!("fft16_avx2_fma.s"),
37    include_str!("ifft16_avx2_fma.s")
38);
39
40#[inline(always)]
41pub(crate) fn as_arr<const SIZE: usize, R: Float + FloatConst>(x: &[R]) -> &[R; SIZE] {
42    debug_assert!(x.len() >= SIZE);
43    unsafe { &*(x.as_ptr() as *const [R; SIZE]) }
44}
45
46#[inline(always)]
47pub(crate) fn as_arr_mut<const SIZE: usize, R: Float + FloatConst>(x: &mut [R]) -> &mut [R; SIZE] {
48    debug_assert!(x.len() >= SIZE);
49    unsafe { &mut *(x.as_mut_ptr() as *mut [R; SIZE]) }
50}
51
52pub struct ReimFFTAvx;
53
54impl ReimDFTExecute<ReimFFTTable<f64>, f64> for ReimFFTAvx {
55    #[inline(always)]
56    fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
57        unsafe {
58            fft_avx2_fma(table.m(), table.omg(), data);
59        }
60    }
61}
62
63pub struct ReimIFFTAvx;
64
65impl ReimDFTExecute<ReimIFFTTable<f64>, f64> for ReimIFFTAvx {
66    #[inline(always)]
67    fn reim_dft_execute(table: &ReimIFFTTable<f64>, data: &mut [f64]) {
68        unsafe {
69            ifft_avx2_fma(table.m(), table.omg(), data);
70        }
71    }
72}