/******************************************************************************
* Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are not permitted.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
******************************************************************************/
#pragma once
/**
* \file
* Utilities for interacting with the opaque CUDA __half type
*/
#include <stdint.h>
#include <cuda_fp16.h>
#include <iosfwd>
#include <cub/util_type.cuh>
/******************************************************************************
* half_t
******************************************************************************/
/**
* Host-based fp16 data type compatible and convertible with __half
*/
struct half_t
{
uint16_t __x;
/// Constructor from __half
__host__ __device__ __forceinline__
half_t(const __half &other)
{
__x = reinterpret_cast<const uint16_t&>(other);
}
/// Constructor from integer
__host__ __device__ __forceinline__
half_t(int a)
{
*this = half_t(float(a));
}
/// Default constructor
__host__ __device__ __forceinline__
half_t() : __x(0)
{}
/// Constructor from float
__host__ __device__ __forceinline__
half_t(float a)
{
// Stolen from Norbert Juffa
uint32_t ia = *reinterpret_cast<uint32_t*>(&a);
uint16_t ir;
ir = (ia >> 16) & 0x8000;
if ((ia & 0x7f800000) == 0x7f800000)
{
if ((ia & 0x7fffffff) == 0x7f800000)
{
ir |= 0x7c00; /* infinity */
}
else
{
ir = 0x7fff; /* canonical NaN */
}
}
else if ((ia & 0x7f800000) >= 0x33000000)
{
int32_t shift = (int32_t) ((ia >> 23) & 0xff) - 127;
if (shift > 15)
{
ir |= 0x7c00; /* infinity */
}
else
{
ia = (ia & 0x007fffff) | 0x00800000; /* extract mantissa */
if (shift < -14)
{ /* denormal */
ir |= ia >> (-1 - shift);
ia = ia << (32 - (-1 - shift));
}
else
{ /* normal */
ir |= ia >> (24 - 11);
ia = ia << (32 - (24 - 11));
ir = ir + ((14 + shift) << 10);
}
/* IEEE-754 round to nearest of even */
if ((ia > 0x80000000) || ((ia == 0x80000000) && (ir & 1)))
{
ir++;
}
}
}
this->__x = ir;
}
/// Cast to __half
__host__ __device__ __forceinline__
operator __half() const
{
return reinterpret_cast<const __half&>(__x);
}
/// Cast to float
__host__ __device__ __forceinline__
operator float() const
{
// Stolen from Andrew Kerr
int sign = ((this->__x >> 15) & 1);
int exp = ((this->__x >> 10) & 0x1f);
int mantissa = (this->__x & 0x3ff);
uint32_t f = 0;
if (exp > 0 && exp < 31)
{
// normal
exp += 112;
f = (sign << 31) | (exp << 23) | (mantissa << 13);
}
else if (exp == 0)
{
if (mantissa)
{
// subnormal
exp += 113;
while ((mantissa & (1 << 10)) == 0)
{
mantissa <<= 1;
exp--;
}
mantissa &= 0x3ff;
f = (sign << 31) | (exp << 23) | (mantissa << 13);
}
else if (sign)
{
f = 0x80000000; // negative zero
}
else
{
f = 0x0; // zero
}
}
else if (exp == 31)
{
if (mantissa)
{
f = 0x7fffffff; // not a number
}
else
{
f = (0xff << 23) | (sign << 31); // inf
}
}
return *reinterpret_cast<float const *>(&f);
}
/// Get raw storage
__host__ __device__ __forceinline__
uint16_t raw()
{
return this->__x;
}
/// Equality
__host__ __device__ __forceinline__
bool operator ==(const half_t &other)
{
return (this->__x == other.__x);
}
/// Inequality
__host__ __device__ __forceinline__
bool operator !=(const half_t &other)
{
return (this->__x != other.__x);
}
/// Assignment by sum
__host__ __device__ __forceinline__
half_t& operator +=(const half_t &rhs)
{
*this = half_t(float(*this) + float(rhs));
return *this;
}
/// Multiply
__host__ __device__ __forceinline__
half_t operator*(const half_t &other)
{
return half_t(float(*this) * float(other));
}
/// Add
__host__ __device__ __forceinline__
half_t operator+(const half_t &other)
{
return half_t(float(*this) + float(other));
}
/// Less-than
__host__ __device__ __forceinline__
bool operator<(const half_t &other) const
{
return float(*this) < float(other);
}
/// Less-than-equal
__host__ __device__ __forceinline__
bool operator<=(const half_t &other) const
{
return float(*this) <= float(other);
}
/// Greater-than
__host__ __device__ __forceinline__
bool operator>(const half_t &other) const
{
return float(*this) > float(other);
}
/// Greater-than-equal
__host__ __device__ __forceinline__
bool operator>=(const half_t &other) const
{
return float(*this) >= float(other);
}
/// numeric_traits<half_t>::max
__host__ __device__ __forceinline__
static half_t max() {
uint16_t max_word = 0x7BFF;
return reinterpret_cast<half_t&>(max_word);
}
/// numeric_traits<half_t>::lowest
__host__ __device__ __forceinline__
static half_t lowest() {
uint16_t lowest_word = 0xFBFF;
return reinterpret_cast<half_t&>(lowest_word);
}
};
/******************************************************************************
* I/O stream overloads
******************************************************************************/
/// Insert formatted \p half_t into the output stream
std::ostream& operator<<(std::ostream &out, const half_t &x)
{
out << (float)x;
return out;
}
/// Insert formatted \p __half into the output stream
std::ostream& operator<<(std::ostream &out, const __half &x)
{
return out << half_t(x);
}
/******************************************************************************
* Traits overloads
******************************************************************************/
template <>
struct cub::FpLimits<half_t>
{
static __host__ __device__ __forceinline__ half_t Max() { return half_t::max(); }
static __host__ __device__ __forceinline__ half_t Lowest() { return half_t::lowest(); }
};
template <> struct cub::NumericTraits<half_t> : cub::BaseTraits<FLOATING_POINT, true, false, unsigned short, half_t> {};