[MLIR] Add f6E2M3FN type (#107999)

This PR adds `f6E2M3FN` type to mlir.

`f6E2M3FN` 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 S1E2M3. Unlike
IEEE-754 types, there are no infinity or NaN values.

```c
f6E2M3FN
- Exponent bias: 1
- Maximum stored exponent value: 3 (binary 11)
- Maximum unbiased exponent value: 3 - 1 = 2
- Minimum stored exponent value: 1 (binary 01)
- Minimum unbiased exponent value: 1 − 1 = 0
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.00.000
- Max normal number: S.11.111 = ±2^(2) x (1 + 0.875) = ±7.5
- Min normal number: S.01.000 = ±2^(0) = ±1.0
- Max subnormal number: S.00.111 = ±2^(0) x 0.875 = ±0.875
- Min subnormal number: S.00.001 = ±2^(0) x 0.125 = ±0.125
```

Related PRs:
- [PR-94735](llvm/llvm-project#94735) [APFloat]
Add APFloat support for FP6 data types
- [PR-105573](llvm/llvm-project#105573) [MLIR]
Add f6E3M2FN 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 Float6E2M3FN type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E2M3FNTypeGetTypeID(void);
+
+/// Checks whether the given type is an f6E2M3FN type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat6E2M3FN(MlirType type);
+
+/// Creates an f6E2M3FN type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat6E2M3FNTypeGet(MlirContext ctx);
+
 /// Returns the typeID of an Float6E3M2FN type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E3M2FNTypeGetTypeID(void);
 

mlir/lib/Bindings/Python/IRTypes.cpp

@@ -124,6 +124,27 @@ class PyFloatType : public PyConcreteType<PyFloatType> {
   }
 };
 
+/// Floating Point Type subclass - Float6E2M3FNType.
+class PyFloat6E2M3FNType
+    : public PyConcreteType<PyFloat6E2M3FNType, PyFloatType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat6E2M3FN;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat6E2M3FNTypeGetTypeID;
+  static constexpr const char *pyClassName = "Float6E2M3FNType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat6E2M3FNTypeGet(context->get());
+          return PyFloat6E2M3FNType(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float6_e2m3fn type.");
+  }
+};
+
 /// Floating Point Type subclass - Float6E3M2FNType.
 class PyFloat6E3M2FNType
     : public PyConcreteType<PyFloat6E3M2FNType, PyFloatType> {
@@ -901,6 +922,7 @@ void mlir::python::populateIRTypes(py::module &m) {
   PyIntegerType::bind(m);
   PyFloatType::bind(m);
   PyIndexType::bind(m);
+  PyFloat6E2M3FNType::bind(m);
   PyFloat6E3M2FNType::bind(m);
   PyFloat8E4M3FNType::bind(m);
   PyFloat8E5M2Type::bind(m);

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -85,6 +85,18 @@ unsigned mlirFloatTypeGetWidth(MlirType type) {
   return llvm::cast<FloatType>(unwrap(type)).getWidth();
 }
 
+MlirTypeID mlirFloat6E2M3FNTypeGetTypeID() {
+  return wrap(Float6E2M3FNType::getTypeID());
+}
+
+bool mlirTypeIsAFloat6E2M3FN(MlirType type) {
+  return unwrap(type).isFloat6E2M3FN();
+}
+
+MlirType mlirFloat6E2M3FNTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat6E2M3FN(unwrap(ctx)));
+}
+
 MlirTypeID mlirFloat6E3M2FNTypeGetTypeID() {
   return wrap(Float6E3M2FNType::getTypeID());
 }

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

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

mlir/python/mlir/extras/types.py

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