arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft4_avx_fma.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft4_avx_fma.cRaw
 1#include <assert.h>
 2#include <immintrin.h>
 3#include <stdint.h>
 4
 5#include "reim_fft_private.h"
 6
 7__always_inline void reim_ctwiddle_avx_fma(__m128d* ra, __m128d* rb, __m128d* ia, __m128d* ib, const __m128d omre,
 8                                           const __m128d omim) {
 9  // rb * omre - ib * omim;
10  __m128d rprod0 = _mm_mul_pd(*ib, omim);
11  rprod0 = _mm_fmsub_pd(*rb, omre, rprod0);
12
13  // rb * omim + ib * omre;
14  __m128d iprod0 = _mm_mul_pd(*rb, omim);
15  iprod0 = _mm_fmadd_pd(*ib, omre, iprod0);
16
17  *rb = _mm_sub_pd(*ra, rprod0);
18  *ib = _mm_sub_pd(*ia, iprod0);
19  *ra = _mm_add_pd(*ra, rprod0);
20  *ia = _mm_add_pd(*ia, iprod0);
21}
22
23EXPORT void reim_fft4_avx_fma(double* dre, double* dim, const void* ompv) {
24  const double* omp = (const double*)ompv;
25
26  __m128d ra0 = _mm_loadu_pd(dre);
27  __m128d ra2 = _mm_loadu_pd(dre + 2);
28  __m128d ia0 = _mm_loadu_pd(dim);
29  __m128d ia2 = _mm_loadu_pd(dim + 2);
30
31  // 1
32  {
33    // duplicate omegas in precomp?
34    __m128d om = _mm_loadu_pd(omp);
35    __m128d omre = _mm_permute_pd(om, 0);
36    __m128d omim = _mm_permute_pd(om, 3);
37
38    reim_ctwiddle_avx_fma(&ra0, &ra2, &ia0, &ia2, omre, omim);
39  }
40
41  // 2
42  {
43    const __m128d fft4neg = _mm_castsi128_pd(_mm_set_epi64x(UINT64_C(1) << 63, 0));
44    __m128d om = _mm_loadu_pd(omp + 2);      // om: r,i
45    __m128d omim = _mm_permute_pd(om, 1);    // omim: i,r
46    __m128d omre = _mm_xor_pd(om, fft4neg);  // omre: r,-i
47
48    __m128d rb = _mm_unpackhi_pd(ra0, ra2);  // (r0, r1), (r2, r3) -> (r1, r3)
49    __m128d ib = _mm_unpackhi_pd(ia0, ia2);  // (i0, i1), (i2, i3) -> (i1, i3)
50    __m128d ra = _mm_unpacklo_pd(ra0, ra2);  // (r0, r1), (r2, r3) -> (r0, r2)
51    __m128d ia = _mm_unpacklo_pd(ia0, ia2);  // (i0, i1), (i2, i3) -> (i0, i2)
52
53    reim_ctwiddle_avx_fma(&ra, &rb, &ia, &ib, omre, omim);
54
55    ra0 = _mm_unpacklo_pd(ra, rb);
56    ia0 = _mm_unpacklo_pd(ia, ib);
57    ra2 = _mm_unpackhi_pd(ra, rb);
58    ia2 = _mm_unpackhi_pd(ia, ib);
59  }
60
61  // 4
62  _mm_storeu_pd(dre, ra0);
63  _mm_storeu_pd(dre + 2, ra2);
64  _mm_storeu_pd(dim, ia0);
65  _mm_storeu_pd(dim + 2, ia2);
66}