[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 59b7cf266d9632be6b579a8ade840ffffd4ef067 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/7] [mlir][bufferization] handle exit aliases in
 buffer-loop-hoisting

---
 .../Bufferization/Transforms/Passes.td        |   7 +
 .../Transforms/BufferOptimizations.cpp        | 131 ++++++-
 .../Transforms/buffer-loop-hoisting.mlir      | 341 ++++++++++++++++++
 .../Transforms/transform-ops.mlir             |  61 ++++
 4 files changed, 529 insertions(+), 11 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 0638b339cf9dc..f57e9527486d7 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 e03887c8c5910..87d13e1067d3f 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;
@@ -371,14 +467,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-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 7498f7af8c5d1..62df7c10fc0df 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -567,3 +567,344 @@ func.func @loop_unreachable_parent() {
   }
   return
 }
+
+// -----
+
+// 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 27ff9964c2ca2..3cd2af07a9393 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -289,3 +289,64 @@ module attributes {transform.with_named_sequence} {
     return
   }
 }
+
+// -----
+
+// 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 f834b0190dcd1048b55bd2aef467d8cf62d45c6a 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/7] [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 62df7c10fc0df..3cc2becbf2402 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -877,7 +877,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 cf13eb97f98276aaf9a0b518d252eb8afc1350e3 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/7] [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        | 100 +++++-----------
 .../Transforms/buffer-loop-hoisting.mlir      | 111 ++++++++++++------
 3 files changed, 107 insertions(+), 116 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index f57e9527486d7..5af0dd4793ac1 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 87d13e1067d3f..67dc1e2cdf3f8 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)
-    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); }))
+  // 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;
 
-  // 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;
     }
   }
@@ -388,7 +349,7 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
         if (parentOp == scopeOp ||
             parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>() ||
             !isKnownControlFlowInterface(parentOp) ||
-            !state.isLegalPlacement(parentOp))
+            !state.isLegalPlacement(parentOp, scopeOp))
           break;
         // Move to our parent block by notifying the current StateT
         // implementation.
@@ -437,7 +398,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) {
@@ -468,26 +429,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-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 3cc2becbf2402..426ccb82aa3b3 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -710,37 +710,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> {
@@ -760,14 +738,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
@@ -908,3 +878,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
+}

>From 2acd4b0395762adafd1f7a577a4d2ca76e57010b Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Thu, 1 Oct 2026 12:34:13 +0200
Subject: [PATCH 4/7] style: clang format

---
 .../Transforms/BufferOptimizations.cpp            | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 67dc1e2cdf3f8..0ec4a933f3edb 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -72,9 +72,9 @@ static bool isSequentialLoop(Operation *op) {
 /// 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) {
+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()) ||
@@ -398,7 +398,9 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
   }
 
   /// Returns true if the given operation does not represent a loop.
