[Mlir-commits] [mlir] Fix llvm-mlir-use-after-erase findings (PR #210733)

[llvmlistbot at llvm.org]

https://github.com/Harald-R updated https://github.com/llvm/llvm-project/pull/210733

>From 48d79b0b4341d2f96aa33d011f00af0e378eeb01 Mon Sep 17 00:00:00 2001
From: Harald-R <rotuna.razvan at gmail.com>
Date: Mon, 20 Jul 2026 12:52:09 +0300
Subject: [PATCH] Fix llvm-mlir-use-after-erase findings

---
 mlir/lib/Bytecode/Reader/BytecodeReader.cpp   |  4 ++--
 .../MemRef/Transforms/NormalizeMemRefs.cpp    | 20 ++++++++++---------
 .../Transforms/Utils/DialectConversion.cpp    |  3 +--
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 96ead28d0aa40..5c54738dbed96 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -1648,10 +1648,10 @@ class mlir::BytecodeReader::Impl {
           return failure();
         continue;
       }
-      op->dropAllReferences();
-      op->erase();
       lazyLoadableOps.pop_front();
       lazyLoadableOpsMap.erase(op);
+      op->dropAllReferences();
+      op->erase();
     }
     return success();
   }
diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
index fa7991e6c6a80..23591dd045f78 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
@@ -301,14 +301,15 @@ void NormalizeMemRefs::updateFunctionSignature(func::FuncOp funcOp,
         // If it failed (due to escapes for example), bail out.
         // It should never hit this part of the code because it is called by
         // only those functions which are normalizable.
-        newCallOp->erase();
         replacingMemRefUsesFailed = true;
         break;
       }
       returnTypeChanged = true;
     }
-    if (replacingMemRefUsesFailed)
+    if (replacingMemRefUsesFailed) {
+      newCallOp->erase();
       continue;
+    }
     // Replace all uses for other non-memref result types.
     userOp->replaceAllUsesWith(newCallOp);
     userOp->erase();
@@ -456,17 +457,18 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
                                               /*userFilterFn=*/nullptr,
                                               /*allowNonDereferencingOps=*/true,
                                               /*replaceInDeallocOp=*/true))) {
-            newOp->erase();
             replacingMemRefUsesFailed = true;
-            continue;
+            break;
           }
         }
-        if (!replacingMemRefUsesFailed) {
-          // Replace other ops with new op and delete the old op when the
-          // replacement succeeded.
-          op->replaceAllUsesWith(newOp);
-          op->erase();
+        if (replacingMemRefUsesFailed) {
+          newOp->erase();
+          return;
         }
+        // Replace other ops with new op and delete the old op when the
+        // replacement succeeded.
+        op->replaceAllUsesWith(newOp);
+        op->erase();
       }
     }
   });
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index c76e3808d3b37..c8db49c226e43 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -2696,8 +2696,7 @@ LogicalResult OperationLegalizer::legalizeWithFold(Operation *op) {
             "op '" + opName +
             "' folder rollback of IR modifications requested");
       }
-      rewriterImpl.resetState(
-          curState, std::string(op->getName().getStringRef()) + " folder");
+      rewriterImpl.resetState(curState, std::string(opName) + " folder");
       return failure();
     }
   }