arnaucube
arnaucube/poulpyPublic Self-hosted repository
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}