Add bfloat16 support

README.md

@@ -37,8 +37,7 @@ Here are the supported features:
 
 - [x] fp32/64
 - [x] float16
-- [ ] bfloat16
-- [ ] float8
+- [x] bfloat16
 
 ---
 

src/core/include/accumulate_type.h

@@ -15,6 +15,7 @@ template <typename T>
 using acc_type = typename AccumulateType<T>::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:

src/core/include/half.h

@@ -5,19 +5,22 @@
 #include <iosfwd>
 #include <limits>
 #include <ostream>
+#include <cmath>
 #include <bit>
 
 #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<float>(w);
 }
 
-HOST_DEVICE inline uint32_t fp32_to_bits(float f) {
+inline uint32_t fp32_to_bits(float f) {
     return std::bit_cast<uint32_t>(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<uint16_t>((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

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) {

src/core/include/tensor.h

@@ -233,6 +233,7 @@ class Tensor {
     std::vector<Tensor> split(std::vector<int64_t> indices, int64_t dim) const;
 
     Tensor _half() const;
+    Tensor _bfloat16() const;
     Tensor _float() const;
 
     Tensor operator+(const Tensor &other) const;

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);
 }

src/device/launcher_cuda.h

@@ -4,6 +4,7 @@
 #include <device_launch_parameters.h>
 #include <mma.h>
 #include <cuda_fp16.h>
+#include <cuda_bf16.h>
 
 #include <iostream>
 #include <string>
@@ -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<uint16_t *>(&f_);
+}
+
 #include "cutlass/gemm/device/gemm.h"
 
 template <

src/device/utils/sorting_common.h

@@ -219,6 +219,24 @@ struct KeyTraits<dtype::Half> {
     }
 };
 
+template <>
+struct KeyTraits<dtype::BFloat16> {
+    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<dtype::BFloat16 &>(v_de);
+    }
+    HOST_DEVICE static inline int endbit() {
+        return sizeof(Type) << 3;
+    }
+};
+
 template <int N, int CURRENT_VAL = N, int COUNT = 0>
 struct Log2 {
     enum { VALUE = Log2<N, (CURRENT_VAL >> 1), COUNT + 1>::VALUE };

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);
 }

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