1#include <immintrin.h>
  2#include <string.h>
  3
  4#include "cplx_fft_internal.h"
  5#include "cplx_fft_private.h"
  6
  7typedef double D4MEM[4];
  8
  9/**
 10 * @brief complex fft via bfs strategy (for m between 2 and 8)
 11 * @param dat the data to run the algorithm on
 12 * @param omg precomputed tables (must have been filled with fill_omega)
 13 * @param m ring dimension of the FFT (modulo X^m-i)
 14 */
 15void cplx_fft_avx2_fma_bfs_2(D4MEM* dat, const D4MEM** omg, uint32_t m) {
 16  double* data = (double*)dat;
 17  int32_t _2nblock = m >> 1;  // = h in ref code
 18  D4MEM* const finaldd = (D4MEM*)(data + 2 * m);
 19  while (_2nblock >= 2) {
 20    int32_t nblock = _2nblock >> 1;  // =h/2 in ref code
 21    D4MEM* dd = (D4MEM*)data;
 22    do {
 23      const __m256d om = _mm256_load_pd(*omg[0]);
 24      const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
 25      const __m256d omre = _mm256_unpacklo_pd(om, om);
 26      D4MEM* const ddend = (dd + nblock);
 27      D4MEM* ddmid = ddend;
 28      do {
 29        const __m256d b = _mm256_loadu_pd(ddmid[0]);
 30        const __m256d t1 = _mm256_mul_pd(b, omre);
 31        const __m256d barb = _mm256_shuffle_pd(b, b, 5);
 32        const __m256d t2 = _mm256_fmadd_pd(barb, omim, t1);
 33        const __m256d a = _mm256_loadu_pd(dd[0]);
 34        const __m256d newa = _mm256_add_pd(a, t2);
 35        const __m256d newb = _mm256_sub_pd(a, t2);
 36        _mm256_storeu_pd(dd[0], newa);
 37        _mm256_storeu_pd(ddmid[0], newb);
 38        dd += 1;
 39        ddmid += 1;
 40      } while (dd < ddend);
 41      dd += nblock;
 42      *omg += 1;
 43    } while (dd < finaldd);
 44    _2nblock >>= 1;
 45  }
 46  // last iteration when _2nblock == 1
 47  {
 48    D4MEM* dd = (D4MEM*)data;
 49    do {
 50      const __m256d om = _mm256_load_pd(*omg[0]);
 51      const __m256d omre = _mm256_unpacklo_pd(om, om);
 52      const __m256d omim = _mm256_unpackhi_pd(om, om);
 53      const __m256d ab = _mm256_loadu_pd(dd[0]);
 54      const __m256d bb = _mm256_permute4x64_pd(ab, 0b11101110);
 55      const __m256d bbbar = _mm256_permute4x64_pd(ab, 0b10111011);
 56      const __m256d t1 = _mm256_mul_pd(bbbar, omim);
 57      const __m256d t2 = _mm256_fmaddsub_pd(bb, omre, t1);
 58      const __m256d aa = _mm256_permute4x64_pd(ab, 0b01000100);
 59      const __m256d newab = _mm256_add_pd(aa, t2);
 60      _mm256_storeu_pd(dd[0], newab);
 61      dd += 1;
 62      *omg += 1;
 63    } while (dd < finaldd);
 64  }
 65}
 66
 67__always_inline void cplx_twiddle_fft_avx2(int32_t h, D4MEM* data, const void* omg) {
 68  const __m256d om = _mm256_loadu_pd(omg);
 69  const __m256d omim = _mm256_unpackhi_pd(om, om);
 70  const __m256d omre = _mm256_unpacklo_pd(om, om);
 71  D4MEM* d0 = data;
 72  D4MEM* const ddend = d0 + (h >> 1);
 73  D4MEM* d1 = ddend;
 74  do {
 75    const __m256d b = _mm256_loadu_pd(d1[0]);
 76    const __m256d barb = _mm256_shuffle_pd(b, b, 5);
 77    const __m256d t1 = _mm256_mul_pd(barb, omim);
 78    const __m256d t2 = _mm256_fmaddsub_pd(b, omre, t1);
 79    const __m256d a = _mm256_loadu_pd(d0[0]);
 80    const __m256d newa = _mm256_add_pd(a, t2);
 81    const __m256d newb = _mm256_sub_pd(a, t2);
 82    _mm256_storeu_pd(d0[0], newa);
 83    _mm256_storeu_pd(d1[0], newb);
 84    d0 += 1;
 85    d1 += 1;
 86  } while (d0 < ddend);
 87}
 88
 89__always_inline void cplx_bitwiddle_fft_avx2(int32_t h, void* data, const void* powom) {
 90  const __m256d omx = _mm256_loadu_pd(powom);
 91  const __m256d oma = _mm256_permute2f128_pd(omx, omx, 0x00);
 92  const __m256d omb = _mm256_permute2f128_pd(omx, omx, 0x11);
 93  const __m256d omaim = _mm256_unpackhi_pd(oma, oma);
 94  const __m256d omare = _mm256_unpacklo_pd(oma, oma);
 95  const __m256d ombim = _mm256_unpackhi_pd(omb, omb);
 96  const __m256d ombre = _mm256_unpacklo_pd(omb, omb);
 97  D4MEM* d0 = (D4MEM*)data;
 98  D4MEM* const ddend = d0 + (h >> 1);
 99  D4MEM* d1 = ddend;
100  D4MEM* d2 = d0 + h;
101  D4MEM* d3 = d1 + h;
102  __m256d reg0, reg1, reg2, reg3, tmp0, tmp1;
103  do {
104    reg0 = _mm256_loadu_pd(d0[0]);
105    reg1 = _mm256_loadu_pd(d1[0]);
106    reg2 = _mm256_loadu_pd(d2[0]);
107    reg3 = _mm256_loadu_pd(d3[0]);
108    tmp0 = _mm256_shuffle_pd(reg2, reg2, 5);
109    tmp1 = _mm256_shuffle_pd(reg3, reg3, 5);
110    tmp0 = _mm256_mul_pd(tmp0, omaim);
111    tmp1 = _mm256_mul_pd(tmp1, omaim);
112    tmp0 = _mm256_fmaddsub_pd(reg2, omare, tmp0);
113    tmp1 = _mm256_fmaddsub_pd(reg3, omare, tmp1);
114    reg2 = _mm256_sub_pd(reg0, tmp0);
115    reg3 = _mm256_sub_pd(reg1, tmp1);
116    reg0 = _mm256_add_pd(reg0, tmp0);
117    reg1 = _mm256_add_pd(reg1, tmp1);
118    //--------------------------------------
119    tmp0 = _mm256_shuffle_pd(reg1, reg1, 5);
120    tmp1 = _mm256_shuffle_pd(reg3, reg3, 5);
121    tmp0 = _mm256_mul_pd(tmp0, ombim);  //(r,i)
122    tmp1 = _mm256_mul_pd(tmp1, ombre);  //(-i,r)
123    tmp0 = _mm256_fmaddsub_pd(reg1, ombre, tmp0);
124    tmp1 = _mm256_fmsubadd_pd(reg3, ombim, tmp1);
125    reg1 = _mm256_sub_pd(reg0, tmp0);
126    reg3 = _mm256_add_pd(reg2, tmp1);
127    reg0 = _mm256_add_pd(reg0, tmp0);
128    reg2 = _mm256_sub_pd(reg2, tmp1);
129    /////
130    _mm256_storeu_pd(d0[0], reg0);
131    _mm256_storeu_pd(d1[0], reg1);
132    _mm256_storeu_pd(d2[0], reg2);
133    _mm256_storeu_pd(d3[0], reg3);
134    d0 += 1;
135    d1 += 1;
136    d2 += 1;
137    d3 += 1;
138  } while (d0 < ddend);
139}
140
141/**
142 * @brief complex fft via bfs strategy (for m >= 16)
143 * @param dat the data to run the algorithm on
144 * @param omg precomputed tables (must have been filled with fill_omega)
145 * @param m ring dimension of the FFT (modulo X^m-i)
146 */
147void cplx_fft_avx2_fma_bfs_16(D4MEM* dat, const D4MEM** omg, uint32_t m) {
148  double* data = (double*)dat;
149  D4MEM* const finaldd = (D4MEM*)(data + 2 * m);
150  uint32_t mm = m;
151  uint32_t log2m = _mm_popcnt_u32(m - 1);  // log2(m)
152  if (log2m % 2 == 1) {
153    uint32_t h = mm >> 1;
154    cplx_twiddle_fft_avx2(h, dat, **omg);
155    *omg += 1;
156    mm >>= 1;
157  }
158  while (mm > 16) {
159    uint32_t h = mm / 4;
160    for (CPLX* d = (CPLX*)data; d < (CPLX*)finaldd; d += mm) {
161      cplx_bitwiddle_fft_avx2(h, d, (CPLX*)*omg);
162      *omg += 1;
163    }
164    mm = h;
165  }
166  {
167    D4MEM* dd = (D4MEM*)data;
168    do {
169      cplx_fft16_avx_fma(dd, *omg);
170      dd += 8;
171      *omg += 4;
172    } while (dd < finaldd);
173    _mm256_zeroupper();
174  }
175  /*
176  int32_t _2nblock = m >> 1;  // = h in ref code
177  while (_2nblock >= 16) {
178    int32_t nblock = _2nblock >> 1;  // =h/2 in ref code
179    D4MEM* dd = (D4MEM*)data;
180    do {
181      const __m256d om = _mm256_load_pd(*omg[0]);
182      const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
183      const __m256d omre = _mm256_unpacklo_pd(om, om);
184      D4MEM* const ddend = (dd + nblock);
185      D4MEM* ddmid = ddend;
186      do {
187        const __m256d b = _mm256_loadu_pd(ddmid[0]);
188        const __m256d t1 = _mm256_mul_pd(b, omre);
189        const __m256d barb = _mm256_shuffle_pd(b, b, 5);
190        const __m256d t2 = _mm256_fmadd_pd(barb, omim, t1);
191        const __m256d a = _mm256_loadu_pd(dd[0]);
192        const __m256d newa = _mm256_add_pd(a, t2);
193        const __m256d newb = _mm256_sub_pd(a, t2);
194        _mm256_storeu_pd(dd[0], newa);
195        _mm256_storeu_pd(ddmid[0], newb);
196        dd += 1;
197        ddmid += 1;
198      } while (dd < ddend);
199      dd += nblock;
200      *omg += 1;
201    } while (dd < finaldd);
202    _2nblock >>= 1;
203  }
204  // last iteration when _2nblock == 8
205  {
206    D4MEM* dd = (D4MEM*)data;
207    do {
208      cplx_fft16_avx_fma(dd, *omg);
209      dd += 8;
210      *omg += 4;
211    } while (dd < finaldd);
212    _mm256_zeroupper();
213  }
214   */
215}
216
217/**
218 * @brief complex fft via dfs recursion (for m >= 16)
219 * @param dat the data to run the algorithm on
220 * @param omg precomputed tables (must have been filled with fill_omega)
221 * @param m ring dimension of the FFT (modulo X^m-i)
222 */
223void cplx_fft_avx2_fma_rec_16(D4MEM* dat, const D4MEM** omg, uint32_t m) {
224  if (m <= 8) return cplx_fft_avx2_fma_bfs_2(dat, omg, m);
225  if (m <= 2048) return cplx_fft_avx2_fma_bfs_16(dat, omg, m);
226  double* data = (double*)dat;
227  int32_t _2nblock = m >> 1;       // = h in ref code
228  int32_t nblock = _2nblock >> 1;  // =h/2 in ref code
229  D4MEM* dd = (D4MEM*)data;
230  const __m256d om = _mm256_load_pd(*omg[0]);
231  const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
232  const __m256d omre = _mm256_unpacklo_pd(om, om);
233  D4MEM* const ddend = (dd + nblock);
234  D4MEM* ddmid = ddend;
235  do {
236    const __m256d b = _mm256_loadu_pd(ddmid[0]);
237    const __m256d t1 = _mm256_mul_pd(b, omre);
238    const __m256d barb = _mm256_shuffle_pd(b, b, 5);
239    const __m256d t2 = _mm256_fmadd_pd(barb, omim, t1);
240    const __m256d a = _mm256_loadu_pd(dd[0]);
241    const __m256d newa = _mm256_add_pd(a, t2);
242    const __m256d newb = _mm256_sub_pd(a, t2);
243    _mm256_storeu_pd(dd[0], newa);
244    _mm256_storeu_pd(ddmid[0], newb);
245    dd += 1;
246    ddmid += 1;
247  } while (dd < ddend);
248  *omg += 1;
249  cplx_fft_avx2_fma_rec_16(dat, omg, _2nblock);
250  cplx_fft_avx2_fma_rec_16(ddend, omg, _2nblock);
251}
252
253/**
254 * @brief complex fft via best strategy (for m>=1)
255 * @param dat the data to run the algorithm on: m complex numbers
256 * @param omg precomputed tables (must have been filled with fill_omega)
257 * @param m ring dimension of the FFT (modulo X^m-i)
258 */
259EXPORT void cplx_fft_avx2_fma(const CPLX_FFT_PRECOMP* precomp, void* d) {
260  const uint32_t m = precomp->m;
261  const D4MEM* omg = (D4MEM*)precomp->powomegas;
262  if (m <= 1) return;
263  if (m <= 8) return cplx_fft_avx2_fma_bfs_2(d, &omg, m);
264  if (m <= 2048) return cplx_fft_avx2_fma_bfs_16(d, &omg, m);
265  cplx_fft_avx2_fma_rec_16(d, &omg, m);
266}