-  bool isLegalPlacement(Operation *op, Operation *scopeOp) { 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) {
@@ -438,8 +440,9 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
     auto alloc = allocValue.getDefiningOp<memref::AllocOp>();
     if (!loop || !alloc)
       return false;
-    // Moving outside the analysis scope leaves liveness information unavailable.
-    // An unreachable parent block has no dominator node for the placement walk.
+    // 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;

>From 57f210d4301b704b0b2724a68959b0b5b625a2b1 Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Sat, 3 Oct 2026 10:27:33 +0200
Subject: [PATCH 5/7] [mlir][bufferization] Simplify scf.while exit-alias
 hoisting

Also addresses review comments and relies on the previously implemented shared placement checks.
---
 .../Bufferization/Transforms/Passes.td        |  8 ++---
 .../Transforms/BufferOptimizations.cpp        | 30 +++++-----------
 .../Transforms/buffer-loop-hoisting.mlir      | 34 +++++++++++++++----
 .../Transforms/transform-ops.mlir             |  1 -
 4 files changed, 40 insertions(+), 33 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 5af0dd4793ac1..9e0e59ec98648 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -308,10 +308,10 @@ def BufferLoopHoistingPass : Pass<"buffer-loop-hoisting", "func::FuncOp"> {
 
     `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.
+    This extension requires a loop directly in an isolated function body 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 0ec4a933f3edb..bf02036befb92 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -69,11 +69,11 @@ static bool isSequentialLoop(Operation *op) {
   return !op->hasTrait<OpTrait::HasParallelRegion>() && isLoop(op);
 }
 
-/// 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.
+/// Returns true if an allocation 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,
+canHoistFromWhile(scf::WhileOp loop,
                   const BufferViewFlowAnalysis::ValueSetT &aliases) {
   // Further hoisting across an enclosing loop with unmodeled parallel
   // execution could share the allocation across iterations.
@@ -81,11 +81,6 @@ canHoistFromWhile(memref::AllocOp alloc, scf::WhileOp loop,
       !loop->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
     return false;
 
-  // 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;
-
   // 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) {
@@ -349,7 +344,7 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
         if (parentOp == scopeOp ||
             parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>() ||
             !isKnownControlFlowInterface(parentOp) ||
-            !state.isLegalPlacement(parentOp, scopeOp))
+            !state.isLegalPlacement(parentOp))
           break;
         // Move to our parent block by notifying the current StateT
         // implementation.
@@ -398,9 +393,7 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
   }
 
   /// Returns true if the given operation does not represent a loop.
-  bool isLegalPlacement(Operation *op, Operation *scopeOp) {
-    return !isLoop(op);
-  }
+  bool isLegalPlacement(Operation *op) { return !isLoop(op); }
 
   /// Returns true if the given operation should be considered for hoisting.
   static bool shouldHoistOpType(Operation *op) {
@@ -431,22 +424,15 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
   }
 
   /// Returns true if the allocation can be moved across the given loop.
-  bool isLegalPlacement(Operation *op, Operation *scopeOp) {
+  bool isLegalPlacement(Operation *op) {
     if (!isSequentialLoop(op))
       return false;
     if (!dominators->dominates(aliasDominatorBlock, op->getBlock()))
       return true;
     auto loop = dyn_cast<scf::WhileOp>(op);
-    auto alloc = allocValue.getDefiningOp<memref::AllocOp>();
-    if (!loop || !alloc)
-      return false;
-    // 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()))
+    if (!loop || !allocValue.getDefiningOp<memref::AllocOp>())
       return false;
-    return canHoistFromWhile(alloc, loop, aliases);
+    return canHoistFromWhile(loop, aliases);
   }
 
   /// Returns true if the given operation should be considered for hoisting.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 426ccb82aa3b3..304a2bc6633b1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -922,15 +922,14 @@ func.func @no_hoist_while_alloca(%condition: i1) {
 
 // -----
 
-// Conservatively exclude unreachable parent blocks from the fallback.
-// CHECK-LABEL: func @no_hoist_while_unreachable_parent(
+// Hoist into the unreachable parent block, then stop the placement walk there.
+// CHECK-LABEL: func @while_unreachable_parent(
 // CHECK: return
 // CHECK: ^bb1:
-// CHECK-NOT: memref.alloc
-// CHECK: scf.while
-// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.while
 // CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
-func.func @no_hoist_while_unreachable_parent() {
+func.func @while_unreachable_parent() {
   return
 ^dead:
   %false = arith.constant false
@@ -943,3 +942,26 @@ func.func @no_hoist_while_unreachable_parent() {
   }
   return
 }
+
+// -----
+
+// Storing an alias as a value can expose the buffer to later iterations.
+// CHECK-LABEL: func @no_hoist_while_store_alias(
+// CHECK-SAME: %{{.*}}: i1, %[[HOLDER:.*]]: memref<memref<1xindex>>
+// CHECK-NOT: memref.alloc
+// CHECK: scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: ^bb0(%[[CURRENT:.*]]: memref<1xindex>):
+// CHECK-NEXT: memref.store %[[CURRENT]], %[[HOLDER]][]
+func.func @no_hoist_while_store_alias(%condition: i1, %holder: memref<memref<1xindex>>) {
+  %last = scf.while () : () -> memref<1xindex> {
+    %buffer = memref.alloc() : memref<1xindex>
+    scf.condition(%condition) %buffer : memref<1xindex>
+  } do {
+  ^bb0(%current: memref<1xindex>):
+    memref.store %current, %holder[] : memref<memref<1xindex>>
+    scf.yield
+  }
+  return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 3cd2af07a9393..7e2d4bf08a2fc 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -333,7 +333,6 @@ 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 {

>From ae9ded89ea85ff35a84f51bbb0a1f3f002b18cfe Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Sat, 3 Oct 2026 18:40:52 +0200
Subject: [PATCH 6/7] [mlir][bufferization] Allow exit-alias hoisting in
 non-isolated functions

Also keep the alias reference in the loop state instead of the shared base
---
 .../Bufferization/Transforms/Passes.td        |   2 +-
 .../Transforms/BufferOptimizations.cpp        |  29 ++--
 .../Transforms/buffer-loop-hoisting.mlir      | 131 +++++++++++++++---
 .../Transforms/transform-ops.mlir             |  34 +++++
 4 files changed, 164 insertions(+), 32 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 9e0e59ec98648..874e2049d5f80 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -308,7 +308,7 @@ def BufferLoopHoistingPass : Pass<"buffer-loop-hoisting", "func::FuncOp"> {
 
     `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 and
+    This extension requires a loop directly in a function body 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 bf02036befb92..53ba49b25d7fd 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -77,8 +77,7 @@ canHoistFromWhile(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>())
+  if (!isa<FunctionOpInterface>(loop->getParentOp()))
     return false;
 
   // Check all aliases, including loop results: memory effects alone do not rule
@@ -229,15 +228,11 @@ 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,
-      const BufferViewFlowAnalysis::ValueSetT &aliases)
+  BufferAllocationHoistingStateBase(DominanceInfo *dominators, Value allocValue,
+                                    Block *placementBlock)
       : dominators(dominators), allocValue(allocValue),
-        placementBlock(placementBlock), aliases(aliases) {}
+        placementBlock(placementBlock) {}
 };
 
 /// Implements the actual hoisting logic for allocation nodes.
