[Mlir-commits] [mlir] [mlir][bufferization] Hoist allocations through exactly-once regions (PR #224921)

[purnima shrivastava]

https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/224921

>From f240813ff587fe9cff9b20c01f6b3b7ad17ee98c Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Sun, 20 Sep 2026 21:17:47 +0530
Subject: [PATCH] [mlir][bufferization] Hoist allocations through exactly-once
 regions

---
 .../Transforms/BufferOptimizations.cpp        | 39 ++++++---
 .../Transforms/buffer-loop-hoisting.mlir      | 82 +++++++++++++++++++
 2 files changed, 109 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 4f2a4734056b6..a3f1f56923d6f 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,22 @@ 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 = cast<RegionBranchOpInterface>(op);
+
+      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 +398,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..c00909c3b9455 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -519,3 +519,85 @@ 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 {{.*}}, %[[ALLOC]]
+
+// -----
+
+// CHECK-LABEL: func @negative_loop_if_alloc_not_hoisted
+func.func @negative_loop_if_alloc_not_hoisted(%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]]
+
+// -----
+
+// CHECK-LABEL: func @loop_nested_execute_region_hoisted
+func.func @loop_nested_execute_region_hoisted() {
+  %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