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

[Maksim Levental]

================
@@ -720,25 +795,26 @@ size_t PyMlirContext::clearLiveOperations() {
   {
     nb::ft_lock_guard lock(liveOperationsMutex);
     std::swap(operations, liveOperations);
+    for (auto &op : operations)
+      op.second.second->setInvalidLocked();
   }
-  for (auto &op : operations)
-    op.second.second->setInvalid();
   size_t numInvalidated = operations.size();
   return numInvalidated;
 }
 
-void PyMlirContext::clearOperation(MlirOperation op) {
-  PyOperation *py_op;
-  {
-    nb::ft_lock_guard lock(liveOperationsMutex);
-    auto it = liveOperations.find(op.ptr);
-    if (it == liveOperations.end()) {
-      return;
-    }
-    py_op = it->second.second;
-    liveOperations.erase(it);
+void PyMlirContext::clearOperationLocked(MlirOperation op) {
+  auto it = liveOperations.find(op.ptr);
+  if (it == liveOperations.end()) {
+    return;
   }
-  py_op->setInvalid();
+  PyOperation *py_op = it->second.second;
+  py_op->setInvalidLocked();
+  liveOperations.erase(it);
----------------
makslevental wrote:

wouldn't you rather reorder these? the `py_op` isn't actually "invalid" until it's erased from `liveOperations` right?

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