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}