From f2959b6795f00deccb956009fb81b45edc732945 Mon Sep 17 00:00:00 2001 From: xytpai Date: Wed, 30 Jul 2025 11:37:42 +0800 Subject: [PATCH] add bf16 --- README.md | 3 +- src/core/include/accumulate_type.h | 3 + src/core/include/half.h | 120 +++++++++++++++++++++++++---- src/core/include/scalar_type.h | 23 +++--- src/core/include/tensor.h | 1 + src/core/tensor.cpp | 4 + src/device/launcher_cuda.h | 11 +++ src/device/utils/sorting_common.h | 18 +++++ src/register.cpp | 2 + test/test_tensor.py | 6 ++ 10 files changed, 161 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 8d6f178..f997abf 100644 --- a/README.md +++ b/README.md @@ -37,8 +37,7 @@ Here are the supported features: - [x] fp32/64 - [x] float16 -- [ ] bfloat16 -- [ ] float8 +- [x] bfloat16 --- diff --git a/src/core/include/accumulate_type.h b/src/core/include/accumulate_type.h index fb6d441..e3d3dad 100644 --- a/src/core/include/accumulate_type.h +++ b/src/core/include/accumulate_type.h @@ -15,6 +15,7 @@ template using acc_type = typename AccumulateType::type; ACC_TYPE(dtype::Half, float) +ACC_TYPE(dtype::BFloat16, float) ACC_TYPE(float, float) ACC_TYPE(double, double) ACC_TYPE(int8_t, int64_t) @@ -29,6 +30,8 @@ static ScalarType accumulate_type(ScalarType dtype) { switch (dtype) { case ScalarType::Half: return ScalarType::Float; + case ScalarType::BFloat16: + return ScalarType::Float; case ScalarType::Float: return ScalarType::Float; case ScalarType::Double: diff --git a/src/core/include/half.h b/src/core/include/half.h index c61e9ec..eda1f53 100644 --- a/src/core/include/half.h +++ b/src/core/include/half.h @@ -5,19 +5,22 @@ #include #include #include +#include #include #include "launcher.h" -HOST_DEVICE inline float fp32_from_bits(uint32_t w) { +namespace dtype { + +inline float fp32_from_bits(uint32_t w) { return std::bit_cast(w); } -HOST_DEVICE inline uint32_t fp32_to_bits(float f) { +inline uint32_t fp32_to_bits(float f) { return std::bit_cast(f); } -HOST_DEVICE inline float fp16_ieee_to_fp32_value(uint16_t h) { +inline float fp16_ieee_to_fp32_value(uint16_t h) { /* * Extend the half-precision floating-point number to 32 bits and shift to the * upper part of the 32-bit word: @@ -181,28 +184,111 @@ inline uint16_t fp16_ieee_from_fp32_value(float f) { (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)); } -namespace dtype { +inline float f32_from_bits(uint16_t src) { + float res = 0; + uint32_t tmp = src; + tmp <<= 16; + std::memcpy(&res, &tmp, sizeof(tmp)); + return res; +} + +inline uint16_t round_to_nearest_even(float src) { + if (std::isnan(src)) { + return UINT16_C(0x7FC0); + } else { + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) + union { + uint32_t U32; // NOLINT(facebook-hte-BadMemberName) + float F32; // NOLINT(facebook-hte-BadMemberName) + }; + F32 = src; + uint32_t rounding_bias = ((U32 >> 16) & 1) + UINT32_C(0x7FFF); + return static_cast((U32 + rounding_bias) >> 16); + } +} + +class BFloat16; struct alignas(2) Half { unsigned short x; Half() = default; #if defined(__CUDACC__) - inline DEVICE operator float() const { - return half_to_float(x); - } - - inline DEVICE Half(float f) { - x = float_to_half(f); - } + DEVICE inline operator float() const; + DEVICE inline Half(float f); + DEVICE inline operator BFloat16() const; #else - inline HOST_DEVICE operator float() const { - return fp16_ieee_to_fp32_value(x); - } + inline operator float() const; + inline Half(float f); + inline operator BFloat16() const; +#endif +}; - inline HOST_DEVICE Half(float f) { - x = fp16_ieee_from_fp32_value(f); - } +struct alignas(2) BFloat16 { + unsigned short x; + BFloat16() = default; +#if defined(__CUDACC__) + DEVICE inline operator float() const; + DEVICE inline BFloat16(float f); + DEVICE inline operator Half() const; +#else + inline operator float() const; + inline BFloat16(float f); + inline operator Half() const; #endif }; +#if defined(__CUDACC__) + +DEVICE inline Half::operator float() const { + return half_to_float(x); +} + +DEVICE inline Half::Half(float f) { + x = float_to_half(f); +} + +DEVICE inline Half::operator BFloat16() const { + return BFloat16(half_to_float(x)); +} + +DEVICE inline BFloat16::operator float() const { + return bfloat16_to_float(x); +} + +DEVICE inline BFloat16::BFloat16(float f) { + x = float_to_bfloat16(f); +} + +DEVICE inline BFloat16::operator Half() const { + return Half(bfloat16_to_float(x)); +} + +#else + +inline Half::operator float() const { + return fp16_ieee_to_fp32_value(x); +} + +inline Half::Half(float f) { + x = fp16_ieee_from_fp32_value(f); +} + +inline Half::operator BFloat16() const { + return BFloat16(fp16_ieee_to_fp32_value(x)); +} + +inline BFloat16::operator float() const { + return f32_from_bits(x); +} + +inline BFloat16::BFloat16(float f) { + x = round_to_nearest_even(f); +} + +inline BFloat16::operator Half() const { + return Half(f32_from_bits(x)); +} + +#endif + } // namespace dtype diff --git a/src/core/include/scalar_type.h b/src/core/include/scalar_type.h index e6d172a..ad9f74a 100644 --- a/src/core/include/scalar_type.h +++ b/src/core/include/scalar_type.h @@ -6,16 +6,17 @@ #include "half.h" #include "exception.h" -#define FORALL_BASIC_SCALAR_TYPES(_, ...) \ - _(bool, Bool, __VA_ARGS__) /* 0 */ \ - _(uint8_t, Byte, __VA_ARGS__) /* 1 */ \ - _(int8_t, Char, __VA_ARGS__) /* 2 */ \ - _(int16_t, Short, __VA_ARGS__) /* 3 */ \ - _(int, Int, __VA_ARGS__) /* 4 */ \ - _(int64_t, Long, __VA_ARGS__) /* 5 */ \ - _(dtype::Half, Half, __VA_ARGS__) /* 6 */ \ - _(float, Float, __VA_ARGS__) /* 7 */ \ - _(double, Double, __VA_ARGS__) /* 8 */ +#define FORALL_BASIC_SCALAR_TYPES(_, ...) \ + _(bool, Bool, __VA_ARGS__) /* 0 */ \ + _(uint8_t, Byte, __VA_ARGS__) /* 1 */ \ + _(int8_t, Char, __VA_ARGS__) /* 2 */ \ + _(int16_t, Short, __VA_ARGS__) /* 3 */ \ + _(int, Int, __VA_ARGS__) /* 4 */ \ + _(int64_t, Long, __VA_ARGS__) /* 5 */ \ + _(dtype::Half, Half, __VA_ARGS__) /* 6 */ \ + _(dtype::BFloat16, BFloat16, __VA_ARGS__) /* 7 */ \ + _(float, Float, __VA_ARGS__) /* 8 */ \ + _(double, Double, __VA_ARGS__) /* 9 */ enum class ScalarType : int8_t { #define DEFINE_ENUM(_, n, ...) n, @@ -109,7 +110,7 @@ FORALL_BASIC_SCALAR_TYPES(SPECIALIZE_CppTypeToScalarType) #undef SPECIALIZE_CppTypeToScalarType inline bool is_floating_type(ScalarType t) { - return t == ScalarType::Double || t == ScalarType::Float || t == ScalarType::Half; + return t == ScalarType::Double || t == ScalarType::Float || t == ScalarType::Half || t == ScalarType::BFloat16; } inline bool is_unsigned_int_type(ScalarType t) { diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index ccac42a..d12b9b8 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -233,6 +233,7 @@ class Tensor { std::vector split(std::vector indices, int64_t dim) const; Tensor _half() const; + Tensor _bfloat16() const; Tensor _float() const; Tensor operator+(const Tensor &other) const; diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index af3451d..7e1f791 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -291,6 +291,10 @@ Tensor Tensor::_half() const { return gpu::convert(*this, ScalarType::Half); } +Tensor Tensor::_bfloat16() const { + return gpu::convert(*this, ScalarType::BFloat16); +} + Tensor Tensor::_float() const { return gpu::convert(*this, ScalarType::Float); } diff --git a/src/device/launcher_cuda.h b/src/device/launcher_cuda.h index fe314c4..ca9d482 100644 --- a/src/device/launcher_cuda.h +++ b/src/device/launcher_cuda.h @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -523,6 +524,16 @@ inline DEVICE uint16_t float_to_half(float f) { return tmp.u16[0]; } +inline DEVICE float bfloat16_to_float(uint16_t h) { + auto h_ = *reinterpret_cast<__nv_bfloat16 *>(&h); + return __bfloat162float(h_); +} + +inline DEVICE uint16_t float_to_bfloat16(float f) { + auto f_ = __float2bfloat16(f); + return *reinterpret_cast(&f_); +} + #include "cutlass/gemm/device/gemm.h" template < diff --git a/src/device/utils/sorting_common.h b/src/device/utils/sorting_common.h index c7d9cd3..731501d 100644 --- a/src/device/utils/sorting_common.h +++ b/src/device/utils/sorting_common.h @@ -219,6 +219,24 @@ struct KeyTraits { } }; +template <> +struct KeyTraits { + using Type = uint16_t; + HOST_DEVICE static inline Type convert(dtype::BFloat16 v) { + Type x = *((Type *)&v); + Type mask = -((x >> 15)) | 0x8000; + return (x ^ mask); + } + HOST_DEVICE static inline dtype::BFloat16 deconvert(Type v) { + Type mask = ((v >> 15) - 1) | 0x8000; + Type v_de = v ^ mask; + return reinterpret_cast(v_de); + } + HOST_DEVICE static inline int endbit() { + return sizeof(Type) << 3; + } +}; + template struct Log2 { enum { VALUE = Log2> 1), COUNT + 1>::VALUE }; diff --git a/src/register.cpp b/src/register.cpp index c05adf3..9b2869b 100644 --- a/src/register.cpp +++ b/src/register.cpp @@ -68,6 +68,7 @@ PYBIND11_MODULE(kfunca, m) { .value("int", ScalarType::Int) .value("long", ScalarType::Long) .value("half", ScalarType::Half) + .value("bfloat16", ScalarType::BFloat16) .value("float", ScalarType::Float) .value("double", ScalarType::Double) .value("bool", ScalarType::Bool) @@ -207,5 +208,6 @@ PYBIND11_MODULE(kfunca, m) { .def("mean_var", &Tensor::mean_var) .def("norm_stat", &Tensor::norm_stat) .def("half", &Tensor::_half) + .def("bfloat16", &Tensor::_bfloat16) .def("float", &Tensor::_float); } diff --git a/test/test_tensor.py b/test/test_tensor.py index 4516428..873de90 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -151,6 +151,12 @@ def test_convert(self): arr_gpu *= arr_gpu arr_gpu_half *= arr_gpu_half assert_allclose(arr_gpu, arr_gpu_half.float()) + arr = np.random.uniform(-10, 10, size=(2, 3)) + arr_gpu = kfunca.from_numpy(arr, 0) + arr_gpu_bf = arr_gpu.bfloat16() + arr_gpu *= arr_gpu + arr_gpu_bf *= arr_gpu_bf + assert_allclose(arr_gpu, arr_gpu_bf.float(), atol=1e-1, rtol=1e-1) def test_permute(self): arr = np.random.uniform(-10, 10, size=(16, 8, 64, 11)) # 0,1,2,3