[Mlir-commits] [mlir] [MLIR][Python] Call `notifyOperationInserted` while constructing new op in rewrite patterns (PR #163694)

[llvmlistbot at llvm.org]

================
@@ -202,7 +207,15 @@ class PyRewritePatternSet {
           PyMlirContext::forContext(mlirOperationGetContext(op));
       nb::object opView = PyOperation::forOperation(ctx, op)->createOpView();
 
-      nb::object res = f(opView, PyPatternRewriter(rewriter));
+      PyPatternRewriter pyRewriter(rewriter);
+      nb::object listener = nb::cast(pyRewriter.getListener());
+
+      listener.attr("__enter__")();
+      auto exit = llvm::make_scope_exit([listener] {
+        listener.attr("__exit__")(nb::none(), nb::none(), nb::none());
+      });
----------------
PragmaTwice wrote:

I'll refactor this change soon.

> basically add an add overload that takes a listener and sticks it into `*userData` along with the callable?

yeah I can get your point, but we cannot get the listener before we get the rewriter, and only after we enter the callback we can finally get the rewriter. (shown as an example below)

```python
def f(op, rewriter):
  # now we can get the rewriter!

patterns = RewritePatternSet()
patterns.add(arith.SomeOp, f) # we cannot get it here
frozen = patterns.freeze()

apply_and_fold_greedily(op, frozen) # I think the rewriter is constructed inside this method
```

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