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

[Matteo Di Iorio]

================
@@ -371,14 +423,16 @@ 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;
+    auto loop = dyn_cast<scf::WhileOp>(op);
+    if (!loop || !allocValue.getDefiningOp<memref::AllocOp>())
----------------
splicer3 wrote:

Essentially only `memref.alloc` is allowed here. That is because I found that `memref.realloc` also passes the loop-hoisting interface check. Hoisting it could move it before a store to its source in `before`, so the store would end up using the old memref after the realloc. The forward alias checks on the result don’t cover that use.

I think an explicit extension to `memref.alloca` could be harmless, but just dropping this check would also allow that `memref.realloc` case.

https://github.com/llvm/llvm-project/pull/226951