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