[MLIR][APFloat] Add Type definitions for UE5M3 type (#212712)

This commit adds UE5M3 type definitions in MLIR based on the Float8E5M3FNU APFloat type

mlir/include/mlir-c/BuiltinTypes.h

@@ -215,6 +215,18 @@ MLIR_CAPI_EXPORTED MlirType mlirFloat8E8M0FNUTypeGet(MlirContext ctx);
 
 MLIR_CAPI_EXPORTED MlirStringRef mlirFloat8E8M0FNUTypeGetName(void);
 
+/// Returns the typeID of a Float8E5M3FNU type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M3FNUTypeGetTypeID(void);
+
+/// Checks whether the given type is an f8E5M3FNU type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M3FNU(MlirType type);
+
+/// Creates an f8E5M3FNU type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M3FNUTypeGet(MlirContext ctx);
+
+MLIR_CAPI_EXPORTED MlirStringRef mlirFloat8E5M3FNUTypeGetName(void);
+
 /// Returns the typeID of an BFloat16 type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirBFloat16TypeGetTypeID(void);
 

mlir/include/mlir/Bindings/Python/IRTypes.h

@@ -211,6 +211,19 @@ public:
   static void bindDerived(ClassTy &c);
 };
 
+/// Floating Point Type subclass - Float8E5M3FNUType.
+class MLIR_PYTHON_API_EXPORTED PyFloat8E5M3FNUType
+    : public PyConcreteType<PyFloat8E5M3FNUType, PyFloatType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M3FNU;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E5M3FNUTypeGetTypeID;
+  static constexpr const char *pyClassName = "Float8E5M3FNUType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c);
+};
+
 /// Floating Point Type subclass - BF16Type.
 class MLIR_PYTHON_API_EXPORTED PyBF16Type
     : public PyConcreteType<PyBF16Type, PyFloatType> {

mlir/include/mlir/IR/Builders.h

@@ -62,6 +62,7 @@ public:
 
   // Types.
   FloatType getF8E8M0Type();
+  FloatType getF8E5M3FNUType();
   FloatType getF8E4M3FNType();
   FloatType getF8E5M2Type();
   FloatType getBF16Type();

mlir/include/mlir/IR/BuiltinTypes.td

@@ -433,6 +433,26 @@ def Builtin_Float8E8M0FNU : Builtin_FloatType<"Float8E8M0FNU", "f8E8M0FNU"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Float8E5M3FNUType
+//===----------------------------------------------------------------------===//
+
+def Builtin_Float8E5M3FNU : Builtin_FloatType<"Float8E5M3FNU", "f8E5M3FNU"> {
+  let summary = "8-bit unsigned floating point with 5-bit exponent, 3-bit mantissa";
+  let description = [{
+    An 8-bit floating point type with 0 sign bit, 5 bits exponent and 3 bits
+    mantissa. This is not a standard type as defined by IEEE-754, but it
+    follows similar conventions with the following characteristics:
+
+      * bit encoding: S0E5M3
+      * exponent bias: 15
+      * infinities: Not supported
+      * NaNs: Supported with all bits set to 1
+      * Zero: Supported
+      * denormals when exponent is 0
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // BFloat16Type
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/CommonTypeConstraints.td

@@ -377,6 +377,8 @@ def F6E3M2FN : Type<CPred<"::llvm::isa<::mlir::Float6E3M2FNType>($_self)">, "f6E
                BuildableType<"$_builder.getType<Float6E3M2FNType>()">;
 def F8E8M0FNU : Type<CPred<"::llvm::isa<::mlir::Float8E8M0FNUType>($_self)">, "f8E8M0FNU type">,
                 BuildableType<"$_builder.getType<Float8E8M0FNUType>()">;
+def F8E5M3FNU : Type<CPred<"::llvm::isa<::mlir::Float8E5M3FNUType>($_self)">, "f8E5M3FNU type">,
+                BuildableType<"$_builder.getType<Float8E5M3FNUType>()">;
 
 def AnyComplex : Type<CPred<"::llvm::isa<::mlir::ComplexType>($_self)">,
                       "complex-type", "::mlir::ComplexType">;

mlir/lib/AsmParser/TokenKinds.def

@@ -106,6 +106,7 @@ TOK_KEYWORD(f4E2M1FN)
 TOK_KEYWORD(f6E2M3FN)
 TOK_KEYWORD(f6E3M2FN)
 TOK_KEYWORD(f8E8M0FNU)
+TOK_KEYWORD(f8E5M3FNU)
 TOK_KEYWORD(f128)
 TOK_KEYWORD(false)
 TOK_KEYWORD(floordiv)

mlir/lib/AsmParser/TypeParser.cpp

@@ -49,6 +49,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
   case Token::kw_f8E4M3B11FNUZ:
   case Token::kw_f8E3M4:
   case Token::kw_f8E8M0FNU:
+  case Token::kw_f8E5M3FNU:
   case Token::kw_bf16:
   case Token::kw_f16:
   case Token::kw_tf32:
@@ -340,6 +341,9 @@ Type Parser::parseNonFunctionType() {
   case Token::kw_f8E8M0FNU:
     consumeToken(Token::kw_f8E8M0FNU);
     return builder.getType<Float8E8M0FNUType>();
+  case Token::kw_f8E5M3FNU:
+    consumeToken(Token::kw_f8E5M3FNU);
+    return builder.getType<Float8E5M3FNUType>();
   case Token::kw_bf16:
     consumeToken(Token::kw_bf16);
     return builder.getType<BFloat16Type>();

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -239,6 +239,16 @@ void PyFloat8E8M0FNUType::bindDerived(ClassTy &c) {
       nb::arg("context") = nb::none(), "Create a float8_e8m0fnu type.");
 }
 
+void PyFloat8E5M3FNUType::bindDerived(ClassTy &c) {
+  c.def_static(
+      "get",
+      [](DefaultingPyMlirContext context) {
+        MlirType t = mlirFloat8E5M3FNUTypeGet(context->get());
+        return PyFloat8E5M3FNUType(context->getRef(), t);
+      },
+      nb::arg("context") = nb::none(), "Create a float8_e5m3fnu type.");
+}
+
 void PyBF16Type::bindDerived(ClassTy &c) {
   c.def_static(
       "get",
@@ -936,6 +946,7 @@ void populateIRTypes(nb::module_ &m) {
   PyFloat8E5M2FNUZType::bind(m);
   PyFloat8E3M4Type::bind(m);
   PyFloat8E8M0FNUType::bind(m);
+  PyFloat8E5M3FNUType::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyTF32Type::bind(m);

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -265,6 +265,22 @@ MlirStringRef mlirFloat8E8M0FNUTypeGetName(void) {
   return wrap(Float8E8M0FNUType::name);
 }
 
+MlirTypeID mlirFloat8E5M3FNUTypeGetTypeID() {
+  return wrap(Float8E5M3FNUType::getTypeID());
+}
+
+bool mlirTypeIsAFloat8E5M3FNU(MlirType type) {
+  return llvm::isa<Float8E5M3FNUType>(unwrap(type));
+}
+
+MlirType mlirFloat8E5M3FNUTypeGet(MlirContext ctx) {
+  return wrap(Float8E5M3FNUType::get(unwrap(ctx)));
+}
+
+MlirStringRef mlirFloat8E5M3FNUTypeGetName(void) {
+  return wrap(Float8E5M3FNUType::name);
+}
+
 MlirTypeID mlirBFloat16TypeGetTypeID() {
   return wrap(BFloat16Type::getTypeID());
 }

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

@@ -301,7 +301,7 @@ Type LLVMTypeConverter::convertFloatType(FloatType type) const {
   if (isa<Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType, Float8E5M2FNUZType,
           Float8E4M3FNUZType, Float8E4M3B11FNUZType, Float8E3M4Type,
           Float4E2M1FNType, Float6E2M3FNType, Float6E3M2FNType,
-          Float8E8M0FNUType>(type))
+          Float8E8M0FNUType, Float8E5M3FNUType>(type))
     return IntegerType::get(&getContext(), type.getWidth());
 
   // Other floating-point types: A custom type conversion rule must be

mlir/lib/IR/AsmPrinter.cpp

@@ -2805,6 +2805,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) {
       .Case<Float8E4M3B11FNUZType>([&](Type) { os << "f8E4M3B11FNUZ"; })
       .Case<Float8E3M4Type>([&](Type) { os << "f8E3M4"; })
       .Case<Float8E8M0FNUType>([&](Type) { os << "f8E8M0FNU"; })
+      .Case<Float8E5M3FNUType>([&](Type) { os << "f8E5M3FNU"; })
       .Case<BFloat16Type>([&](Type) { os << "bf16"; })
       .Case<Float16Type>([&](Type) { os << "f16"; })
       .Case<FloatTF32Type>([&](Type) { os << "tf32"; })

mlir/lib/IR/Builders.cpp

@@ -34,6 +34,10 @@ Location Builder::getFusedLoc(ArrayRef<Location> locs, Attribute metadata) {
 
 FloatType Builder::getF8E8M0Type() { return Float8E8M0FNUType::get(context); }
 
+FloatType Builder::getF8E5M3FNUType() {
+  return Float8E5M3FNUType::get(context);
+}
+
 FloatType Builder::getF8E4M3FNType() { return Float8E4M3FNType::get(context); }
 
 FloatType Builder::getF8E5M2Type() { return Float8E5M2Type::get(context); }

mlir/lib/IR/BuiltinTypes.cpp

@@ -192,6 +192,7 @@ FLOAT_TYPE_SEMANTICS(Float8E4M3FNUZType, Float8E4M3FNUZ)
 FLOAT_TYPE_SEMANTICS(Float8E4M3B11FNUZType, Float8E4M3B11FNUZ)
 FLOAT_TYPE_SEMANTICS(Float8E3M4Type, Float8E3M4)
 FLOAT_TYPE_SEMANTICS(Float8E8M0FNUType, Float8E8M0FNU)
+FLOAT_TYPE_SEMANTICS(Float8E5M3FNUType, Float8E5M3FNU)
 FLOAT_TYPE_SEMANTICS(BFloat16Type, BFloat)
 FLOAT_TYPE_SEMANTICS(Float16Type, IEEEhalf)
 FLOAT_TYPE_SEMANTICS(FloatTF32Type, FloatTF32)

mlir/python/mlir/extras/types.py

@@ -21,6 +21,7 @@ from ..ir import (
     Float8E4M3Type,
     Float8E5M2Type,
     Float8E8M0FNUType,
+    Float8E5M3FNUType,
     FloatTF32Type,
     FunctionType,
     IndexType,
@@ -84,6 +85,7 @@ f4E2M1FN = lambda: Float4E2M1FNType.get()
 f6E2M3FN = lambda: Float6E2M3FNType.get()
 f6E3M2FN = lambda: Float6E3M2FNType.get()
 f8E8M0FNU = lambda: Float8E8M0FNUType.get()
+f8E5M3FNU = lambda: Float8E5M3FNUType.get()
 
 none = lambda: NoneType.get()
 

mlir/test/IR/attribute.mlir

@@ -80,6 +80,10 @@ func.func @float_attrs_pass() {
     // CHECK: float_attr = 2.000000e+00 : f8E8M0FNU
     float_attr = 2. : f8E8M0FNU
   } : () -> ()
+  "test.float_attrs"() {
+    // CHECK: float_attr = 2.000000e+00 : f8E5M3FNU
+    float_attr = 2. : f8E5M3FNU
+  } : () -> ()
   "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f16
     float_attr = 2. : f16

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -99,6 +99,9 @@ llvm.mlir.global internal @f8E4M3B11FNUZ_global_as_i8(1.5 : f8E4M3B11FNUZ) : i8
 // CHECK: @f8E8M0FNU_global_as_i8 = internal global i8 127
 llvm.mlir.global internal @f8E8M0FNU_global_as_i8(1.0 : f8E8M0FNU) : i8
 
+// CHECK: @f8E5M3FNU_global_as_i8 = internal global i8 120
+llvm.mlir.global internal @f8E5M3FNU_global_as_i8(1.0 : f8E5M3FNU) : i8
+
 // CHECK: @bf16_global_as_i16 = internal global i16 16320
 llvm.mlir.global internal @bf16_global_as_i16(1.5 : bf16) : i16
 

mlir/test/python/ir/builtin_types.py

@@ -135,6 +135,8 @@ def testFloatTypeSubclasses():
     # CHECK: True
     print(isinstance(Type.parse("f8E8M0FNU", ctx), FloatType))
     # CHECK: True
+    print(isinstance(Type.parse("f8E5M3FNU", ctx), FloatType))
+    # CHECK: True
     print(isinstance(Type.parse("f16", ctx), FloatType))
     # CHECK: True
     print(isinstance(Type.parse("bf16", ctx), FloatType))
@@ -277,6 +279,8 @@ def testFloatType():
         print("float:", Float8E4M3B11FNUZType.get())
         # CHECK: float: f8E8M0FNU
         print("float:", Float8E8M0FNUType.get())
+        # CHECK: float: f8E5M3FNU
+        print("float:", Float8E5M3FNUType.get())
         # CHECK: float: bf16
         print("float:", BF16Type.get())
         # CHECK: float: f16
@@ -676,6 +680,7 @@ def testTypeIDs():
             (Float8E4M3B11FNUZType, Float8E4M3B11FNUZType.get()),
             (Float8E5M2FNUZType, Float8E5M2FNUZType.get()),
             (Float8E8M0FNUType, Float8E8M0FNUType.get()),
+            (Float8E5M3FNUType, Float8E5M3FNUType.get()),
             (BF16Type, BF16Type.get()),
             (F16Type, F16Type.get()),
             (F32Type, F32Type.get()),
@@ -706,6 +711,7 @@ def testTypeIDs():
         # CHECK: Float8E4M3B11FNUZType(f8E4M3B11FNUZ)
         # CHECK: Float8E5M2FNUZType(f8E5M2FNUZ)
         # CHECK: Float8E8M0FNUType(f8E8M0FNU)
+        # CHECK: Float8E5M3FNUType(f8E5M3FNU)
         # CHECK: BF16Type(bf16)
         # CHECK: F16Type(f16)
         # CHECK: F32Type(f32)
@@ -858,6 +864,9 @@ def testConcreteTypesRoundTrip():
         # CHECK: Float8E8M0FNUType
         # CHECK: Float8E8M0FNUType(f8E8M0FNU)
         print_downcasted(Float8E8M0FNUType.get())
+        # CHECK: Float8E5M3FNUType
+        # CHECK: Float8E5M3FNUType(f8E5M3FNU)
+        print_downcasted(Float8E5M3FNUType.get())
         # CHECK: BF16Type
         # CHECK: BF16Type(bf16)
         print_downcasted(BF16Type.get())

mlir/utils/lldb-scripts/mlirDataFormatters.py

@@ -61,6 +61,7 @@ builtin_attr_type_mnemonics = {
     "mlir::Float8E4M3B11FNUZType": '"f8E4M3B11FNUZ"',
     "mlir::Float8E3M4Type": '"f8E3M4"',
     "mlir::Float8E8M0FNUType": '"f8E8M0FNU"',
+    "mlir::Float8E5M3FNUType": '"f8E5M3FNU"',
     "mlir::BFloat16Type": '"bf16"',
     "mlir::Float16Type": '"f16"',
     "mlir::FloatTF32Type": '"tf32"',

mlir/utils/tree-sitter-mlir/grammar.js

@@ -232,7 +232,7 @@ const common = {
   float_type : $ => token(
       choice('f16', 'f32', 'f64', 'f80', 'f128', 'bf16', 'f8E3M4', 'f8E4M3FN',
              'f8E4M3', 'f8E5M2', 'f4E2M1FN', 'f6E2M3FN', 'f6E3M2FN',
-             'f8E8M0FNU')),
+             'f8E8M0FNU', 'f8E5M3FNU')),
   index_type : $ => token('index'),
   none_type : $ => token('none'),
   complex_type : $ => seq(token('complex'), '<', $._prim_type, '>'),