@@ -376,7 +371,11 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
 /// that hoists allocations into dominator blocks while keeping them inside of
 /// loops.
 struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
-  using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
+  BufferAllocationHoistingState(DominanceInfo *dominators, Value allocValue,
+                                Block *placementBlock,
+                                const BufferViewFlowAnalysis::ValueSetT &)
+      : BufferAllocationHoistingStateBase(dominators, allocValue,
+                                          placementBlock) {}
 
   /// Computes the upper bound for the placement block search.
   Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
@@ -410,11 +409,19 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
 /// A state implementation compatible with the `BufferAllocationHoisting` class
 /// that hoists allocations out of loops.
 struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
-  using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
+  /// The forward alias closure of the current allocation.
+  const BufferViewFlowAnalysis::ValueSetT &aliases;
 
   /// Remembers the dominator block of all aliases.
   Block *aliasDominatorBlock = nullptr;
 
+  BufferAllocationLoopHoistingState(
+      DominanceInfo *dominators, Value allocValue, Block *placementBlock,
+      const BufferViewFlowAnalysis::ValueSetT &aliases)
+      : BufferAllocationHoistingStateBase(dominators, allocValue,
+                                          placementBlock),
+        aliases(aliases) {}
+
   /// Computes the upper bound for the placement block search.
   Block *computeUpperBound(Block *dominatorBlock, Block *dependencyBlock) {
     aliasDominatorBlock = dominatorBlock;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 304a2bc6633b1..61ba3e5d4e15a 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -570,35 +570,95 @@ func.func @loop_unreachable_parent() {
 
 // -----
 
-// An exit alias, invariant extent and unrelated carried buffer permit hoisting.
+// An exit alias does not carry the allocation back to before.
 // CHECK-LABEL: func @while_exit_alias(
-// CHECK-SAME: %{{.*}}: index, %[[SIZE:.*]]: index,
-//      CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
-// CHECK-NEXT: %[[LAST:.*]]:2 = scf.while
+//      CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
 // CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
-//      CHECK: scf.condition{{.*}} %[[ALLOC]], %{{.*}}
-//      CHECK: ^bb0(%[[CURRENT:.*]]: memref<?xindex>, %[[KEEP:.*]]: memref<1xindex>):
+//      CHECK: scf.condition{{.*}} %[[ALLOC]]
+//      CHECK: ^bb0(%[[CURRENT:.*]]: 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 {
+//      CHECK: memref.load %[[LAST]]
+func.func @while_exit_alias(%n: index) -> 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>
+  %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, %keep : memref<?xindex>, memref<1xindex>
+    scf.condition(%continue) %buffer : memref<1xindex>
   } do {
-  ^bb0(%current: memref<?xindex>, %keep: memref<1xindex>):
-    %value = memref.load %current[%c0] : memref<?xindex>
+  ^bb0(%current: memref<1xindex>):
+    %value = memref.load %current[%c0] : memref<1xindex>
     %next = arith.addi %value, %c1 : index
-    scf.yield %next, %keep : index, memref<1xindex>
+    scf.yield %next : index
+  }
+  %result = memref.load %last[%c0] : memref<1xindex>
+  return %result : index
+}
+
+// -----
+
+// A loop-invariant allocation extent permits hoisting with an exit alias.
+// CHECK-LABEL: func @while_exit_alias_invariant_extent(
+// CHECK-SAME: %{{.*}}: i1, %[[SIZE:.*]]: index)
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// CHECK-NEXT: %{{.*}} = scf.while
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+func.func @while_exit_alias_invariant_extent(%condition: i1, %size: index) {
+  %last = scf.while () : () -> memref<?xindex> {
+    %buffer = memref.alloc(%size) : memref<?xindex>
+    scf.condition(%condition) %buffer : memref<?xindex>
+  } do {
+  ^bb0(%current: memref<?xindex>):
+    scf.yield
+  }
+  return
+}
+
+// -----
+
+// An unrelated buffer may be carried to before while the allocation is hoisted.
+// CHECK-LABEL: func @while_exit_alias_unrelated_carried(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %{{.*}}:2 = scf.while
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]], %{{.*}}
+// CHECK: ^bb0(%{{.*}}: memref<1xindex>, %[[KEEP:.*]]: memref<1xindex>):
+// CHECK-NEXT: scf.yield %[[KEEP]]
+func.func @while_exit_alias_unrelated_carried(%condition: i1, %other: memref<1xindex>) {
+  %last:2 = scf.while (%keep = %other)
+      : (memref<1xindex>) -> (memref<1xindex>, memref<1xindex>) {
+    %buffer = memref.alloc() : memref<1xindex>
+    scf.condition(%condition) %buffer, %keep : memref<1xindex>, memref<1xindex>
+  } do {
+  ^bb0(%current: memref<1xindex>, %keep: memref<1xindex>):
+    scf.yield %keep : memref<1xindex>
   }
-  %view = memref.subview %last#0[0] [1] [1] : memref<?xindex> to memref<1xindex>
+  return
+}
+
+// -----
+
+// Casts and subviews of the loop result are included in the alias check.
+// CHECK-LABEL: func @while_exit_alias_view(
+// CHECK: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: %[[VIEW:.*]] = memref.subview %[[LAST]]
+// CHECK-NEXT: %[[CAST:.*]] = memref.cast %[[VIEW]]
+// CHECK-NEXT: %{{.*}} = memref.load %[[CAST]]
+func.func @while_exit_alias_view(%condition: i1) -> index {
+  %c0 = arith.constant 0 : index
+  %last = scf.while () : () -> memref<2xindex> {
+    %buffer = memref.alloc() : memref<2xindex>
+    memref.store %c0, %buffer[%c0] : memref<2xindex>
+    scf.condition(%condition) %buffer : memref<2xindex>
+  } do {
+  ^bb0(%current: memref<2xindex>):
+    scf.yield
+  }
+  %view = memref.subview %last[0] [1] [1] : memref<2xindex> to memref<1xindex>
   %cast = memref.cast %view : memref<1xindex> to memref<?xindex>
   %result = memref.load %cast[%c0] : memref<?xindex>
   return %result : index
@@ -661,6 +721,37 @@ func.func @no_hoist_while_exit_capture(%condition: i1) {
 
 // -----
 
+// A loop result used in a nested region is outside the supported alias uses.
+// CHECK-LABEL: func @no_hoist_while_exit_nested_region(
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT: scf.condition{{.*}} %[[ALLOC]]
+// CHECK: scf.if
+// CHECK-NEXT: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: scf.yield %[[VALUE]] : index
+func.func @no_hoist_while_exit_nested_region(%condition: i1, %read: 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 = scf.if %read -> (index) {
+    %value = memref.load %last[%c0] : memref<1xindex>
+    scf.yield %value : index
+  } else {
+    scf.yield %c0 : index
+  }
+  return %result : index
+}
+
+// -----
+
 // Hoisting would require the loaded value to be invariant across iterations.
 // CHECK-LABEL: func @no_hoist_while_invariant_load(
 //  CHECK-NOT: memref.alloc
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 7e2d4bf08a2fc..a1a8f1ca139a1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -349,3 +349,37 @@ module attributes {transform.with_named_sequence} {
     return %result : index
   }
 }
+
+// -----
+
+// Targeting the module still keeps the allocation inside the non-isolated
+// function, even when the allocation size is captured from outside it.
+// CHECK-LABEL: module attributes
+// CHECK: %[[SIZE:.*]] = arith.constant 4 : index
+// CHECK-NOT: memref.alloc
+// CHECK: test.conversion_func_op @buffer_loop_hoisting_nonisolated_function(
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[SIZE]])
+// 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}) {
+    transform.bufferization.buffer_loop_hoisting %root : !transform.any_op
+    transform.yield
+  }
+  %size = arith.constant 4 : index
+  test.conversion_func_op @buffer_loop_hoisting_nonisolated_function(%condition: i1) {
+    %c0 = arith.constant 0 : index
+    %last = scf.while () : () -> memref<?xindex> {
+      %buffer = memref.alloc(%size) : memref<?xindex>
+      memref.store %c0, %buffer[%c0] : memref<?xindex>
+      scf.condition(%condition) %buffer : memref<?xindex>
+    } do {
+    ^bb0(%current: memref<?xindex>):
+      scf.yield
+    }
+    %result = memref.load %last[%c0] : memref<?xindex>
+    "test.return"() : () -> ()
+  }
+}

