Intoducing bfloat16 type (#21522)

izdeby · facebook-github-bot · commit e72b617eb556 · 2019-07-09T21:14:10.000-07:00
Summary: Pull Request resolved: #21522 ghimport-source-id: 4803f19 Test Plan: Imported from OSS Differential Revision: D15819369 Pulled By: izdeby fbshipit-source-id: 46408dc316a5c4dc644a736dc42da2422b34bcb9
diff --git a/c10/test/util/bfloat16_test.cpp b/c10/test/util/bfloat16_test.cpp
@@ -0,0 +1,69 @@
+#include <c10/util/BFloat16.h>
+#include <gtest/gtest.h>
+
+namespace {
+  float float_from_bytes(
+      uint32_t sign,
+      uint32_t exponent,
+      uint32_t fraction
+  ) {
+      uint32_t bytes;
+      bytes = 0;
+      bytes |= sign;
+      bytes <<= 8;
+      bytes |= exponent;
+      bytes <<= 23;
+      bytes |= fraction;
+
+      float res;
+      std::memcpy(&res, &bytes, sizeof(res));
+      return res;
+  }
+
+  TEST(BFloat16Conversion, FloatToBFloat16AndBack) {
+    float in[100];
+    for (int i = 0; i < 100; ++i) {
+      in[i] = i + 1.25;
+    }
+
+    c10::BFloat16 bfloats[100];
+    float out[100];
+
+    for (int i = 0; i < 100; ++i) {
+      bfloats[i].x = c10::detail::bits_from_f32(in[i]);
+      out[i] = c10::detail::f32_from_bits(bfloats[i].x);
+
+      // The relative error should be less than 1/(2^7) since bfloat16
+      // has 7 bits mantissa.
+      EXPECT_LE(fabs(out[i] - in[i]) / in[i], 1.0 / 128);
+    }
+  }
+
+  TEST(BFloat16Conversion, NaN) {
+    float inNaN = float_from_bytes(0, 0xFF, 0x7FFFFF);
+    EXPECT_TRUE(std::isnan(inNaN));
+
+    c10::BFloat16 a = c10::BFloat16(inNaN);
+    float out = c10::detail::f32_from_bits(a.x);
+
+    EXPECT_TRUE(std::isnan(out));
+  }
+
+  TEST(BFloat16Conversion, Inf) {
+    float inInf = float_from_bytes(0, 0xFF, 0);
+    EXPECT_TRUE(std::isinf(inInf));
+
+    c10::BFloat16 a = c10::BFloat16(inInf);
+    float out = c10::detail::f32_from_bits(a.x);
+
+    EXPECT_TRUE(std::isinf(out));
+  }
+
+  TEST(BFloat16Conversion, SmallestDenormal) {
+    float in =  std::numeric_limits<float>::denorm_min(); // The smallest non-zero subnormal number
+    c10::BFloat16 a = c10::BFloat16(in);
+    float out = c10::detail::f32_from_bits(a.x);
+
+    EXPECT_FLOAT_EQ(in, out);
+  }
+} // namespace
diff --git a/c10/util/BFloat16-inl.h b/c10/util/BFloat16-inl.h
@@ -0,0 +1,17 @@
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+namespace c10 {
+
+/// Constructors
+inline C10_HOST_DEVICE BFloat16::BFloat16(float value) {
+  x = detail::bits_from_f32(value);
+}
+
+/// Implicit conversions
+inline C10_HOST_DEVICE BFloat16::operator float() const {
+  return detail::f32_from_bits(x);
+}
+
+} // namespace c10
diff --git a/c10/util/BFloat16.h b/c10/util/BFloat16.h
@@ -0,0 +1,66 @@
+#pragma once
+
+// Defines the bloat16 type (brain floating-point). This representation uses
+// 1 bit for the sign, 8 bits for the exponent and 7 bits for the mantissa.
+
+#include <c10/macros/Macros.h>
+#include <cmath>
+#include <cstring>
+
+namespace c10 {
+
+namespace detail {
+  inline C10_HOST_DEVICE float f32_from_bits(uint16_t src) {
+    float res = 0;
+    uint32_t tmp = src;
+    tmp <<= 16;
+
+#ifdef __HIP_PLATFORM_HCC__
+    float* tempRes;
+
+    // We should be using memcpy in order to respect the strict aliasing rule
+    // but it fails in the HIP environment.
+    tempRes = reinterpret_cast<float*>(&tmp);
+    res = *tempRes;
+#else
+    std::memcpy(&res, &tmp, sizeof(tmp));
+#endif
+
+    return res;
+  }
+
+  inline C10_HOST_DEVICE uint16_t bits_from_f32(float src) {
+    uint32_t res = 0;
+
+#ifdef __HIP_PLATFORM_HCC__
+    // We should be using memcpy in order to respect the strict aliasing rule
+    // but it fails in the HIP environment.
+    uint32_t* tempRes = reinterpret_cast<uint32_t*>(&src);
+    res = *tempRes;
+#else
+    std::memcpy(&res, &src, sizeof(res));
+#endif
+
+    return res >> 16;
+  }
+} // namespace detail
+
+struct alignas(2) BFloat16 {
+  uint16_t x;
+
+  // HIP wants __host__ __device__ tag, CUDA does not
+#ifdef __HIP_PLATFORM_HCC__
+  C10_HOST_DEVICE BFloat16() = default;
+#else
+  BFloat16() = default;
+
+#endif
+
+  explicit inline C10_HOST_DEVICE BFloat16(float value);
+  explicit inline C10_HOST_DEVICE operator float() const;
+};
+
+} // namespace c10
+
+
+#include <c10/util/BFloat16-inl.h>
diff --git a/c10/util/typeid.cpp b/c10/util/typeid.cpp
@@ -85,6 +85,7 @@ CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(28, at::Half*)
 CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(29, c10::qint8)
 CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(30, c10::quint8)
 CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(31, c10::qint32)
-CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(32, _CaffeHighestPreallocatedTypeId)
+CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(32, at::BFloat16)
+CAFFE_DEFINE_PREALLOCATED_KNOWN_TYPE(33, _CaffeHighestPreallocatedTypeId)
 
 } // namespace caffe2
diff --git a/c10/util/typeid.h b/c10/util/typeid.h
@@ -27,6 +27,7 @@
 #include <c10/util/qint32.h>
 #include <c10/util/qint8.h>
 #include <c10/util/quint8.h>
+#include <c10/util/BFloat16.h>
 
 /*
  * TypeIdentifier is a small type containing an id.
@@ -629,6 +630,7 @@ CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(28, at::Half*)
 CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(29, c10::qint8)
 CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(30, c10::quint8)
 CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(31, c10::qint32)
-CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(32, _CaffeHighestPreallocatedTypeId)
+CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(32, at::BFloat16)
+CAFFE_DECLARE_PREALLOCATED_KNOWN_TYPE(33, _CaffeHighestPreallocatedTypeId)
 
 } // namespace caffe2
