1#include <immintrin.h>
  2
  3#include "cplx_fft_internal.h"
  4#include "cplx_fft_private.h"
  5
  6typedef double D2MEM[2];
  7
  8EXPORT void cplx_fftvec_addmul_sse(const CPLX_FFTVEC_ADDMUL_PRECOMP* precomp, void* r, const void* a, const void* b) {
  9  const uint32_t m = precomp->m;
 10  const D2MEM* aa = (D2MEM*)a;
 11  const D2MEM* bb = (D2MEM*)b;
 12  D2MEM* rr = (D2MEM*)r;
 13  const D2MEM* const aend = aa + m;
 14  do {
 15    /*
 16BEGIN_TEMPLATE
 17const __m128d ari% = _mm_loadu_pd(aa[%]);
 18const __m128d bri% = _mm_loadu_pd(bb[%]);
 19const __m128d rri% = _mm_loadu_pd(rr[%]);
 20const __m128d bir% = _mm_shuffle_pd(bri%,bri%, 5);
 21const __m128d aii% = _mm_shuffle_pd(ari%,ari%, 15);
 22const __m128d pro% = _mm_fmaddsub_pd(aii%,bir%,rri%);
 23const __m128d arr% = _mm_shuffle_pd(ari%,ari%, 0);
 24const __m128d res% = _mm_fmaddsub_pd(arr%,bri%,pro%);
 25_mm_storeu_pd(rr[%],res%);
 26rr += @; // ONCE
 27aa += @; // ONCE
 28bb += @; // ONCE
 29END_TEMPLATE
 30     */
 31    // BEGIN_INTERLEAVE 2
 32    const __m128d ari0 = _mm_loadu_pd(aa[0]);
 33    const __m128d ari1 = _mm_loadu_pd(aa[1]);
 34    const __m128d bri0 = _mm_loadu_pd(bb[0]);
 35    const __m128d bri1 = _mm_loadu_pd(bb[1]);
 36    const __m128d rri0 = _mm_loadu_pd(rr[0]);
 37    const __m128d rri1 = _mm_loadu_pd(rr[1]);
 38    const __m128d bir0 = _mm_shuffle_pd(bri0, bri0, 0b01);
 39    const __m128d bir1 = _mm_shuffle_pd(bri1, bri1, 0b01);
 40    const __m128d aii0 = _mm_shuffle_pd(ari0, ari0, 0b11);
 41    const __m128d aii1 = _mm_shuffle_pd(ari1, ari1, 0b11);
 42    const __m128d pro0 = _mm_fmaddsub_pd(aii0, bir0, rri0);
 43    const __m128d pro1 = _mm_fmaddsub_pd(aii1, bir1, rri1);
 44    const __m128d arr0 = _mm_shuffle_pd(ari0, ari0, 0b00);
 45    const __m128d arr1 = _mm_shuffle_pd(ari1, ari1, 0b00);
 46    const __m128d res0 = _mm_fmaddsub_pd(arr0, bri0, pro0);
 47    const __m128d res1 = _mm_fmaddsub_pd(arr1, bri1, pro1);
 48    _mm_storeu_pd(rr[0], res0);
 49    _mm_storeu_pd(rr[1], res1);
 50    rr += 2;  // ONCE
 51    aa += 2;  // ONCE
 52    bb += 2;  // ONCE
 53              // END_INTERLEAVE
 54  } while (aa < aend);
 55}
 56
 57#if 0
 58EXPORT void cplx_fftvec_mul_fma(const CPLX_FFTVEC_MUL_PRECOMP* precomp, void* r, const void* a, const void* b) {
 59  const uint32_t m = precomp->m;
 60  const double(*aa)[4] = (double(*)[4])a;
 61  const double(*bb)[4] = (double(*)[4])b;
 62  double(*rr)[4] = (double(*)[4])r;
 63  const double(*const aend)[4] = aa + (m >> 1);
 64  do {
 65    /*
 66BEGIN_TEMPLATE
 67const __m256d ari% = _mm256_loadu_pd(aa[%]);
 68const __m256d bri% = _mm256_loadu_pd(bb[%]);
 69const __m256d bir% = _mm256_shuffle_pd(bri%,bri%, 5); // conj of b
 70const __m256d aii% = _mm256_shuffle_pd(ari%,ari%, 15); // im of a
 71const __m256d pro% = _mm256_mul_pd(aii%,bir%);
 72const __m256d arr% = _mm256_shuffle_pd(ari%,ari%, 0); // rr of a
 73const __m256d res% = _mm256_fmaddsub_pd(arr%,bri%,pro%);
 74_mm256_storeu_pd(rr[%],res%);
 75END_TEMPLATE
 76     */
 77    // BEGIN_INTERLEAVE 4
 78const __m256d ari0 = _mm256_loadu_pd(aa[0]);
 79const __m256d ari1 = _mm256_loadu_pd(aa[1]);
 80const __m256d ari2 = _mm256_loadu_pd(aa[2]);
 81const __m256d ari3 = _mm256_loadu_pd(aa[3]);
 82const __m256d bri0 = _mm256_loadu_pd(bb[0]);
 83const __m256d bri1 = _mm256_loadu_pd(bb[1]);
 84const __m256d bri2 = _mm256_loadu_pd(bb[2]);
 85const __m256d bri3 = _mm256_loadu_pd(bb[3]);
 86const __m256d bir0 = _mm256_shuffle_pd(bri0,bri0, 5); // conj of b
 87const __m256d bir1 = _mm256_shuffle_pd(bri1,bri1, 5); // conj of b
 88const __m256d bir2 = _mm256_shuffle_pd(bri2,bri2, 5); // conj of b
 89const __m256d bir3 = _mm256_shuffle_pd(bri3,bri3, 5); // conj of b
 90const __m256d aii0 = _mm256_shuffle_pd(ari0,ari0, 15); // im of a
 91const __m256d aii1 = _mm256_shuffle_pd(ari1,ari1, 15); // im of a
 92const __m256d aii2 = _mm256_shuffle_pd(ari2,ari2, 15); // im of a
 93const __m256d aii3 = _mm256_shuffle_pd(ari3,ari3, 15); // im of a
 94const __m256d pro0 = _mm256_mul_pd(aii0,bir0);
 95const __m256d pro1 = _mm256_mul_pd(aii1,bir1);
 96const __m256d pro2 = _mm256_mul_pd(aii2,bir2);
 97const __m256d pro3 = _mm256_mul_pd(aii3,bir3);
 98const __m256d arr0 = _mm256_shuffle_pd(ari0,ari0, 0); // rr of a
 99const __m256d arr1 = _mm256_shuffle_pd(ari1,ari1, 0); // rr of a
100const __m256d arr2 = _mm256_shuffle_pd(ari2,ari2, 0); // rr of a
101const __m256d arr3 = _mm256_shuffle_pd(ari3,ari3, 0); // rr of a
102const __m256d res0 = _mm256_fmaddsub_pd(arr0,bri0,pro0);
103const __m256d res1 = _mm256_fmaddsub_pd(arr1,bri1,pro1);
104const __m256d res2 = _mm256_fmaddsub_pd(arr2,bri2,pro2);
105const __m256d res3 = _mm256_fmaddsub_pd(arr3,bri3,pro3);
106_mm256_storeu_pd(rr[0],res0);
107_mm256_storeu_pd(rr[1],res1);
108_mm256_storeu_pd(rr[2],res2);
109_mm256_storeu_pd(rr[3],res3);
110    // END_INTERLEAVE
111    rr += 4;
112    aa += 4;
113    bb += 4;
114  } while (aa < aend);
115}
116
117EXPORT void cplx_fftvec_add_fma(uint32_t m, void* r, const void* a, const void* b) {
118  const double(*aa)[4] = (double(*)[4])a;
119  const double(*bb)[4] = (double(*)[4])b;
120  double(*rr)[4] = (double(*)[4])r;
121  const double(*const aend)[4] = aa + (m >> 1);
122  do {
123    /*
124BEGIN_TEMPLATE
125const __m256d ari% = _mm256_loadu_pd(aa[%]);
126const __m256d bri% = _mm256_loadu_pd(bb[%]);
127const __m256d res% = _mm256_add_pd(ari%,bri%);
128_mm256_storeu_pd(rr[%],res%);
129END_TEMPLATE
130     */
131    // BEGIN_INTERLEAVE 4
132const __m256d ari0 = _mm256_loadu_pd(aa[0]);
133const __m256d ari1 = _mm256_loadu_pd(aa[1]);
134const __m256d ari2 = _mm256_loadu_pd(aa[2]);
135const __m256d ari3 = _mm256_loadu_pd(aa[3]);
136const __m256d bri0 = _mm256_loadu_pd(bb[0]);
137const __m256d bri1 = _mm256_loadu_pd(bb[1]);
138const __m256d bri2 = _mm256_loadu_pd(bb[2]);
139const __m256d bri3 = _mm256_loadu_pd(bb[3]);
140const __m256d res0 = _mm256_add_pd(ari0,bri0);
141const __m256d res1 = _mm256_add_pd(ari1,bri1);
142const __m256d res2 = _mm256_add_pd(ari2,bri2);
143const __m256d res3 = _mm256_add_pd(ari3,bri3);
144_mm256_storeu_pd(rr[0],res0);
145_mm256_storeu_pd(rr[1],res1);
146_mm256_storeu_pd(rr[2],res2);
147_mm256_storeu_pd(rr[3],res3);
148    // END_INTERLEAVE
149    rr += 4;
150    aa += 4;
151    bb += 4;
152  } while (aa < aend);
153}
154
155EXPORT void cplx_fftvec_sub2_to_fma(uint32_t m, void* r, const void* a, const void* b) {
156  const double(*aa)[4] = (double(*)[4])a;
157  const double(*bb)[4] = (double(*)[4])b;
158  double(*rr)[4] = (double(*)[4])r;
159  const double(*const aend)[4] = aa + (m >> 1);
160  do {
161    /*
162BEGIN_TEMPLATE
163const __m256d ari% = _mm256_loadu_pd(aa[%]);
164const __m256d bri% = _mm256_loadu_pd(bb[%]);
165const __m256d sum% = _mm256_add_pd(ari%,bri%);
166const __m256d rri% = _mm256_loadu_pd(rr[%]);
167const __m256d res% = _mm256_sub_pd(rri%,sum%);
168_mm256_storeu_pd(rr[%],res%);
169END_TEMPLATE
170     */
171    // BEGIN_INTERLEAVE 4
172const __m256d ari0 = _mm256_loadu_pd(aa[0]);
173const __m256d ari1 = _mm256_loadu_pd(aa[1]);
174const __m256d ari2 = _mm256_loadu_pd(aa[2]);
175const __m256d ari3 = _mm256_loadu_pd(aa[3]);
176const __m256d bri0 = _mm256_loadu_pd(bb[0]);
177const __m256d bri1 = _mm256_loadu_pd(bb[1]);
178const __m256d bri2 = _mm256_loadu_pd(bb[2]);
179const __m256d bri3 = _mm256_loadu_pd(bb[3]);
180const __m256d sum0 = _mm256_add_pd(ari0,bri0);
181const __m256d sum1 = _mm256_add_pd(ari1,bri1);
182const __m256d sum2 = _mm256_add_pd(ari2,bri2);
183const __m256d sum3 = _mm256_add_pd(ari3,bri3);
184const __m256d rri0 = _mm256_loadu_pd(rr[0]);
185const __m256d rri1 = _mm256_loadu_pd(rr[1]);
186const __m256d rri2 = _mm256_loadu_pd(rr[2]);
187const __m256d rri3 = _mm256_loadu_pd(rr[3]);
188const __m256d res0 = _mm256_sub_pd(rri0,sum0);
189const __m256d res1 = _mm256_sub_pd(rri1,sum1);
190const __m256d res2 = _mm256_sub_pd(rri2,sum2);
191const __m256d res3 = _mm256_sub_pd(rri3,sum3);
192_mm256_storeu_pd(rr[0],res0);
193_mm256_storeu_pd(rr[1],res1);
194_mm256_storeu_pd(rr[2],res2);
195_mm256_storeu_pd(rr[3],res3);
196    // END_INTERLEAVE
197    rr += 4;
198    aa += 4;
199    bb += 4;
200  } while (aa < aend);
201}
202
203EXPORT void cplx_fftvec_copy_fma(uint32_t m, void* r, const void* a) {
204  const double(*aa)[4] = (double(*)[4])a;
205  double(*rr)[4] = (double(*)[4])r;
206  const double(*const aend)[4] = aa + (m >> 1);
207  do {
208    /*
209BEGIN_TEMPLATE
210const __m256d ari% = _mm256_loadu_pd(aa[%]);
211_mm256_storeu_pd(rr[%],ari%);
212END_TEMPLATE
213     */
214    // BEGIN_INTERLEAVE 4
215const __m256d ari0 = _mm256_loadu_pd(aa[0]);
216const __m256d ari1 = _mm256_loadu_pd(aa[1]);
217const __m256d ari2 = _mm256_loadu_pd(aa[2]);
218const __m256d ari3 = _mm256_loadu_pd(aa[3]);
219_mm256_storeu_pd(rr[0],ari0);
220_mm256_storeu_pd(rr[1],ari1);
221_mm256_storeu_pd(rr[2],ari2);
222_mm256_storeu_pd(rr[3],ari3);
223    // END_INTERLEAVE
224    rr += 4;
225    aa += 4;
226  } while (aa < aend);
227}
228
229EXPORT void cplx_fftvec_twiddle_fma(uint32_t m, void* a, void* b, const void* omg) {
230  double(*aa)[4] = (double(*)[4])a;
231  double(*bb)[4] = (double(*)[4])b;
232  const double(*const aend)[4] = aa + (m >> 1);
233  const __m256d om = _mm256_loadu_pd(omg);
234  const __m256d omrr = _mm256_shuffle_pd(om, om, 0);
235  const __m256d omii = _mm256_shuffle_pd(om, om, 15);
236  do {
237    /*
238BEGIN_TEMPLATE
239const __m256d bri% = _mm256_loadu_pd(bb[%]);
240const __m256d bir% = _mm256_shuffle_pd(bri%,bri%,5);
241__m256d p% = _mm256_mul_pd(bir%,omii);
242p% = _mm256_fmaddsub_pd(bri%,omrr,p%);
243const __m256d ari% = _mm256_loadu_pd(aa[%]);
244_mm256_storeu_pd(aa[%],_mm256_add_pd(ari%,p%));
245_mm256_storeu_pd(bb[%],_mm256_sub_pd(ari%,p%));
246END_TEMPLATE
247     */
248    // BEGIN_INTERLEAVE 4
249const __m256d bri0 = _mm256_loadu_pd(bb[0]);
250const __m256d bri1 = _mm256_loadu_pd(bb[1]);
251const __m256d bri2 = _mm256_loadu_pd(bb[2]);
252const __m256d bri3 = _mm256_loadu_pd(bb[3]);
253const __m256d bir0 = _mm256_shuffle_pd(bri0,bri0,5);
254const __m256d bir1 = _mm256_shuffle_pd(bri1,bri1,5);
255const __m256d bir2 = _mm256_shuffle_pd(bri2,bri2,5);
256const __m256d bir3 = _mm256_shuffle_pd(bri3,bri3,5);
257__m256d p0 = _mm256_mul_pd(bir0,omii);
258__m256d p1 = _mm256_mul_pd(bir1,omii);
259__m256d p2 = _mm256_mul_pd(bir2,omii);
260__m256d p3 = _mm256_mul_pd(bir3,omii);
261p0 = _mm256_fmaddsub_pd(bri0,omrr,p0);
262p1 = _mm256_fmaddsub_pd(bri1,omrr,p1);
263p2 = _mm256_fmaddsub_pd(bri2,omrr,p2);
264p3 = _mm256_fmaddsub_pd(bri3,omrr,p3);
265const __m256d ari0 = _mm256_loadu_pd(aa[0]);
266const __m256d ari1 = _mm256_loadu_pd(aa[1]);
267const __m256d ari2 = _mm256_loadu_pd(aa[2]);
268const __m256d ari3 = _mm256_loadu_pd(aa[3]);
269_mm256_storeu_pd(aa[0],_mm256_add_pd(ari0,p0));
270_mm256_storeu_pd(aa[1],_mm256_add_pd(ari1,p1));
271_mm256_storeu_pd(aa[2],_mm256_add_pd(ari2,p2));
272_mm256_storeu_pd(aa[3],_mm256_add_pd(ari3,p3));
273_mm256_storeu_pd(bb[0],_mm256_sub_pd(ari0,p0));
274_mm256_storeu_pd(bb[1],_mm256_sub_pd(ari1,p1));
275_mm256_storeu_pd(bb[2],_mm256_sub_pd(ari2,p2));
276_mm256_storeu_pd(bb[3],_mm256_sub_pd(ari3,p3));
277    // END_INTERLEAVE
278    bb += 4;
279    aa += 4;
280  } while (aa < aend);
281}
282
283EXPORT void cplx_fftvec_innerprod_avx2_fma(const CPLX_FFTVEC_INNERPROD_PRECOMP* precomp, const int32_t ellbar,
284                                   const uint64_t lda, const uint64_t ldb,
285                                   void* r, const void* a, const void* b) {
286  const uint32_t m = precomp->m;
287  const uint32_t blk = precomp->blk;
288  const uint32_t nblocks = precomp->nblocks;
289  const CPLX* aa = (CPLX*)a;
290  const CPLX* bb = (CPLX*)b;
291  CPLX* rr = (CPLX*)r;
292  const uint64_t ldda = lda >> 4;  // in CPLX
293  const uint64_t lddb = ldb >> 4;
294  if (m==0) {
295    memset(r, 0, m*sizeof(CPLX));
296    return;
297  }
298  for (uint32_t k=0; k<nblocks; ++k) {
299    const uint64_t offset = k*blk;
300    const CPLX* aaa = aa+offset;
301    const CPLX* bbb = bb+offset;
302    CPLX *rrr = rr+offset;
303    cplx_fftvec_mul_fma(&precomp->mul_func, rrr, aaa, bbb);
304    for (int32_t i=1; i<ellbar; ++i) {
305      cplx_fftvec_addmul_fma(&precomp->addmul_func, rrr, aaa + i * ldda, bbb + i * lddb);
306    }
307  }
308}
309#endif