[mlir][python] Expose fp8 types with pybind.

[zhangqiaorjc]

Summary

Expose fp8 types with pybind.

[zhangqiaorjc]

Ran clang-format. Removed a blank line.

[stellaraccident]

Thank you! Do you need help landing this patch or do you have commit access?

[zhangqiaorjc]

@stellaraccident could you help me land it? this is my first patch, so i don't have commit access. thanks!

[mehdi_amini]

Can we cover this in tests?

[zhangqiaorjc]

Added a test in test/python/ir/builtin_types.py.

[zhangqiaorjc]

Thank you @mehdi_amini !

mlir/lib/Bindings/Python/IRTypes.cpp

c.def_static(
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirIndexTypeGet(context->get());
           return PyIndexType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a index type.");
   }
 };
 
+/// Floating Point Type subclass - Float8E4M3FNType.
+class PyFloat8E4M3FNType : public PyConcreteType<PyFloat8E4M3FNType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FN;
+  static constexpr const char *pyClassName = "Float8E4M3FNType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat8E4M3FNTypeGet(context->get());
+          return PyFloat8E4M3FNType(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float8_e4m3fn type.");
+  }
+};
+
+/// Floating Point Type subclass - Float8M5E2Type.
+class PyFloat8E5M2Type : public PyConcreteType<PyFloat8E5M2Type> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2;
+  static constexpr const char *pyClassName = "Float8E5M2Type";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat8E5M2TypeGet(context->get());
+          return PyFloat8E5M2Type(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float8_e5m2 type.");
+  }
+};
+
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsABF16;
   static constexpr const char *pyClassName = "BF16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
public:
   }
 };
 
 } // namespace
 
 void mlir::python::populateIRTypes(py::module &m) {
   PyIntegerType::bind(m);
   PyIndexType::bind(m);
+  PyFloat8E4M3FNType::bind(m);
+  PyFloat8E5M2Type::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyF32Type::bind(m);
   PyF64Type::bind(m);
   PyNoneType::bind(m);
   PyComplexType::bind(m);
   PyShapedType::bind(m);
   PyVectorType::bind(m);
   PyRankedTensorType::bind(m);
   PyUnrankedTensorType::bind(m);
   PyMemRefType::bind(m);
   PyUnrankedMemRefType::bind(m);
   PyTupleType::bind(m);
   PyFunctionType::bind(m);
   PyOpaqueType::bind(m);
 }

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

__all__ = [
     "DenseIntElementsAttr",
     "Dialect",
     "DialectDescriptor",
     "Dialects",
     "Diagnostic",
     "DiagnosticHandler",
     "DiagnosticSeverity",
     "DictAttr",
+    "Float8E4M3FNType",
+    "Float8E5M2Type",
     "F16Type",
     "F32Type",
     "F64Type",
     "FlatSymbolRefAttr",
     "FloatAttr",
     "FunctionType",
     "IndexType",
     "InferTypeOpInterface",
class DictAttr(Attribute):
     @overload
     def __getitem__(self, arg0: str) -> Attribute: ...
     @overload
     def __getitem__(self, arg0: int) -> NamedAttribute: ...
     def __len__(self) -> int: ...
     @property
     def type(self) -> Type: ...
 
+class Float8E4M3FNType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    @staticmethod
+    def get(*args, **kwargs) -> Float8E4M3FNType: ...
+    @staticmethod
+    def isinstance(arg: Any) -> bool: ...
+
+class Float8E5M2Type(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    @staticmethod
+    def get(*args, **kwargs) -> Float8E5M2Type: ...
+    @staticmethod
+    def isinstance(arg: Any) -> bool: ...
+
 # TODO: Auto-generated. Audit and fix.
 class F16Type(Type):
     def __init__(self, cast_from_type: Type) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> F16Type: ...
     @staticmethod
     def isinstance(arg: Any) -> bool: ...
 

mlir/test/python/ir/builtin_types.py

with Context() as ctx:
     # CHECK: index type: index
     print("index type:", IndexType.get())
 
 
 # CHECK-LABEL: TEST: testFloatType
 @run
 def testFloatType():
   with Context():
+    # CHECK: float: f8E4M3FN
+    print("float:", Float8E4M3FNType.get())
+    # CHECK: float: f8E5M2
+    print("float:", Float8E5M2Type.get())
     # CHECK: float: bf16
     print("float:", BF16Type.get())
     # CHECK: float: f16
     print("float:", F16Type.get())
     # CHECK: float: f32
     print("float:", F32Type.get())
     # CHECK: float: f64
     print("float:", F64Type.get())