[Mlir-commits] [mlir] [mlir][bufferization] Fix bug when moving dealloc operation (PR #229296)
Pranav Alavandi
llvmlistbot at llvm.org
Mon Oct 12 19:49:36 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/3] [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 a41d1543ca480e769fefdc60537eb83b65618c27 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/3] Fixes from PR comment
---
.../Transforms/OptimizeAllocationLiveness.cpp | 13 +++++++----
.../optimize-allocation-liveness.mlir | 23 +++++++++++++++++++
2 files changed, 31 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 18e87db44d177..902d24e186ef5 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -107,10 +107,10 @@ struct 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 +119,7 @@ struct OptimizeAllocationLiveness
if (!definingOp)
continue;
- if (!target || domInfo.properlyDominates(target, definingOp))
+ if (!target || domInfo->properlyDominates(target, definingOp))
target = definingOp;
}
@@ -128,7 +128,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 +143,7 @@ struct OptimizeAllocationLiveness
return;
BufferViewFlowAnalysis analysis = BufferViewFlowAnalysis(func);
+ auto domInfo = DominanceInfo();
func.walk([&](MemoryEffectOpInterface memEffectOp) -> WalkResult {
if (!hasMemoryAllocEffect(memEffectOp))
@@ -190,7 +192,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
+}
>From a176e3706de62fc1470a48c32f1722e81b658d49 Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Tue, 13 Oct 2026 13:47:48 +1100
Subject: [PATCH 3/3] Add docstring + comments to explain
findValidInsertionPoint
---
.../Transforms/OptimizeAllocationLiveness.cpp | 21 +++++++++++++++++++
1 file changed, 21 insertions(+)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 902d24e186ef5..41ba73ff29949 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -107,6 +107,13 @@ struct OptimizeAllocationLiveness
public:
OptimizeAllocationLiveness() = default;
+ /// When moving a dealloc operation to be closer to its companion alloc we
+ /// have to consider the following. The dealloc should be:
+ /// - After all the operand definitions that the dealloc uses
+ /// - Before any of ops that consume the results produced by the dealloc
+ /// To do this we first ensure the "target" position is after all the operands
+ /// Then we ensure this target dominates all consumers of deallocs results. If
+ /// not no valid position exists and return nullptr.
static Operation *findValidInsertionPoint(DominanceInfo *domInfo,
Operation *allocOp,
Operation *deallocOp,
@@ -114,6 +121,12 @@ struct OptimizeAllocationLiveness
auto *target = lastUser;
+ // Start at the last user of the buffer, then move later if an operand of
+ // the dealloc is defined after it. E.g. target starts at %x but moves to
+ // %w:
+ // %x = memref.load %a[%c0] // lastUser
+ // %w = gpu.wait async [%d] // target
+ // %done = gpu.dealloc async [%w] %a : memref<8xf32> // deallocOp
for (Value operand : deallocOp->getOperands()) {
auto *definingOp = operand.getDefiningOp();
if (!definingOp)
@@ -126,6 +139,14 @@ struct OptimizeAllocationLiveness
if (!target)
target = allocOp;
+ // Ensure target dominates all result users. Here the realloc is the
+ // deallocOp since it frees %a, and it also produces %b. The first loop
+ // leaves target at the memref.dealloc (the last user of the aliased
+ // buffer), but that op uses %b, so placing the realloc after it is invalid:
+ // %a = memref.alloc() : memref<8xf32>
+ // %b = memref.realloc %a : memref<8xf32> to memref<16xf32> // deallocOp
+ // memref.dealloc %b : memref<16xf32> // target
+ // nullptr is returned as a result.
for (Value result : deallocOp->getResults()) {
for (Operation *user : result.getUsers()) {
if (user == target ||
More information about the Mlir-commits
mailing list