[Mlir-commits] [mlir] 2c41964 - [mlir][bufferization] Remove memref.alloc-only check for while exit aliases (#232834)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 18 23:21:45 PDT 2026
Author: Matteo Di Iorio
Date: 2026-10-19T15:21:40+09:00
New Revision: 2c419645e808a6667eee9dbb16c828f77e2fd2f0
URL: https://github.com/llvm/llvm-project/commit/2c419645e808a6667eee9dbb16c828f77e2fd2f0
DIFF: https://github.com/llvm/llvm-project/commit/2c419645e808a6667eee9dbb16c828f77e2fd2f0.diff
LOG: [mlir][bufferization] Remove memref.alloc-only check for while exit aliases (#232834)
The `scf.while` exit-alias path added in #226951 currently restricts
allocations to `memref.alloc`. With the realloc interface model fixed in
#231608, I think we can remove that restriction.
This allows any eligible allocation op that supports
`HoistingKind::Loop` to use that path, which at this time means an
extension to `memref.alloca`.
Also remove the allocation value from the hoisting state since only that
check used it.
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 954bc67ce7bf6..a0a585e6c977c 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -306,13 +306,15 @@ def BufferLoopHoistingPass : Pass<"buffer-loop-hoisting", "func::FuncOp"> {
This pass implements an approach to aggressively move allocations upwards
out of loop nests. It does not move allocations into common dominators.
- `memref.alloc` operations directly in the `before` region of `scf.while`
- may also be hoisted when their buffers are forwarded through loop results.
+ Allocations directly in the `before` region of `scf.while` may also be
+ hoisted when their buffers are forwarded through loop results. They must
+ advertise `HoistingKind::Loop` through `AllocationOpInterface`.
This extension requires a loop strictly inside the transformation target.
Its parent must stop further loop hoisting: it must either be a non-loop
operation or have the `HasParallelRegion` trait. No alias may be carried
back to `before`, and all alias uses, including uses of loop results, must
- permit reusing the allocation across iterations.
+ permit reusing the allocation across iterations. Stack allocations remain
+ within their original automatic allocation scope.
}];
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index f7bc9a70b0864..37444edccbf56 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -225,17 +225,13 @@ struct BufferAllocationHoistingStateBase {
/// A pointer to the current dominance info.
DominanceInfo *dominators;
- /// The current allocation value.
- Value allocValue;
-
/// The current placement block (if any).
Block *placementBlock;
/// Initializes the state base.
- BufferAllocationHoistingStateBase(DominanceInfo *dominators, Value allocValue,
+ BufferAllocationHoistingStateBase(DominanceInfo *dominators,
Block *placementBlock)
- : dominators(dominators), allocValue(allocValue),
- placementBlock(placementBlock) {}
+ : dominators(dominators), placementBlock(placementBlock) {}
};
/// Implements the actual hoisting logic for allocation nodes.
@@ -271,8 +267,7 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
Block *dominatorBlock =
findCommonDominator(allocValue, resultAliases, dominators);
// Init the initial hoisting state.
- StateT state(&dominators, allocValue, allocValue.getParentBlock(),
- resultAliases);
+ StateT state(&dominators, allocValue.getParentBlock(), resultAliases);
// Check for additional allocation dependencies to compute an upper bound
// for hoisting.
Block *dependencyBlock = nullptr;
@@ -374,11 +369,10 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
/// that hoists allocations into dominator blocks while keeping them inside of
/// loops.
struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
- BufferAllocationHoistingState(DominanceInfo *dominators, Value allocValue,
+ BufferAllocationHoistingState(DominanceInfo *dominators,
Block *placementBlock,
const BufferViewFlowAnalysis::ValueSetT &)
- : BufferAllocationHoistingStateBase(dominators, allocValue,
- placementBlock) {}
+ : BufferAllocationHoistingStateBase(dominators, placementBlock) {}
/// Computes the upper bound for the placement block search.
Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
@@ -419,10 +413,9 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
Block *aliasDominatorBlock = nullptr;
BufferAllocationLoopHoistingState(
- DominanceInfo *dominators, Value allocValue, Block *placementBlock,
+ DominanceInfo *dominators, Block *placementBlock,
const BufferViewFlowAnalysis::ValueSetT &aliases)
- : BufferAllocationHoistingStateBase(dominators, allocValue,
- placementBlock),
+ : BufferAllocationHoistingStateBase(dominators, placementBlock),
aliases(aliases) {}
/// Computes the upper bound for the placement block search.
@@ -440,12 +433,7 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
if (!dominators->dominates(aliasDominatorBlock, op->getBlock()))
return true;
auto loop = dyn_cast<scf::WhileOp>(op);
- // The alias checks do not cover uses of allocation operands. Hoisting
- // memref.realloc before a store to its source would leave a use of the
- // old memref after reallocation.
- if (!loop || !allocValue.getDefiningOp<memref::AllocOp>())
- return false;
- return canHoistFromWhile(loop, aliases);
+ return loop && canHoistFromWhile(loop, aliases);
}
/// Returns true if the given operation should be considered for hoisting.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 807ee95d65b4f..4a7fbd297112f 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -1149,13 +1149,12 @@ func.func @no_hoist_while_capture(%condition: i1) {
// -----
-// The exit-alias fallback does not support stack allocations.
-// CHECK-LABEL: func @no_hoist_while_alloca(
-// CHECK-NOT: memref.alloca
-// CHECK: scf.while
-// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloca()
+// Hoisting preserves the function's automatic allocation scope.
+// CHECK-LABEL: func @hoist_while_alloca(
+// CHECK: %[[ALLOC:.*]] = memref.alloca()
+// CHECK-NEXT: scf.while
// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
-func.func @no_hoist_while_alloca(%condition: i1) {
+func.func @hoist_while_alloca(%condition: i1) {
%last = scf.while () : () -> memref<1xindex> {
%buffer = memref.alloca() : memref<1xindex>
scf.condition(%condition) %buffer : memref<1xindex>
@@ -1168,6 +1167,140 @@ func.func @no_hoist_while_alloca(%condition: i1) {
// -----
+// Hoist an alloca whose dynamic size is loop-invariant, even when an alias
+// produced by a subview and a cast is forwarded to both loop results.
+// CHECK-LABEL: func.func @alloca_exit_dynamic(
+// CHECK-SAME: %{{.*}}: index, %[[SIZE:.*]]: index)
+// CHECK-NOT: scf.while
+// CHECK: %[[BUF:.*]] = memref.alloca(%[[SIZE]]) : memref<?xi32>
+// CHECK-NEXT: %[[LAST:.*]]:2 = scf.while
+// CHECK: memref.store %{{.*}}, %[[BUF]][%{{.*}}]
+// CHECK-NEXT: %[[VIEW:.*]] = memref.subview %[[BUF]]
+// CHECK-NEXT: %[[CAST:.*]] = memref.cast %[[VIEW]]
+// CHECK: scf.condition(%{{.*}}) %[[CAST]], %[[CAST]]
+// CHECK: } do {
+// CHECK-NEXT: ^bb0(%[[FIRST:.*]]: memref<?xi32>, %{{.*}}: memref<?xi32>):
+// CHECK-NEXT: %{{.*}} = memref.load %[[FIRST]]
+// CHECK: %[[VALUE:.*]] = memref.load %[[LAST]]#1
+// CHECK-NEXT: return %[[VALUE]] : i32
+func.func @alloca_exit_dynamic(%n: index, %size: index) -> i32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last:2 = scf.while (%i = %c0) : (index) -> (memref<?xi32>, memref<?xi32>) {
+ %buffer = memref.alloca(%size) : memref<?xi32>
+ %v = arith.index_cast %i : index to i32
+ memref.store %v, %buffer[%c0] : memref<?xi32>
+ %view = memref.subview %buffer[0] [1] [1] : memref<?xi32> to memref<1xi32>
+ %cast = memref.cast %view : memref<1xi32> to memref<?xi32>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %cast, %cast : memref<?xi32>, memref<?xi32>
+ } do {
+ ^bb0(%first: memref<?xi32>, %second: memref<?xi32>):
+ %value = memref.load %first[%c0] : memref<?xi32>
+ %i = arith.index_cast %value : i32 to index
+ %next = arith.addi %i, %c1 : index
+ scf.yield %next : index
+ }
+ %result = memref.load %last#1[%c0] : memref<?xi32>
+ return %result : i32
+}
+
+// -----
+
+// Keep the alloca in the after region: hoisting would make the store
+// overwrite the previous iteration's buffer before it is read.
+// CHECK-LABEL: func.func @after_backedge(
+// CHECK-NOT: memref.alloca()
+// CHECK: scf.while
+// CHECK-NOT: memref.alloca()
+// CHECK: } do {
+// CHECK: %[[FRESH:.*]] = memref.alloca()
+// CHECK: memref.store %[[NEXT:.*]], %[[FRESH]]
+// CHECK: scf.yield %[[FRESH]], %[[NEXT]]
+func.func @after_backedge(%seed: memref<1xindex>, %out: memref<2xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %result:2 = scf.while (%old = %seed, %i = %c0)
+ : (memref<1xindex>, index) -> (memref<1xindex>, index) {
+ %continue = arith.cmpi slt, %i, %c2 : index
+ scf.condition(%continue) %old, %i : memref<1xindex>, index
+ } do {
+ ^bb0(%previous: memref<1xindex>, %i: index):
+ %fresh = memref.alloca() : memref<1xindex>
+ %next = arith.addi %i, %c1 : index
+ memref.store %next, %fresh[%c0] : memref<1xindex>
+ %oldValue = memref.load %previous[%c0] : memref<1xindex>
+ memref.store %oldValue, %out[%i] : memref<2xindex>
+ scf.yield %fresh, %next : memref<1xindex>, index
+ }
+ return
+}
+
+// -----
+
+// Hoist before the while without crossing the explicit allocation scope.
+// CHECK-LABEL: func.func @alloca_exit_in_scope(
+// CHECK-NOT: memref.alloca()
+// CHECK: %[[RESULT:.*]] = memref.alloca_scope -> (i32) {
+// CHECK-NEXT: %[[BUF:.*]] = memref.alloca() : memref<1xi32>
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store %{{.*}}, %[[BUF]]
+// CHECK-NEXT: scf.condition(%{{.*}}) %[[BUF]]
+// CHECK: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: memref.alloca_scope.return %[[VALUE]] : i32
+// CHECK: return %[[RESULT]] : i32
+func.func @alloca_exit_in_scope(%condition: i1, %value: i32) -> i32 {
+ %c0 = arith.constant 0 : index
+ %result = memref.alloca_scope -> i32 {
+ %last = scf.while () : () -> memref<1xi32> {
+ %buffer = memref.alloca() : memref<1xi32>
+ memref.store %value, %buffer[%c0] : memref<1xi32>
+ scf.condition(%condition) %buffer : memref<1xi32>
+ } do {
+ ^bb0(%current: memref<1xi32>):
+ scf.yield
+ }
+ %loaded = memref.load %last[%c0] : memref<1xi32>
+ memref.alloca_scope.return %loaded : i32
+ }
+ return %result : i32
+}
+
+// -----
+
+// Keep the hoisted alloca private to each parallel iteration.
+// CHECK-LABEL: func.func @alloca_exit_in_parallel(
+// CHECK-NOT: memref.alloca()
+// CHECK: scf.parallel (%[[I:.*]]) =
+// CHECK-NEXT: %[[BUF:.*]] = memref.alloca() : memref<1xindex>
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store %[[I]], %[[BUF]]
+// CHECK-NEXT: scf.condition(%{{.*}}) %[[BUF]]
+// CHECK: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: memref.store %[[VALUE]], %{{.*}}[%[[I]]]
+func.func @alloca_exit_in_parallel(%condition: i1, %out: memref<2xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloca() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %value = memref.load %last[%c0] : memref<1xindex>
+ memref.store %value, %out[%i] : memref<2xindex>
+ scf.reduce
+ }
+ return
+}
+
+// -----
+
// Hoisting realloc before the store would invalidate the store's memref.
// The loop exits immediately to avoid reusing the source after reallocation.
// CHECK-LABEL: func @no_hoist_while_realloc(
More information about the Mlir-commits
mailing list