[MLIR] Add f6E3M2FN type (#105573).

`f6E3M2FN` type is proposed in [OpenCompute MX Specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
It defines a 6-bit floating point number with bit layout S1E3M2. Unlike IEEE-754 types, there are no infinity or NaN values.

```c
f6E3M2FN
- Exponent bias: 3
- Maximum stored exponent value: 7 (binary 111)
- Maximum unbiased exponent value: 7 - 3 = 4
- Minimum stored exponent value: 1 (binary 001)
- Minimum unbiased exponent value: 1 − 3 = −2
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.000.00
- Max normal number: S.111.11 = ±2^(4) x (1 + 0.75) = ±28
- Min normal number: S.001.00 = ±2^(-2) = ±0.25
- Max subnormal number: S.000.11 = ±2^(-2) x 0.75 = ±0.1875
- Min subnormal number: S.000.01 = ±2^(-2) x 0.25 = ±0.0625
```

Related PRs:
- [PR-94735](#94735) [APFloat] Add APFloat support for FP6 data types
- [PR-97118](#97118) [MLIR] Add f8E4M3 type - was used as a template for this PR

llvm/unittests/ADT/APFloatTest.cpp

@@ -2084,16 +2084,17 @@ TEST(APFloatTest, getSmallestNormalized) {
   EXPECT_FALSE(test.isDenormal());
   EXPECT_TRUE(test.bitwiseIsEqual(expected));
   EXPECT_TRUE(test.isSmallestNormalized());
+
   test = APFloat::getSmallestNormalized(APFloat::Float6E3M2FN(), false);
   expected = APFloat(APFloat::Float6E3M2FN(), "0x1p-2");
-
-  test = APFloat::getSmallestNormalized(APFloat::Float4E2M1FN(), false);
-  expected = APFloat(APFloat::Float4E2M1FN(), "0x1p0");
   EXPECT_FALSE(test.isNegative());
   EXPECT_TRUE(test.isFiniteNonZero());
   EXPECT_FALSE(test.isDenormal());
   EXPECT_TRUE(test.bitwiseIsEqual(expected));
   EXPECT_TRUE(test.isSmallestNormalized());
+
+  test = APFloat::getSmallestNormalized(APFloat::Float4E2M1FN(), false);
+  expected = APFloat(APFloat::Float4E2M1FN(), "0x1p0");
   EXPECT_FALSE(test.isNegative());
   EXPECT_TRUE(test.isFiniteNonZero());
   EXPECT_FALSE(test.isDenormal());

mlir/include/mlir-c/BuiltinTypes.h

@@ -149,6 +149,16 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E3M4(MlirType type);
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E3M4TypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float6E3M2FN type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E3M2FNTypeGetTypeID(void);
+
+/// Checks whether the given type is an f6E3M2FN type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat6E3M2FN(MlirType type);
+
+/// Creates an f6E3M2FN type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat6E3M2FNTypeGet(MlirContext ctx);
+
 /// Returns the typeID of an BFloat16 type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirBFloat16TypeGetTypeID(void);
 

mlir/include/mlir/IR/Builders.h

@@ -67,6 +67,7 @@ class Builder {
   FloatType getFloat8E4M3FNUZType();
   FloatType getFloat8E4M3B11FNUZType();
   FloatType getFloat8E3M4Type();
+  FloatType getFloat6E3M2FNType();
   FloatType getBF16Type();
   FloatType getF16Type();
   FloatType getTF32Type();

mlir/include/mlir/IR/BuiltinTypes.h

@@ -67,6 +67,7 @@ class FloatType : public Type {
   static FloatType getFloat8E4M3FNUZ(MLIRContext *ctx);
   static FloatType getFloat8E4M3B11FNUZ(MLIRContext *ctx);
   static FloatType getFloat8E3M4(MLIRContext *ctx);
+  static FloatType getFloat6E3M2FN(MLIRContext *ctx);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
@@ -415,9 +416,9 @@ inline bool BaseMemRefType::isValidElementType(Type type) {
 inline bool FloatType::classof(Type type) {
   return llvm::isa<Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType,
                    Float8E5M2FNUZType, Float8E4M3FNUZType,
-                   Float8E4M3B11FNUZType, Float8E3M4Type, BFloat16Type,
-                   Float16Type, FloatTF32Type, Float32Type, Float64Type,
-                   Float80Type, Float128Type>(type);
+                   Float8E4M3B11FNUZType, Float8E3M4Type, Float6E3M2FNType,
+                   BFloat16Type, Float16Type, FloatTF32Type, Float32Type,
+                   Float64Type, Float80Type, Float128Type>(type);
 }
 
 inline FloatType FloatType::getFloat8E5M2(MLIRContext *ctx) {
@@ -448,6 +449,10 @@ inline FloatType FloatType::getFloat8E3M4(MLIRContext *ctx) {
   return Float8E3M4Type::get(ctx);
 }
 
+inline FloatType FloatType::getFloat6E3M2FN(MLIRContext *ctx) {
+  return Float6E3M2FNType::get(ctx);
+}
+
 inline FloatType FloatType::getBF16(MLIRContext *ctx) {
   return BFloat16Type::get(ctx);
 }

mlir/include/mlir/IR/BuiltinTypes.td

@@ -233,6 +233,24 @@ def Builtin_Float8E3M4 : Builtin_FloatType<"Float8E3M4", "f8E3M4"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Float6E3M2FNType
+
+def Builtin_Float6E3M2FN : Builtin_FloatType<"Float6E3M2FN", "f6E3M2FN"> {
+  let summary = "6-bit floating point with 3 bits exponent and 2 bit mantissa";
+  let description = [{
+    An 6-bit floating point type with 1 sign bit, 3 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: S1E3M2
+      * exponent bias: 3
+      * infinities: Not supported
+      * NaNs: Not supported
+      * denormals when exponent is 0
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // BFloat16Type
 

mlir/include/mlir/IR/CommonTypeConstraints.td

@@ -344,6 +344,8 @@ def F8E5M2FNUZ : Type<CPred<"$_self.isFloat8E5M2FNUZ()">, "f8E5M2FNUZ type">,
                  BuildableType<"$_builder.getFloat8E5M2FNUZType()">;
 def F8E3M4 : Type<CPred<"$_self.isFloat8E3M4()">, "f8E3M4 type">,
              BuildableType<"$_builder.getFloat8E3M4Type()">;
+def F6E3M2FN : Type<CPred<"$_self.isFloat6E3M2FN()">, "f6E3M2FN type">,
+               BuildableType<"$_builder.getFloat6E3M2FNType()">;
 
 def AnyComplex : Type<CPred<"::llvm::isa<::mlir::ComplexType>($_self)">,
                       "complex-type", "::mlir::ComplexType">;

mlir/include/mlir/IR/Types.h

@@ -132,6 +132,7 @@ class Type {
   bool isFloat8E4M3FNUZ() const;
   bool isFloat8E4M3B11FNUZ() const;
   bool isFloat8E3M4() const;
+  bool isFloat6E3M2FN() const;
   bool isBF16() const;
   bool isF16() const;
   bool isTF32() const;

mlir/lib/AsmParser/TokenKinds.def

@@ -101,6 +101,7 @@ TOK_KEYWORD(f8E5M2FNUZ)
 TOK_KEYWORD(f8E4M3FNUZ)
 TOK_KEYWORD(f8E4M3B11FNUZ)
 TOK_KEYWORD(f8E3M4)
+TOK_KEYWORD(f6E3M2FN)
 TOK_KEYWORD(f128)
 TOK_KEYWORD(false)
 TOK_KEYWORD(floordiv)

mlir/lib/AsmParser/TypeParser.cpp

@@ -46,6 +46,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
   case Token::kw_f8E4M3FNUZ:
   case Token::kw_f8E4M3B11FNUZ:
   case Token::kw_f8E3M4:
+  case Token::kw_f6E3M2FN:
   case Token::kw_bf16:
   case Token::kw_f16:
   case Token::kw_tf32:
@@ -324,6 +325,9 @@ Type Parser::parseNonFunctionType() {
   case Token::kw_f8E3M4:
     consumeToken(Token::kw_f8E3M4);
     return builder.getFloat8E3M4Type();
+  case Token::kw_f6E3M2FN:
+    consumeToken(Token::kw_f6E3M2FN);
+    return builder.getFloat6E3M2FNType();
   case Token::kw_bf16:
     consumeToken(Token::kw_bf16);
     return builder.getBF16Type();

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -268,6 +268,27 @@ class PyFloat8E3M4Type : public PyConcreteType<PyFloat8E3M4Type, PyFloatType> {
   }
 };
 
+/// Floating Point Type subclass - Float6E3M2FNType.
+class PyFloat6E3M2FNType
+    : public PyConcreteType<PyFloat6E3M2FNType, PyFloatType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat6E3M2FN;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat6E3M2FNTypeGetTypeID;
+  static constexpr const char *pyClassName = "Float6E3M2FNType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat6E3M2FNTypeGet(context->get());
+          return PyFloat6E3M2FNType(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float6_e3m2fn type.");
+  }
+};
+
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type, PyFloatType> {
 public:
@@ -887,6 +908,7 @@ void mlir::python::populateIRTypes(py::module &m) {
   PyFloat8E4M3B11FNUZType::bind(m);
   PyFloat8E5M2FNUZType::bind(m);
   PyFloat8E3M4Type::bind(m);
+  PyFloat6E3M2FNType::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyTF32Type::bind(m);

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -169,6 +169,18 @@ MlirType mlirFloat8E3M4TypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E3M4(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat6E3M2FNTypeGetTypeID() {
+  return wrap(Float6E3M2FNType::getTypeID());
+}
+
+bool mlirTypeIsAFloat6E3M2FN(MlirType type) {
+  return unwrap(type).isFloat6E3M2FN();
+}
+
+MlirType mlirFloat6E3M2FNTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat6E3M2FN(unwrap(ctx)));
+}
+
 MlirTypeID mlirBFloat16TypeGetTypeID() {
   return wrap(BFloat16Type::getTypeID());
 }

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

@@ -249,7 +249,8 @@ Type LLVMTypeConverter::convertIntegerType(IntegerType type) const {
 Type LLVMTypeConverter::convertFloatType(FloatType type) const {
   if (type.isFloat8E5M2() || type.isFloat8E4M3() || type.isFloat8E4M3FN() ||
       type.isFloat8E5M2FNUZ() || type.isFloat8E4M3FNUZ() ||
-      type.isFloat8E4M3B11FNUZ() || type.isFloat8E3M4())
+      type.isFloat8E4M3B11FNUZ() || type.isFloat8E3M4() ||
+      type.isFloat6E3M2FN())
     return IntegerType::get(&getContext(), type.getWidth());
   return type;
 }

mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp

@@ -61,6 +61,7 @@ static std::optional<FloatType> parseFloatType(MLIRContext *ctx,
       .Case("f8E5M2FNUZ", b.getFloat8E5M2FNUZType())
       .Case("f8E4M3FNUZ", b.getFloat8E4M3FNUZType())
       .Case("f8E3M4", b.getFloat8E3M4Type())
+      .Case("f6E3M2FN", b.getFloat6E3M2FNType())
       .Case("bf16", b.getBF16Type())
       .Case("f16", b.getF16Type())
       .Case("f32", b.getF32Type())

mlir/lib/IR/AsmPrinter.cpp

@@ -2582,6 +2582,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) {
       .Case<Float8E4M3FNUZType>([&](Type) { os << "f8E4M3FNUZ"; })
       .Case<Float8E4M3B11FNUZType>([&](Type) { os << "f8E4M3B11FNUZ"; })
       .Case<Float8E3M4Type>([&](Type) { os << "f8E3M4"; })
+      .Case<Float6E3M2FNType>([&](Type) { os << "f6E3M2FN"; })
       .Case<BFloat16Type>([&](Type) { os << "bf16"; })
       .Case<Float16Type>([&](Type) { os << "f16"; })
       .Case<FloatTF32Type>([&](Type) { os << "tf32"; })

mlir/lib/IR/Builders.cpp

@@ -62,6 +62,10 @@ FloatType Builder::getFloat8E3M4Type() {
   return FloatType::getFloat8E3M4(context);
 }
 
+FloatType Builder::getFloat6E3M2FNType() {
+  return FloatType::getFloat6E3M2FN(context);
+}
+
 FloatType Builder::getBF16Type() { return FloatType::getBF16(context); }
 
 FloatType Builder::getF16Type() { return FloatType::getF16(context); }

mlir/lib/IR/BuiltinTypes.cpp

@@ -91,6 +91,8 @@ IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
 //===----------------------------------------------------------------------===//
 
 unsigned FloatType::getWidth() {
+  if (llvm::isa<Float6E3M2FNType>(*this))
+    return 6;
   if (llvm::isa<Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType,
                 Float8E5M2FNUZType, Float8E4M3FNUZType, Float8E4M3B11FNUZType,
                 Float8E3M4Type>(*this))
@@ -124,6 +126,8 @@ const llvm::fltSemantics &FloatType::getFloatSemantics() {
     return APFloat::Float8E4M3B11FNUZ();
   if (llvm::isa<Float8E3M4Type>(*this))
     return APFloat::Float8E3M4();
+  if (llvm::isa<Float6E3M2FNType>(*this))
+    return APFloat::Float6E3M2FN();
   if (llvm::isa<BFloat16Type>(*this))
     return APFloat::BFloat();
   if (llvm::isa<Float16Type>(*this))

mlir/lib/IR/MLIRContext.cpp

@@ -228,6 +228,7 @@ class MLIRContextImpl {
   Float8E4M3FNUZType f8E4M3FNUZTy;
   Float8E4M3B11FNUZType f8E4M3B11FNUZTy;
   Float8E3M4Type f8E3M4Ty;
+  Float6E3M2FNType f6E3M2FNTy;
   BFloat16Type bf16Ty;
   Float16Type f16Ty;
   FloatTF32Type tf32Ty;
@@ -320,6 +321,7 @@ MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
   impl->f8E4M3FNUZTy = TypeUniquer::get<Float8E4M3FNUZType>(this);
   impl->f8E4M3B11FNUZTy = TypeUniquer::get<Float8E4M3B11FNUZType>(this);
   impl->f8E3M4Ty = TypeUniquer::get<Float8E3M4Type>(this);
+  impl->f6E3M2FNTy = TypeUniquer::get<Float6E3M2FNType>(this);
   impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this);
   impl->f16Ty = TypeUniquer::get<Float16Type>(this);
   impl->tf32Ty = TypeUniquer::get<FloatTF32Type>(this);
@@ -1034,6 +1036,9 @@ Float8E4M3B11FNUZType Float8E4M3B11FNUZType::get(MLIRContext *context) {
 Float8E3M4Type Float8E3M4Type::get(MLIRContext *context) {
   return context->getImpl().f8E3M4Ty;
 }
+Float6E3M2FNType Float6E3M2FNType::get(MLIRContext *context) {
+  return context->getImpl().f6E3M2FNTy;
+}
 BFloat16Type BFloat16Type::get(MLIRContext *context) {
   return context->getImpl().bf16Ty;
 }

mlir/lib/IR/Types.cpp

@@ -47,6 +47,7 @@ bool Type::isFloat8E4M3B11FNUZ() const {
   return llvm::isa<Float8E4M3B11FNUZType>(*this);
 }
 bool Type::isFloat8E3M4() const { return llvm::isa<Float8E3M4Type>(*this); }
+bool Type::isFloat6E3M2FN() const { return llvm::isa<Float6E3M2FNType>(*this); }
 bool Type::isBF16() const { return llvm::isa<BFloat16Type>(*this); }
 bool Type::isF16() const { return llvm::isa<Float16Type>(*this); }
 bool Type::isTF32() const { return llvm::isa<FloatTF32Type>(*this); }

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

@@ -120,6 +120,7 @@ __all__ = [
     "F32Type",
     "F64Type",
     "FlatSymbolRefAttr",
+    "Float6E3M2FNType",
     "Float8E3M4Type",
     "Float8E4M3B11FNUZType",
     "Float8E4M3FNType",
@@ -1539,6 +1540,19 @@ class FlatSymbolRefAttr(Attribute):
         Returns the value of the FlatSymbolRef attribute as a string
         """
 
+class Float6E3M2FNType(FloatType):
+    static_typeid: ClassVar[TypeID]
+    @staticmethod
+    def get(context: Optional[Context] = None) -> Float6E3M2FNType:
+        """
+        Create a float6_e3m2fn type.
+        """
+    @staticmethod
+    def isinstance(other: Type) -> bool: ...
+    def __init__(self, cast_from_type: Type) -> None: ...
+    @property
+    def typeid(self) -> TypeID: ...
+
 class Float8E3M4Type(FloatType):
     static_typeid: ClassVar[TypeID]
     @staticmethod

mlir/python/mlir/extras/types.py

@@ -12,6 +12,7 @@
     F16Type,
     F32Type,
     F64Type,
+    Float6E3M2FNType,
     Float8E3M4Type,
     Float8E4M3B11FNUZType,
     Float8E4M3FNType,
@@ -74,6 +75,7 @@ def ui(width):
 f8E4M3FN = lambda: Float8E4M3FNType.get()
 f8E4M3B11FNUZ = lambda: Float8E4M3B11FNUZType.get()
 f8E3M4 = lambda: Float8E3M4Type.get()
+f6E3M2FN = lambda: Float6E3M2FNType.get()
 
 none = lambda: NoneType.get()
 

mlir/test/IR/attribute.mlir

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

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -39,6 +39,9 @@ llvm.mlir.global internal constant @string_const("foobar") : !llvm.array<6 x i8>
 // CHECK: @int_global_undef = internal global i64 undef
 llvm.mlir.global internal @int_global_undef() : i64
 
+// CHECK: @f6E3M2FN_global_as_i6 = internal global i6 14
+llvm.mlir.global internal @f6E3M2FN_global_as_i6(1.5 : f6E3M2FN) : i6
+
 // CHECK: @f8E3M4_global_as_i8 = internal global i8 56
 llvm.mlir.global internal @f8E3M4_global_as_i8(1.5 : f8E3M4) : i8
 

mlir/test/python/ir/builtin_types.py

@@ -113,6 +113,8 @@ def testTypeIsInstance():
 def testFloatTypeSubclasses():
     ctx = Context()
     # CHECK: True
+    print(isinstance(Type.parse("f6E3M2FN", ctx), FloatType))
+    # CHECK: True
     print(isinstance(Type.parse("f8E3M4", ctx), FloatType))
     # CHECK: True
     print(isinstance(Type.parse("f8E4M3", ctx), FloatType))
@@ -233,6 +235,8 @@ def testIndexType():
 @run
 def testFloatType():
     with Context():
+        # CHECK: float: f6E3M2FN
+        print("float:", Float6E3M2FNType.get())
         # CHECK: float: f8E3M4
         print("float:", Float8E3M4Type.get())
         # CHECK: float: f8E4M3
@@ -609,6 +613,7 @@ def testTypeIDs():
         types = [
             (IntegerType, IntegerType.get_signless(16)),
             (IndexType, IndexType.get()),
+            (Float6E3M2FNType, Float6E3M2FNType.get()),
             (Float8E3M4Type, Float8E3M4Type.get()),
             (Float8E4M3Type, Float8E4M3Type.get()),
             (Float8E4M3FNType, Float8E4M3FNType.get()),
@@ -634,6 +639,7 @@ def testTypeIDs():
 
         # CHECK: IntegerType(i16)
         # CHECK: IndexType(index)
+        # CHECK: Float6E3M2FNType(f6E3M2FN)
         # CHECK: Float8E3M4Type(f8E3M4)
         # CHECK: Float8E4M3Type(f8E4M3)
         # CHECK: Float8E4M3FNType(f8E4M3FN)
@@ -713,6 +719,9 @@ def print_downcasted(typ):
         # CHECK: F64Type
         # CHECK: F64Type(f64)
         print_downcasted(F64Type.get())
+        # CHECK: Float6E3M2FNType
+        # CHECK: Float6E3M2FNType(f6E3M2FN)
+        print_downcasted(Float6E3M2FNType.get())
         # CHECK: Float8E3M4Type
         # CHECK: Float8E3M4Type(f8E3M4)
         print_downcasted(Float8E3M4Type.get())