[MLIR] Add f4E2M1FN type (#108877)

This PR adds `f4E2M1FN` type to mlir.

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

```c
f4E2M1FN
- 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.0
- Max normal number: S.11.1 = ±2^(2) x (1 + 0.5) = ±6.0
- Min normal number: S.01.0 = ±2^(0) = ±1.0
- Min subnormal number: S.00.1 = ±2^(0) x 0.5 = ±0.5
```

Related PRs:
- [PR-95392](llvm/llvm-project#95392) [APFloat]
Add APFloat support for FP4 data type
- [PR-105573](llvm/llvm-project#105573) [MLIR]
Add f6E3M2FN type - was used as a template for this PR
- [PR-107999](llvm/llvm-project#107999) [MLIR]
Add f6E2M3FN type

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 Float4E2M1FN type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat4E2M1FNTypeGetTypeID(void);
+
+/// Checks whether the given type is an f4E2M1FN type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat4E2M1FN(MlirType type);
+
+/// Creates an f4E2M1FN type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat4E2M1FNTypeGet(MlirContext ctx);
+
 /// Returns the typeID of an Float6E2M3FN type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E2M3FNTypeGetTypeID(void);
 

mlir/lib/Bindings/Python/IRTypes.cpp

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

mlir/lib/CAPI/IR/BuiltinTypes.cpp

@@ -85,6 +85,18 @@ unsigned mlirFloatTypeGetWidth(MlirType type) {
   return llvm::cast<FloatType>(unwrap(type)).getWidth();
 }
 
+MlirTypeID mlirFloat4E2M1FNTypeGetTypeID() {
+  return wrap(Float4E2M1FNType::getTypeID());
+}
+
+bool mlirTypeIsAFloat4E2M1FN(MlirType type) {
+  return unwrap(type).isFloat4E2M1FN();
+}
+
+MlirType mlirFloat4E2M1FNTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat4E2M1FN(unwrap(ctx)));
+}
+
 MlirTypeID mlirFloat6E2M3FNTypeGetTypeID() {
   return wrap(Float6E2M3FNType::getTypeID());
 }

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

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

mlir/python/mlir/extras/types.py

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