math.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2021, Lawrence Livermore National Security, LLC and LvArray contributors.
 * All rights reserved.
 * See the LICENSE file for details.
 * SPDX-License-Identifier: (BSD-3-Clause)
 */

#pragma once

// Source includes
#include "LvArrayConfig.hpp"
#include "Macros.hpp"

// System includes
#include <cmath>
#include <type_traits>

#if defined( LVARRAY_USE_CUDA )
  #include <cuda_fp16.h>
#endif

namespace LvArray
{

namespace math
{

namespace internal
{

template< typename T, typename U >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
T convert( T const, U const u )
{ return u; }

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
int numValues( T const )
{ return 1; }

template< typename T >
struct SingleType
{
  using type = T;
};

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
SingleType< T > getFirst( T const x )
{ return x; }

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
SingleType< T > getSecond( T const x )
{ return x; }

template< typename T >
LVARRAY_HOST_DEVICE inline constexpr
T lessThan( T const x, T const y )
{ return __hlt( x, y ); }

#if defined( LVARRAY_USE_CUDA )

template< typename U >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
__half convert( __half const, U const u )
{ return __float2half_rn( u ); }

template< typename U >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
__half2 convert( __half2 const, U const u )
{ return __float2half2_rn( u ); }

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
__half2 convert( __half2 const, __half const u )
{
#if defined( LVARRAY_DEVICE_COMPILE )
  return __half2half2( u );
#else
  return __float2half2_rn( u );
#endif
}

template< typename U, typename V >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
__half2 convert( __half2 const, U const u, V const v )
{ return __floats2half2_rn( u, v ); }

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
__half2 convert( __half2 const, __half const u, __half const v )
{
#if defined( LVARRAY_DEVICE_COMPILE )
  return __halves2half2( u, v );
#else
  return __floats2half2_rn( u, v );
#endif
}

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
int numValues( __half2 const & )
{ return 2; }

template<>
struct SingleType< __half2 >
{
  using type = __half;
};

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half getFirst( __half2 const x )
{ return __low2half( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half getSecond( __half2 const x )
{ return __high2half( x ); }

#endif

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half lessThan( __half const x, __half const y )
{ return __hlt( x, y ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 lessThan( __half2 const x, __half2 const y )
{ return __hlt2( x, y ); }
#endif

} // namespace internal


template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
int numValues()
{ return internal::numValues( T() ); }

template< typename T >
using SingleType = typename internal::SingleType< T >::type;

template< typename T, typename U >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
T convert( U const u )
{ return internal::convert( T(), u ); }

template< typename T, typename U, typename V >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
T convert( U const u, V const v )
{ return internal::convert( T(), u, v ); }

template< typename T >
LVARRAY_DEVICE LVARRAY_FORCE_INLINE
SingleType< T > getFirst( T const x )
{ return internal::getFirst( x ); }

template< typename T >
LVARRAY_DEVICE LVARRAY_FORCE_INLINE
SingleType< T > getSecond( T const x )
{ return internal::getSecond( x ); }

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
std::enable_if_t< std::is_arithmetic< T >::value, T >
max( T const a, T const b )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::max( a, b );
#else
  return std::max( a, b );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half max( __half const a, __half const b )
{
#if defined(LVARRAY_USE_CUDA) && CUDART_VERSION > 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
  return __hmax( a, b );
#elif defined(LVARRAY_USE_HIP)
  return __hgt( a, b ) ? a : b;
#else
  return a > b ? a : b;
#endif
}

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 max( __half2 const a, __half2 const b )
{
#if defined(LVARRAY_USE_CUDA) && CUDART_VERSION > 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
  return __hmax2( a, b );
#else
  __half2 const aFactor = __hge2( a, b );
  __half2 const bFactor = convert< __half2 >( 1 ) - aFactor;
  return a * aFactor + bFactor * b;
#endif
}

#endif


template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
std::enable_if_t< std::is_arithmetic< T >::value, T >
min( T const a, T const b )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::min( a, b );
#else
  return std::min( a, b );
#endif
}

#if defined( LVARRAY_USE_CUDA )

LVARRAY_DEVICE
LVARRAY_FORCE_INLINE
__half min( __half const a, __half const b )
{
#if CUDART_VERSION > 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
  return __hmin( a, b );
#else
  return a < b ? a : b;
#endif
}

LVARRAY_DEVICE
LVARRAY_FORCE_INLINE
__half2 min( __half2 const a, __half2 const b )
{
#if CUDART_VERSION > 11000 && (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
  return __hmin2( a, b );
#else
  __half2 const aFactor = __hle2( a, b );
  __half2 const bFactor = convert< __half2 >( 1 ) - aFactor;
  return a * aFactor + bFactor * b;
#endif
}

#endif

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
T abs( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::abs( x );
#else
  return std::abs( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half abs( __half const x )
{
#if CUDART_VERSION > 11000
  return __habs( x );
#else
  return x > __half( 0 ) ? x : -x;
#endif
}

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 abs( __half2 const x )
{
#if CUDART_VERSION > 11000
  return __habs2( x );
#else
  return x - __hle2( x, convert< __half2 >( 0 ) ) * ( x + x );
#endif
}

#endif

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE constexpr
T square( T const x )
{ return x * x; }



LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float sqrt( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::sqrtf( x );
#else
  return std::sqrt( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double sqrt( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::sqrt( double( x ) );
#else
  return std::sqrt( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half sqrt( __half const x )
{ return ::hsqrt( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 sqrt( __half2 const x )
{ return ::h2sqrt( x ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float invSqrt( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::rsqrtf( x );
#else
  return 1 / std::sqrt( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double invSqrt( T const x )
{
#if defined( LVARRAY_DEVICE_COMPILE )
  return ::rsqrt( double( x ) );
#else
  return 1 / std::sqrt( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half invSqrt( __half const x )
{ return ::hrsqrt( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 invSqrt( __half2 const x )
{ return ::h2rsqrt( x ); }

#endif



LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float sin( float const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::sinf( theta );
#else
  return std::sin( theta );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double sin( T const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::sin( double( theta ) );
#else
  return std::sin( theta );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half sin( __half const theta )
{ return ::hsin( theta ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 sin( __half2 const theta )
{ return ::h2sin( theta ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float cos( float const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::cosf( theta );
#else
  return std::cos( theta );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double cos( T const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::cos( double( theta ) );
#else
  return std::cos( theta );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half cos( __half const theta )
{ return ::hcos( theta ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 cos( __half2 const theta )
{ return ::h2cos( theta ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
void sincos( float const theta, float & sinTheta, float & cosTheta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  #if defined(LVARRAY_USE_CUDA)
  ::sincos( theta, &sinTheta, &cosTheta );
  #elif defined(LVARRAY_USE_HIP)
  ::sincosf( theta, &sinTheta, &cosTheta );
  #endif
#else
  sinTheta = std::sin( theta );
  cosTheta = std::cos( theta );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
void sincos( double const theta, double & sinTheta, double & cosTheta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  ::sincos( theta, &sinTheta, &cosTheta ); // hip and cuda versions both use double
#else
  sinTheta = std::sin( theta );
  cosTheta = std::cos( theta );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
void sincos( T const theta, double & sinTheta, double & cosTheta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  double s, c;
  ::sincos( theta, &s, &c );
  sinTheta = s;
  cosTheta = c;
#else
  sinTheta = std::sin( theta );
  cosTheta = std::cos( theta );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
void sincos( __half const theta, __half & sinTheta, __half & cosTheta )
{
  sinTheta = ::hsin( theta );
  cosTheta = ::hcos( theta );
}

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
void sincos( __half2 const theta, __half2 & sinTheta, __half2 & cosTheta )
{
  sinTheta = ::h2sin( theta );
  cosTheta = ::h2cos( theta );
}

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float tan( float const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::tanf( theta );
#else
  return std::tan( theta );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double tan( T const theta )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::tan( double( theta ) );
#else
  return std::tan( theta );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half tan( __half const theta )
{
  __half s, c;
  sincos( theta, s, c );
  return s / c;
}

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 tan( __half2 const theta )
{
  __half2 s, c;
  sincos( theta, s, c );
  return s / c;
}

#endif



namespace internal
{

template< typename T >
LVARRAY_DEVICE LVARRAY_FORCE_INLINE
T asinImpl( T const x )
{
  T const negate = lessThan( x, math::convert< T >( 0 ) );
  T const absX = abs( x );

  T ret = math::convert< T >( -0.0187293 ) * absX + math::convert< T >( 0.0742610 );
  ret = ret * absX - math::convert< T >( 0.2121144 );
  ret = ret * absX + math::convert< T >( 1.5707288 );
  ret = math::convert< T >( 3.14159265358979 * 0.5 ) - ret * sqrt( math::convert< T >( 1 ) - absX );
  ret = ret - negate * ( ret + ret );
  T const smallAngle = lessThan( absX, math::convert< T >( 1.7e-1 ) );
  return smallAngle * x + ( math::convert< T >( 1 ) - smallAngle ) * ret;
}

template< typename T >
LVARRAY_DEVICE LVARRAY_FORCE_INLINE
T acosImpl( T const x )
{
  T const negate = lessThan( x, math::convert< T >( 0 ) );
  T const absX = abs( x );

  T ret = math::convert< T >( -0.0187293 ) * absX + math::convert< T >( 0.0742610 );
  ret = ret * absX - math::convert< T >( 0.2121144 );
  ret = ret * absX + math::convert< T >( 1.5707288 );
  ret = ret * sqrt( math::convert< T >( 1 ) - absX );
  ret = ret - negate * ( ret + ret );
  return negate * math::convert< T >( 3.14159265358979 ) + ret;
}

template< typename T >
LVARRAY_DEVICE
LVARRAY_FORCE_INLINE
T atan2Impl( T const y, T const x )
{
  T const absX = abs( x );
  T const absY = abs( y );
  T const ratio = min( absX, absY ) / max( absX, absY );
  T const ratio2 = ratio * ratio;

  T ret = math::convert< T >( -0.013480470 ) * ratio2 + math::convert< T >( 0.057477314 );
  ret = ret * ratio2 - math::convert< T >( 0.121239071 );
  ret = ret * ratio2 + math::convert< T >( 0.195635925 );
  ret = ret * ratio2 - math::convert< T >( 0.332994597 );
  ret = ret * ratio2 + math::convert< T >( 0.999995630 );
  ret = ret * ratio;

  // For single values the following works out to:
  // ret = absX < absY ? 1.570796327 - ret : ret;
  // ret = x < 0 ? 3.141592654 - ret : ret;
  // ret = y < 0 ? -ret : ret;
  ret = internal::lessThan( absX, absY ) * ( math::convert< T >( 1.570796327 ) - ret - ret ) + ret;
  ret = internal::lessThan( x, math::convert< T >( 0 ) ) * ( math::convert< T >( 3.141592654 ) - ret - ret ) + ret;
  ret = ret - internal::lessThan( y, math::convert< T >( 0 ) ) * ( ret + ret );

  return ret;
}

} // namespace internal

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float asin( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::asinf( x );
#else
  return std::asin( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double asin( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::asin( double( x ) );
#else
  return std::asin( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half asin( __half const x )
{ return internal::asinImpl( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 asin( __half2 const x )
{ return internal::asinImpl( x ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float acos( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::acosf( x );
#else
  return std::acos( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double acos( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::acos( double( x ) );
#else
  return std::acos( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half acos( __half const x )
{ return internal::acosImpl( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 acos( __half2 const x )
{ return internal::acosImpl( x ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float atan2( float const y, float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::atan2f( y, x );
#else
  return std::atan2( y, x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double atan2( T const y, T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::atan2( double( y ), double( x ) );
#else
  return std::atan2( y, x );
#endif
}

#if defined( LVARRAY_USE_CUDA )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half atan2( __half const y, __half const x )
{ return internal::atan2Impl( y, x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 atan2( __half2 const y, __half2 const x )
{ return internal::atan2Impl( y, x ); }

#endif



LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float exp( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::expf( x );
#else
  return std::exp( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double exp( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::exp( double( x ) );
#else
  return std::exp( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half exp( __half const x )
{ return ::hexp( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 exp( __half2 const x )
{ return ::h2exp( x ); }

#endif

LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
float log( float const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::logf( x );
#else
  return std::log( x );
#endif
}

template< typename T >
LVARRAY_HOST_DEVICE LVARRAY_FORCE_INLINE
double log( T const x )
{
#if defined(LVARRAY_DEVICE_COMPILE)
  return ::log( double( x ) );
#else
  return std::log( x );
#endif
}

#if defined( LVARRAY_USE_DEVICE )

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half log( __half const x )
{ return ::hlog( x ); }

LVARRAY_DEVICE LVARRAY_FORCE_INLINE
__half2 log( __half2 const x )
{ return ::h2log( x ); }

#endif


} // namespace math
} // namespace LvArray