[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 9cd5aedced66d8d94cc1573b97f6cd2d13fbae0a 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

---
 .../Dialect/MemRef/Transforms/NormalizeMemRefs.cpp |  2 +-
 mlir/lib/Transforms/Utils/DialectConversion.cpp    |  3 +--
 .../test/Dialect/MemRef/normalize-memrefs-ops.mlir | 14 ++++++++++++++
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
index fa7991e6c6a80..d5e0dace3c775 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
@@ -458,7 +458,7 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
                                               /*replaceInDeallocOp=*/true))) {
             newOp->erase();
             replacingMemRefUsesFailed = true;
-            continue;
+            break;
           }
         }
         if (!replacingMemRefUsesFailed) {
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();
     }
   }
diff --git a/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir b/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
index 344da4e5e2462..0a9204d4573c9 100644
--- a/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
+++ b/mlir/test/Dialect/MemRef/normalize-memrefs-ops.mlir
@@ -127,6 +127,20 @@ func.func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x
     // CHECK-NEXT: return %[[v1]], %[[v2]] : memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>
 }
 
+// Test with op_norm_ret, where results are consumed by a non-normalizable op.
+
+// CHECK-LABEL: test_norm_ret_denorm
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x32xf32>)
+func.func @test_norm_ret_denorm(%arg0: memref<1x16x14x14xf32, #map_tile>) {
+    %0, %1 = "test.op_norm_ret"(%arg0) : (memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>)
+    // CHECK: %[[v0:.*]], %[[v1:.*]] = "test.op_norm_ret"(%[[ARG0]]) : (memref<1x16x1x1x32x32xf32>) -> (memref<1x16x14x14xf32, #[[MAP:.*]]>, memref<1x16x14x14xf32, #[[MAP]]>)
+    "test.op_nonnorm"(%0, %0) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>) -> ()
+    // CHECK: "test.op_nonnorm"(%[[v0]], %[[v0]]) : (memref<1x16x14x14xf32, #[[MAP]]>, memref<1x16x14x14xf32, #[[MAP]]>) -> ()
+    memref.dealloc %1 : memref<1x16x14x14xf32, #map_tile>
+    // CHECK: memref.dealloc %[[v1]] : memref<1x16x14x14xf32, #[[MAP]]>
+    return
+}
+
 // Test with an arbitrary op that references the function symbol.
 
 "test.op_funcref"() {func = @test_norm_mix} : () -> ()