Add Float8E4M3FN type to MLIR.

[reedwm]

Summary

The paper https://arxiv.org/abs/2209.05433 introduces two new FP8 dtypes: E5M2 (called Float8E5M2 in LLVM) and E4M3 (called Float8E4M3FN in LLVM). Support for Float8E5M2 in APFloat and MLIR was added in https://reviews.llvm.org/D133823. Support for Float8E4M3FN in APFloat was added in https://reviews.llvm.org/D137760. This change adds Float8E4M3FN to MLIR as well.

There is an RFC for adding the FP8 dtypes here: https://discourse.llvm.org/t/rfc-add-apfloat-and-mlir-type-support-for-fp8-e5m2/65279.

This change is identical to the MLIR changes in the patch that added Float8E5M2, except that Float8E4M3FN is added instead.

[stellaraccident]

Thanks - this is great!

mlir/include/mlir-c/BuiltinTypes.h

 
 /// Checks whether the given type is an f8E5M2 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M2(MlirType type);
 
 /// Creates an f8E5M2 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M2TypeGet(MlirContext ctx);
 
+/// Checks whether the given type is an f8E4M3FN type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3FN(MlirType type);
+
+/// Creates an f8E4M3FN type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx);
+
 /// Checks whether the given type is a bf16 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsABF16(MlirType type);
 
 /// Creates a bf16 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirBF16TypeGet(MlirContext ctx);
 
 /// Checks whether the given type is an f16 type.

mlir/include/mlir/IR/Builders.h

public:
 
   // Locations.
   Location getUnknownLoc();
   Location getFusedLoc(ArrayRef<Location> locs,
                        Attribute metadata = Attribute());
 
   // Types.
   FloatType getFloat8E5M2Type();
+  FloatType getFloat8E4M3FNType();
   FloatType getBF16Type();
   FloatType getF16Type();
   FloatType getF32Type();
   FloatType getF64Type();
   FloatType getF80Type();
   FloatType getF128Type();
 
   IndexType getIndexType();

mlir/include/mlir/IR/BuiltinTypes.h

public:
   // Convenience factories.
   static FloatType getBF16(MLIRContext *ctx);
   static FloatType getF16(MLIRContext *ctx);
   static FloatType getF32(MLIRContext *ctx);
   static FloatType getF64(MLIRContext *ctx);
   static FloatType getF80(MLIRContext *ctx);
   static FloatType getF128(MLIRContext *ctx);
   static FloatType getFloat8E5M2(MLIRContext *ctx);
+  static FloatType getFloat8E4M3FN(MLIRContext *ctx);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Return the bitwidth of this float type.
   unsigned getWidth();
 
   /// Return the width of the mantissa of this type.

 
 inline bool BaseMemRefType::isValidElementType(Type type) {
   return type.isIntOrIndexOrFloat() ||
          type.isa<ComplexType, MemRefType, VectorType, UnrankedMemRefType>() ||
          type.isa<MemRefElementTypeInterface>();
 }
 
 inline bool FloatType::classof(Type type) {
-  return type.isa<Float8E5M2Type, BFloat16Type, Float16Type, Float32Type,
-                  Float64Type, Float80Type, Float128Type>();
+  return type.isa<Float8E5M2Type, Float8E4M3FNType, BFloat16Type, Float16Type,
+                  Float32Type, Float64Type, Float80Type, Float128Type>();
 }
 
 inline FloatType FloatType::getFloat8E5M2(MLIRContext *ctx) {
   return Float8E5M2Type::get(ctx);
 }
 
+inline FloatType FloatType::getFloat8E4M3FN(MLIRContext *ctx) {
+  return Float8E4M3FNType::get(ctx);
+}
+
 inline FloatType FloatType::getBF16(MLIRContext *ctx) {
   return BFloat16Type::get(ctx);
 }
 
 inline FloatType FloatType::getF16(MLIRContext *ctx) {
   return Float16Type::get(ctx);
 }
 

mlir/include/mlir/IR/BuiltinTypes.td

