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}