[Mlir-commits] [mlir] [mlir][bufferization] Support multi-result alloc ops in OptimizeAllocationLiveness (PR #189527)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Mohamed Deraz Nasr (MDerazNasr)

<details>
<summary>Changes</summary>

The pass previously bailed out entirely when an op produced more than one allocated result (guarded by `allocationResults.size() != 1`). This meant none of the results had their dealloc moved, even if each had a perfectly valid single dealloc in the same block.

This patch removes that guard and replaces the single-result path with a loop over all allocation results. Each result is handled independently: if a result has no single dealloc in the same block it is skipped (`continue`), otherwise its dealloc is moved to right after its last user.

To exercise this path a new test op `test.alloc_with_two_memref_results` is added (both results carry `MemAlloc` effects) along with a lit test that verifies each dealloc ends up immediately after the last use of its corresponding result.

---
Full diff: https://github.com/llvm/llvm-project/pull/189527.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp (+32-35) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir (+20) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+8) 


``````````diff
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();
     });
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
+}
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
 //===----------------------------------------------------------------------===//

``````````

</details>


https://github.com/llvm/llvm-project/pull/189527