[Mlir-commits] [mlir] b3e6487 - [mlir][DenseElementsAttr] Fix storage size for bfloat16 when parsing from hex.

[River Riddle]

Author: River Riddle
Date: 2020-02-25T15:00:32-08:00
New Revision: b3e6487f02ac98b6c4a501f5da10c867bd98f769

URL: https://github.com/llvm/llvm-project/commit/b3e6487f02ac98b6c4a501f5da10c867bd98f769
DIFF: https://github.com/llvm/llvm-project/commit/b3e6487f02ac98b6c4a501f5da10c867bd98f769.diff

LOG: [mlir][DenseElementsAttr] Fix storage size for bfloat16 when parsing from hex.

Summary: bfloat16 is stored internally as a double, so we can't direct use Type::getIntOrFloatBitWidth.

Differential Revision: https://reviews.llvm.org/D75133

Added: 
    

Modified: 
    mlir/lib/Parser/Parser.cpp
    mlir/test/IR/dense-elements-hex.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 03c676f905d4..9064ff92650d 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -2118,8 +2118,12 @@ DenseElementsAttr TensorLiteralParser::getHexAttr(llvm::SMLoc loc,
 
   // Check that the size of the hex data correpsonds to the size of the type, or
   // a splat of the type.
+  // TODO: bf16 is currently stored as a double, this should be removed when
+  // APFloat properly supports it.
+  int64_t elementWidth =
+      elementType.isBF16() ? 64 : elementType.getIntOrFloatBitWidth();
   if (static_cast<int64_t>(data.size() * CHAR_BIT) !=
-      (type.getNumElements() * elementType.getIntOrFloatBitWidth())) {
+      (type.getNumElements() * elementWidth)) {
     p.emitError(loc) << "elements hex data size is invalid for provided type: "
                      << type;
     return nullptr;

diff  --git a/mlir/test/IR/dense-elements-hex.mlir b/mlir/test/IR/dense-elements-hex.mlir
index b3f9539e0111..f2ce955357fb 100644
--- a/mlir/test/IR/dense-elements-hex.mlir
+++ b/mlir/test/IR/dense-elements-hex.mlir
@@ -7,6 +7,9 @@
 // CHECK: dense<[1.000000e+01, 5.000000e+00]> : tensor<2xf64>
 "foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
 
+// CHECK: dense<[1.000000e+01, 5.000000e+00]> : tensor<2xbf16>
+"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2xbf16>} : () -> ()
+
 // -----
 
 // expected-error at +1 {{elements hex string should start with '0x'}}