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