[Mlir-commits] [mlir] [MLIR:Python] Fix race on PyOperations. (PR #139721)

[Maksim Levental]

================
@@ -512,6 +512,49 @@ def _original_test_create_module_with_consts(self):
                 arith.constant(dtype, py_values[2])
 
 
+    def test_check_pyoperation_race(self):
+        # Regression test for a race where:
+        # * one thread is in the process of destroying a PyOperation,
+        # * while simultaneously another thread looks up the PyOperation is
+        #   the liveOperations map and attempts to increase its reference count.
+        # It is illegal to attempt to revive an object that is in the process of
+        # being deleted, and this was producing races and heap use-after-frees.
+        num_workers = 40
+        num_runs = 20
+
+        barrier = threading.Barrier(num_workers)
+
+        def walk_operations(op):
+            _ = op.operation.name
+            for region in op.operation.regions:
+                for block in region:
+                    for op in block:
+                        walk_operations(op)
+
+        with Context():
+            mlir_module = Module.parse(
+                """
+    module @m {
+    func.func public @main(%arg0: tensor<f32>) -> (tensor<f32>) {
+        return %arg0 : tensor<f32>
+    }
+    }
+                """
+            )
----------------
makslevental wrote:

ultra-nit

```suggestion
            mlir_module = Module.parse(
                dedent(
                    """
                    module @m {
                      func.func public @main(%arg0: tensor<f32>) -> (tensor<f32>) {
                        return %arg0 : tensor<f32>
                      }
                    }
                    """
                )
            )
```

will need `from textwrap import dedent` at the top

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