1use std::arch::x86_64::{
  2    __m128d, __m256d, _mm_load_pd, _mm256_add_pd, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd,
  3    _mm256_permute2f128_pd, _mm256_set_m128d, _mm256_storeu_pd, _mm256_sub_pd, _mm256_unpackhi_pd, _mm256_unpacklo_pd,
  4};
  5
  6use crate::cpu_fft64_avx::reim::{as_arr, as_arr_mut};
  7
  8#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
  9#[target_feature(enable = "avx2,fma")]
 10pub(crate) fn fft_avx2_fma(m: usize, omg: &[f64], data: &mut [f64]) {
 11    if m < 16 {
 12        use poulpy_hal::reference::fft64::reim::fft_ref;
 13
 14        fft_ref(m, omg, data);
 15        return;
 16    }
 17
 18    assert!(data.len() == 2 * m);
 19    let (re, im) = data.split_at_mut(m);
 20
 21    if m == 16 {
 22        fft16_avx2_fma(
 23            as_arr_mut::<16, f64>(re),
 24            as_arr_mut::<16, f64>(im),
 25            as_arr::<16, f64>(omg),
 26        )
 27    } else if m <= 2048 {
 28        fft_bfs_16_avx2_fma(m, re, im, omg, 0);
 29    } else {
 30        fft_rec_16_avx2_fma(m, re, im, omg, 0);
 31    }
 32}
 33
 34unsafe extern "sysv64" {
 35    unsafe fn fft16_avx2_fma_asm(re: *mut f64, im: *mut f64, omg: *const f64);
 36}
 37
 38#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 39#[target_feature(enable = "avx2,fma")]
 40fn fft16_avx2_fma(re: &mut [f64; 16], im: &mut [f64; 16], omg: &[f64; 16]) {
 41    unsafe {
 42        fft16_avx2_fma_asm(re.as_mut_ptr(), im.as_mut_ptr(), omg.as_ptr());
 43    }
 44}
 45
 46#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 47#[target_feature(enable = "avx2,fma")]
 48fn fft_rec_16_avx2_fma(m: usize, re: &mut [f64], im: &mut [f64], omg: &[f64], mut pos: usize) -> usize {
 49    if m <= 2048 {
 50        return fft_bfs_16_avx2_fma(m, re, im, omg, pos);
 51    };
 52
 53    let h: usize = m >> 1;
 54    twiddle_fft_avx2_fma(h, re, im, *as_arr::<2, f64>(&omg[pos..]));
 55    pos += 2;
 56    pos = fft_rec_16_avx2_fma(h, re, im, omg, pos);
 57    pos = fft_rec_16_avx2_fma(h, &mut re[h..], &mut im[h..], omg, pos);
 58    pos
 59}
 60
 61#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
 62#[target_feature(enable = "avx2,fma")]
 63fn fft_bfs_16_avx2_fma(m: usize, re: &mut [f64], im: &mut [f64], omg: &[f64], mut pos: usize) -> usize {
 64    let log_m: usize = (usize::BITS - (m - 1).leading_zeros()) as usize;
 65    let mut mm: usize = m;
 66
 67    if !log_m.is_multiple_of(2) {
 68        let h: usize = mm >> 1;
 69        twiddle_fft_avx2_fma(h, re, im, *as_arr::<2, f64>(&omg[pos..]));
 70        pos += 2;
 71        mm = h
 72    }
 73
 74    while mm > 16 {
 75        let h: usize = mm >> 2;
 76        for off in (0..m).step_by(mm) {
 77            bitwiddle_fft_avx2_fma(
 78                h,
 79                &mut re[off..],
 80                &mut im[off..],
 81                as_arr::<4, f64>(&omg[pos..]),
 82            );
 83
 84            pos += 4;
 85        }
 86        mm = h
 87    }
 88
 89    for off in (0..m).step_by(16) {
 90        fft16_avx2_fma(
 91            as_arr_mut::<16, f64>(&mut re[off..]),
 92            as_arr_mut::<16, f64>(&mut im[off..]),
 93            as_arr::<16, f64>(&omg[pos..]),
 94        );
 95
 96        pos += 16;
 97    }
 98
 99    pos
100}
101
102#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
103#[target_feature(enable = "avx2,fma")]
104fn twiddle_fft_avx2_fma(h: usize, re: &mut [f64], im: &mut [f64], omg: [f64; 2]) {
105    unsafe {
106        let omx: __m128d = _mm_load_pd(omg.as_ptr());
107        let omra: __m256d = _mm256_set_m128d(omx, omx);
108        let omi: __m256d = _mm256_unpackhi_pd(omra, omra);
109        let omr: __m256d = _mm256_unpacklo_pd(omra, omra);
110        let mut r0: *mut f64 = re.as_mut_ptr();
111        let mut r1: *mut f64 = re.as_mut_ptr().add(h);
112        let mut i0: *mut f64 = im.as_mut_ptr();
113        let mut i1: *mut f64 = im.as_mut_ptr().add(h);
114
115        for _ in (0..h).step_by(4) {
116            let mut ur0: __m256d = _mm256_loadu_pd(r0);
117            let mut ur1: __m256d = _mm256_loadu_pd(r1);
118            let mut ui0: __m256d = _mm256_loadu_pd(i0);
119            let mut ui1: __m256d = _mm256_loadu_pd(i1);
120            let mut tra: __m256d = _mm256_mul_pd(omi, ui1);
121            let mut tia: __m256d = _mm256_mul_pd(omi, ur1);
122
123            tra = _mm256_fmsub_pd(omr, ur1, tra);
124            tia = _mm256_fmadd_pd(omr, ui1, tia);
125            ur1 = _mm256_sub_pd(ur0, tra);
126            ui1 = _mm256_sub_pd(ui0, tia);
127            ur0 = _mm256_add_pd(ur0, tra);
128            ui0 = _mm256_add_pd(ui0, tia);
129
130            _mm256_storeu_pd(r0, ur0);
131            _mm256_storeu_pd(r1, ur1);
132            _mm256_storeu_pd(i0, ui0);
133            _mm256_storeu_pd(i1, ui1);
134
135            r0 = r0.add(4);
136            r1 = r1.add(4);
137            i0 = i0.add(4);
138            i1 = i1.add(4);
139        }
140    }
141}
142
143#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
144#[target_feature(enable = "avx2,fma")]
145fn bitwiddle_fft_avx2_fma(h: usize, re: &mut [f64], im: &mut [f64], omg: &[f64; 4]) {
146    unsafe {
147        let mut r0: *mut f64 = re.as_mut_ptr();
148        let mut r1: *mut f64 = re.as_mut_ptr().add(h);
149        let mut r2: *mut f64 = re.as_mut_ptr().add(2 * h);
150        let mut r3: *mut f64 = re.as_mut_ptr().add(3 * h);
151        let mut i0: *mut f64 = im.as_mut_ptr();
152        let mut i1: *mut f64 = im.as_mut_ptr().add(h);
153        let mut i2: *mut f64 = im.as_mut_ptr().add(2 * h);
154        let mut i3: *mut f64 = im.as_mut_ptr().add(3 * h);
155        let om0: __m256d = _mm256_loadu_pd(omg.as_ptr());
156        let omb: __m256d = _mm256_permute2f128_pd(om0, om0, 0x11);
157        let oma: __m256d = _mm256_permute2f128_pd(om0, om0, 0x00);
158        let omai: __m256d = _mm256_unpackhi_pd(oma, oma);
159        let omar: __m256d = _mm256_unpacklo_pd(oma, oma);
160        let ombi: __m256d = _mm256_unpackhi_pd(omb, omb);
161        let ombr: __m256d = _mm256_unpacklo_pd(omb, omb);
162        for _ in (0..h).step_by(4) {
163            let mut ur0: __m256d = _mm256_loadu_pd(r0);
164            let mut ur1: __m256d = _mm256_loadu_pd(r1);
165            let mut ur2: __m256d = _mm256_loadu_pd(r2);
166            let mut ur3: __m256d = _mm256_loadu_pd(r3);
167            let mut ui0: __m256d = _mm256_loadu_pd(i0);
168            let mut ui1: __m256d = _mm256_loadu_pd(i1);
169            let mut ui2: __m256d = _mm256_loadu_pd(i2);
170            let mut ui3: __m256d = _mm256_loadu_pd(i3);
171
172            let mut tra: __m256d = _mm256_mul_pd(omai, ui2);
173            let mut trb: __m256d = _mm256_mul_pd(omai, ui3);
174            let mut tia: __m256d = _mm256_mul_pd(omai, ur2);
175            let mut tib: __m256d = _mm256_mul_pd(omai, ur3);
176            tra = _mm256_fmsub_pd(omar, ur2, tra);
177            trb = _mm256_fmsub_pd(omar, ur3, trb);
178            tia = _mm256_fmadd_pd(omar, ui2, tia);
179            tib = _mm256_fmadd_pd(omar, ui3, tib);
180            ur2 = _mm256_sub_pd(ur0, tra);
181            ur3 = _mm256_sub_pd(ur1, trb);
182            ui2 = _mm256_sub_pd(ui0, tia);
183            ui3 = _mm256_sub_pd(ui1, tib);
184            ur0 = _mm256_add_pd(ur0, tra);
185            ur1 = _mm256_add_pd(ur1, trb);
186            ui0 = _mm256_add_pd(ui0, tia);
187            ui1 = _mm256_add_pd(ui1, tib);
188
189            tra = _mm256_mul_pd(ombi, ui1);
190            trb = _mm256_mul_pd(ombr, ui3);
191            tia = _mm256_mul_pd(ombi, ur1);
192            tib = _mm256_mul_pd(ombr, ur3);
193            tra = _mm256_fmsub_pd(ombr, ur1, tra);
194            trb = _mm256_fmadd_pd(ombi, ur3, trb);
195            tia = _mm256_fmadd_pd(ombr, ui1, tia);
196            tib = _mm256_fmsub_pd(ombi, ui3, tib);
197            ur1 = _mm256_sub_pd(ur0, tra);
198            ur3 = _mm256_add_pd(ur2, trb);
199            ui1 = _mm256_sub_pd(ui0, tia);
200            ui3 = _mm256_add_pd(ui2, tib);
201            ur0 = _mm256_add_pd(ur0, tra);
202            ur2 = _mm256_sub_pd(ur2, trb);
203            ui0 = _mm256_add_pd(ui0, tia);
204            ui2 = _mm256_sub_pd(ui2, tib);
205
206            _mm256_storeu_pd(r0, ur0);
207            _mm256_storeu_pd(r1, ur1);
208            _mm256_storeu_pd(r2, ur2);
209            _mm256_storeu_pd(r3, ur3);
210            _mm256_storeu_pd(i0, ui0);
211            _mm256_storeu_pd(i1, ui1);
212            _mm256_storeu_pd(i2, ui2);
213            _mm256_storeu_pd(i3, ui3);
214
215            r0 = r0.add(4);
216            r1 = r1.add(4);
217            r2 = r2.add(4);
218            r3 = r3.add(4);
219            i0 = i0.add(4);
220            i1 = i1.add(4);
221            i2 = i2.add(4);
222            i3 = i3.add(4);
223        }
224    }
225}
226
227#[test]
228fn test_fft_avx2_fma() {
229    use super::*;
230
231    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
232    #[target_feature(enable = "avx2,fma")]
233    fn internal(log_m: usize) {
234        use poulpy_hal::reference::fft64::reim::ReimFFTRef;
235
236        let m = 1 << log_m;
237
238        let table: ReimFFTTable<f64> = ReimFFTTable::<f64>::new(m);
239
240        let mut values_0: Vec<f64> = vec![0f64; m << 1];
241        let scale: f64 = 1.0f64 / m as f64;
242        values_0
243            .iter_mut()
244            .enumerate()
245            .for_each(|(i, x)| *x = (i + 1) as f64 * scale);
246
247        let mut values_1: Vec<f64> = vec![0f64; m << 1];
248        values_1
249            .iter_mut()
250            .zip(values_0.iter())
251            .for_each(|(y, x)| *y = *x);
252
253        ReimFFTAvx::reim_dft_execute(&table, &mut values_0);
254        ReimFFTRef::reim_dft_execute(&table, &mut values_1);
255
256        let max_diff: f64 = 1.0 / ((1u64 << (53 - log_m - 1)) as f64);
257
258        for i in 0..m * 2 {
259            let diff: f64 = (values_0[i] - values_1[i]).abs();
260            assert!(
261                diff <= max_diff,
262                "{} -> {}-{} = {}",
263                i,
264                values_0[i],
265                values_1[i],
266                diff
267            )
268        }
269    }
270
271    if std::is_x86_feature_detected!("avx2") {
272        for log_m in 0..16 {
273            unsafe { internal(log_m) }
274        }
275    } else {
276        eprintln!("skipping: CPU lacks avx2");
277    }
278}