avx512d_types.h

Go to the documentation of this file.

#ifndef AMPGEN_AVXd_TYPES
#define AMPGEN_AVXd_TYPES 1
 
#include <immintrin.h>
#include <array>
#include <iostream>
#include <complex>
#include <omp.h>
#include <cmath>
 
namespace AmpGen {
  namespace AVX512d {
#define stl_fallback(x)                                                                                                                                        \
  inline real_v x(const real_v &v) {                                                                                                                           \
    auto a = v.to_array();                                                                                                                                     \
    return real_v(std::x(a[0]), std::x(a[1]), std::x(a[2]), std::x(a[3]), std::x(a[4]), std::x(a[5]), std::x(a[6]), std::x(a[7]));                             \
  }
 
    struct real_v {
      __m512d data;
      static constexpr unsigned size = 8;
      typedef double scalar_type;
      real_v() = default;
      real_v(__m512d data) : data(data) {}
      real_v(const double &f) : data(_mm512_set1_pd(f)) {}
      real_v(const double &x0, const double &x1, const double &x2, const double &x3, const double &x4, const double &x5, const double &x6, const double &x7) {
        double tmp[8] = {x0, x1, x2, x3, x4, x5, x6, x7};
        data = _mm512_loadu_pd(tmp);
      }
      real_v(const double *f) : data(_mm512_loadu_pd(f)) {}
      void store(double *ptr) const { _mm512_storeu_pd(ptr, data); }
      std::array<double, 8> to_array() const {
        std::array<double, 8> b;
        store(&b[0]);
        return b;
      }
      double at(const unsigned i) const { return to_array()[i]; }
      operator __m512d() const { return data; }
    };
 
    inline real_v operator+(const real_v &lhs, const real_v &rhs) { return _mm512_add_pd(lhs, rhs); }
    inline real_v operator-(const real_v &lhs, const real_v &rhs) { return _mm512_sub_pd(lhs, rhs); }
    inline real_v operator*(const real_v &lhs, const real_v &rhs) { return _mm512_mul_pd(lhs, rhs); }
    inline real_v operator/(const real_v &lhs, const real_v &rhs) { return _mm512_div_pd(lhs, rhs); }
    inline real_v operator-(const real_v &x) { return -1.f * x; }
    inline real_v operator&(const real_v &lhs, const real_v &rhs) { return _mm512_and_pd(lhs, rhs); }
    inline real_v operator|(const real_v &lhs, const real_v &rhs) { return _mm512_or_pd(lhs, rhs); }
    inline real_v operator^(const real_v &lhs, const real_v &rhs) { return _mm512_xor_pd(lhs, rhs); }
    inline real_v operator+=(real_v &lhs, const real_v &rhs) {
      lhs = lhs + rhs;
      return lhs;
    }
    inline real_v operator-=(real_v &lhs, const real_v &rhs) {
      lhs = lhs - rhs;
      return lhs;
    }
    inline real_v operator*=(real_v &lhs, const real_v &rhs) {
      lhs = lhs * rhs;
      return lhs;
    }
    inline real_v operator/=(real_v &lhs, const real_v &rhs) {
      lhs = lhs / rhs;
      return lhs;
    }
    inline real_v operator&&(const real_v &lhs, const real_v &rhs) { return _mm512_and_pd(lhs, rhs); }
    inline real_v operator||(const real_v &lhs, const real_v &rhs) { return _mm512_or_pd(lhs, rhs); }
    inline real_v operator!(const real_v &x) { return x ^ _mm512_castsi512_pd(_mm512_set1_epi32(-1)); }
    inline __mmask8 operator<(const real_v &lhs, const real_v &rhs) { return _mm512_cmp_pd_mask(lhs, rhs, _CMP_LT_OS); }
    inline __mmask8 operator>(const real_v &lhs, const real_v &rhs) { return _mm512_cmp_pd_mask(lhs, rhs, _CMP_GT_OS); }
    inline __mmask8 operator==(const real_v &lhs, const real_v &rhs) { return _mm512_cmp_pd_mask(lhs, rhs, _CMP_EQ_OS); }
    inline real_v sqrt(const real_v &v) { return _mm512_sqrt_pd(v); }
    inline real_v abs(const real_v &v) { return _mm512_andnot_pd(_mm512_set1_pd(-0.), v); }
    // inline real_v sin( const real_v& v )  { return sin512_pd(v) ; }
    // inline real_v cos( const real_v& v )  { return cos512_pd(v) ; }
    // inline real_v tan( const real_v& v )  { real_v s; real_v c; sincos512_pd(v, (__m512*)&s, (__m512*)&c) ; return s/c; }
    // inline real_v exp( const real_v& v )  { return exp512_ps(v) ; }
    inline real_v select(const __mmask8 &mask, const real_v &a, const real_v &b) { return _mm512_mask_mov_pd(b, mask, a); }
    inline real_v select(const bool &mask, const real_v &a, const real_v &b) { return mask ? a : b; }
    inline real_v sign(const real_v &v) { return select(v > 0., +1., -1.); }
    inline real_v atan2(const real_v &y, const real_v &x) {
      std::array<double, 8> bx{x.to_array()}, by{y.to_array()};
      return real_v(std::atan2(by[0], bx[0]), std::atan2(by[1], bx[1]), std::atan2(by[2], bx[2]), std::atan2(by[3], bx[3]), std::atan2(by[4], bx[4]),
                    std::atan2(by[5], bx[5]), std::atan2(by[6], bx[6]), std::atan2(by[7], bx[7]));
    }
    inline __m512i double_to_int(const real_v &x) {
      auto xr = _mm512_roundscale_pd(x, _MM_FROUND_TO_ZERO);
      // based on:  https://stackoverflow.com/questions/41144668/how-to-efficiently-perform-double-int64-conversions-with-sse-avx
      return _mm512_sub_epi64(_mm512_castpd_si512(_mm512_add_pd(xr, _mm512_set1_pd(0x0018000000000000))),
                              _mm512_castpd_si512(_mm512_set1_pd(0x0018000000000000)));
    }
    inline real_v gather(const double *base_addr, const real_v &offsets) { return _mm512_i64gather_pd(double_to_int(offsets), base_addr, sizeof(double)); }
 
