[Mlir-commits] [mlir] [MLIR][Python] Expose the insertion point of pattern rewriter (PR #161001)

[llvmlistbot at llvm.org]

================
@@ -143,7 +143,21 @@ class PyFrozenRewritePatternSet {
 
 /// Create the `mlir.rewrite` here.
 void mlir::python::populateRewriteSubmodule(nb::module_ &m) {
-  nb::class_<MlirPatternRewriter>(m, "PatternRewriter");
+  nb::class_<MlirPatternRewriter>(m, "PatternRewriter")
+      .def("ip", [](MlirPatternRewriter rewriter) {
+        MlirRewriterBase base = mlirPatternRewriterAsBase(rewriter);
+        MlirBlock block = mlirRewriterBaseGetInsertionBlock(base);
+        MlirOperation op = mlirRewriterBaseGetOperationAfterInsertion(base);
+        MlirOperation owner = mlirBlockGetParentOperation(block);
+        auto ctx = PyMlirContext::forContext(mlirRewriterBaseGetContext(base))
+                       ->getRef();
+        if (mlirOperationIsNull(op)) {
+          auto parent = PyOperation::forOperation(ctx, owner);
+          return PyInsertionPoint(PyBlock(parent, block));
+        }
+
+        return PyInsertionPoint(*PyOperation::forOperation(ctx, op).get());
+      });
----------------
PragmaTwice wrote:

Not sure if there is a good way to cast `Mlir*` CAPI types into `Py*` C++ classes. It seems that here we don't need to care too much about lifetime of blocks/operations (as long as the insertion point does not escape from the scope of the rewrite callback). 🤔 

I'll try to define something like `class PyPatternRewriter` and see if that makes the code cleaner.

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