def Builtin_Float8E5M2 : Builtin_FloatType<"Float8E5M2"> {
   let description = [{
     An 8-bit floating point type with 1 sign bit, 5 bits exponent and 2 bits
     mantissa. This is not a standard type as defined by IEEE-754, but it
     follows similar conventions with the following characteristics:
 
       * bit encoding: S1E5M2
       * exponent bias: 15
       * infinities: supported with exponent set to all 1s and mantissa 0s
-      * NaNs: supported with exponent bits set to all 1s and mantissa of 
+      * NaNs: supported with exponent bits set to all 1s and mantissa of
         (01, 10, or 11)
       * denormals when exponent is 0
 
     Described in: https://arxiv.org/abs/2209.05433
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Float8E4M3FNType
+
+def Builtin_Float8E4M3FN : Builtin_FloatType<"Float8E4M3FN"> {
+  let summary = "8-bit floating point with 3 bit mantissa";
+  let description = [{
+    An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits
+    mantissa. This is not a standard type as defined by IEEE-754, but it follows
+    similar conventions, with the exception that there are no infinity values
+    and only two NaN representations. This type has the following
+    characteristics:
+
+      * bit encoding: S1E4M3
+      * exponent bias: 7
+      * infinities: Not supported
+      * NaNs: supported with exponent bits and mantissa bits set to all 1s
+      * denormals when exponent is 0
+
+    Described in: https://arxiv.org/abs/2209.05433
+  }];
+}
 
 //===----------------------------------------------------------------------===//
 // BFloat16Type
 
 def Builtin_BFloat16 : Builtin_FloatType<"BFloat16"> {
   let summary = "bfloat16 floating-point type";
 }
 

mlir/include/mlir/IR/Types.h

public:
 
   /// Get the dialect this type is registered to.
   Dialect &getDialect() const { return impl->getAbstractType().getDialect(); }
 
   // Convenience predicates.  This is only for floating point types,
   // derived types should use isa/dyn_cast.
   bool isIndex() const;
   bool isFloat8E5M2() const;
+  bool isFloat8E4M3FN() const;
   bool isBF16() const;
   bool isF16() const;
   bool isF32() const;
   bool isF64() const;
   bool isF80() const;
   bool isF128() const;
 
   /// Return true if this is an integer type with the specified width.

mlir/lib/AsmParser/TokenKinds.def

 TOK_KEYWORD(complex)
 TOK_KEYWORD(dense)
 TOK_KEYWORD(dense_resource)
 TOK_KEYWORD(f16)
 TOK_KEYWORD(f32)
 TOK_KEYWORD(f64)
 TOK_KEYWORD(f80)
 TOK_KEYWORD(f8E5M2)
+TOK_KEYWORD(f8E4M3FN)
 TOK_KEYWORD(f128)
 TOK_KEYWORD(false)
 TOK_KEYWORD(floordiv)
 TOK_KEYWORD(for)
 TOK_KEYWORD(func)
 TOK_KEYWORD(index)
 TOK_KEYWORD(loc)
 TOK_KEYWORD(max)

mlir/lib/AsmParser/TypeParser.cpp

OptionalParseResult Parser::parseOptionalType(Type &type) {
   case Token::l_paren:
   case Token::kw_memref:
   case Token::kw_tensor:
   case Token::kw_complex:
   case Token::kw_tuple:
   case Token::kw_vector:
   case Token::inttype:
   case Token::kw_f8E5M2:
+  case Token::kw_f8E4M3FN:
   case Token::kw_bf16:
   case Token::kw_f16:
   case Token::kw_f32:
   case Token::kw_f64:
   case Token::kw_f80:
   case Token::kw_f128:
   case Token::kw_index:
   case Token::kw_none:
case Token::inttype: {
     consumeToken(Token::inttype);
     return IntegerType::get(getContext(), *width, signSemantics);
   }
 
   // float-type
   case Token::kw_f8E5M2:
     consumeToken(Token::kw_f8E5M2);
     return builder.getFloat8E5M2Type();
+  case Token::kw_f8E4M3FN:
+    consumeToken(Token::kw_f8E4M3FN);
+    return builder.getFloat8E4M3FNType();
   case Token::kw_bf16:
     consumeToken(Token::kw_bf16);
     return builder.getBF16Type();
   case Token::kw_f16:
     consumeToken(Token::kw_f16);
     return builder.getF16Type();
   case Token::kw_f32:
     consumeToken(Token::kw_f32);

mlir/lib/CAPI/IR/BuiltinTypes.cpp

 bool mlirTypeIsAFloat8E5M2(MlirType type) {
   return unwrap(type).isFloat8E5M2();
 }
 
 MlirType mlirFloat8E5M2TypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E5M2(unwrap(ctx)));
 }
 
+bool mlirTypeIsAFloat8E4M3FN(MlirType type) {
+  return unwrap(type).isFloat8E4M3FN();
+}
+
+MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat8E4M3FN(unwrap(ctx)));
+}
+
 bool mlirTypeIsABF16(MlirType type) { return unwrap(type).isBF16(); }
 
 MlirType mlirBF16TypeGet(MlirContext ctx) {
   return wrap(FloatType::getBF16(unwrap(ctx)));
 }
 
 bool mlirTypeIsAF16(MlirType type) { return unwrap(type).isF16(); }
 

