[Mlir-commits] [mlir] [mlir][bufferization] handle exit aliases in buffer-loop-hoisting (PR #226951)

[Matteo Di Iorio]

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/2] [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 &region : loop->getRegions())
+    if (!region.hasOneBlock() || region.front().empty() ||
+        !dominance.hasSSADominance(&region) ||
+        !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/2] [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>