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