wehub-resource-sync
69 lines
1.3 KiB
C++
69 lines
1.3 KiB
C++
/**
|
|
* Copyright (c) 2023 by Contributors
|
|
* @file dgl/runtime/ndarray.h
|
|
* @brief BFloat16 CPU header
|
|
*/
|
|
#ifndef DGL_RUNTIME_BFLOAT16_H_
|
|
#define DGL_RUNTIME_BFLOAT16_H_
|
|
|
|
#include <cmath>
|
|
|
|
class BFloat16 {
|
|
uint16_t val;
|
|
|
|
public:
|
|
constexpr BFloat16() : val(0) {}
|
|
// Disable lint "explicit" warning, since implicit usage on constructor is
|
|
// expected.
|
|
BFloat16(float f) { // NOLINT
|
|
if (std::isnan(f)) {
|
|
val = 0x7FC0;
|
|
} else {
|
|
union {
|
|
uint16_t iraw16[2];
|
|
uint32_t iraw32;
|
|
float f32;
|
|
};
|
|
|
|
f32 = f;
|
|
const uint32_t rounding_bias = 0x00007FFF + (iraw16[1] & 0x1);
|
|
val = static_cast<uint16_t>((iraw32 + rounding_bias) >> 16);
|
|
}
|
|
}
|
|
static constexpr BFloat16 Min() {
|
|
BFloat16 min;
|
|
min.val = 0xFF80;
|
|
return min;
|
|
}
|
|
|
|
static constexpr BFloat16 Max() {
|
|
BFloat16 max;
|
|
max.val = 0x7F80;
|
|
return max;
|
|
}
|
|
|
|
BFloat16& operator-=(const float& rhs) {
|
|
float lhs = (*this);
|
|
(*this) = lhs - rhs;
|
|
return *this;
|
|
}
|
|
|
|
BFloat16& operator+=(const float& rhs) {
|
|
float lhs = (*this);
|
|
(*this) = lhs + rhs;
|
|
return *this;
|
|
}
|
|
|
|
operator float() const {
|
|
union {
|
|
float f;
|
|
uint16_t raw[2];
|
|
};
|
|
raw[0] = 0;
|
|
raw[1] = val;
|
|
return f;
|
|
}
|
|
};
|
|
|
|
#endif // DGL_RUNTIME_BFLOAT16_H_
|