From a0d8c51e66ed75bb55a3ae43730cae89e37f6f4f Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov@arm.com>
Date: Sun, 9 Nov 2025 22:03:33 +0000
Subject: [PATCH 1/2] [mlir]: Add handling of escaped memrefs to
 erase_dead_alloc_and_stores

transform.memref.erase_dead_alloc_and_stores will not delete escaped
allocations.

Change-Id: I0fce45a19e08e80fc2c7c2598c7e702865d37fd3
---
 mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp |  3 +++
 mlir/test/Dialect/MemRef/transform-ops.mlir   | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index 6200366cded29..8dc48c67265f4 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -140,6 +140,9 @@ static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {
   std::vector<Operation *> opUses;
   for (OpOperand &use : op->getUses()) {
     Operation *useOp = use.getOwner();
+    // Use escaped the scope
+    if (useOp->hasTrait<mlir::OpTrait::ReturnLike>())
+      return false;
     if (isa<memref::DeallocOp>(useOp) ||
         (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&
          !mlir::hasEffect<MemoryEffects::Read>(useOp)) ||
diff --git a/mlir/test/Dialect/MemRef/transform-ops.mlir b/mlir/test/Dialect/MemRef/transform-ops.mlir
index 3b37c62fcb28e..6e130912c47e9 100644
--- a/mlir/test/Dialect/MemRef/transform-ops.mlir
+++ b/mlir/test/Dialect/MemRef/transform-ops.mlir
@@ -306,6 +306,23 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: func.func @dead_alloc_escaped
+func.func @dead_alloc_escaped() -> memref<8x64xf32, 3> {
+  // CHECK: %{{.+}} = memref.alloc
+  %0 = memref.alloc() : memref<8x64xf32, 3>
+  return %0 : memref<8x64xf32, 3>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    transform.memref.erase_dead_alloc_and_stores %0 : (!transform.any_op) -> ()
+    transform.yield
+  }
+}
+
+// -----
+
 // CHECK-LABEL: func.func @dead_alloc
 func.func @dead_alloc() {
   // CHECK-NOT: %{{.+}} = memref.alloc

From 39247a8ecc849233534f41dfd316ddc224395ec1 Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov@arm.com>
Date: Mon, 10 Nov 2025 10:33:15 +0000
Subject: [PATCH 2/2] Addressed upstream comment

Change-Id: I02495c6cf2d67fcf7e0d266c8b51c661e759a3f4
---
 mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index 8dc48c67265f4..e6adcde72ad66 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -141,7 +141,7 @@ static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {
   for (OpOperand &use : op->getUses()) {
     Operation *useOp = use.getOwner();
     // Use escaped the scope
-    if (useOp->hasTrait<mlir::OpTrait::ReturnLike>())
+    if (useOp->mightHaveTrait<OpTrait::IsTerminator>())
       return false;
     if (isa<memref::DeallocOp>(useOp) ||
         (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&