arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft8_avx_fma.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft8_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(__m256d* ra, __m256d* rb, __m256d* ia, __m256d* ib, const __m256d omre,
 8                                           const __m256d omim) {
 9  // rb * omre - ib * omim;
10  __m256d rprod0 = _mm256_mul_pd(*ib, omim);
11  rprod0 = _mm256_fmsub_pd(*rb, omre, rprod0);
12
13  // rb * omim + ib * omre;
14  __m256d iprod0 = _mm256_mul_pd(*rb, omim);
15  iprod0 = _mm256_fmadd_pd(*ib, omre, iprod0);
16
17  *rb = _mm256_sub_pd(*ra, rprod0);
18  *ib = _mm256_sub_pd(*ia, iprod0);
19  *ra = _mm256_add_pd(*ra, rprod0);
20  *ia = _mm256_add_pd(*ia, iprod0);
21}
22
23EXPORT void reim_fft8_avx_fma(double* dre, double* dim, const void* ompv) {
24  const double* omp = (const double*)ompv;
25
26  __m256d ra0 = _mm256_loadu_pd(dre);
27  __m256d ra4 = _mm256_loadu_pd(dre + 4);
28  __m256d ia0 = _mm256_loadu_pd(dim);
29  __m256d ia4 = _mm256_loadu_pd(dim + 4);
30
31  // 1
32  {
33    // duplicate omegas in precomp?
34    __m128d omri = _mm_loadu_pd(omp);
35    __m256d omriri = _mm256_set_m128d(omri, omri);
36    __m256d omi = _mm256_permute_pd(omriri, 15);
37    __m256d omr = _mm256_permute_pd(omriri, 0);
38
39    reim_ctwiddle_avx_fma(&ra0, &ra4, &ia0, &ia4, omr, omi);
40  }
41
42  // 2
43  {
44    const __m128d fft8neg = _mm_castsi128_pd(_mm_set_epi64x(UINT64_C(1) << 63, 0));
45    __m128d omri = _mm_loadu_pd(omp + 2);             // r,i
46    __m128d omrmi = _mm_xor_pd(omri, fft8neg);        // r,-i
47    __m256d omrirmi = _mm256_set_m128d(omrmi, omri);  // r,i,r,-i
48    __m256d omi = _mm256_permute_pd(omrirmi, 3);      // i,i,r,r
49    __m256d omr = _mm256_permute_pd(omrirmi, 12);     // r,r,-i,-i
50
51    __m256d rb = _mm256_permute2f128_pd(ra0, ra4, 0x31);
52    __m256d ib = _mm256_permute2f128_pd(ia0, ia4, 0x31);
53    __m256d ra = _mm256_permute2f128_pd(ra0, ra4, 0x20);
54    __m256d ia = _mm256_permute2f128_pd(ia0, ia4, 0x20);
55
56    reim_ctwiddle_avx_fma(&ra, &rb, &ia, &ib, omr, omi);
57
58    ra0 = _mm256_permute2f128_pd(ra, rb, 0x20);
59    ra4 = _mm256_permute2f128_pd(ra, rb, 0x31);
60    ia0 = _mm256_permute2f128_pd(ia, ib, 0x20);
61    ia4 = _mm256_permute2f128_pd(ia, ib, 0x31);
62  }
63
64  // 3
65  {
66    const __m256d fft8neg2 = _mm256_castsi256_pd(_mm256_set_epi64x(UINT64_C(1) << 63, UINT64_C(1) << 63, 0, 0));
67    __m256d om = _mm256_loadu_pd(omp + 4);            // r0,r1,i0,i1
68    __m256d omi = _mm256_permute2f128_pd(om, om, 1);  // i0,i1,r0,r1
69    __m256d omr = _mm256_xor_pd(om, fft8neg2);        // r0,r1,-i0,-i1
70
71    __m256d rb = _mm256_unpackhi_pd(ra0, ra4);
72    __m256d ib = _mm256_unpackhi_pd(ia0, ia4);
73    __m256d ra = _mm256_unpacklo_pd(ra0, ra4);
74    __m256d ia = _mm256_unpacklo_pd(ia0, ia4);
75
76    reim_ctwiddle_avx_fma(&ra, &rb, &ia, &ib, omr, omi);
77
78    ra4 = _mm256_unpackhi_pd(ra, rb);
79    ia4 = _mm256_unpackhi_pd(ia, ib);
80    ra0 = _mm256_unpacklo_pd(ra, rb);
81    ia0 = _mm256_unpacklo_pd(ia, ib);
82  }
83
84  // 4
85  _mm256_storeu_pd(dre, ra0);
86  _mm256_storeu_pd(dre + 4, ra4);
87  _mm256_storeu_pd(dim, ia0);
88  _mm256_storeu_pd(dim + 4, ia4);
89}