[Mlir-commits] [mlir] f8be54b - [mlir][bufferization] Fix placement walk crashes and invalid captures (#228050)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 2 19:25:36 PDT 2026
Author: Matteo Di Iorio
Date: 2026-10-03T02:25:29Z
New Revision: f8be54bc80fc677344e6639493894d32a12f0c47
URL: https://github.com/llvm/llvm-project/commit/f8be54bc80fc677344e6639493894d32a12f0c47
DIFF: https://github.com/llvm/llvm-project/commit/f8be54bc80fc677344e6639493894d32a12f0c47.diff
LOG: [mlir][bufferization] Fix placement walk crashes and invalid captures (#228050)
`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.
Co-authored-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Co-authored-by: Matthias Springer <me at m-sp.org>
Added:
Modified:
mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
Removed:
################################################################################
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
+ }
+}
More information about the Mlir-commits
mailing list