[mlir][python] Clear PyOperations instead of invalidating them. (#70044)

`PyOperations` are Python-level handles to `Operation *` instances. When
the latter are modified by C++, the former need to be invalidated.
#69746 implements such invalidation mechanism by setting all
`PyReferences` to `invalid`. However, that is not enough: they also need
to be removed from the `liveOperations` map since other parts of the
code (such as `PyOperation::createDetached`) assume that that map only
contains valid refs.

This is required to actually solve the issue in #69730.

1 files changed, 23 insertions, 2 deletions

diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index e7f79ddc7511..0face028b73f 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -176,6 +176,14 @@ def testRunPipelineError():
 @run
 def testPostPassOpInvalidation():
     with Context() as ctx:
+        log_op_count = lambda: log("live ops:", ctx._get_live_operation_count())
+
+        # CHECK: invalidate_ops=False
+        log("invalidate_ops=False")
+
+        # CHECK: live ops: 0
+        log_op_count()
+
         module = ModuleOp.parse(
             """
           module {
@@ -188,8 +196,8 @@ def testPostPassOpInvalidation():
         """
         )
 
-        # CHECK: invalidate_ops=False
-        log("invalidate_ops=False")
+        # CHECK: live ops: 1
+        log_op_count()
 
         outer_const_op = module.body.operations[0]
         # CHECK: %[[VAL0:.*]] = arith.constant 10 : i64
@@ -206,6 +214,9 @@ def testPostPassOpInvalidation():
         # CHECK: %[[VAL1]] = arith.constant 10 : i64
         log(inner_const_op)
 
+        # CHECK: live ops: 4
+        log_op_count()
+
         PassManager.parse("builtin.module(canonicalize)").run(
             module, invalidate_ops=False
         )
@@ -222,6 +233,9 @@ def testPostPassOpInvalidation():
         # CHECK: invalidate_ops=True
         log("invalidate_ops=True")
 
+        # CHECK: live ops: 4
+        log_op_count()
+
         module = ModuleOp.parse(
             """
           module {
@@ -237,7 +251,14 @@ def testPostPassOpInvalidation():
         func_op = module.body.operations[1]
         inner_const_op = func_op.body.blocks[0].operations[0]
 
+        # CHECK: live ops: 4
+        log_op_count()
+
         PassManager.parse("builtin.module(canonicalize)").run(module)
+
+        # CHECK: live ops: 1
+        log_op_count()
+
         try:
             log(func_op)
         except RuntimeError as e: