[Mlir-commits] [mlir] [mlir][bufferization] Hoist allocations through exactly-once regions (PR #224921)
purnima shrivastava
llvmlistbot at llvm.org
Sun Sep 20 08:18:55 PDT 2026
https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/224921
>From bf2fd8bc329c49ab85baade5e9a3224a39a90882 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sun, 20 Sep 2026 17:33:06 +0530
Subject: [PATCH 1/2] [mlir][bufferization] Hoist allocations through
exactly-once regions
---
.../Transforms/BufferOptimizations.cpp | 41 +++++++++++++------
.../Transforms/buffer-loop-hoisting.mlir | 27 ++++++++++++
2 files changed, 56 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 4f2a4734056b6..97fbbb77a740e 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -288,12 +288,12 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
// implementation, break the walk loop. Furthermore, we should not move
// allocations out of unknown region-based control-flow operations.
if (!isKnownControlFlowInterface(parentOp) ||
- !state.isLegalPlacement(parentOp))
+ !state.isLegalPlacement(parentOp, currentBlock->getParent()))
break;
// Move to our parent block by notifying the current StateT
// implementation.
currentBlock = parentBlock;
- state.recordMoveToParent(currentBlock);
+ state.recordMoveToParent(parentOp, currentBlock);
}
}
// Return the finally determined placement block.
@@ -337,7 +337,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, Region *) { return !isLoop(op); }
/// Returns true if the given operation should be considered for hoisting.
static bool shouldHoistOpType(Operation *op) {
@@ -348,7 +348,9 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
void recordMoveToDominator(Block *block) { placementBlock = block; }
/// Sets the current placement block to the given block.
- void recordMoveToParent(Block *block) { recordMoveToDominator(block); }
+ void recordMoveToParent(Operation *, Block *block) {
+ recordMoveToDominator(block);
+ }
};
/// A state implementation compatible with the `BufferAllocationHoisting` class
@@ -367,12 +369,24 @@ 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.
- bool isLegalPlacement(Operation *op) {
+ /// Returns true if the placement search can cross the given parent operation.
+ /// Non-loop regions may be crossed if the current region is executed exactly
+ /// once. Loops must have sequential execution semantics and must not violate
+ /// the alias back-edge constraint.
+ bool isLegalPlacement(Operation *op, Region *region) {
+ if (!isLoop(op)) {
+ auto branchOp = dyn_cast<RegionBranchOpInterface>(op);
+ if (!branchOp)
+ return false;
+
+ SmallVector<Attribute> operands(op->getNumOperands(), nullptr);
+ SmallVector<InvocationBounds> bounds;
+ branchOp.getRegionInvocationBounds(operands, bounds);
+
+ const InvocationBounds &bound = bounds[region->getRegionNumber()];
+ return bound.getLowerBound() == 1 && bound.getUpperBound() == 1;
+ }
+
return isSequentialLoop(op) &&
!dominators->dominates(aliasDominatorBlock, op->getBlock());
}
@@ -386,8 +400,11 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
/// operations out of loops only.
void recordMoveToDominator(Block *block) {}
- /// Sets the current placement block to the given block.
- void recordMoveToParent(Block *block) { placementBlock = block; }
+ /// Updates the placement block when crossing a sequential loop.
+ void recordMoveToParent(Operation *parentOp, Block *block) {
+ if (isSequentialLoop(parentOp))
+ placementBlock = block;
+ }
};
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index ee7571b2c8466..b3131966570c7 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -519,3 +519,30 @@ func.func @hoist_alloca(
// CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
// CHECK-NEXT: {{.*}} = scf.for
+
+// -----
+
+// CHECK-LABEL: func @loop_execute_region
+func.func @loop_execute_region(%output: memref<8xindex>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+
+ scf.for %i = %c0 to %c8 step %c1 {
+ %value = scf.execute_region -> index {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ %loaded = memref.load %buffer[%c0] : memref<1xindex>
+ scf.yield %loaded : index
+ }
+ memref.store %value, %output[%i] : memref<8xindex>
+ }
+
+ return
+}
+
+// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1xindex>
+// CHECK-NEXT: scf.for
+// CHECK: scf.execute_region
+// CHECK-NOT: memref.alloc()
+// CHECK: memref.store
>From 5b41a81b01d685af2826671acf0a902b29750e97 Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sun, 20 Sep 2026 20:46:03 +0530
Subject: [PATCH 2/2] [mlir][bufferization] Address loop hoisting review
comments
---
.../Transforms/BufferOptimizations.cpp | 4 +-
.../Transforms/buffer-loop-hoisting.mlir | 61 ++++++++++++++++++-
2 files changed, 61 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 97fbbb77a740e..a3f1f56923d6f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -375,9 +375,7 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
/// the alias back-edge constraint.
bool isLegalPlacement(Operation *op, Region *region) {
if (!isLoop(op)) {
- auto branchOp = dyn_cast<RegionBranchOpInterface>(op);
- if (!branchOp)
- return false;
+ auto branchOp = cast<RegionBranchOpInterface>(op);
SmallVector<Attribute> operands(op->getNumOperands(), nullptr);
SmallVector<InvocationBounds> bounds;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index b3131966570c7..90c28fcd13998 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -545,4 +545,63 @@ func.func @loop_execute_region(%output: memref<8xindex>) {
// CHECK-NEXT: scf.for
// CHECK: scf.execute_region
// CHECK-NOT: memref.alloc()
-// CHECK: memref.store
+// CHECK: memref.store {{.*}}, %[[ALLOC]]
+
+// -----
+
+// Conditional regions must not be crossed when hoisting allocations.
+
+// CHECK-LABEL: func @loop_if_alloc
+func.func @loop_if_alloc(%condition: i1) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.if %condition {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ }
+ }
+
+ return
+}
+
+// CHECK-NOT: memref.alloc()
+// CHECK: scf.for
+// CHECK-NEXT: scf.if
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc() : memref<1xindex>
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+
+// -----
+
+// Hoist through multiple regions that each execute exactly once.
+
+// CHECK-LABEL: func @loop_nested_execute_region
+func.func @loop_nested_execute_region() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+
+ scf.for %i = %c0 to %c8 step %c1 {
+ scf.execute_region {
+ scf.execute_region {
+ %buffer = memref.alloc() : memref<1xindex>
+ memref.store %i, %buffer[%c0] : memref<1xindex>
+ scf.yield
+ }
+ scf.yield
+ }
+ }
+
+ return
+}
+
+// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1xindex>
+// CHECK-NEXT: scf.for
+// CHECK-NEXT: scf.execute_region
+// CHECK-NEXT: scf.execute_region
+// CHECK-NOT: memref.alloc()
+// CHECK: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NOT: memref.alloc()
+// CHECK: return
More information about the Mlir-commits
mailing list