    inline void frexp(const real_v &value, real_v &mant, real_v &exponent) {
      auto arg_as_int = _mm512_castpd_si512(value);
      static const real_v offset(4503599627370496.0 + 1022.0);                // 2^52 + 1022.0
      static const __m512i pow2_52_i = _mm512_set1_epi64(0x4330000000000000); // *reinterpret_cast<const uint64_t*>(&pow2_52_d);
      auto b = _mm512_srl_epi64(arg_as_int, _mm_cvtsi32_si128(52));
      auto c = _mm512_or_si512(b, pow2_52_i);
      exponent = real_v(_mm512_castsi512_pd(c)) - offset;
      mant
        = _mm512_castsi512_pd(_mm512_or_si512(_mm512_and_si512(arg_as_int, _mm512_set1_epi64(0x000FFFFFFFFFFFFFll)), _mm512_set1_epi64(0x3FE0000000000000ll)));
    }
 
    inline real_v fmadd(const real_v &a, const real_v &b, const real_v &c) { return _mm512_fmadd_pd(a, b, c); }
    inline real_v log(const real_v &arg) {
      static const real_v corr = 0.693147180559945286226764;
      static const real_v CL15 = 0.148197055177935105296783;
      static const real_v CL13 = 0.153108178020442575739679;
      static const real_v CL11 = 0.181837339521549679055568;
      static const real_v CL9 = 0.22222194152736701733275;
      static const real_v CL7 = 0.285714288030134544449368;
      static const real_v CL5 = 0.399999999989941956712869;
      static const real_v CL3 = 0.666666666666685503450651;
      static const real_v CL1 = 2.0;
      real_v mant, exponent;
      frexp(arg, mant, exponent);
      auto x = (mant - 1.) / (mant + 1.);
      auto x2 = x * x;
      auto p = fmadd(CL15, x2, CL13);
      p = fmadd(p, x2, CL11);
      p = fmadd(p, x2, CL9);
      p = fmadd(p, x2, CL7);
      p = fmadd(p, x2, CL5);
      p = fmadd(p, x2, CL3);
      p = fmadd(p, x2, CL1);
      p = fmadd(p, x, corr * exponent);
      return p;
    }
    stl_fallback(exp) stl_fallback(tan) stl_fallback(sin) stl_fallback(cos) inline real_v remainder(const real_v &a, const real_v &b) {
      return a - real_v(_mm512_roundscale_pd(a / b, _MM_FROUND_TO_NEG_INF)) * b;
    }
    inline real_v fmod(const real_v &a, const real_v &b) {
      auto r = remainder(abs(a), abs(b));
      return select(a > 0., r, -r);
    }
 
