diff --git a/cpp/include/raft/util/cudart_utils.hpp b/cpp/include/raft/util/cudart_utils.hpp --- a/cpp/include/raft/util/cudart_utils.hpp +++ b/cpp/include/raft/util/cudart_utils.hpp @@ -21,6 +21,7 @@ #include +#include #include #include @@ -456,4 +457,27 @@ constexpr inline auto upper_bound( return static_cast(__half_constexpr{0x7c00u}); } +/** + * This is a hack to allow constexpr definition of `bfloat16` constants. + * + * Same reasoning as for `half`: CUDA’s `__nv_bfloat16` has no constexpr constructor. + */ +struct __bfloat16_constexpr : __nv_bfloat16 { // NOLINT + constexpr explicit inline __bfloat16_constexpr(uint16_t u) : __nv_bfloat16() { __x = u; } +}; + +template <> +constexpr inline auto lower_bound<__nv_bfloat16>() -> __nv_bfloat16 +{ + // Negative infinity in bfloat16 (sign=1, exp=all ones, mantissa=0) + return static_cast<__nv_bfloat16>(__bfloat16_constexpr{0xff80u}); +} + +template <> +constexpr inline auto upper_bound<__nv_bfloat16>() -> __nv_bfloat16 +{ + // Positive infinity in bfloat16 (sign=0, exp=all ones, mantissa=0) + return static_cast<__nv_bfloat16>(__bfloat16_constexpr{0x7f80u}); +} + } // namespace raft