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

[rriddle]

Summary

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

[lucyrfox]

LGTM, thanks for the quick fix, River!

mlir/lib/Parser/Parser.cpp

DenseElementsAttr TensorLiteralParser::getHexAttr(llvm::SMLoc loc,
   }
 
   std::string data;
   if (parseElementAttrHexValues(p, hexStorage.getValue(), data))
     return nullptr;
 
   // 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;
   }
 
   return DenseElementsAttr::getFromRawBuffer(
       type, ArrayRef<char>(data.data(), data.size()), /*isSplatBuffer=*/false);
 }

mlir/test/IR/dense-elements-hex.mlir

 // RUN: mlir-opt %s -verify-diagnostics -split-input-file -mlir-print-elementsattrs-with-hex-if-larger=1 | FileCheck %s --check-prefix=HEX
 // RUN: mlir-opt %s -verify-diagnostics -split-input-file | FileCheck %s
 
 // HEX: dense<"0x00000000000024400000000000001440"> : tensor<2xf64>
 "foo.op"() {dense.attr = dense<[10.0, 5.0]> : tensor<2xf64>} : () -> ()
 
 // 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@+1 {{elements hex string should start with '0x'}}
 "foo.op"() {dense.attr = dense<"00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
 
 // -----
 
 // expected-error@+1 {{elements hex string only contains hex digits}}