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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Matteo Di Iorio (splicer3)

<details>
<summary>Changes</summary>

As discussed in #<!-- -->223944, `buffer-loop-hoisting` rejects an allocation in `scf.while` even though the buffer is never passed back to the region containing the allocation. The allocation and the forwarded `after` argument belong to different regions, so their common dominator is the block containing the loop, which in turn makes `isLegalPlacement` return false and reject the hoist.

This PR adds a fallback for when that existing check rejects the hoist. It keeps the full alias set and rejects hoisting if an alias reaches an argument of the allocation's region. It also checks uses of all aliases, including loop results, for capture, deallocation and explicit observations of allocation identity.

For now, the fallback only handles `memref.alloc` in multi-region loops directly inside an isolated function body. Each region must have one SSA block, successor inputs must be fully described by operand forwarding, and the allocation's region must execute on entry. The function-parent restriction avoids exposing the allocation to further hoisting through enclosing loops with incorrectly classified execution semantics, as in #<!-- -->226090.
The shared placement walk also stops at the transform target, isolation boundaries and unreachable enclosing blocks. I've added regression tests to cover those boundaries as well.

There are some points I'd like feedback on, eventually:

1. The terminator check assumes that successor operands are used only for forwarding. This is consistent with the handling in `remove-dead-values` (#<!-- -->173505), which can replace those operands with poison when all their successor inputs are dead.

2. Absence of `HasParallelRegion` does not establish sequential execution, and the trait currently requires the operation to have one region. The function-parent restriction avoids the enclosing-loop problem, but doesn't establish how the candidate loop itself executes.

3. `scf.while` does provide the required execution and forwarding semantics, so if needed the scope can be reduced to that.

Validation:
Release build with assertions, 111 selected lit tests passed with one expected failure.

Fixes #<!-- -->223944.

AI assistance: Codex, mainly for rapid prototyping (very early implementations of what I discussed in the issue and other potential fixes) and codebase exploration. The end result was manually reviewed.

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


5 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+7) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp (+131-17) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir (+17) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir (+341) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir (+61) 


``````````diff
The server is unavailable at this time. Please wait a few minutes before you try again.
``````````

</details>


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