[mlir] Fix representation of BF16 constants

[dcaballe]

Summary

This patch is a follow-up on https://reviews.llvm.org/D81127

BF16 constants were represented as 64-bit floating point values due to the lack
of support for BF16 in APFloat. APFloat was recently extended to support
BF16 so this patch is fixing the BF16 constant representation to be 16-bit.

[dcaballe]

mlir/test/IR/parser.mlir

1095

This test is actually not passing right now because the output values are not in hex format. This is what I get:

func @bfloat16_special_values() {
  %cst = constant 1.272040e-29 : bf16
  %cst_0 = constant -1.272040e-29 : bf16
  %cst_1 = constant 1.893270e-29 : bf16
  %cst_2 = constant -1.893270e-29 : bf16
  %cst_3 = constant 1.262180e-29 : bf16
  %cst_4 = constant -1.262180e-29 : bf16
  return
}

Not sure if this means that something else needs to be changed and I missed it, or this is just the way it should be printed now.

mlir/unittests/IR/AttributeTest.cpp

146

FloatAttrSplat was already resting BF16 using an attribute. I decided to change that test to F32 and then use the BF16 attribute for this. I couldn't find a way to pass a 16-bit value (not attribute) to testSplat that worked for BF16. I tried with int16_t value = ... but it crashes.

[rriddle]

Thanks for doing this! Was just thinking about this today.

mlir/test/IR/parser.mlir

1095

That leads me to believe that you have the wrong nan/inf values. The printer checks for that when deciding whether to use hex or not, unless something is wrong for the APFloat hex formatter for bfloat.

https://github.com/llvm/llvm-project/blob/d76e62fdb7a93d9a33f642b6b528f2562cc3c3f4/mlir/lib/IR/AsmPrinter.cpp#L1103

Can you manually print a nan or inf to make sure that your hex is correct/being parsed correctly?

https://github.com/llvm/llvm-project/blob/d76e62fdb7a93d9a33f642b6b528f2562cc3c3f4/mlir/lib/IR/AsmPrinter.cpp#L1137

mlir/unittests/IR/AttributeTest.cpp

146

That is fine, and expected. Literal types are checked to match, and BF16 doesn't have a c++ literal type.

[dcaballe]

If no more comments, I'll land this after https://reviews.llvm.org/D81302

mlir/test/IR/parser.mlir

1095

Thanks! The values are OK. It's a bug in LLVM: https://reviews.llvm.org/D81302

mlir/lib/IR/AttributeDetail.h

bool operator==(const KeyTy &key) const {
     return key.first == getType() && key.second.bitwiseIsEqual(getValue());
   }
   static unsigned hashKey(const KeyTy &key) {
     return llvm::hash_combine(key.first, llvm::hash_value(key.second));
   }
 
   /// Construct a key with a type and double.
   static KeyTy getKey(Type type, double value) {
-    // Treat BF16 as double because it is not supported in LLVM's APFloat.
-    // TODO(b/121118307): add BF16 support to APFloat?
-    if (type.isBF16() || type.isF64())
+    if (type.isF64())
       return KeyTy(type, APFloat(value));
 
     // This handles, e.g., F16 because there is no APFloat constructor for it.
     bool unused;
     APFloat val(value);
     val.convert(type.cast<FloatType>().getFloatSemantics(),
                 APFloat::rmNearestTiesToEven, &unused);
     return KeyTy(type, val);

 // Elements Attributes
 //===----------------------------------------------------------------------===//
 
 /// Return the bit width which DenseElementsAttr should use for this type.
 inline size_t getDenseElementBitWidth(Type eltType) {
   // Align the width for complex to 8 to make storage and interpretation easier.
   if (ComplexType comp = eltType.dyn_cast<ComplexType>())
     return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) * 2;
-  // FIXME(b/121118307): using 64 bits for BF16 because it is currently stored
-  // with double semantics.
-  if (eltType.isBF16())
-    return 64;
   if (eltType.isIndex())
     return IndexType::kInternalStorageBitWidth;
   return eltType.getIntOrFloatBitWidth();
 }
 
 /// An attribute representing a reference to a dense vector or tensor object.
 struct DenseElementsAttributeStorage : public AttributeStorage {
 public:

mlir/lib/IR/StandardTypes.cpp

unsigned FloatType::getWidth() {
   default:
     llvm_unreachable("unexpected type");
   }
 }
 
 /// Returns the floating semantics for the given type.
 const llvm::fltSemantics &FloatType::getFloatSemantics() {
   if (isBF16())
-    // Treat BF16 like a double. This is unfortunate but BF16 fltSemantics is
-    // not defined in LLVM.
-    // TODO(jpienaar): add BF16 to LLVM? fltSemantics are internal to APFloat.cc
-    // else one could add it.
-    //  static const fltSemantics semBF16 = {127, -126, 8, 16};
-    return APFloat::IEEEdouble();
+    return APFloat::BFloat();
   if (isF16())
     return APFloat::IEEEhalf();
   if (isF32())
     return APFloat::IEEEsingle();
   if (isF64())
     return APFloat::IEEEdouble();
   llvm_unreachable("non-floating point type used");
 }

