[Mlir-commits] [mlir] 860259c - [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (#226951)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 13 00:51:49 PDT 2026
Author: Matteo Di Iorio
Date: 2026-10-13T09:51:44+02:00
New Revision: 860259ca3c72d0c62f3c21248d29d81cf96ea5fc
URL: https://github.com/llvm/llvm-project/commit/860259ca3c72d0c62f3c21248d29d81cf96ea5fc
DIFF: https://github.com/llvm/llvm-project/commit/860259ca3c72d0c62f3c21248d29d81cf96ea5fc.diff
LOG: [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (#226951)
As discussed in #223944, `buffer-loop-hoisting` rejects an allocation in
`scf.while` even though the buffer is never passed back to the region
containing the allocation. The allocation and the forwarded `after`
argument belong to different regions, so their common dominator is the
block containing the loop, which in turn makes `isLegalPlacement` return
false and reject the hoist.
This PR adds a fallback for `memref.alloc` directly in the `before`
region of `scf.while` when the existing common-dominator check rejects
the hoist. It keeps the full alias set and rejects hoisting if an alias
is carried back to `before`. It also checks uses of all aliases,
including loop results, for capture, deallocation and explicit
observations of allocation identity. The fallback requires the loop's
parent to stop further loop hoisting: it must either be a non-loop
operation or have the `HasParallelRegion` trait. This restriction avoids
exposing the allocation to further hoisting through enclosing loops with
incorrectly classified execution semantics, as in #226090.
Validation:
Release build with assertions: `check-mlir` completed with 4072 passed
and one expected failure.
Fixes #223944.
AI assistance: Codex, mainly for rapid prototyping (very early
implementations of what I discussed in the issue and other potential
fixes) and codebase exploration. The end result was manually reviewed.
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
mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 0638b339cf9dc..954bc67ce7bf6 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -305,6 +305,14 @@ def BufferLoopHoistingPass : Pass<"buffer-loop-hoisting", "func::FuncOp"> {
let description = [{
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.
+ 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.
}];
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index e03887c8c5910..f7bc9a70b0864 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -18,7 +18,9 @@
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Pass/Pass.h"
@@ -66,6 +68,53 @@ static bool isSequentialLoop(Operation *op) {
return !op->hasTrait<OpTrait::HasParallelRegion>() && isLoop(op);
}
+/// Returns true if an allocation can be reused across iterations when its
+/// aliases have a common dominator outside the while loop.
+/// `aliases` must contain the complete forward alias set of the allocation.
+static bool
+canHoistFromWhile(scf::WhileOp loop,
+ const BufferViewFlowAnalysis::ValueSetT &aliases) {
+ // Require a parent that stops further loop hoisting. Otherwise, the
+ // dominator-based path could move the allocation across an enclosing loop
+ // with unmodeled parallel execution.
+ if (isSequentialLoop(loop->getParentOp()))
+ return false;
+
+ // Check all aliases, including loop results: memory effects alone do not rule
+ // out capture or uses that depend on allocation identity or lifetime.
+ for (Value alias : aliases) {
+ if (!isa<BaseMemRefType>(alias.getType()))
+ return false;
+ for (OpOperand &use : alias.getUses()) {
+ Operation *user = use.getOwner();
+ if (user->getParentOp() != loop && user->getBlock() != loop->getBlock())
+ return false;
+ if (auto load = dyn_cast<memref::LoadOp>(user)) {
+ // Reusing the allocation across iterations can invalidate invariant
+ // loads if
diff erent iterations store
diff erent values to the same
+ // location.
+ if (load.getInvariant())
+ return false;
+ continue;
+ }
+ if (auto store = dyn_cast<memref::StoreOp>(user)) {
+ if (store.getMemref() != alias)
+ return false;
+ continue;
+ }
+ // ViewLikeOpInterface describes aliases but does not guarantee that
+ // changing allocation identity preserves the view's access semantics.
+ if (isa<memref::CastOp, memref::SubViewOp>(user))
+ continue;
+ // scf.condition only forwards memrefs to after or to the loop results.
+ // Reject scf.yield to keep aliases out of the next iteration.
+ if (user != loop.getConditionOp().getOperation())
+ return false;
+ }
+ }
+ return true;
+}
+
/// Returns true if the given operation implements the AllocationOpInterface
/// and it supports the dominate block hoisting.
static bool allowAllocDominateBlockHoisting(Operation *op) {
@@ -222,7 +271,8 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
Block *dominatorBlock =
findCommonDominator(allocValue, resultAliases, dominators);
// Init the initial hoisting state.
- StateT state(&dominators, allocValue, allocValue.getParentBlock());
+ StateT state(&dominators, allocValue, allocValue.getParentBlock(),
+ resultAliases);
// Check for additional allocation dependencies to compute an upper bound
// for hoisting.
Block *dependencyBlock = nullptr;
@@ -324,7 +374,11 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
/// that hoists allocations into dominator blocks while keeping them inside of
/// loops.
struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
- using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
+ BufferAllocationHoistingState(DominanceInfo *dominators, Value allocValue,
+ Block *placementBlock,
+ const BufferViewFlowAnalysis::ValueSetT &)
+ : BufferAllocationHoistingStateBase(dominators, allocValue,
+ placementBlock) {}
/// Computes the upper bound for the placement block search.
Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
@@ -358,11 +412,19 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
/// A state implementation compatible with the `BufferAllocationHoisting` class
/// that hoists allocations out of loops.
struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
- using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
+ /// The forward alias closure of the current allocation.
+ const BufferViewFlowAnalysis::ValueSetT &aliases;
/// Remembers the dominator block of all aliases.
Block *aliasDominatorBlock = nullptr;
+ BufferAllocationLoopHoistingState(
+ DominanceInfo *dominators, Value allocValue, Block *placementBlock,
+ const BufferViewFlowAnalysis::ValueSetT &aliases)
+ : BufferAllocationHoistingStateBase(dominators, allocValue,
+ placementBlock),
+ aliases(aliases) {}
+
/// Computes the upper bound for the placement block search.
Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
aliasDominatorBlock = dominatorBlock;
@@ -371,14 +433,19 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
return dependencyBlock ? dependencyBlock : nullptr;
}
- /// Returns true if the given operation represents a loop with sequential
- /// execution semantics and one of the aliases caused the
- /// `aliasDominatorBlock` to be "above" the block of the given loop operation.
- /// If this is the case, it indicates that the allocation is passed via a back
- /// edge.
+ /// Returns true if the allocation can be moved across the given loop.
bool isLegalPlacement(Operation *op) {
- return isSequentialLoop(op) &&
- !dominators->dominates(aliasDominatorBlock, op->getBlock());
+ if (!isSequentialLoop(op))
+ return false;
+ 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);
}
/// 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 7498f7af8c5d1..93083b2ca3cdd 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -567,3 +567,610 @@ func.func @loop_unreachable_parent() {
}
return
}
+
+// -----
+
+// An exit alias does not carry the allocation back to before.
+// CHECK-LABEL: func @while_exit_alias(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<1xindex>):
+// CHECK-NEXT: %{{.*}} = memref.load %[[CURRENT]]
+// CHECK: memref.load %[[LAST]]
+func.func @while_exit_alias(%n: index) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last = scf.while (%i = %c0) : (index) -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ %value = memref.load %current[%c0] : memref<1xindex>
+ %next = arith.addi %value, %c1 : index
+ scf.yield %next : index
+ }
+ %result = memref.load %last[%c0] : memref<1xindex>
+ return %result : index
+}
+
+// -----
+
+// A loop-invariant allocation extent permits hoisting with an exit alias.
+// CHECK-LABEL: func @while_exit_alias_invariant_extent(
+// CHECK-SAME: %{{.*}}: i1, %[[SIZE:.*]]: index)
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// CHECK-NEXT: %{{.*}} = scf.while
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func @while_exit_alias_invariant_extent(%condition: i1, %size: index) {
+ %last = scf.while () : () -> memref<?xindex> {
+ %buffer = memref.alloc(%size) : memref<?xindex>
+ scf.condition(%condition) %buffer : memref<?xindex>
+ } do {
+ ^bb0(%current: memref<?xindex>):
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// An unrelated buffer may be carried to before while the allocation is hoisted.
+// CHECK-LABEL: func @while_exit_alias_unrelated_carried(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %{{.*}}:2 = scf.while
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]], %{{.*}}
+// CHECK: ^bb0(%{{.*}}: memref<1xindex>, %[[KEEP:.*]]: memref<1xindex>):
+// CHECK-NEXT: scf.yield %[[KEEP]]
+func.func @while_exit_alias_unrelated_carried(%condition: i1, %other: memref<1xindex>) {
+ %last:2 = scf.while (%keep = %other)
+ : (memref<1xindex>) -> (memref<1xindex>, memref<1xindex>) {
+ %buffer = memref.alloc() : memref<1xindex>
+ scf.condition(%condition) %buffer, %keep : memref<1xindex>, memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>, %keep: memref<1xindex>):
+ scf.yield %keep : memref<1xindex>
+ }
+ return
+}
+
+// -----
+
+// Casts and subviews of the loop result are included in the alias check.
+// CHECK-LABEL: func @while_exit_alias_view(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: %[[VIEW:.*]] = memref.subview %[[LAST]]
+// CHECK-NEXT: %[[CAST:.*]] = memref.cast %[[VIEW]]
+// CHECK-NEXT: %{{.*}} = memref.load %[[CAST]]
+func.func @while_exit_alias_view(%condition: i1) -> index {
+ %c0 = arith.constant 0 : index
+ %last = scf.while () : () -> memref<2xindex> {
+ %buffer = memref.alloc() : memref<2xindex>
+ memref.store %c0, %buffer[%c0] : memref<2xindex>
+ scf.condition(%condition) %buffer : memref<2xindex>
+ } do {
+ ^bb0(%current: memref<2xindex>):
+ scf.yield
+ }
+ %view = memref.subview %last[0] [1] [1] : memref<2xindex> to memref<1xindex>
+ %cast = memref.cast %view : memref<1xindex> to memref<?xindex>
+ %result = memref.load %cast[%c0] : memref<?xindex>
+ return %result : index
+}
+
+// -----
+
+// A view passed back to the allocation region prevents hoisting.
+// CHECK-LABEL: func @no_hoist_while_carried_view(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: scf.condition{{.*}} %{{.*}}, %[[ALLOC]]
+// CHECK: ^bb0(%{{.*}}: index, %[[CURRENT:.*]]: memref<1xindex>):
+// CHECK: %[[CAST:.*]] = memref.cast %[[CURRENT]]
+// CHECK: %[[VIEW:.*]] = memref.subview %[[CAST]]
+// CHECK: scf.yield %{{.*}}, %[[VIEW]]
+func.func @no_hoist_while_carried_view(%n: index, %seed: memref<1xindex>) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %result:2 = scf.while (%i = %c0, %old = %seed)
+ : (index, memref<1xindex>) -> (index, memref<1xindex>) {
+ %fresh = memref.alloc() : memref<1xindex>
+ memref.store %i, %fresh[%c0] : memref<1xindex>
+ %previous = memref.load %old[%c0] : memref<1xindex>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %previous, %fresh : index, memref<1xindex>
+ } do {
+ ^bb0(%previous: index, %current: memref<1xindex>):
+ %value = memref.load %current[%c0] : memref<1xindex>
+ %next = arith.addi %value, %c1 : index
+ %cast = memref.cast %current : memref<1xindex> to memref<?xindex>
+ %view = memref.subview %cast[0] [1] [1] : memref<?xindex> to memref<1xindex>
+ scf.yield %next, %view : index, memref<1xindex>
+ }
+ return %result#0 : index
+}
+
+// -----
+
+// Uses of the loop result must also be checked for capture.
+// CHECK-LABEL: func @no_hoist_while_exit_capture(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: call @capture(%[[LAST]])
+func.func private @capture(memref<1xindex>)
+func.func @no_hoist_while_exit_capture(%condition: i1) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ func.call @capture(%last) : (memref<1xindex>) -> ()
+ return
+}
+
+// -----
+
+// A loop result used in a nested region is outside the supported alias uses.
+// CHECK-LABEL: func @no_hoist_while_exit_nested_region(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: scf.if
+// CHECK-NEXT: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: scf.yield %[[VALUE]] : index
+func.func @no_hoist_while_exit_nested_region(%condition: i1, %read: i1) -> index {
+ %c0 = arith.constant 0 : index
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %result = scf.if %read -> (index) {
+ %value = memref.load %last[%c0] : memref<1xindex>
+ scf.yield %value : index
+ } else {
+ scf.yield %c0 : index
+ }
+ return %result : index
+}
+
+// -----
+
+// Hoisting would require the loaded value to be invariant across iterations.
+// CHECK-LABEL: func @no_hoist_while_invariant_load(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<1xindex>):
+// CHECK: memref.load %[[CURRENT]]{{.*}} invariant(true)
+func.func @no_hoist_while_invariant_load(%n: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last = scf.while (%i = %c0) : (index) -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ %value = memref.load %current[%c0] invariant(true) : memref<1xindex>
+ %next = arith.addi %value, %c1 : index
+ scf.yield %next : index
+ }
+ return
+}
+
+// -----
+
+// An allocation freed in the loop cannot be reused by later iterations.
+// CHECK-LABEL: func @no_hoist_while_dealloc(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<1xindex>):
+// CHECK-NEXT: memref.dealloc %[[CURRENT]]
+func.func @no_hoist_while_dealloc(%condition: i1) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ memref.dealloc %current : memref<1xindex>
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// A memory-effect-free pointer observation can expose allocation identity.
+// CHECK-LABEL: func @no_hoist_while_pointer_escape(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[BUF:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[BUF]]
+// CHECK-NEXT: %[[ADDRESS:.*]] = memref.extract_aligned_pointer_as_index %[[BUF]]
+// CHECK-NEXT: memref.store %[[ADDRESS]]
+func.func @no_hoist_while_pointer_escape(%addresses: memref<2xindex>) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last = scf.while (%i = %c0) : (index) -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ %address = memref.extract_aligned_pointer_as_index %buffer : memref<1xindex> -> index
+ memref.store %address, %addresses[%i] : memref<2xindex>
+ %continue = arith.cmpi slt, %i, %c1 : index
+ scf.condition(%continue) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ %value = memref.load %current[%c0] : memref<1xindex>
+ %next = arith.addi %value, %c1 : index
+ scf.yield %next : index
+ }
+ %result = memref.load %last[%c0] : memref<1xindex>
+ return %result : index
+}
+
+// -----
+
+// Hoist before the while without speculating the allocation out of the branch.
+// CHECK-LABEL: func @while_exit_alias_if_parent(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.if
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @while_exit_alias_if_parent(%run: i1, %condition: i1, %out: memref<index>) {
+ %c0 = arith.constant 0 : index
+ scf.if %run {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %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[] : memref<index>
+ }
+ return
+}
+
+// -----
+
+// The non-loop parent stops further hoisting after crossing the while.
+// CHECK-LABEL: func @while_exit_alias_execute_region_parent(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.execute_region
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @while_exit_alias_execute_region_parent(%condition: i1, %out: memref<index>) {
+ %c0 = arith.constant 0 : index
+ scf.execute_region {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %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[] : memref<index>
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// Hoisting from the while must preserve a private allocation per parallel
+// iteration.
+// CHECK-LABEL: func @while_exit_alias_parallel_parent(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.parallel
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @while_exit_alias_parallel_parent(%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.alloc() : 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 into acc.loop would allow a subsequent hoist to share the allocation
+// between parallel iterations. Keep it inside scf.while.
+// CHECK-LABEL: func @no_hoist_while_parallel_acc_parent(
+// CHECK-NOT: memref.alloc
+// CHECK: acc.parallel
+// CHECK-NEXT: acc.loop
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @no_hoist_while_parallel_acc_parent(%out: memref<2xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %false = arith.constant false
+ acc.parallel {
+ acc.loop gang vector control(%i : index) = (%c0 : index)
+ to (%c2 : index) step (%c1 : index) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ scf.condition(%false) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %value = memref.load %last[%c0] : memref<1xindex>
+ memref.store %value, %out[%i] : memref<2xindex>
+ acc.yield
+ } independent
+ acc.yield
+ }
+ return
+}
+
+// -----
+
+// The fallback requires a parent that stops further loop hoisting.
+// An enclosing sequential loop does not provide that boundary.
+// CHECK-LABEL: func @no_hoist_while_nested_for(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.for
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @no_hoist_while_nested_for(%n: index, %condition: i1, %out: memref<?xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : 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<?xindex>
+ }
+ return
+}
+
+// -----
+
+// A loop-variant allocation extent prevents hoisting even if the result only
+// leaves through an exit edge.
+// CHECK-LABEL: func @no_hoist_while_dynamic_size(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[SIZE:.*]] = arith.addi
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @no_hoist_while_dynamic_size(%n: index) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last = scf.while (%i = %c0) : (index) -> memref<?xindex> {
+ %size = arith.addi %i, %c1 : index
+ %buffer = memref.alloc(%size) : memref<?xindex>
+ memref.store %i, %buffer[%c0] : memref<?xindex>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %buffer : memref<?xindex>
+ } do {
+ ^bb0(%current: memref<?xindex>):
+ %value = memref.load %current[%c0] : memref<?xindex>
+ %next = arith.addi %value, %c1 : index
+ scf.yield %next : index
+ }
+ %value = memref.load %last[%c0] : memref<?xindex>
+ return %value : index
+}
+
+// -----
+
+// A write to another buffer does not prevent the dominator-based hoist.
+// CHECK-LABEL: func @hoist_with_unrelated_loop_effect(
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: test.store_with_a_loop_region
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT: %[[VALUE:.*]] = memref.load %[[ALLOC]]
+// CHECK-NEXT: memref.store %[[VALUE]]
+func.func @hoist_with_unrelated_loop_effect(%out: memref<f32>, %value: f32) {
+ test.store_with_a_loop_region %out <store_before_region = true> {
+ %buffer = memref.alloc() : memref<f32>
+ memref.store %value, %buffer[] : memref<f32>
+ %loaded = memref.load %buffer[] : memref<f32>
+ memref.store %loaded, %out[] : memref<f32>
+ test.store_with_a_region_terminator
+ } : memref<f32>
+ return
+}
+
+// -----
+
+// Preserve dominator-based hoisting from sequential OpenACC loops.
+// CHECK-LABEL: func @hoist_sequential_acc_loop(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: acc.loop
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT: %[[VALUE:.*]] = memref.load %[[ALLOC]]
+// CHECK-NEXT: memref.store %[[VALUE]]
+func.func @hoist_sequential_acc_loop(%n: index, %out: memref<?xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ acc.loop control(%i : index) = (%c0 : index) to (%n : index) step (%c1 : index) {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ %value = memref.load %buffer[%c0] : memref<1xindex>
+ memref.store %value, %out[%i] : memref<?xindex>
+ acc.yield
+ } seq
+ return
+}
+
+// -----
+
+// A call may capture an allocation even when no alias is carried to before.
+// CHECK-LABEL: func @no_hoist_while_capture(
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: call @capture(%[[ALLOC]])
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func private @capture(memref<1xindex>)
+func.func @no_hoist_while_capture(%condition: i1) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ func.call @capture(%buffer) : (memref<1xindex>) -> ()
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// 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()
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func @no_hoist_while_alloca(%condition: i1) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloca() : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ 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(
+// CHECK-SAME: %[[SOURCE:.*]]: memref<1xi32>
+// CHECK-NOT: memref.realloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store %{{.*}}, %[[SOURCE]][%{{.*}}]
+// CHECK-NEXT: %[[NEW:.*]] = memref.realloc %[[SOURCE]]
+// CHECK-NEXT: scf.condition{{.*}} %[[NEW]]
+// CHECK: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: return %[[VALUE]]
+func.func @no_hoist_while_realloc(%source: memref<1xi32>) -> i32 {
+ %c0 = arith.constant 0 : index
+ %c7 = arith.constant 7 : i32
+ %false = arith.constant false
+ %last = scf.while () : () -> memref<2xi32> {
+ memref.store %c7, %source[%c0] : memref<1xi32>
+ %new = memref.realloc %source : memref<1xi32> to memref<2xi32>
+ scf.condition(%false) %new : memref<2xi32>
+ } do {
+ ^bb0(%current: memref<2xi32>):
+ scf.yield
+ }
+ %value = memref.load %last[%c0] : memref<2xi32>
+ return %value : i32
+}
+
+// -----
+
+// Hoist into the unreachable parent block, then stop the placement walk there.
+// CHECK-LABEL: func @while_unreachable_parent(
+// CHECK: return
+// CHECK: ^bb1:
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.while
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func @while_unreachable_parent() {
+ return
+^dead:
+ %false = arith.constant false
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ scf.condition(%false) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// Storing an alias as a value can expose the buffer to later iterations.
+// CHECK-LABEL: func @no_hoist_while_store_alias(
+// CHECK-SAME: %{{.*}}: i1, %[[HOLDER:.*]]: memref<memref<1xindex>>
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<1xindex>):
+// CHECK-NEXT: memref.store %[[CURRENT]], %[[HOLDER]][]
+func.func @no_hoist_while_store_alias(%condition: i1, %holder: memref<memref<1xindex>>) {
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ memref.store %current, %holder[] : memref<memref<1xindex>>
+ scf.yield
+ }
+ return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 27ff9964c2ca2..a1a8f1ca139a1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -289,3 +289,97 @@ module attributes {transform.with_named_sequence} {
return
}
}
+
+// -----
+
+// Hoisting must stay within the operation targeted by the transform.
+// CHECK-LABEL: func @buffer_loop_hoisting_while_scope(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+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(%condition: i1) -> index {
+ %c0 = arith.constant 0 : index
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %result = memref.load %last[%c0] : memref<1xindex>
+ return %result : index
+ }
+}
+
+// -----
+
+// Targeting the enclosing function permits hoisting before the loop.
+// CHECK-LABEL: func @buffer_loop_hoisting_function_scope(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+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(%condition: i1) -> index {
+ %c0 = arith.constant 0 : index
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c0, %buffer[%c0] : memref<1xindex>
+ scf.condition(%condition) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %result = memref.load %last[%c0] : memref<1xindex>
+ return %result : index
+ }
+}
+
+// -----
+
+// Targeting the module still keeps the allocation inside the non-isolated
+// function, even when the allocation size is captured from outside it.
+// CHECK-LABEL: module attributes
+// CHECK: %[[SIZE:.*]] = arith.constant 4 : index
+// CHECK-NOT: memref.alloc
+// CHECK: test.conversion_func_op @buffer_loop_hoisting_nonisolated_function(
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ transform.bufferization.buffer_loop_hoisting %root : !transform.any_op
+ transform.yield
+ }
+ %size = arith.constant 4 : index
+ test.conversion_func_op @buffer_loop_hoisting_nonisolated_function(%condition: i1) {
+ %c0 = arith.constant 0 : index
+ %last = scf.while () : () -> memref<?xindex> {
+ %buffer = memref.alloc(%size) : memref<?xindex>
+ memref.store %c0, %buffer[%c0] : memref<?xindex>
+ scf.condition(%condition) %buffer : memref<?xindex>
+ } do {
+ ^bb0(%current: memref<?xindex>):
+ scf.yield
+ }
+ %result = memref.load %last[%c0] : memref<?xindex>
+ "test.return"() : () -> ()
+ }
+}
More information about the Mlir-commits
mailing list