[Mlir-commits] [mlir] [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (PR #226951)
Matteo Di Iorio
llvmlistbot at llvm.org
Thu Oct 1 03:15:09 PDT 2026
https://github.com/splicer3 updated https://github.com/llvm/llvm-project/pull/226951
>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 1/3] [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 b2608fc5b574c..6100e3c71e790 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 4f2a4734056b6..fd5f6cb81d2b9 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 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..ad612c04a25e8 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 4e897a951bc6a..9cf0fd5744347 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
+ }
+}
>From d579015ada4da6fad85882e1c91ddc82b42aefb1 Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Mon, 28 Sep 2026 13:24:40 +0200
Subject: [PATCH 2/3] [mlir][bufferization] fix property syntax in
buffer-loop-hoisting test
---
.../Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index ad612c04a25e8..8b91eae09ba5f 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -853,7 +853,7 @@ func.func @no_hoist_while_dynamic_size(%n: index) -> index {
// 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} {
+ 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>
>From 492535bdb3b860f76feb8affdf210e52de00f10e Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Thu, 1 Oct 2026 12:14:13 +0200
Subject: [PATCH 3/3] [mlir][bufferization] Restrict exit-alias hoisting to
scf.while
Use scf.while semantics for the exit-alias fallback while preserving the
existing hoisting path for other loops.
Also split the shared placement fixes into a separate change.
---
.../Bufferization/Transforms/Passes.td | 12 +-
.../Transforms/BufferOptimizations.cpp | 117 +++++-------------
.../Transforms/buffer-hoisting.mlir | 17 ---
.../Transforms/buffer-loop-hoisting.mlir | 111 +++++++++++------
4 files changed, 113 insertions(+), 144 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 6100e3c71e790..947b90948fe6f 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -306,12 +306,12 @@ 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 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.
+ `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 directly in an isolated function body, in a
+ reachable block and strictly inside the transformation target. 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 fd5f6cb81d2b9..3c62cf52158c4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -18,12 +18,12 @@
#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/FunctionInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
-#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
@@ -69,57 +69,23 @@ 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.
+/// Returns true if an allocation in the before region 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(
+ memref::AllocOp alloc, scf::WhileOp loop,
+ const BufferViewFlowAnalysis::ValueSetT &aliases) {
// 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)
+ // The single-block before region executes on entry, even if the condition
+ // is initially false. The after region may be skipped.
+ if (alloc->getBlock() != loop.getBeforeBody())
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) {
@@ -127,11 +93,10 @@ canHoistFromMultiRegionLoop(Value alloc, Operation *loop,
return false;
for (OpOperand &use : alias.getUses()) {
Operation *user = use.getOwner();
- if (user->getNumRegions() != 0 ||
- (user->getParentOp() != loop && user->getBlock() != loop->getBlock()))
+ if (user->getParentOp() != loop && user->getBlock() != loop->getBlock())
return false;
if (auto load = dyn_cast<memref::LoadOp>(user)) {
- if (load.getMemref() != alias || load.getInvariant())
+ if (load.getInvariant())
return false;
continue;
}
@@ -140,17 +105,13 @@ canHoistFromMultiRegionLoop(Value alloc, Operation *loop,
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;
+ if (isa<memref::CastOp, memref::SubViewOp>(user))
continue;
- }
- if (!isa<RegionBranchTerminatorOpInterface>(user) ||
- user->getParentOp() != loop || !mapping.contains(&use))
+ // 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;
}
}
@@ -364,11 +325,6 @@ 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;
@@ -383,13 +339,13 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
currentBlock = idom->getBlock();
state.recordMoveToDominator(currentBlock);
} else {
- // 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))
+ // 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. Also, we should not move
+ // allocations out of unknown region-based control-flow operations.
+ if (!isKnownControlFlowInterface(parentOp) ||
+ !state.isLegalPlacement(parentOp, scopeOp))
break;
// Move to our parent block by notifying the current StateT
// implementation.
@@ -438,7 +394,7 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
}
/// Returns true if the given operation does not represent a loop.
- bool isLegalPlacement(Operation *op) { return !isLoop(op); }
+ bool isLegalPlacement(Operation *op, Operation *scopeOp) { return !isLoop(op); }
/// Returns true if the given operation should be considered for hoisting.
static bool shouldHoistOpType(Operation *op) {
@@ -469,26 +425,21 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
}
/// Returns true if the allocation can be moved across the given loop.
- bool isLegalPlacement(Operation *op) {
+ bool isLegalPlacement(Operation *op, Operation *scopeOp) {
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>())
+ auto loop = dyn_cast<scf::WhileOp>(op);
+ auto alloc = allocValue.getDefiningOp<memref::AllocOp>();
+ if (!loop || !alloc)
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);
+ // Moving outside the analysis scope leaves liveness information unavailable.
+ // An unreachable parent block has no dominator node for the placement walk.
+ if (!scopeOp->isProperAncestor(op) ||
+ !dominators->isReachableFromEntry(op->getBlock()))
+ return false;
+ return canHoistFromWhile(alloc, loop, aliases);
}
/// 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 827fcbb0a156b..4c07df951e280 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir
@@ -802,20 +802,3 @@ 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 8b91eae09ba5f..c4ba8f13e9d68 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -686,37 +686,15 @@ func.func @no_hoist_while_dealloc(%condition: i1) {
// -----
-// 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 {
+// 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> {
@@ -736,14 +714,6 @@ func.func @pointer_escape(%addresses: memref<2xindex>) -> index {
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
@@ -884,3 +854,68 @@ func.func @hoist_sequential_acc_loop(%n: index, %out: memref<?xindex>) {
} 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
+}
+
+// -----
+
+// Conservatively exclude unreachable parent blocks from the fallback.
+// CHECK-LABEL: func @no_hoist_while_unreachable_parent(
+// CHECK: return
+// CHECK: ^bb1:
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func @no_hoist_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
+}
More information about the Mlir-commits
mailing list