    inline std::ostream &operator<<(std::ostream &os, const real_v &obj) {
      auto buffer = obj.to_array();
      for(unsigned i = 0; i != 8; ++i) os << buffer[i] << " ";
      return os;
    }
 
    using complex_v = std::complex<real_v>;
    inline complex_v operator+(const complex_v &lhs, const real_v &rhs) { return complex_v(lhs.real() + rhs, lhs.imag()); }
    inline complex_v operator-(const complex_v &lhs, const real_v &rhs) { return complex_v(lhs.real() - rhs, lhs.imag()); }
    inline complex_v operator*(const complex_v &lhs, const real_v &rhs) { return complex_v(lhs.real() * rhs, lhs.imag() * rhs); }
    inline complex_v operator/(const complex_v &lhs, const real_v &rhs) { return complex_v(lhs.real() / rhs, lhs.imag() / rhs); }
    inline complex_v operator+(const real_v &lhs, const complex_v &rhs) { return complex_v(lhs + rhs.real(), rhs.imag()); }
    inline complex_v operator-(const real_v &lhs, const complex_v &rhs) { return complex_v(lhs - rhs.real(), -rhs.imag()); }
    inline complex_v operator*(const real_v &lhs, const complex_v &rhs) { return complex_v(lhs * rhs.real(), lhs * rhs.imag()); }
    inline complex_v operator/(const real_v &lhs, const complex_v &rhs) {
      return complex_v(lhs * rhs.real(), -lhs * rhs.imag()) / (rhs.real() * rhs.real() + rhs.imag() * rhs.imag());
    }
    inline real_v abs(const complex_v &v) { return sqrt(v.real() * v.real() + v.imag() * v.imag()); }
    inline real_v norm(const complex_v &v) { return (v.real() * v.real() + v.imag() * v.imag()); }
    inline complex_v select(const __mmask8 &mask, const complex_v &a, const complex_v &b) {
      return complex_v(select(mask, a.real(), b.real()), select(mask, a.imag(), b.imag()));
    }
    inline complex_v select(const __mmask8 &mask, const real_v &a, const complex_v &b) {
      return complex_v(select(mask, a, b.real()), select(mask, 0.f, b.imag()));
    }
    inline complex_v select(const __mmask8 &mask, const complex_v &a, const real_v &b) {
      return complex_v(select(mask, a.real(), b), select(mask, a.imag(), 0.f));
    }
    inline complex_v select(const bool &mask, const complex_v &a, const complex_v &b) { return mask ? a : b; }
    inline complex_v exp(const complex_v &v) {
      auto [s, c] = sincos(v.imag());
      return exp(v.real()) * complex_v(c, s);
    }
    inline complex_v sqrt(const complex_v &v) {
      auto r = abs(v);
      return complex_v(sqrt(0.5 * (r + v.real())), sign(v.imag()) * sqrt(0.5 * (r - v.real())));
    }
    inline complex_v log(const complex_v &v) { return complex_v(0.5 * log(norm(v)), atan2(v.imag(), v.real())); }
 
    inline std::ostream &operator<<(std::ostream &os, const complex_v &obj) { return os << "( " << obj.real() << ") (" << obj.imag() << ")"; }
#pragma omp declare reduction(+ : real_v : omp_out = omp_out + omp_in)
#pragma omp declare reduction(+ : complex_v : omp_out = omp_out + omp_in)
 
  }
}
 
#endif