mlir/lib/Parser/Parser.cpp

if (!type.isa<FloatType>())
     return (emitError("floating point value not valid for specified type"),
             nullptr);
   return FloatAttr::get(type, isNegative ? -val.getValue() : val.getValue());
 }
 
 /// Construct a float attribute bitwise equivalent to the integer literal.
 static Optional<APFloat> buildHexadecimalFloatLiteral(Parser *p, FloatType type,
                                                       uint64_t value) {
-  // FIXME: bfloat is currently stored as a double internally because it doesn't
-  // have valid APFloat semantics.
-  if (type.isF64() || type.isBF16())
+  if (type.isF64())
     return APFloat(type.getFloatSemantics(), APInt(/*numBits=*/64, value));
 
   APInt apInt(type.getWidth(), value);
   if (apInt != value) {
     p->emitError("hexadecimal float constant out of range for type");
     return llvm::None;
   }
   return APFloat(type.getFloatSemantics(), apInt);
if (!token.is(Token::floatliteral))
       return p.emitError()
              << "expected floating-point elements, but parsed integer";
 
     // Build the float values from tokens.
     auto val = token.getFloatingPointValue();
     if (!val.hasValue())
       return p.emitError("floating point value too large for attribute");
 
-    // Treat BF16 as double because it is not supported in LLVM's APFloat.
     APFloat apVal(isNegative ? -*val : *val);
-    if (!eltTy.isBF16() && !eltTy.isF64()) {
+    if (!eltTy.isF64()) {
       bool unused;
       apVal.convert(eltTy.getFloatSemantics(), APFloat::rmNearestTiesToEven,
                     &unused);
     }
     floatValues.push_back(apVal);
   }
   return success();
 }

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

 // RUN: mlir-opt -allow-unregistered-dialect %s -verify-diagnostics -split-input-file -mlir-print-elementsattrs-with-hex-if-larger=1 | FileCheck %s --check-prefix=HEX
 // RUN: mlir-opt -allow-unregistered-dialect %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<2xcomplex<f64>>
 "foo.op"() {dense.attr = dense<"0x0000000000002440000000000000144000000000000024400000000000001440"> : tensor<2xcomplex<f64>>} : () -> ()
 
 // CHECK: dense<[1.000000e+01, 5.000000e+00]> : tensor<2xbf16>
-"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2xbf16>} : () -> ()
+"foo.op"() {dense.attr = dense<"0x2041A040"> : 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}}
 "foo.op"() {dense.attr = dense<"0x0000000000002440000000000000144X"> : tensor<2xf64>} : () -> ()
 
 // -----
 
 // expected-error@+1 {{elements hex data size is invalid for provided type}}
 "foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<4xf64>} : () -> ()

mlir/test/IR/parser.mlir

func @f64_special_values() {
   // Check that values that can't be represented with the default format, use
   // hex instead.
   // CHECK: constant 0xC1CDC00000000000 : f64
   %6 = constant 0xC1CDC00000000000 : f64
 
   return
 }
 
