1#include "immintrin.h"
  2#include "reim_fft_internal.h"
  3#include "reim_fft_private.h"
  4
  5void reim_invtwiddle_ifft_ref(uint64_t h, double* re, double* im, double om[2]);
  6void reim_invbitwiddle_ifft_ref(uint64_t h, double* re, double* im, double om[4]);
  7
  8__always_inline void reim_invtwiddle_ifft_avx2_fma(uint32_t h, double* re, double* im, double om[2]) {
  9  const __m128d omx = _mm_load_pd(om);
 10  const __m256d omra = _mm256_set_m128d(omx, omx);
 11  const __m256d omi = _mm256_unpackhi_pd(omra, omra);
 12  const __m256d omr = _mm256_unpacklo_pd(omra, omra);
 13  double* r0 = re;
 14  double* r1 = re + h;
 15  double* i0 = im;
 16  double* i1 = im + h;
 17  for (uint32_t i = 0; i < h; i += 4) {
 18    __m256d ur0 = _mm256_loadu_pd(r0 + i);
 19    __m256d ur1 = _mm256_loadu_pd(r1 + i);
 20    __m256d ui0 = _mm256_loadu_pd(i0 + i);
 21    __m256d ui1 = _mm256_loadu_pd(i1 + i);
 22    __m256d tra = _mm256_sub_pd(ur0, ur1);
 23    __m256d tia = _mm256_sub_pd(ui0, ui1);
 24    ur0 = _mm256_add_pd(ur0, ur1);
 25    ui0 = _mm256_add_pd(ui0, ui1);
 26    ur1 = _mm256_mul_pd(omi, tia);
 27    ui1 = _mm256_mul_pd(omi, tra);
 28    ur1 = _mm256_fmsub_pd(omr, tra, ur1);
 29    ui1 = _mm256_fmadd_pd(omr, tia, ui1);
 30    _mm256_storeu_pd(r0 + i, ur0);
 31    _mm256_storeu_pd(r1 + i, ur1);
 32    _mm256_storeu_pd(i0 + i, ui0);
 33    _mm256_storeu_pd(i1 + i, ui1);
 34  }
 35}
 36
 37__always_inline void reim_invbitwiddle_ifft_avx2_fma(uint32_t h, double* re, double* im, double om[4]) {
 38  // reim_invbitwiddle_ifft_ref(h,re,im, om);
 39  double* const r0 = re;
 40  double* const r1 = re + h;
 41  double* const r2 = re + 2 * h;
 42  double* const r3 = re + 3 * h;
 43  double* const i0 = im;
 44  double* const i1 = im + h;
 45  double* const i2 = im + 2 * h;
 46  double* const i3 = im + 3 * h;
 47  const __m256d om0 = _mm256_loadu_pd(om);
 48  const __m256d omb = _mm256_permute2f128_pd(om0, om0, 0x11);
 49  const __m256d oma = _mm256_permute2f128_pd(om0, om0, 0x00);
 50  const __m256d omai = _mm256_unpackhi_pd(oma, oma);
 51  const __m256d omar = _mm256_unpacklo_pd(oma, oma);
 52  const __m256d ombi = _mm256_unpackhi_pd(omb, omb);
 53  const __m256d ombr = _mm256_unpacklo_pd(omb, omb);
 54  for (uint32_t i = 0; i < h; i += 4) {
 55    __m256d ur0 = _mm256_loadu_pd(r0 + i);
 56    __m256d ur1 = _mm256_loadu_pd(r1 + i);
 57    __m256d ur2 = _mm256_loadu_pd(r2 + i);
 58    __m256d ur3 = _mm256_loadu_pd(r3 + i);
 59    __m256d ui0 = _mm256_loadu_pd(i0 + i);
 60    __m256d ui1 = _mm256_loadu_pd(i1 + i);
 61    __m256d ui2 = _mm256_loadu_pd(i2 + i);
 62    __m256d ui3 = _mm256_loadu_pd(i3 + i);
 63    __m256d tra, trb, tia, tib;
 64    //------ twiddles 2
 65    tra = _mm256_sub_pd(ur0, ur1);
 66    trb = _mm256_sub_pd(ur2, ur3);
 67    tia = _mm256_sub_pd(ui0, ui1);
 68    tib = _mm256_sub_pd(ui2, ui3);
 69    ur0 = _mm256_add_pd(ur0, ur1);
 70    ur2 = _mm256_add_pd(ur2, ur3);
 71    ui0 = _mm256_add_pd(ui0, ui1);
 72    ui2 = _mm256_add_pd(ui2, ui3);
 73    ur1 = _mm256_mul_pd(omai, tia);
 74    ur3 = _mm256_mul_pd(omar, tib);
 75    ui1 = _mm256_mul_pd(omai, tra);
 76    ui3 = _mm256_mul_pd(omar, trb);
 77    ur1 = _mm256_fmsub_pd(omar, tra, ur1);
 78    ur3 = _mm256_fmadd_pd(omai, trb, ur3);
 79    ui1 = _mm256_fmadd_pd(omar, tia, ui1);
 80    ui3 = _mm256_fmsub_pd(omai, tib, ui3);
 81    //------ twiddles 1
 82    tra = _mm256_sub_pd(ur0, ur2);
 83    trb = _mm256_sub_pd(ur1, ur3);
 84    tia = _mm256_sub_pd(ui0, ui2);
 85    tib = _mm256_sub_pd(ui1, ui3);
 86    ur0 = _mm256_add_pd(ur0, ur2);
 87    ur1 = _mm256_add_pd(ur1, ur3);
 88    ui0 = _mm256_add_pd(ui0, ui2);
 89    ui1 = _mm256_add_pd(ui1, ui3);
 90    ur2 = _mm256_mul_pd(ombi, tia);
 91    ur3 = _mm256_mul_pd(ombi, tib);
 92    ui2 = _mm256_mul_pd(ombi, tra);
 93    ui3 = _mm256_mul_pd(ombi, trb);
 94    ur2 = _mm256_fmsub_pd(ombr, tra, ur2);
 95    ur3 = _mm256_fmsub_pd(ombr, trb, ur3);
 96    ui2 = _mm256_fmadd_pd(ombr, tia, ui2);
 97    ui3 = _mm256_fmadd_pd(ombr, tib, ui3);
 98    ///---
 99    _mm256_storeu_pd(r0 + i, ur0);
100    _mm256_storeu_pd(r1 + i, ur1);
101    _mm256_storeu_pd(r2 + i, ur2);
102    _mm256_storeu_pd(r3 + i, ur3);
103    _mm256_storeu_pd(i0 + i, ui0);
104    _mm256_storeu_pd(i1 + i, ui1);
105    _mm256_storeu_pd(i2 + i, ui2);
106    _mm256_storeu_pd(i3 + i, ui3);
107  }
108}
109
110void reim_ifft_bfs_16_avx2_fma(uint32_t m, double* re, double* im, double** omg) {
111  uint32_t log2m = _mm_popcnt_u32(m - 1);  // log2(m);
112  for (uint32_t off = 0; off < m; off += 16) {
113    reim_ifft16_avx_fma(re + off, im + off, *omg);
114    *omg += 16;
115  }
116  uint32_t h = 16;
117  uint32_t ms2 = m >> 1;  // m/2
118  while (h < ms2) {
119    uint32_t mm = h << 2;
120    for (uint32_t off = 0; off < m; off += mm) {
121      reim_invbitwiddle_ifft_avx2_fma(h, re + off, im + off, *omg);
122      *omg += 4;
123    }
124    h = mm;
125  }
126  if (log2m & 1) {
127    // if (h!=ms2) abort(); // bug
128    //  do the first twiddle iteration normally
129    reim_invtwiddle_ifft_avx2_fma(h, re, im, *omg);
130    *omg += 2;
131    // h = m;
132  }
133}
134
135void reim_ifft_rec_16_avx2_fma(uint32_t m, double* re, double* im, double** omg) {
136  if (m <= 2048) return reim_ifft_bfs_16_avx2_fma(m, re, im, omg);
137  const uint32_t h = m >> 1;  // m / 2;
138  reim_ifft_rec_16_avx2_fma(h, re, im, omg);
139  reim_ifft_rec_16_avx2_fma(h, re + h, im + h, omg);
140  reim_invtwiddle_ifft_avx2_fma(h, re, im, *omg);
141  *omg += 2;
142}
143
144void reim_ifft_avx2_fma(const REIM_IFFT_PRECOMP* precomp, double* dat) {
145  const int32_t m = precomp->m;
146  double* omg = precomp->powomegas;
147  double* re = dat;
148  double* im = dat + m;
149  if (m <= 16) {
150    switch (m) {
151      case 1:
152        return;
153      case 2:
154        return reim_ifft2_ref(re, im, omg);
155      case 4:
156        return reim_ifft4_avx_fma(re, im, omg);
157      case 8:
158        return reim_ifft8_avx_fma(re, im, omg);
159      case 16:
160        return reim_ifft16_avx_fma(re, im, omg);
161      default:
162        abort();  // m is not a power of 2
163    }
164  }
165  if (m <= 2048) return reim_ifft_bfs_16_avx2_fma(m, re, im, &omg);
166  return reim_ifft_rec_16_avx2_fma(m, re, im, &omg);
167}