Flush denormals to +/- 0 when converting float to bfloat16.

tensorflower-gardener · tensorflower-gardener · commit b04c4e0e4338 · 2020-03-19T20:10:36.000-07:00
PiperOrigin-RevId: 301948798
Change-Id: Ic24b699b2e23683d3710d7abb4317833df252af0
diff --git a/tensorflow/core/framework/bfloat16_test.cc b/tensorflow/core/framework/bfloat16_test.cc
@@ -23,6 +23,35 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+TEST(Bfloat16Test, ZeroRepresentations) {
+  ASSERT_EQ(bfloat16{0.0f}, bfloat16{0.0f});
+  ASSERT_EQ(bfloat16{-0.0f}, bfloat16{0.0f});
+  ASSERT_EQ(bfloat16{-0.0f}, bfloat16{-0.0f});
+  ASSERT_EQ(bfloat16{0.0f}.value, 0x0000);
+  ASSERT_EQ(bfloat16{-0.0f}.value, 0x8000);
+}
+
+TEST(Bfloat16Test, FlushDenormalsToZero) {
+  for (float denorm = -std::numeric_limits<float>::denorm_min();
+       denorm < std::numeric_limits<float>::denorm_min();
+       denorm = std::nextafterf(denorm, 1.0f)) {
+    bfloat16 bf_trunc = bfloat16::truncate_to_bfloat16(denorm);
+    ASSERT_EQ(float{bf_trunc}, 0.0f);
+    if (std::signbit(denorm)) {
+      ASSERT_EQ(bf_trunc.value, 0x8000) << denorm;
+    } else {
+      ASSERT_EQ(bf_trunc.value, 0x0000) << denorm;
+    }
+    bfloat16 bf_round = bfloat16::round_to_bfloat16(denorm);
+    ASSERT_EQ(float{bf_round}, 0.0f);
+    if (std::signbit(denorm)) {
+      ASSERT_EQ(bf_round.value, 0x8000) << denorm;
+    } else {
+      ASSERT_EQ(bf_round.value, 0x0000) << denorm;
+    }
+  }
+}
+
 TEST(Bfloat16Test, DefaultValueIsZero) {
   EXPECT_EQ(0.0f, static_cast<float>(bfloat16()));
 }
@@ -65,6 +94,7 @@ TEST_P(Bfloat16Test, TruncateTest) {
     EXPECT_TRUE(std::isnan(float(truncated)) || std::isinf(float(truncated)));
     return;
   }
+
   EXPECT_EQ(GetParam().expected_truncation, float(truncated));
 
   bfloat16 rounded = bfloat16::round_to_bfloat16((GetParam().input));
@@ -114,14 +144,16 @@ INSTANTIATE_TEST_SUITE_P(
             BinaryToFloat(0, 0b10000000, 0b1001000, 0b1000000000000000),
             BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000),
             BinaryToFloat(0, 0b10000000, 0b1001000, 0b0000000000000000)},
+        // The following two floats are denormals and will be flushed
+        // to zero.
         Bfloat16TestParam{
             BinaryToFloat(0, 0b00000000, 0b1001000, 0b1000000000000000),
-            BinaryToFloat(0, 0b00000000, 0b1001000, 0b0000000000000000),
-            BinaryToFloat(0, 0b00000000, 0b1001000, 0b0000000000000000)},
+            BinaryToFloat(0, 0b00000000, 0b0000000, 0b0000000000000000),
+            BinaryToFloat(0, 0b00000000, 0b0000000, 0b0000000000000000)},
         Bfloat16TestParam{
             BinaryToFloat(0, 0b00000000, 0b1111111, 0b1100000000000000),
-            BinaryToFloat(0, 0b00000000, 0b1111111, 0b0000000000000000),
-            BinaryToFloat(0, 0b00000001, 0b0000000, 0b0000000000000000)}));
+            BinaryToFloat(0, 0b00000000, 0b0000000, 0b0000000000000000),
+            BinaryToFloat(0, 0b00000000, 0b0000000, 0b0000000000000000)}));
 
 TEST(Bfloat16Test, Conversion) {
   float a[100];
diff --git a/tensorflow/core/lib/bfloat16/bfloat16.h b/tensorflow/core/lib/bfloat16/bfloat16.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <cmath>
 #include <complex>
 #include <iostream>
+#include <limits>
 
 #include "tensorflow/core/platform/byte_order.h"
 
@@ -53,6 +54,10 @@ struct bfloat16 {
     if (float_isnan(v)) {
       output.value = NAN_VALUE;
       return output;
+    } else if (std::fabs(v) < std::numeric_limits<float>::min()) {
+      // Flush denormal to +/- 0.
+      output.value = std::signbit(v) ? 0x8000 : 0;
+      return output;
     }
     const uint16_t* p = reinterpret_cast<const uint16_t*>(&v);
 #if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
@@ -196,6 +201,9 @@ struct bfloat16 {
       // qNaN magic: All exponent bits set + most significant bit of fraction
       // set.
       output.value = 0x7fc0;
+    } else if (std::fabs(v) < std::numeric_limits<float>::min()) {
+      // Flush denormal to +/- 0.0
+      output.value = std::signbit(v) ? 0x8000 : 0;
     } else {
       // Fast rounding algorithm that rounds a half value to nearest even. This
       // reduces expected error when we convert a large number of floats. Here