-// FIXME: bfloat16 currently uses f64 as a storage format. This test should be
-// changed when that gets fixed.
 // CHECK-LABEL: @bfloat16_special_values
 func @bfloat16_special_values() {
   // bfloat16 signaling NaNs.
-  // CHECK: constant 0x7FF0000000000001 : bf16
-  %0 = constant 0x7FF0000000000001 : bf16
-  // CHECK: constant 0x7FF8000000000000 : bf16
-  %1 = constant 0x7FF8000000000000 : bf16
+  // CHECK: constant 0x7F81 : bf16
+  %0 = constant 0x7F81 : bf16
+  // CHECK: constant 0xFF81 : bf16
+  %1 = constant 0xFF81 : bf16
 
   // bfloat16 quiet NaNs.
-  // CHECK: constant 0x7FF0000001000000 : bf16
-  %2 = constant 0x7FF0000001000000 : bf16
-  // CHECK: constant 0xFFF0000001000000 : bf16
-  %3 = constant 0xFFF0000001000000 : bf16
+  // CHECK: constant 0x7FC0 : bf16
+  %2 = constant 0x7FC0 : bf16
+  // CHECK: constant 0xFFC0 : bf16
+  %3 = constant 0xFFC0 : bf16
 
   // bfloat16 positive infinity.
-  // CHECK: constant 0x7FF0000000000000 : bf16
-  %4 = constant 0x7FF0000000000000 : bf16
+  // CHECK: constant 0x7F80 : bf16
+  %4 = constant 0x7F80 : bf16
   // bfloat16 negative infinity.
-  // CHECK: constant 0xFFF0000000000000 : bf16
-  %5 = constant 0xFFF0000000000000 : bf16
+  // CHECK: constant 0xFF80 : bf16
+  %5 = constant 0xFF80 : bf16
 
   return
 }
 
 // We want to print floats in exponential notation with 6 significant digits,
 // but it may lead to precision loss when parsing back, in which case we print
 // the decimal form instead.
 // CHECK-LABEL: @f32_potential_precision_loss()

 
 // This tests the behavior
 func @pretty_names() {
   // Simple case, should parse and print as %x being an implied 'name'
   // attribute.
   %x = test.string_attr_pretty_name
   // CHECK: %x = test.string_attr_pretty_name
   // CHECK-NOT: attributes
-  
   // This specifies an explicit name, which should override the result.
   %YY = test.string_attr_pretty_name attributes { names = ["y"] }
   // CHECK: %y = test.string_attr_pretty_name
   // CHECK-NOT: attributes
-  
   // Conflicts with the 'y' name, so need an explicit attribute.
   %0 = "test.string_attr_pretty_name"() { names = ["y"]} : () -> i32
   // CHECK: %y_0 = test.string_attr_pretty_name attributes {names = ["y"]}
 
   // Name contains a space.
   %1 = "test.string_attr_pretty_name"() { names = ["space name"]} : () -> i32
   // CHECK: %space_name = test.string_attr_pretty_name attributes {names = ["space name"]}
 

mlir/test/Target/llvmir.mlir

-// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
 
 // CHECK: @i32_global = internal global i32 42
 llvm.mlir.global internal @i32_global(42: i32) : !llvm.i32
 
 // CHECK: @i32_const = internal constant i53 52
 llvm.mlir.global internal constant @i32_const(52: i53) : !llvm.i53
 
 // CHECK: @int_global_array = internal global [3 x i32] [i32 62, i32 62, i32 62]
llvm.func @passthrough() attributes {passthrough = ["noinline", ["alignstack", "4"], "null_pointer_is_valid", ["foo", "bar"]]} {
   llvm.return
 }
 
 // CHECK: attributes #[[ATTR_GROUP]] = {
 // CHECK-DAG: noinline
 // CHECK-DAG: alignstack=4
 // CHECK-DAG: null_pointer_is_valid
 // CHECK-DAG: "foo"="bar"
+
+// -----
+
+// CHECK-LABEL: @constant_bf16
+llvm.func @constant_bf16() -> !llvm<"bfloat"> {
+  %0 = llvm.mlir.constant(1.000000e+01 : bf16) : !llvm<"bfloat">
+  llvm.return %0 : !llvm<"bfloat">
+}
+
+// CHECK: ret bfloat 0xR4120
+

mlir/unittests/IR/AttributeTest.cpp

TEST(DenseSplatTest, F64Splat) {
   FloatType floatTy = FloatType::getF64(&context);
   double value = 10.0;
 
   testSplat(floatTy, APFloat(value));
 }
 
 TEST(DenseSplatTest, FloatAttrSplat) {
   MLIRContext context;
-  FloatType floatTy = FloatType::getBF16(&context);
+  FloatType floatTy = FloatType::getF32(&context);
   Attribute value = FloatAttr::get(floatTy, 10.0);
 
   testSplat(floatTy, value);
 }
 
 TEST(DenseSplatTest, BF16Splat) {
   MLIRContext context;
   FloatType floatTy = FloatType::getBF16(&context);
-  // Note: We currently use double to represent bfloat16.
-  double value = 10.0;
+  Attribute value = FloatAttr::get(floatTy, 10.0);
 
   testSplat(floatTy, value);
 }
 
 TEST(DenseSplatTest, StringSplat) {
   MLIRContext context;
   Type stringType =
       OpaqueType::get(Identifier::get("test", &context), "string", &context);