mlir/lib/IR/AsmPrinter.cpp

 void AsmPrinter::Impl::printTypeImpl(Type type) {
   TypeSwitch<Type>(type)
       .Case<OpaqueType>([&](OpaqueType opaqueTy) {
         printDialectSymbol(os, "!", opaqueTy.getDialectNamespace(),
                            opaqueTy.getTypeData());
       })
       .Case<IndexType>([&](Type) { os << "index"; })
       .Case<Float8E5M2Type>([&](Type) { os << "f8E5M2"; })
+      .Case<Float8E4M3FNType>([&](Type) { os << "f8E4M3FN"; })
       .Case<BFloat16Type>([&](Type) { os << "bf16"; })
       .Case<Float16Type>([&](Type) { os << "f16"; })
       .Case<Float32Type>([&](Type) { os << "f32"; })
       .Case<Float64Type>([&](Type) { os << "f64"; })
       .Case<Float80Type>([&](Type) { os << "f80"; })
       .Case<Float128Type>([&](Type) { os << "f128"; })
       .Case<IntegerType>([&](IntegerType integerTy) {
         if (integerTy.isSigned())

mlir/lib/IR/Builders.cpp

 //===----------------------------------------------------------------------===//
 // Types.
 //===----------------------------------------------------------------------===//
 
 FloatType Builder::getFloat8E5M2Type() {
   return FloatType::getFloat8E5M2(context);
 }
 
+FloatType Builder::getFloat8E4M3FNType() {
+  return FloatType::getFloat8E4M3FN(context);
+}
+
 FloatType Builder::getBF16Type() { return FloatType::getBF16(context); }
 
 FloatType Builder::getF16Type() { return FloatType::getF16(context); }
 
 FloatType Builder::getF32Type() { return FloatType::getF32(context); }
 
 FloatType Builder::getF64Type() { return FloatType::getF64(context); }
 

mlir/lib/IR/BuiltinTypes.cpp

IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
   return IntegerType::get(getContext(), scale * getWidth(), getSignedness());
 }
 
 //===----------------------------------------------------------------------===//
 // Float Type
 //===----------------------------------------------------------------------===//
 
 unsigned FloatType::getWidth() {
-  if (isa<Float8E5M2Type>())
+  if (isa<Float8E5M2Type, Float8E4M3FNType>())
     return 8;
   if (isa<Float16Type, BFloat16Type>())
     return 16;
   if (isa<Float32Type>())
     return 32;
   if (isa<Float64Type>())
     return 64;
   if (isa<Float80Type>())
     return 80;
   if (isa<Float128Type>())
     return 128;
   llvm_unreachable("unexpected float type");
 }
 
 /// Returns the floating semantics for the given type.
 const llvm::fltSemantics &FloatType::getFloatSemantics() {
   if (isa<Float8E5M2Type>())
     return APFloat::Float8E5M2();
+  if (isa<Float8E4M3FNType>())
+    return APFloat::Float8E4M3FN();
   if (isa<BFloat16Type>())
     return APFloat::BFloat();
   if (isa<Float16Type>())
     return APFloat::IEEEhalf();
   if (isa<Float32Type>())
     return APFloat::IEEEsingle();
   if (isa<Float64Type>())
     return APFloat::IEEEdouble();

mlir/lib/IR/MLIRContext.cpp

#endif
   // Type uniquing
   //===--------------------------------------------------------------------===//
 
   DenseMap<TypeID, AbstractType *> registeredTypes;
   StorageUniquer typeUniquer;
 
   /// Cached Type Instances.
   Float8E5M2Type f8E5M2Ty;
+  Float8E4M3FNType f8E4M3FNTy;
   BFloat16Type bf16Ty;
   Float16Type f16Ty;
   Float32Type f32Ty;
   Float64Type f64Ty;
   Float80Type f80Ty;
   Float128Type f128Ty;
   IndexType indexTy;
   IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
   getOrLoadDialect<BuiltinDialect>();
 
   // Initialize several common attributes and types to avoid the need to lock
   // the context when accessing them.
 
   //// Types.
   /// Floating-point Types.
   impl->f8E5M2Ty = TypeUniquer::get<Float8E5M2Type>(this);
+  impl->f8E4M3FNTy = TypeUniquer::get<Float8E4M3FNType>(this);
   impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this);
   impl->f16Ty = TypeUniquer::get<Float16Type>(this);
   impl->f32Ty = TypeUniquer::get<Float32Type>(this);
   impl->f64Ty = TypeUniquer::get<Float64Type>(this);
   impl->f80Ty = TypeUniquer::get<Float80Type>(this);
   impl->f128Ty = TypeUniquer::get<Float128Type>(this);
   /// Index Type.
   impl->indexTy = TypeUniquer::get<IndexType>(this);

 
 /// Returns the storage uniquer used for constructing type storage instances.
 /// This should not be used directly.
 StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }
 
 Float8E5M2Type Float8E5M2Type::get(MLIRContext *context) {
   return context->getImpl().f8E5M2Ty;
 }
+Float8E4M3FNType Float8E4M3FNType::get(MLIRContext *context) {
+  return context->getImpl().f8E4M3FNTy;
+}
 BFloat16Type BFloat16Type::get(MLIRContext *context) {
   return context->getImpl().bf16Ty;
 }
 Float16Type Float16Type::get(MLIRContext *context) {
   return context->getImpl().f16Ty;
 }
 Float32Type Float32Type::get(MLIRContext *context) {
   return context->getImpl().f32Ty;

mlir/lib/IR/Types.cpp

 
 //===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//
 
 MLIRContext *Type::getContext() const { return getDialect().getContext(); }
 
 bool Type::isFloat8E5M2() const { return isa<Float8E5M2Type>(); }
+bool Type::isFloat8E4M3FN() const { return isa<Float8E4M3FNType>(); }
 bool Type::isBF16() const { return isa<BFloat16Type>(); }
 bool Type::isF16() const { return isa<Float16Type>(); }
 bool Type::isF32() const { return isa<Float32Type>(); }
 bool Type::isF64() const { return isa<Float64Type>(); }
 bool Type::isF80() const { return isa<Float80Type>(); }
 bool Type::isF128() const { return isa<Float128Type>(); }
 
 bool Type::isIndex() const { return isa<IndexType>(); }

mlir/test/IR/attribute.mlir

 //===----------------------------------------------------------------------===//
 
 func.func @float_attrs_pass() {
   "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f8E5M2
     float_attr = 2. : f8E5M2
   } : () -> ()
   "test.float_attrs"() {
+    // CHECK: float_attr = 2.000000e+00 : f8E4M3FN
+    float_attr = 2. : f8E4M3FN
+  } : () -> ()
+  "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f16
     float_attr = 2. : f16
   } : () -> ()
   "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : bf16
     float_attr = 2. : bf16
   } : () -> ()
   "test.float_attrs"() {