[Mlir-commits] [mlir] [mlir][bufferization] Skip escaping allocation users (PR #208979)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: lianjinfeng2003 (mygitljf)

<details>
<summary>Changes</summary>

I made the liveness optimization bail out when an allocation user cannot be safely represented in the allocation block. In that case the pass leaves the dealloc in place instead of trying to compare or move across CFG boundaries.
Fixes #<!-- -->205982 

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp (+2) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
index 28ee5b8e32b99..cfa6fca9d601a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp
@@ -149,6 +149,8 @@ struct OptimizeAllocationLiveness
 
           // find the ancestor of user that is in the same block as the allocOp.
           auto *topUser = allocOp->getBlock()->findAncestorOpInBlock(*user);
+          if (!topUser || topUser->hasTrait<OpTrait::IsTerminator>())
+            return WalkResult::advance();
           if (!lastUser || happensBefore(lastUser, topUser)) {
             lastUser = topUser;
           }
diff --git a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
index 63d33e3a88bed..4ba39d84d184f 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir
@@ -234,3 +234,25 @@ func.func private @test_alloc_with_multiple_results() -> () {
   memref.dealloc %alloc2 : memref<64xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func.func @allocation_alias_in_successor_block(
+// CHECK:         %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT:    test.buffer_based
+// CHECK-NEXT:    memref.dealloc %[[ALLOC]]
+// CHECK-NEXT:    cf.br ^bb3(%[[ALLOC]] : memref<2xf32>)
+func.func @allocation_alias_in_successor_block(
+    %arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+  cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  cf.br ^bb3(%arg1 : memref<2xf32>)
+^bb2:
+  %alloc = memref.alloc() : memref<2xf32>
+  test.buffer_based in(%arg1 : memref<2xf32>) out(%alloc : memref<2xf32>)
+  memref.dealloc %alloc : memref<2xf32>
+  cf.br ^bb3(%alloc : memref<2xf32>)
+^bb3(%0: memref<2xf32>):
+  test.copy(%0, %arg2) : (memref<2xf32>, memref<2xf32>)
+  return
+}

``````````

</details>


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