[Mlir-commits] [mlir] [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (PR #226951)
Matteo Di Iorio
llvmlistbot at llvm.org
Mon Sep 28 03:39:47 PDT 2026
https://github.com/splicer3 created https://github.com/llvm/llvm-project/pull/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 draft adds a fallback for when that existing check rejects the hoist. It keeps the full alias set and rejects hoisting if an alias reaches an argument of the allocation's region. It also checks uses of all aliases, including loop results, for capture, deallocation and explicit observations of allocation identity. I'm making this a draft because it does have some assumptions listed below.
For now, the fallback only handles `memref.alloc` in multi-region loops directly inside an isolated function body. Each region must have one SSA block, successor inputs must be fully described by operand forwarding, and the allocation's region must execute on entry. The function-parent restriction avoids exposing the allocation to further hoisting through enclosing loops with incorrectly classified execution semantics, as in #226090.
The shared placement walk also stops at the transform target, isolation boundaries and unreachable enclosing blocks. I've added regression tests to cover those boundaries as well.
Now, there are some caveats:
1. The terminator check assumes that successor operands are used only for forwarding. This is consistent with the handling in `remove-dead-values` (#173505), which can replace those operands with poison when all their successor inputs are dead.
2. Absence of `HasParallelRegion` does not establish sequential execution, and the trait currently requires the operation to have one region. The function-parent restriction avoids the enclosing-loop problem, but doesn't establish how the candidate loop itself executes.
3. `scf.while` provides the required execution and forwarding semantics and I can reduce the scope to just that. However, there is also the possibility of extending interface guarantees, but I guess that would need some more discussion and a RFC?
Validation:
Release build with assertions, 111 selected lit tests passed with 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.
>From c93687fc5ba1e1419118a4488b2e16a3413807ca Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Mon, 28 Sep 2026 11:29:15 +0200
Subject: [PATCH] [mlir][bufferization] handle exit aliases in
buffer-loop-hoisting
---
.../Bufferization/Transforms/Passes.td | 7 +
.../Transforms/BufferOptimizations.cpp | 148 +++++++-
.../Transforms/buffer-hoisting.mlir | 17 +
.../Transforms/buffer-loop-hoisting.mlir | 341 ++++++++++++++++++
.../Transforms/transform-ops.mlir | 61 ++++
5 files changed, 557 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index b2608fc5b574c6..6100e3c71e790f 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -305,6 +305,13 @@ 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 may be hoisted from multi-region loops even when
+ their buffers are forwarded through loop results. This extension is limited
+ to loops directly in an isolated function body, with single-block SSA regions
+ and complete operand forwarding through `RegionBranchOpInterface`. The
+ allocation's region must execute on entry, its aliases must not reach that
+ region through a backedge, and their uses must permit reusing the allocation.
}];
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 4f2a4734056b63..fd5f6cb81d2b9f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -19,8 +19,11 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
@@ -66,6 +69,94 @@ static bool isSequentialLoop(Operation *op) {
return !op->hasTrait<OpTrait::HasParallelRegion>() && isLoop(op);
}
+/// Checks additional hoisting restrictions when the common-dominator check
+/// cannot rule out aliases carried across loop iterations.
+///
+/// Assumes sequential region execution and that region terminators only forward
+/// buffer operands.
+static bool
+canHoistFromMultiRegionLoop(Value alloc, Operation *loop,
+ const BufferViewFlowAnalysis::ValueSetT &aliases,
+ DominanceInfo &dominance) {
+ // Restrict this check to loops directly in an isolated function body.
+ // Further hoisting across an enclosing loop with unmodeled parallel
+ // execution could share the allocation across iterations.
+ if (!isa<FunctionOpInterface>(loop->getParentOp()) ||
+ !loop->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
+ return false;
+
+ auto branch = dyn_cast<RegionBranchOpInterface>(loop);
+ Region *allocRegion = alloc.getParentRegion();
+ if (!branch || !alloc.getDefiningOp<memref::AllocOp>() ||
+ allocRegion->getParentOp() != loop)
+ return false;
+
+ for (Region ®ion : loop->getRegions())
+ if (!region.hasOneBlock() || region.front().empty() ||
+ !dominance.hasSSADominance(®ion) ||
+ !isa<RegionBranchTerminatorOpInterface>(region.front().getTerminator()))
+ return false;
+ // Hoisting must not introduce an allocation on a path that skips its region.
+ SmallVector<RegionSuccessor> entries;
+ branch.getSuccessorRegions(RegionBranchPoint::parent(), entries);
+ if (entries.size() != 1 || !entries.front().isRegion() ||
+ entries.front().getSuccessor() != allocRegion)
+ return false;
+
+ // An alias at the allocation region's entry may refer to an allocation from
+ // an earlier iteration. Loop results alone do not imply such a backedge.
+ if (llvm::any_of(allocRegion->getArguments(),
+ [&](Value arg) { return aliases.contains(arg); }))
+ return false;
+
+ // Alias analysis cannot trace inputs that are not forwarded from operands.
+ for (RegionBranchPoint point : branch.getAllRegionBranchPoints()) {
+ SmallVector<RegionSuccessor> successors;
+ branch.getSuccessorRegions(point, successors);
+ for (RegionSuccessor successor : successors)
+ if (!branch.getNonSuccessorInputs(successor).empty())
+ return false;
+ }
+ RegionBranchSuccessorMapping mapping;
+ branch.getSuccessorOperandInputMapping(mapping);
+
+ // 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->getNumRegions() != 0 ||
+ (user->getParentOp() != loop && user->getBlock() != loop->getBlock()))
+ return false;
+ if (auto load = dyn_cast<memref::LoadOp>(user)) {
+ if (load.getMemref() != alias || load.getInvariant())
+ return false;
+ continue;
+ }
+ if (auto store = dyn_cast<memref::StoreOp>(user)) {
+ if (store.getMemref() != alias)
+ return false;
+ continue;
+ }
+ if (!isMemoryEffectFree(user))
+ return false;
+ // ViewLikeOpInterface describes aliases but does not guarantee that
+ // changing allocation identity preserves the view's access semantics.
+ if (isa<memref::CastOp, memref::SubViewOp>(user)) {
+ if (!aliases.contains(user->getResult(0)))
+ return false;
+ continue;
+ }
+ if (!isa<RegionBranchTerminatorOpInterface>(user) ||
+ user->getParentOp() != loop || !mapping.contains(&use))
+ 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) {
@@ -182,11 +273,15 @@ struct BufferAllocationHoistingStateBase {
/// The current placement block (if any).
Block *placementBlock;
+ /// The forward alias closure of the current allocation.
+ const BufferViewFlowAnalysis::ValueSetT &aliases;
+
/// Initializes the state base.
- BufferAllocationHoistingStateBase(DominanceInfo *dominators, Value allocValue,
- Block *placementBlock)
+ BufferAllocationHoistingStateBase(
+ DominanceInfo *dominators, Value allocValue, Block *placementBlock,
+ const BufferViewFlowAnalysis::ValueSetT &aliases)
: dominators(dominators), allocValue(allocValue),
- placementBlock(placementBlock) {}
+ placementBlock(placementBlock), aliases(aliases) {}
};
/// Implements the actual hoisting logic for allocation nodes.
@@ -222,7 +317,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;
@@ -268,6 +364,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 +383,12 @@ 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) ||
+ // No suitable immediate dominator is available. Only cross known
+ // control-flow operations within the analysis scope, and preserve
+ // isolation and the placement constraints imposed by StateT.
+ if (parentOp == scopeOp ||
+ parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>() ||
+ !isKnownControlFlowInterface(parentOp) ||
!state.isLegalPlacement(parentOp))
break;
// Move to our parent block by notifying the current StateT
@@ -367,14 +468,27 @@ 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;
+ // Only nested operations may affect addressable memory. Effects on
+ // non-addressable resources cannot access the allocated buffer.
+ if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+ return false;
+ if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ effectInterface.getEffects(effects);
+ if (llvm::any_of(effects,
+ [](const MemoryEffects::EffectInstance &effect) {
+ return effect.getResource()->isAddressable();
+ }))
+ return false;
+ }
+ return op->getNumRegions() > 1 &&
+ canHoistFromMultiRegionLoop(allocValue, op, aliases, *dominators);
}
/// Returns true if the given operation should be considered for hoisting.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
index 4c07df951e2806..827fcbb0a156b8 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 ee04564f22f6ba..ad612c04a25e83 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -543,3 +543,344 @@ func.func @hoist_alloca(
// CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: {{.*}} = scf.for
+
+// -----
+
+// An exit alias, invariant extent and unrelated carried buffer permit hoisting.
+// CHECK-LABEL: func @while_exit_alias(
+// CHECK-SAME: %{{.*}}: index, %[[SIZE:.*]]: index,
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// CHECK-NEXT: %[[LAST:.*]]:2 = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]], %{{.*}}
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<?xindex>, %[[KEEP:.*]]: memref<1xindex>):
+// CHECK-NEXT: %{{.*}} = memref.load %[[CURRENT]]
+// CHECK: scf.yield %{{.*}}, %[[KEEP]]
+// CHECK: %[[VIEW:.*]] = memref.subview %[[LAST]]#0
+// CHECK: %[[CAST:.*]] = memref.cast %[[VIEW]]
+// CHECK: memref.load %[[CAST]]
+func.func @while_exit_alias(%n: index, %size: index, %other: memref<1xindex>) -> index {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %last:2 = scf.while (%i = %c0, %keep = %other)
+ : (index, memref<1xindex>) -> (memref<?xindex>, memref<1xindex>) {
+ %buffer = memref.alloc(%size) : memref<?xindex>
+ memref.store %i, %buffer[%c0] : memref<?xindex>
+ %continue = arith.cmpi slt, %i, %n : index
+ scf.condition(%continue) %buffer, %keep : memref<?xindex>, memref<1xindex>
+ } do {
+ ^bb0(%current: memref<?xindex>, %keep: memref<1xindex>):
+ %value = memref.load %current[%c0] : memref<?xindex>
+ %next = arith.addi %value, %c1 : index
+ scf.yield %next, %keep : index, memref<1xindex>
+ }
+ %view = memref.subview %last#0[0] [1] [1] : memref<?xindex> 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
+}
+
+// -----
+
+// 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
+}
+
+// -----
+
+// Stop at an unreachable enclosing block without querying its dominator node.
+// CHECK-LABEL: func @while_unreachable_parent(
+// CHECK: return
+// CHECK: ^bb1:
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: memref.load %[[LAST]]
+func.func @while_unreachable_parent() -> index {
+ %zero = arith.constant 0 : index
+ return %zero : index
+^dead:
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %false = arith.constant false
+ %last = scf.while () : () -> memref<1xindex> {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %c1, %buffer[%c0] : memref<1xindex>
+ scf.condition(%false) %buffer : memref<1xindex>
+ } do {
+ ^bb0(%current: memref<1xindex>):
+ scf.yield
+ }
+ %result = memref.load %last[%c0] : memref<1xindex>
+ return %result : index
+}
+
+// -----
+
+// A memory-effect-free pointer observation can expose allocation identity.
+func.func @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
+}
+
+// CHECK-LABEL: func.func @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]]
+
+// -----
+
+// 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
+}
+
+// -----
+
+// Hoisting with exit aliases requires a loop directly in a function body,
+// even when the enclosing loop is sequential.
+// 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 attributes {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
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 4e897a951bc6a9..9cf0fd57443476 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -236,3 +236,64 @@ func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
%0 = bufferization.alloc_tensor() : tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
+
+// -----
+
+// 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.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
+ }
+}
More information about the Mlir-commits
mailing list