[Mlir-commits] [mlir] [mlir][memref] Move dealloc for unused memref to after alloc (PR #227975)

[Pranav Alavandi]

https://github.com/pranavalavandi updated https://github.com/llvm/llvm-project/pull/227975

>From c048f1a1819f68cd5209943e992b98d4b936099e Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Thu, 1 Oct 2026 16:25:30 +1000
Subject: [PATCH 1/2] [mlir][memref] Move dealloc for unused memref to after
 alloc

---
 .../Transforms/OptimizeAllocationLiveness.cpp     |  3 +++
 .../Transforms/optimize-allocation-liveness.mlir  | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 28ee5b8e32b9913..3d4b8d2adb9117c 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -155,6 +155,9 @@ struct OptimizeAllocationLiveness
         }
       }
       if (lastUser == nullptr) {
+        LDBG() << "Last user not found ";
+        LDBG() << "Moved dealloc op after alloc: " << *allocOp;
+        deallocOp->moveAfter(allocOp);
         return WalkResult::advance();
       }
       LDBG() << "Last user found: " << *lastUser;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
index b8fcccda8833aef..07620859cc0bb0a 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -234,3 +234,18 @@ func.func private @test_alloc_with_multiple_results() -> () {
   memref.dealloc %alloc2 : memref<64xf32>
   return
 }
+
+// -----
+func.func private @work()
+// CHECK-LABEL:   func.func private @test_alloc_with_no_uses() {
+// CHECK:           %[[ALLOC_0:.*]] = memref.alloc() : memref<1048576xi32>
+// CHECK:           memref.dealloc %[[ALLOC_0]] : memref<1048576xi32>
+// CHECK:           call @work() : () -> ()
+// CHECK:           return
+// CHECK:         }
+func.func private @test_alloc_with_no_uses() -> () {
+    %alloc = memref.alloc() : memref<1048576xi32>
+    call @work() : () -> ()
+    memref.dealloc %alloc : memref<1048576xi32>
+    return
+}

>From 06a57354de29f2936d17a4a02870f365b3fbd42a Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Thu, 1 Oct 2026 18:36:52 +1000
Subject: [PATCH 2/2] fix bug where dealloc operands were not considered

---
 .../Transforms/OptimizeAllocationLiveness.cpp | 35 +++++++++++------
 .../optimize-allocation-liveness.mlir         | 38 +++++++++++++++++++
 2 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 3d4b8d2adb9117c..7634f92c6281400 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -154,20 +154,33 @@ struct OptimizeAllocationLiveness
           }
         }
       }
-      if (lastUser == nullptr) {
-        LDBG() << "Last user not found ";
-        LDBG() << "Moved dealloc op after alloc: " << *allocOp;
-        deallocOp->moveAfter(allocOp);
-        return WalkResult::advance();
+
+      Operation *lastDeallocOperandDef = nullptr;
+      for (auto operand : deallocOp->getOperands()) {
+        auto *operandDefOp = operand.getDefiningOp();
+        if (!operandDefOp)
+          continue;
+        auto *topUser =
+            allocOp->getBlock()->findAncestorOpInBlock(*operandDefOp);
+        if (!lastDeallocOperandDef ||
+            happensBefore(lastDeallocOperandDef, topUser)) {
+          lastDeallocOperandDef = topUser;
+        }
       }
-      LDBG() << "Last user found: " << *lastUser;
-      assert(lastUser->getBlock() == allocOp->getBlock());
-      assert(lastUser->getBlock() == deallocOp->getBlock());
-      // Move the dealloc op after the last user.
-      deallocOp->moveAfter(lastUser);
-      LDBG() << "Moved dealloc op after: " << *lastUser;
 
+      Operation *target = lastUser;
+      if (!target || (lastDeallocOperandDef &&
+                      happensBefore(target, lastDeallocOperandDef)))
+        target = lastDeallocOperandDef;
+
+      if (!target)
+        target = allocOp;
+
+      LDBG() << "Moving dealloc op after: " << *target;
+
+      deallocOp->moveAfter(target);
       return WalkResult::advance();
+
     });
   }
 };
diff --git a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
index 07620859cc0bb0a..6bb62d556feef31 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -249,3 +249,41 @@ func.func private @test_alloc_with_no_uses() -> () {
     memref.dealloc %alloc : memref<1048576xi32>
     return
 }
+
+// -----
+// CHECK-LABEL:   func.func @test_alloc_with_dealloc_operand(
+// CHECK-SAME:      %[[ARG0:.*]]: !gpu.async.token) {
+// CHECK:           %[[ALLOC_0:.*]] = gpu.alloc  () : memref<8xf32>
+// CHECK:           %[[WAIT_0:.*]] = gpu.wait async {{\[}}%[[ARG0]]]
+// CHECK:           %[[DEALLOC_0:.*]] = gpu.dealloc async {{\[}}%[[WAIT_0]]] %[[ALLOC_0]] : memref<8xf32>
+// CHECK:           gpu.wait {{\[}}%[[DEALLOC_0]]]
+// CHECK:           return
+// CHECK:         }
+func.func @test_alloc_with_dealloc_operand(%d: !gpu.async.token) {
+  %a = gpu.alloc () : memref<8xf32>
+  %w = gpu.wait async [%d]
+  %done = gpu.dealloc async [%w] %a : memref<8xf32>
+  gpu.wait [%done]
+  return
+}
+
+// -----
+// CHECK-LABEL:   func.func @test_alloc_with_dealloc_operand_and_load(
+// CHECK-SAME:      %[[ARG0:.*]]: !gpu.async.token) {
+// CHECK:           %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[ALLOC_0:.*]] = gpu.alloc  () : memref<8xf32>
+// CHECK:           %[[LOAD_0:.*]] = memref.load %[[ALLOC_0]]{{\[}}%[[CONSTANT_0]]] : memref<8xf32>
+// CHECK:           %[[WAIT_0:.*]] = gpu.wait async {{\[}}%[[ARG0]]]
+// CHECK:           %[[DEALLOC_0:.*]] = gpu.dealloc async {{\[}}%[[WAIT_0]]] %[[ALLOC_0]] : memref<8xf32>
+// CHECK:           gpu.wait {{\[}}%[[DEALLOC_0]]]
+// CHECK:           return
+// CHECK:         }
+func.func @test_alloc_with_dealloc_operand_and_load(%d: !gpu.async.token) {
+  %c0 = arith.constant 0 : index
+  %a = gpu.alloc () : memref<8xf32>
+  %x = memref.load %a[%c0] : memref<8xf32>
+  %w = gpu.wait async [%d]
+  %done = gpu.dealloc async [%w] %a : memref<8xf32>
+  gpu.wait [%done]
+  return
+}