avx2f_types.h

Go to the documentation of this file.

#ifndef AMPGEN_AVX_TYPES
#define AMPGEN_AVX_TYPES 1
 
#include <immintrin.h>
#include <array>
#include <iostream>
#include <complex>
#include <omp.h>
#include <cmath>
#include "AmpGen/Complex.h"
 
#if USE_MVEC 
  extern "C" void    _ZGVdN8vvv_sincos(__m256 x, __m256i ptrs, __m256i ptrc);
#define libmvec_alias(F, O) \
  extern "C" __m256 _ZGVcN8v_##F(__m256 x);             \
  inline real_v O( const real_v& v ){ return _ZGVcN8v_##F(v) ; }
#else
#define libmvec_alias(F, O) \
  inline real_v O( const real_v& v ){ auto arr = v.to_ptr(); return real_v(  \
      std::F(arr[0]), std::F(arr[1]), std::F(arr[2]), std::F(arr[3]),          \
      std::F(arr[4]), std::F(arr[5]), std::F(arr[6]), std::F(arr[7]) ) ; }
#endif
 
namespace AmpGen {
  namespace AVX2f {
    struct real_v {
      __m256 data;
      static constexpr unsigned size = 8 ;
      typedef float scalar_type;
      real_v() = default; 
      real_v(__m256 data ) : data(data) {}
      real_v(const int& f ) : real_v(scalar_type(f)) {}
      real_v(const scalar_type& f ) : data( _mm256_set1_ps(f) ) {}
      real_v(const double& f )      : data( _mm256_set1_ps( scalar_type(f) )) {}
      explicit real_v(const scalar_type* f )      : data( _mm256_loadu_ps( f ) ) {}
      real_v(const scalar_type& x0, const scalar_type& x1, const scalar_type& x2, const scalar_type& x3,
          const scalar_type& x4, const scalar_type& x5, const scalar_type& x6, const scalar_type& x7)
      {
        data = _mm256_set_ps(x7,x6,x5,x4,x3,x2,x1,x0); 
      }
 
      void store( scalar_type* ptr ) const { _mm256_storeu_ps( ptr, data ); }
      std::array<scalar_type, 8> to_array() const { std::array<scalar_type, 8> b; store( &b[0] ); return b; }
      const scalar_type* to_ptr() const { return reinterpret_cast<const scalar_type*>( &data ) ; }
            scalar_type* to_ptr() { return reinterpret_cast<scalar_type*>( &data ) ; }
      scalar_type at(const unsigned i) const { return to_ptr()[i] ; }       
      operator __m256() const { return data ; } 
      inline real_v operator+=(const real_v& rhs ); 
      inline real_v operator-=(const real_v& rhs ); 
      inline real_v operator*=(const real_v& rhs ); 
      inline real_v operator/=(const real_v& rhs ); 
      inline __m256i to_int() const { return _mm256_cvtps_epi32(data); }
    };
    
    inline real_v operator+( const real_v& lhs, const real_v& rhs ) { return _mm256_add_ps(lhs, rhs); }
    inline real_v operator-( const real_v& lhs, const real_v& rhs ) { return _mm256_sub_ps(lhs, rhs); }
    inline real_v operator*( const real_v& lhs, const real_v& rhs ) { return _mm256_mul_ps(lhs, rhs); }
    inline real_v operator/( const real_v& lhs, const real_v& rhs ) { return _mm256_div_ps(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 _mm256_and_ps( lhs, rhs ); }
    inline real_v operator|( const real_v& lhs, const real_v& rhs ) { return _mm256_or_ps( lhs, rhs ); }
    inline real_v operator^( const real_v& lhs, const real_v& rhs ) { return _mm256_xor_ps( lhs, rhs ); }
    inline real_v operator&&( const real_v& lhs, const real_v& rhs ) { return _mm256_and_ps( lhs, rhs ); }
    inline real_v operator||( const real_v& lhs, const real_v& rhs ) { return _mm256_or_ps( lhs, rhs ); }
    inline real_v operator!( const real_v& x ) { return x ^ _mm256_castsi256_ps( _mm256_set1_epi32( -1 ) ); }
    inline real_v operator<( const real_v& lhs, const real_v& rhs ) { return _mm256_cmp_ps( lhs, rhs, _CMP_LT_OS ); }
    inline real_v operator>( const real_v& lhs, const real_v& rhs ) { return _mm256_cmp_ps( lhs, rhs, _CMP_GT_OS ); }
    inline real_v operator<=( const real_v& lhs, const real_v& rhs ) { return _mm256_cmp_ps( lhs, rhs, _CMP_LE_OS ); }
    inline real_v operator>=( const real_v& lhs, const real_v& rhs ) { return _mm256_cmp_ps( lhs, rhs, _CMP_GE_OS ); }
    inline real_v operator==( const real_v& lhs, const real_v& rhs ){ return _mm256_cmp_ps( lhs, rhs, _CMP_EQ_OS ); }
    inline real_v sqrt( const real_v& v ) { return _mm256_sqrt_ps(v); } 
    inline real_v real_v::operator+=(const real_v& rhs ){ *this = *this + rhs; return *this; }
    inline real_v real_v::operator-=(const real_v& rhs ){ *this = *this - rhs; return *this; }
    inline real_v real_v::operator*=(const real_v& rhs ){ *this = *this * rhs; return *this; }
    inline real_v real_v::operator/=(const real_v& rhs ){ *this = *this / rhs; return *this; }
    libmvec_alias(sinf, sin)
    libmvec_alias(cosf, cos)
    libmvec_alias(expf, exp)
    libmvec_alias(logf, log)
    inline std::array<int32_t, real_v::size> store( const __m256i& v )
    {
      alignas(32) std::array<int32_t, real_v::size> rt; 
       _mm256_store_si256( (__m256i*)&rt[0], v);
       return rt; 
    }
    
    inline void sincos( const real_v& v, real_v& s, real_v& c )
    { 
      s = sin(v);
      c = cos(v);
    }
    inline std::pair<real_v, real_v> sincos( const real_v& v )
    {
      std::pair<real_v, real_v> rt;
      sincos( v, rt.first, rt.second );
      return rt; 
    }
    inline real_v tan( const real_v& v )
    {
      auto [s,c] = sincos( v );
      return s / c ;  
    }
    
    inline real_v abs   ( const real_v& v ) { return v & _mm256_castsi256_ps( _mm256_set1_epi32( 0x7FFFFFFF ) ); }
    inline real_v select(const real_v& mask, const real_v& a, const real_v& b ) { return _mm256_blendv_ps( b, a, mask ); }
    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 fmadd ( const real_v& a, const real_v& b, const real_v& c ) { return _mm256_fmadd_ps(a, b, c); }
    inline real_v remainder( const real_v& a, const real_v& b ){ return a - real_v(_mm256_round_ps(a/b, _MM_FROUND_TO_NEG_INF)) * b; }
    inline real_v atan2( const real_v& y, const real_v& x ){ 
      const auto* bx = x.to_ptr();
      const auto* by = y.to_ptr();
      real_v rt; 
      for( unsigned i = 0 ; i != real_v::size ; ++i ) rt.to_ptr()[i] = std::atan2( by[i] , bx[i] );
      return rt; 
    }
    inline real_v gather( const double* base_addr, const real_v& offsets)
    {
      std::array<float, real_v::size> tmp;
      auto ptr = store( offsets.to_int() );
   //   int32_t* ptr = (int32_t*)(&ints); 
      for( int i = 0 ; i != real_v::size; ++i ) tmp[i] = real_v::scalar_type( base_addr[ptr[i]] );
      return real_v( tmp.data() ); 
    }
    
    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 != real_v::size; ++i ) os << buffer[i] << " ";
      return os; 
    }
 
    using complex_v = Complex<real_v>; 
    inline complex_v select(const real_v& 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 real_v& 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 real_v& 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; }
    #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