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!(include_str!("fft16_avx2_fma.s"), include_str!("ifft16_avx2_fma.s"));
36
37#[inline(always)]
38pub(crate) fn as_arr<const SIZE: usize, R: Float + FloatConst>(x: &[R]) -> &[R; SIZE] {
39    debug_assert!(x.len() >= SIZE);
40    unsafe { &*(x.as_ptr() as *const [R; SIZE]) }
41}
42
43#[inline(always)]
44pub(crate) fn as_arr_mut<const SIZE: usize, R: Float + FloatConst>(x: &mut [R]) -> &mut [R; SIZE] {
45    debug_assert!(x.len() >= SIZE);
46    unsafe { &mut *(x.as_mut_ptr() as *mut [R; SIZE]) }
47}
48
49pub struct ReimFFTAvx;
50
51impl ReimDFTExecute<ReimFFTTable<f64>, f64> for ReimFFTAvx {
52    #[inline(always)]
53    fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
54        unsafe {
55            fft_avx2_fma(table.m(), table.omg(), data);
56        }
57    }
58}
59
60pub struct ReimIFFTAvx;
61
62impl ReimDFTExecute<ReimIFFTTable<f64>, f64> for ReimIFFTAvx {
63    #[inline(always)]
64    fn reim_dft_execute(table: &ReimIFFTTable<f64>, data: &mut [f64]) {
65        unsafe {
66            ifft_avx2_fma(table.m(), table.omg(), data);
67        }
68    }
69}