[Mlir-commits] [mlir] [mlir][bufferization] Support multi-result alloc ops in OptimizeAllocationLiveness (PR #189527)
Mohamed Deraz Nasr
llvmlistbot at llvm.org
Mon Mar 30 21:43:03 PDT 2026
https://github.com/MDerazNasr updated https://github.com/llvm/llvm-project/pull/189527
>From 79f6c96558e9f5e7fec8debdd79838a3808ac3ca Mon Sep 17 00:00:00 2001
From: Mohamed Deraz Nasr <mderaznasr at gmail.com>
Date: Tue, 31 Mar 2026 00:34:59 -0400
Subject: [PATCH 1/4] [mlir][bufferization] Support multi-result alloc ops in
OptimizeAllocationLiveness
---
.../Transforms/OptimizeAllocationLiveness.cpp | 67 +++++++++----------
1 file changed, 32 insertions(+), 35 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 28ee5b8e32b99..6b1e8f396cd7f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -122,47 +122,44 @@ struct OptimizeAllocationLiveness
LDBG() << "Checking alloc op: " << allocOp;
SmallVector<OpResult> allocationResults = collectAllocations(allocOp);
- // Multiple allocations from a single op are not considered here yet.
- if (allocationResults.size() != 1)
- return WalkResult::advance();
- OpResult allocResult = allocationResults[0];
- LDBG() << "On allocation result: " << allocResult;
+ for (OpResult allocResult : allocationResults) {
+ LDBG() << "On allocation result: " << allocResult;
- auto *deallocOp = findUserWithFreeSideEffect(allocResult);
- if (!deallocOp || (deallocOp->getBlock() != allocOp->getBlock())) {
- // The pass handles allocations that have a single dealloc op in the
- // same block. We also should not hoist the dealloc op out of
- // conditionals.
- return WalkResult::advance();
- }
+ auto *deallocOp = findUserWithFreeSideEffect(allocResult);
+ if (!deallocOp || (deallocOp->getBlock() != allocOp->getBlock())) {
+ // Skip results whose dealloc is in a different block. Moving
+ // across blocks could hoist the dealloc out of conditionals.
+ continue;
+ }
- Operation *lastUser = nullptr;
- const BufferViewFlowAnalysis::ValueSetT &deps =
- analysis.resolve(allocResult);
- for (auto dep : llvm::make_early_inc_range(deps)) {
- for (auto *user : dep.getUsers()) {
- // We are looking for a non dealloc op user.
- // check if user is the dealloc op itself.
- if (user == deallocOp)
- continue;
-
- // find the ancestor of user that is in the same block as the allocOp.
- auto *topUser = allocOp->getBlock()->findAncestorOpInBlock(*user);
- if (!lastUser || happensBefore(lastUser, topUser)) {
- lastUser = topUser;
+ Operation *lastUser = nullptr;
+ const BufferViewFlowAnalysis::ValueSetT &deps =
+ analysis.resolve(allocResult);
+ for (auto dep : llvm::make_early_inc_range(deps)) {
+ for (auto *user : dep.getUsers()) {
+ // We are looking for a non dealloc op user.
+ // check if user is the dealloc op itself.
+ if (user == deallocOp)
+ continue;
+
+ // find the ancestor of user that is in the same block as the allocOp.
+ auto *topUser = allocOp->getBlock()->findAncestorOpInBlock(*user);
+ if (!lastUser || happensBefore(lastUser, topUser)) {
+ lastUser = topUser;
+ }
}
}
+ if (lastUser == nullptr)
+ continue;
+
+ 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;
}
- if (lastUser == nullptr) {
- return WalkResult::advance();
- }
- 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;
return WalkResult::advance();
});
>From 5cd16f436d8ac1a9a6a21963a9aab1b06516d5a5 Mon Sep 17 00:00:00 2001
From: Mohamed Deraz Nasr <mderaznasr at gmail.com>
Date: Tue, 31 Mar 2026 00:35:10 -0400
Subject: [PATCH 2/4] [mlir][bufferization] Support multi-result alloc ops in
OptimizeAllocationLiveness
---
mlir/test/lib/Dialect/Test/TestOps.td | 8 ++++++++
1 file changed, 8 insertions(+)
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 4c9e6b3fe9e45..18ef5ce3d58c0 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3887,6 +3887,14 @@ def TestAllocWithMultipleResults : TEST_Op<"alloc_with_multiple_results"> {
}];
}
+def TestAllocWithTwoMemRefResults : TEST_Op<"alloc_with_two_memref_results"> {
+ let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref0,
+ Res<AnyMemRef, "", [MemAlloc]>:$memref1);
+ let assemblyFormat = [{
+ attr-dict `:` type($memref0) `,` type($memref1)
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Test Ops bufferization
//===----------------------------------------------------------------------===//
>From 2ec5619108073cea4c85bd50d5aa36c7c507d97c Mon Sep 17 00:00:00 2001
From: Mohamed Deraz Nasr <mderaznasr at gmail.com>
Date: Tue, 31 Mar 2026 00:35:28 -0400
Subject: [PATCH 3/4] [mlir][bufferization] Support multi-result alloc ops in
OptimizeAllocationLiveness
---
.../optimize-allocation-liveness.mlir | 20 +++++++++++++++++++
1 file changed, 20 insertions(+)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
index 63d33e3a88bed..f810a70389e65 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -234,3 +234,23 @@ func.func private @test_alloc_with_multiple_results() -> () {
memref.dealloc %alloc2 : memref<64xf32>
return
}
+
+// -----
+
+// CHECK-LABEL: func.func private @test_two_alloc_results
+// CHECK: %[[a0:.*]], %[[a1:.*]] = test.alloc_with_two_memref_results
+// CHECK-NEXT: memref.load %[[a0]]
+// CHECK-NEXT: memref.dealloc %[[a0]]
+// CHECK: memref.store
+// CHECK-NEXT: memref.dealloc %[[a1]]
+
+// Op produces two allocated results. Each dealloc should move to right
+// after the last use of its corresponding result.
+func.func private @test_two_alloc_results(%c0: index) {
+ %alloc0, %alloc1 = test.alloc_with_two_memref_results : memref<32xf32>, memref<32xf32>
+ %val = memref.load %alloc0[%c0] : memref<32xf32>
+ memref.store %val, %alloc1[%c0] : memref<32xf32>
+ memref.dealloc %alloc0 : memref<32xf32>
+ memref.dealloc %alloc1 : memref<32xf32>
+ return
+}
>From 02ef918486eb19d37232a38227af9da488e686f5 Mon Sep 17 00:00:00 2001
From: Mohamed Deraz Nasr <mderaznasr at gmail.com>
Date: Tue, 31 Mar 2026 00:42:53 -0400
Subject: [PATCH 4/4] use alloc0/alloc1 names in CHECK patterns
---
.../Transforms/optimize-allocation-liveness.mlir | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
index f810a70389e65..b62a91af196a3 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -238,11 +238,11 @@ func.func private @test_alloc_with_multiple_results() -> () {
// -----
// CHECK-LABEL: func.func private @test_two_alloc_results
-// CHECK: %[[a0:.*]], %[[a1:.*]] = test.alloc_with_two_memref_results
-// CHECK-NEXT: memref.load %[[a0]]
-// CHECK-NEXT: memref.dealloc %[[a0]]
+// CHECK: %[[alloc0:.*]], %[[alloc1:.*]] = test.alloc_with_two_memref_results
+// CHECK-NEXT: memref.load %[[alloc0]]
+// CHECK-NEXT: memref.dealloc %[[alloc0]]
// CHECK: memref.store
-// CHECK-NEXT: memref.dealloc %[[a1]]
+// CHECK-NEXT: memref.dealloc %[[alloc1]]
// Op produces two allocated results. Each dealloc should move to right
// after the last use of its corresponding result.
More information about the Mlir-commits
mailing list