[Mlir-commits] [mlir] [mlir][bufferization] Fix placement walk crashes and invalid captures (PR #228050)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 05:06:10 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Matteo Di Iorio (splicer3)
<details>
<summary>Changes</summary>
`findPlacementBlock`, which is shared by `buffer-hoisting` and `buffer-loop-hoisting`, can move an allocation outside the operation used to build the liveness analysis. For example, when `transform.bufferization.buffer_loop_hoisting` targets a `scf.while`, moving the allocation into the enclosing function block causes the liveness lookup to crash.
The walk can also cross an `IsIsolatedFromAbove` operation and introduce an illegal capture, or reach an unreachable enclosing block and dereference a missing dominator node. The initial reachability check on the allocation's block doesn't cover that second case.
This PR keeps allocations inside the analysis target and isolated regions, and stops the walk before querying a dominator node for an unreachable block. Movement between blocks within the same region remains allowed.
I've also added regression tests for these cases, plus one that targets the enclosing function and checks that the allocation still moves out of the loop. On the base, the scope and unreachable-parent tests crash, while the isolation test fails verification.
Changes were made by splitting them out of #<!-- -->226951, as suggested during [review](https://github.com/llvm/llvm-project/pull/226951#pullrequestreview-5372976164). The failures also reproduce without the exit-alias fallback presented in that PR.
Validation:
Release build with assertions on macOS, `check-mlir` completed with 4052 passed and one expected failure.
---
Full diff: https://github.com/llvm/llvm-project/pull/228050.diff
4 Files Affected:
- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+10-6)
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir (+17)
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir (+24)
- (modified) mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir (+53)
``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 4f2a4734056b6..e03887c8c5910 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -268,6 +268,11 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
(parentBlock = parentOp->getBlock()) &&
(!upperBound ||
dominators.properlyDominates(upperBound, currentBlock))) {
+ // A reachable nested region may have an unreachable enclosing block,
+ // which has no node in the dominator tree.
+ if (!dominators.isReachableFromEntry(currentBlock))
+ break;
+
// Try to find an immediate dominator and check whether the parent block
// is above the immediate dominator (if any).
DominanceInfoNode *idom = nullptr;
@@ -282,12 +287,11 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
currentBlock = idom->getBlock();
state.recordMoveToDominator(currentBlock);
} else {
- // We have to move to our parent block since an immediate dominator does
- // either not exist or is above our parent block. If we cannot move to
- // our parent operation due to constraints given by the StateT
- // implementation, break the walk loop. Furthermore, we should not move
- // allocations out of unknown region-based control-flow operations.
- if (!isKnownControlFlowInterface(parentOp) ||
+ // The target's parent block is outside the liveness analysis. Keep
+ // allocations within the target and preserve isolation boundaries.
+ if (parentOp == scopeOp ||
+ parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>() ||
+ !isKnownControlFlowInterface(parentOp) ||
!state.isLegalPlacement(parentOp))
break;
// Move to our parent block by notifying the current StateT
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
index 4c07df951e280..827fcbb0a156b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
@@ -802,3 +802,20 @@ func.func @unreachable_alloc(%arg0: f32) {
// CHECK: memref.alloc
return
}
+
+// -----
+
+// Hoisting must not introduce a capture into an isolated region.
+// CHECK-LABEL: func @no_hoist_isolated_region(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[RESULT:.*]] = "test.isolated_region_branch"
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: "test.isolated_region_yield"(%[[ALLOC]])
+// CHECK: return %[[RESULT]]
+func.func @no_hoist_isolated_region() -> memref<1xi32> {
+ %result = "test.isolated_region_branch"() ({
+ %buffer = memref.alloc() : memref<1xi32>
+ "test.isolated_region_yield"(%buffer) : (memref<1xi32>) -> ()
+ }) : () -> memref<1xi32>
+ return %result : memref<1xi32>
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index ee04564f22f6b..7498f7af8c5d1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -543,3 +543,27 @@ func.func @hoist_alloca(
// CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: {{.*}} = scf.for
+
+// -----
+
+// A nested entry block can be reachable even when its enclosing block is not.
+// CHECK-LABEL: func @loop_unreachable_parent(
+// CHECK: return
+// CHECK: ^bb1:
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+func.func @loop_unreachable_parent() {
+ return
+^dead:
+ %c0 = arith.constant 0 : index
+ %false = arith.constant false
+ scf.while : () -> () {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %buffer[%c0] : memref<1xindex>
+ scf.condition(%false)
+ } do {
+ scf.yield
+ }
+ return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 4e897a951bc6a..27ff9964c2ca2 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -236,3 +236,56 @@ func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
%0 = bufferization.alloc_tensor() : tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
+
+// -----
+
+// The loop's parent block is outside the analysis rooted at the target loop.
+// CHECK-LABEL: func @buffer_loop_hoisting_while_scope_no_exit_alias(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %loop = transform.structured.match ops{["scf.while"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.bufferization.buffer_loop_hoisting %loop : !transform.any_op
+ transform.yield
+ }
+ func.func @buffer_loop_hoisting_while_scope_no_exit_alias(%condition: i1, %value: index) {
+ %c0 = arith.constant 0 : index
+ scf.while : () -> () {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %value, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition)
+ } do {
+ scf.yield
+ }
+ return
+ }
+}
+
+// -----
+
+// The same loop can be crossed when the function is the analysis target.
+// CHECK-LABEL: func @buffer_loop_hoisting_function_scope_no_exit_alias(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.bufferization.buffer_loop_hoisting %func : !transform.any_op
+ transform.yield
+ }
+ func.func @buffer_loop_hoisting_function_scope_no_exit_alias(%condition: i1, %value: index) {
+ %c0 = arith.constant 0 : index
+ scf.while : () -> () {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %value, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition)
+ } do {
+ scf.yield
+ }
+ return
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/228050
More information about the Mlir-commits
mailing list