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

[llvmlistbot at llvm.org]

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

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 

>From 66a6e3ca26f80461c1209f9caf5388cef97629f0 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sun, 12 Jul 2026 13:27:38 +0000
Subject: [PATCH] [mlir][bufferization] Skip escaping allocation users

---
 .../Transforms/OptimizeAllocationLiveness.cpp |  2 ++
 .../optimize-allocation-liveness.mlir         | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+)

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
+}