poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_conversions_avx2_fma.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_conversions_avx2_fma.cRaw
1#include <immintrin.h>
2#include <stdio.h>
3
4#include "cplx_fft_internal.h"
5#include "cplx_fft_private.h"
6
7typedef int32_t I8MEM[8];
8typedef double D4MEM[4];
9
10__always_inline void cplx_from_any_fma(uint64_t m, void* r, const int32_t* x, const __m256i C, const __m256d R) {
11 const __m256i S = _mm256_set1_epi32(0x80000000);
12 const I8MEM* inre = (I8MEM*)(x);
13 const I8MEM* inim = (I8MEM*)(x + m);
14 D4MEM* out = (D4MEM*)r;
15 const uint64_t ms8 = m / 8;
16 for (uint32_t i = 0; i < ms8; ++i) {
17 __m256i rea = _mm256_loadu_si256((__m256i*)inre[0]);
18 __m256i ima = _mm256_loadu_si256((__m256i*)inim[0]);
19 rea = _mm256_add_epi32(rea, S);
20 ima = _mm256_add_epi32(ima, S);
21 __m256i tmpa = _mm256_unpacklo_epi32(rea, ima);
22 __m256i tmpc = _mm256_unpackhi_epi32(rea, ima);
23 __m256i cpla = _mm256_permute2x128_si256(tmpa, tmpc, 0x20);
24 __m256i cplc = _mm256_permute2x128_si256(tmpa, tmpc, 0x31);
25 tmpa = _mm256_unpacklo_epi32(cpla, C);
26 __m256i tmpb = _mm256_unpackhi_epi32(cpla, C);
27 tmpc = _mm256_unpacklo_epi32(cplc, C);
28 __m256i tmpd = _mm256_unpackhi_epi32(cplc, C);
29 cpla = _mm256_permute2x128_si256(tmpa, tmpb, 0x20);
30 __m256i cplb = _mm256_permute2x128_si256(tmpa, tmpb, 0x31);
31 cplc = _mm256_permute2x128_si256(tmpc, tmpd, 0x20);
32 __m256i cpld = _mm256_permute2x128_si256(tmpc, tmpd, 0x31);
33 __m256d dcpla = _mm256_sub_pd(_mm256_castsi256_pd(cpla), R);
34 __m256d dcplb = _mm256_sub_pd(_mm256_castsi256_pd(cplb), R);
35 __m256d dcplc = _mm256_sub_pd(_mm256_castsi256_pd(cplc), R);
36 __m256d dcpld = _mm256_sub_pd(_mm256_castsi256_pd(cpld), R);
37 _mm256_storeu_pd(out[0], dcpla);
38 _mm256_storeu_pd(out[1], dcplb);
39 _mm256_storeu_pd(out[2], dcplc);
40 _mm256_storeu_pd(out[3], dcpld);
41 inre += 1;
42 inim += 1;
43 out += 4;
44 }
45}
46
47EXPORT void cplx_from_znx32_avx2_fma(const CPLX_FROM_ZNX32_PRECOMP* precomp, void* r, const int32_t* x) {
48 // note: the hex code of 2^31 + 2^52 is 0x4330000080000000
49 const __m256i C = _mm256_set1_epi32(0x43300000);
50 const __m256d R = _mm256_set1_pd((INT64_C(1) << 31) + (INT64_C(1) << 52));
51 // double XX = INT64_C(1) + (INT64_C(1)<<31) + (INT64_C(1)<<52);
52 // printf("\n\n%016lx\n", *(uint64_t*)&XX);
53 // abort();
54 const uint64_t m = precomp->m;
55 cplx_from_any_fma(m, r, x, C, R);
56}
57
58EXPORT void cplx_from_tnx32_avx2_fma(const CPLX_FROM_TNX32_PRECOMP* precomp, void* r, const int32_t* x) {
59 // note: the hex code of 2^-1 + 2^30 is 0x4130000080000000
60 const __m256i C = _mm256_set1_epi32(0x41300000);
61 const __m256d R = _mm256_set1_pd(0.5 + (INT64_C(1) << 20));
62 // double XX = (double)(INT64_C(1) + (INT64_C(1)<<31) + (INT64_C(1)<<52))/(INT64_C(1)<<32);
63 // printf("\n\n%016lx\n", *(uint64_t*)&XX);
64 // abort();
65 const uint64_t m = precomp->m;
66 cplx_from_any_fma(m, r, x, C, R);
67}
68
69EXPORT void cplx_to_tnx32_avx2_fma(const CPLX_TO_TNX32_PRECOMP* precomp, int32_t* r, const void* x) {
70 const __m256d R = _mm256_set1_pd((0.5 + (INT64_C(3) << 19)) * precomp->divisor);
71 const __m256i MASK = _mm256_set1_epi64x(0xFFFFFFFFUL);
72 const __m256i S = _mm256_set1_epi32(0x80000000);
73 // const __m256i IDX = _mm256_set_epi32(0,4,1,5,2,6,3,7);
74 const __m256i IDX = _mm256_set_epi32(7, 3, 6, 2, 5, 1, 4, 0);
75 const uint64_t m = precomp->m;
76 const uint64_t ms8 = m / 8;
77 I8MEM* outre = (I8MEM*)r;
78 I8MEM* outim = (I8MEM*)(r + m);
79 const D4MEM* in = x;
80 // Note: this formula will only work if abs(in) < 2^32
81 for (uint32_t i = 0; i < ms8; ++i) {
82 __m256d cpla = _mm256_loadu_pd(in[0]);
83 __m256d cplb = _mm256_loadu_pd(in[1]);
84 __m256d cplc = _mm256_loadu_pd(in[2]);
85 __m256d cpld = _mm256_loadu_pd(in[3]);
86 __m256i icpla = _mm256_castpd_si256(_mm256_add_pd(cpla, R));
87 __m256i icplb = _mm256_castpd_si256(_mm256_add_pd(cplb, R));
88 __m256i icplc = _mm256_castpd_si256(_mm256_add_pd(cplc, R));
89 __m256i icpld = _mm256_castpd_si256(_mm256_add_pd(cpld, R));
90 icpla = _mm256_or_si256(_mm256_and_si256(icpla, MASK), _mm256_slli_epi64(icplb, 32));
91 icplc = _mm256_or_si256(_mm256_and_si256(icplc, MASK), _mm256_slli_epi64(icpld, 32));
92 icpla = _mm256_xor_si256(icpla, S);
93 icplc = _mm256_xor_si256(icplc, S);
94 __m256i re = _mm256_unpacklo_epi64(icpla, icplc);
95 __m256i im = _mm256_unpackhi_epi64(icpla, icplc);
96 re = _mm256_permutevar8x32_epi32(re, IDX);
97 im = _mm256_permutevar8x32_epi32(im, IDX);
98 _mm256_storeu_si256((__m256i*)outre[0], re);
99 _mm256_storeu_si256((__m256i*)outim[0], im);
100 outre += 1;
101 outim += 1;
102 in += 4;
103 }
104}