[Mlir-commits] [mlir] [mlir][bufferization] Fix bug when moving dealloc operation (PR #229296)
Pranav Alavandi
llvmlistbot at llvm.org
Sun Oct 11 22:41:24 PDT 2026
https://github.com/pranavalavandi updated https://github.com/llvm/llvm-project/pull/229296
>From 25b80ce6a57bded0cff6af6d350b593996cb9148 Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Tue, 6 Oct 2026 13:25:06 +1100
Subject: [PATCH 1/2] [mlir][bufferization] fix bug when moving dealloc
operation
---
.../Transforms/OptimizeAllocationLiveness.cpp | 41 ++++++++++++++++--
.../optimize-allocation-liveness.mlir | 43 +++++++++++++++++++
2 files changed, 81 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 28ee5b8e32b99..18e87db44d177 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -106,6 +107,34 @@ struct OptimizeAllocationLiveness
public:
OptimizeAllocationLiveness() = default;
+ static Operation *findValidInsertionPoint(Operation *allocOp,
+ Operation *deallocOp,
+ Operation *lastUser) {
+ auto domInfo = DominanceInfo();
+
+ auto *target = lastUser;
+
+ for (Value operand : deallocOp->getOperands()) {
+ auto *definingOp = operand.getDefiningOp();
+ if (!definingOp)
+ continue;
+
+ if (!target || domInfo.properlyDominates(target, definingOp))
+ target = definingOp;
+ }
+
+ if (!target)
+ target = allocOp;
+
+ for (Value result : deallocOp->getResults()) {
+ for (Operation *user : result.getUsers()) {
+ if (user == target || !domInfo.properlyDominates(target, user))
+ return nullptr;
+ }
+ }
+ return target;
+ }
+
void runOnOperation() override {
func::FuncOp func = getOperation();
@@ -160,9 +189,15 @@ struct OptimizeAllocationLiveness
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 = findValidInsertionPoint(allocOp, deallocOp, lastUser);
+ if (!target) {
+ LDBG() << "No valid insertion point found; skipping move";
+ return WalkResult::advance();
+ }
+
+ deallocOp->moveAfter(target);
+ LDBG() << "Moved dealloc op after: " << *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 b8fcccda8833a..ef3a1bb0552ca 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -234,3 +234,46 @@ func.func private @test_alloc_with_multiple_results() -> () {
memref.dealloc %alloc2 : memref<64xf32>
return
}
+
+// -----
+// CHECK-LABEL: func.func @test_dealloc_blocked_by_own_result_user() {
+// CHECK: %[[ALLOC_0:.*]] = memref.alloc() : memref<8xf32>
+// CHECK: %[[REALLOC_0:.*]] = memref.realloc %[[ALLOC_0]] : memref<8xf32> to memref<16xf32>
+// CHECK: memref.dealloc %[[REALLOC_0]] : memref<16xf32>
+// CHECK: return
+// CHECK: }
+
+// This test ensures that the dealloc does not move above its operand definitions
+func.func @test_dealloc_blocked_by_own_result_user() {
+ %a = memref.alloc() : memref<8xf32>
+ %b = memref.realloc %a : memref<8xf32> to memref<16xf32>
+ memref.dealloc %b : memref<16xf32>
+ return
+}
+
+// -----
+// CHECK-LABEL: func.func @test_gpu_dealloc_moved_after_async_token(
+// CHECK-SAME: %[[ARG0:.*]]: !gpu.async.token) {
+// CHECK: %[[ALLOC_0:.*]] = gpu.alloc () : memref<8xf32>
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// 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: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 2 : index
+// CHECK: gpu.wait {{\[}}%[[DEALLOC_0]]]
+// CHECK: return
+// CHECK: }
+
+// This test ensures that async token operand is respected when moving the dealloc
+func.func @test_gpu_dealloc_moved_after_async_token(%d: !gpu.async.token) {
+ %a = gpu.alloc () : memref<8xf32>
+ %c0 = arith.constant 0 : index
+ %x = memref.load %a[%c0] : memref<8xf32>
+ %w = gpu.wait async [%d]
+ %y = arith.constant 1 : index
+ %z = arith.constant 2 : index
+ %done = gpu.dealloc async [%w] %a : memref<8xf32>
+ gpu.wait [%done]
+ return
+}
>From b1531dc67d8997e60bfc8a80c6a673002cfef13a Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Mon, 12 Oct 2026 16:39:26 +1100
Subject: [PATCH 2/2] Fixes from PR comment
---
.../Transforms/OptimizeAllocationLiveness.cpp | 14 +++++++----
.../optimize-allocation-liveness.mlir | 23 +++++++++++++++++++
2 files changed, 32 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 18e87db44d177..844d36fce4d0d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -104,13 +104,14 @@ collectAllocations(MemoryEffectOpInterface allocOp) {
struct OptimizeAllocationLiveness
: public bufferization::impl::OptimizeAllocationLivenessPassBase<
OptimizeAllocationLiveness> {
+
public:
OptimizeAllocationLiveness() = default;
- static Operation *findValidInsertionPoint(Operation *allocOp,
+ static Operation *findValidInsertionPoint(DominanceInfo *domInfo,
+ Operation *allocOp,
Operation *deallocOp,
Operation *lastUser) {
- auto domInfo = DominanceInfo();
auto *target = lastUser;
@@ -119,7 +120,7 @@ struct OptimizeAllocationLiveness
if (!definingOp)
continue;
- if (!target || domInfo.properlyDominates(target, definingOp))
+ if (!target || domInfo->properlyDominates(target, definingOp))
target = definingOp;
}
@@ -128,7 +129,8 @@ struct OptimizeAllocationLiveness
for (Value result : deallocOp->getResults()) {
for (Operation *user : result.getUsers()) {
- if (user == target || !domInfo.properlyDominates(target, user))
+ if (user == target ||
+ !domInfo->properlyDominates(target, user, /*enclosingOpOk=*/false))
return nullptr;
}
}
@@ -142,6 +144,7 @@ struct OptimizeAllocationLiveness
return;
BufferViewFlowAnalysis analysis = BufferViewFlowAnalysis(func);
+ auto domInfo = DominanceInfo();
func.walk([&](MemoryEffectOpInterface memEffectOp) -> WalkResult {
if (!hasMemoryAllocEffect(memEffectOp))
@@ -190,7 +193,8 @@ struct OptimizeAllocationLiveness
assert(lastUser->getBlock() == allocOp->getBlock());
assert(lastUser->getBlock() == deallocOp->getBlock());
- Operation *target = findValidInsertionPoint(allocOp, deallocOp, lastUser);
+ Operation *target =
+ findValidInsertionPoint(&domInfo, allocOp, deallocOp, lastUser);
if (!target) {
LDBG() << "No valid insertion point found; skipping move";
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 ef3a1bb0552ca..02313cab3333e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -277,3 +277,26 @@ func.func @test_gpu_dealloc_moved_after_async_token(%d: !gpu.async.token) {
gpu.wait [%done]
return
}
+
+// -----
+// CHECK-LABEL: func.func @nested_result(
+// CHECK-SAME: %[[ARG0:.*]]: i1) {
+// CHECK: %[[ALLOC_0:.*]] = memref.alloc() : memref<8xf32>
+// CHECK: %[[REALLOC_0:.*]] = memref.realloc %[[ALLOC_0]] : memref<8xf32> to memref<16xf32>
+// CHECK: scf.if %[[ARG0]] {
+// CHECK: memref.dealloc %[[REALLOC_0]] : memref<16xf32>
+// CHECK: } else {
+// CHECK: memref.dealloc %[[REALLOC_0]] : memref<16xf32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+func.func @nested_result(%cond: i1) {
+ %a = memref.alloc() : memref<8xf32>
+ %b = memref.realloc %a : memref<8xf32> to memref<16xf32>
+ scf.if %cond {
+ memref.dealloc %b : memref<16xf32>
+ } else {
+ memref.dealloc %b : memref<16xf32>
+ }
+ return
+}
More information about the Mlir-commits
mailing list