[MLIR] Add f6E3M2FN type (#105573)

This PR adds `f6E3M2FN` type to mlir.

`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](llvm/llvm-project#94735) [APFloat]
Add APFloat support for FP6 data types
- [PR-97118](llvm/llvm-project#97118) [MLIR] Add
f8E4M3 type - was used as a template for this PR

Author: sergey-kozub

mlir/include/mlir-c/BuiltinTypes.h

@@ -79,6 +79,16 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat(MlirType type);
 /// Returns the bitwidth of a floating-point type.
 MLIR_CAPI_EXPORTED unsigned mlirFloatTypeGetWidth(MlirType type);
 
+/// 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 Float8E5M2 type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M2TypeGetTypeID(void);
 

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -124,6 +124,27 @@ class PyFloatType : public PyConcreteType<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 - Float8E4M3FNType.
 class PyFloat8E4M3FNType
     : public PyConcreteType<PyFloat8E4M3FNType, PyFloatType> {
@@ -880,6 +901,7 @@ void mlir::python::populateIRTypes(py::module &m) {
   PyIntegerType::bind(m);
   PyFloatType::bind(m);
   PyIndexType::bind(m);
+  PyFloat6E3M2FNType::bind(m);
   PyFloat8E4M3FNType::bind(m);
   PyFloat8E5M2Type::bind(m);
   PyFloat8E4M3Type::bind(m);

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -85,6 +85,18 @@ unsigned mlirFloatTypeGetWidth(MlirType type) {
   return llvm::cast<FloatType>(unwrap(type)).getWidth();
 }
 
+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 mlirFloat8E5M2TypeGetTypeID() {
   return wrap(Float8E5M2Type::getTypeID());
 }

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()