[Mlir-commits] [mlir] [MLIR][Python] Add `convert_type` API for TypeConverter (PR #183561)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

<details>
<summary>Changes</summary>

This PR adds the  `convert_type` API for `TypeConverter`.

---
Full diff: https://github.com/llvm/llvm-project/pull/183561.diff


4 Files Affected:

- (modified) mlir/include/mlir-c/Rewrite.h (+4) 
- (modified) mlir/lib/Bindings/Python/Rewrite.cpp (+12-1) 
- (modified) mlir/lib/CAPI/Transforms/Rewrite.cpp (+5) 
- (modified) mlir/test/python/rewrite.py (+9) 


``````````diff
diff --git a/mlir/include/mlir-c/Rewrite.h b/mlir/include/mlir-c/Rewrite.h
index b4f93fd5a9b78..5e952edad23cb 100644
--- a/mlir/include/mlir-c/Rewrite.h
+++ b/mlir/include/mlir-c/Rewrite.h
@@ -546,6 +546,10 @@ mlirTypeConverterAddConversion(MlirTypeConverter typeConverter,
                                MlirTypeConverterConversionCallback convertType,
                                void *userData);
 
+/// Convert the given type using the given TypeConverter.
+MLIR_CAPI_EXPORTED MlirType
+mlirTypeConverterConvertType(MlirTypeConverter typeConverter, MlirType type);
+
 //===----------------------------------------------------------------------===//
 /// ConversionPattern API
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp
index e370552c00a9a..6c414e1a4c023 100644
--- a/mlir/lib/Bindings/Python/Rewrite.cpp
+++ b/mlir/lib/Bindings/Python/Rewrite.cpp
@@ -101,6 +101,15 @@ class PyTypeConverter {
         convert.ptr());
   }
 
+  nb::typed<nb::object, std::optional<PyType>> convertType(PyType &type) {
+    MlirType converted = mlirTypeConverterConvertType(typeConverter, type);
+    if (mlirTypeIsNull(converted))
+      return nb::none();
+    return PyType(PyMlirContext::forContext(mlirTypeGetContext(converted)),
+                  converted)
+        .maybeDownCast();
+  }
+
   MlirTypeConverter get() { return typeConverter; }
 
 private:
@@ -621,7 +630,9 @@ void populateRewriteSubmodule(nb::module_ &m) {
   nb::class_<PyTypeConverter>(m, "TypeConverter")
       .def(nb::init<>(), "Create a new TypeConverter.")
       .def("add_conversion", &PyTypeConverter::addConversion, "convert"_a,
-           nb::keep_alive<0, 1>(), "Register a type conversion function.");
+           nb::keep_alive<0, 1>(), "Register a type conversion function.")
+      .def("convert_type", &PyTypeConverter::convertType, "type"_a,
+           "Convert the given type. Returns None if conversion fails.");
 
   //----------------------------------------------------------------------------
   // Mapping of the PDLResultList and PDLModule
diff --git a/mlir/lib/CAPI/Transforms/Rewrite.cpp b/mlir/lib/CAPI/Transforms/Rewrite.cpp
index 5900f08ae1730..a7e43254767ad 100644
--- a/mlir/lib/CAPI/Transforms/Rewrite.cpp
+++ b/mlir/lib/CAPI/Transforms/Rewrite.cpp
@@ -590,6 +590,11 @@ void mlirTypeConverterAddConversion(
           });
 }
 
+MlirType mlirTypeConverterConvertType(MlirTypeConverter typeConverter,
+                                      MlirType type) {
+  return wrap(unwrap(typeConverter)->convertType(unwrap(type)));
+}
+
 //===----------------------------------------------------------------------===//
 /// ConversionPattern API
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/python/rewrite.py b/mlir/test/python/rewrite.py
index 1a9bcc87a3bda..35d88833e69e0 100644
--- a/mlir/test/python/rewrite.py
+++ b/mlir/test/python/rewrite.py
@@ -318,3 +318,12 @@ def convert_muli(op, adaptor, type_converter, rewriter):
             # CHECK: caught exception: partial conversion failed
             # CHECK: failed to legalize unresolved materialization
             print("caught exception:", e)
+
+        t1 = converter.convert_type(IntegerType.get_signless(64))
+        # CHECK: IntType
+        print(type(t1))
+        # CHECK: !smt.int
+        print(str(t1))
+        t2 = converter.convert_type(F32Type.get())
+        # CHECK: None
+        print(t2)

``````````

</details>


https://github.com/llvm/llvm-project/pull/183561