>From c3365fa5e3703cd3559c4c015326a3f6c45004ad Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.diiorio2 at studio.unibo.it>
Date: Mon, 5 Oct 2026 07:50:15 +0200
Subject: [PATCH 7/7] [mlir][bufferization] Document and test the realloc
 hoisting restriction

---
 .../Transforms/BufferOptimizations.cpp        |  3 ++
 .../Transforms/buffer-loop-hoisting.mlir      | 29 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 53ba49b25d7fd..5c94baffa76a4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -437,6 +437,9 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
     if (!dominators->dominates(aliasDominatorBlock, op->getBlock()))
       return true;
     auto loop = dyn_cast<scf::WhileOp>(op);
+    // The alias checks do not cover uses of allocation operands. Hoisting
+    // memref.realloc before a store to its source would leave a use of the
+    // old memref after reallocation.
     if (!loop || !allocValue.getDefiningOp<memref::AllocOp>())
       return false;
     return canHoistFromWhile(loop, aliases);
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 61ba3e5d4e15a..1ed4bf2c34691 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -1013,6 +1013,35 @@ func.func @no_hoist_while_alloca(%condition: i1) {
 
 // -----
 
+// Hoisting realloc before the store would invalidate the store's memref.
+// The loop exits immediately to avoid reusing the source after reallocation.
+// CHECK-LABEL: func @no_hoist_while_realloc(
+// CHECK-SAME: %[[SOURCE:.*]]: memref<1xi32>
+// CHECK-NOT: memref.realloc
+// CHECK: %[[LAST:.*]] = scf.while
+// CHECK-NEXT: memref.store %{{.*}}, %[[SOURCE]][%{{.*}}]
+// CHECK-NEXT: %[[NEW:.*]] = memref.realloc %[[SOURCE]]
+// CHECK-NEXT: scf.condition{{.*}} %[[NEW]]
+// CHECK: %[[VALUE:.*]] = memref.load %[[LAST]]
+// CHECK-NEXT: return %[[VALUE]]
+func.func @no_hoist_while_realloc(%source: memref<1xi32>) -> i32 {
+  %c0 = arith.constant 0 : index
+  %c7 = arith.constant 7 : i32
+  %false = arith.constant false
+  %last = scf.while () : () -> memref<2xi32> {
+    memref.store %c7, %source[%c0] : memref<1xi32>
+    %new = memref.realloc %source : memref<1xi32> to memref<2xi32>
+    scf.condition(%false) %new : memref<2xi32>
+  } do {
+  ^bb0(%current: memref<2xi32>):
+    scf.yield
+  }
+  %value = memref.load %last[%c0] : memref<2xi32>
+  return %value : i32
+}
+
+// -----
+
 // Hoist into the unreachable parent block, then stop the placement walk there.
 // CHECK-LABEL: func @while_unreachable_parent(
 // CHECK: return