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