llvm / llvm/llvm-project

[mlir][VectorToSCF] alloca for the transfer staging buffer is placed inside the enclosing scf.for

Open
#223,187 4 comments 0 reactions 0 assignees View on GitHub
mlir
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

`--convert-vector-to-scf` creates the staging buffer for a multi-dimensional
`vector.transfer_read`/`transfer_write` with `memref.alloca`, and places it in
the body of the enclosing `scf.for` rather than at function scope. In a loop
nest the allocation is then reached once per iteration.

### Reproducer

`repro.mlir` — three nested `scf.for` around one vectorized contraction, which
is what a tiled and vectorized `linalg.batch_matmul` looks like after
bufferization. Buffers are function arguments, so no bufferization pass and no
out-of-tree dialect is involved.

```mlir
func.func @tiled_contraction(%A: memref<1x120x320xf32>,
%B: memref<1x320x320xf32>,
%C: memref<1x120x320xf32>) {
%c0 = arith.constant 0 : index
%c8 = arith.constant 8 : index
%c32 = arith.constant 32 : index
%c120 = arith.constant 120 : index
%c320 = arith.constant 320 : index
%cst = arith.constant 0.000000e+00 : f32

scf.for %i = %c0 to %c120 step %c8 {
scf.for %j = %c0 to %c320 step %c8 {
scf.for %k = %c0 to %c320 step %c32 {
%sa = memref.subview %A[0, %i, %k] [1, 8, 32] [1, 1, 1]
: memref<1x120x320xf32>
to memref<1x8x32xf32, strided<[38400, 320, 1], offset: ?>>
%sb = memref.subview %B[0, %k, %j] [1, 32, 8] [1, 1, 1]
: memref<1x320x320xf32>
to memref<1x32x8xf32, strided<[102400, 320, 1], offset: ?>>
%sc = memref.subview %C[0, %i, %j] [1, 8, 8] [1, 1, 1]
: memref<1x120x320xf32>
to memref<1x8x8xf32, strided<[38400, 320, 1], offset: ?>>

%va = vector.transfer_read %sa[%c0, %c0, %c0], %cst
{in_bounds = [true, true, true]}
: memref<1x8x32xf32, strided<[38400, 320, 1], offset: ?>>,
vector<1x8x32xf32>
%vb = vector.transfer_read %sb[%c0, %c0, %c0], %cst
{in_bounds = [true, true, true]}
: memref<1x32x8xf32, strided<[102400, 320, 1], offset: ?>>,
vector<1x32x8xf32>
%vc = vector.transfer_read %sc[%c0, %c0, %c0], %cst
{in_bounds = [true, true, true]}
: memref<1x8x8xf32, strided<[38400, 320, 1], offset: ?>>,
vector<1x8x8xf32>

%r = vector.contract {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"],
kind = #vector.kind}
%va, %vb, %vc
: vector<1x8x32xf32>, vector<1x32x8xf32> into vector<1x8x8xf32>

vector.transfer_write %r, %sc[%c0, %c0, %c0]
{in_bounds = [true, true, true]}
: vector<1x8x8xf32>,
memref<1x8x8xf32, strided<[38400, 320, 1], offset: ?>>
}
}
}
return
}
```

```
mlir-opt repro.mlir --convert-vector-to-scf
```

### Actual output

```mlir
scf.for %arg3 = %c0 to %c120 step %c8 {
scf.for %arg4 = %c0 to %c320 step %c8 {
scf.for %arg5 = %c0 to %c320 step %c32 {
%alloca = memref.alloca() : memref>
%alloca_0 = memref.alloca() : memref>
%alloca_1 = memref.alloca() : memref>
%alloca_2 = memref.alloca() : memref>
...
```

All four allocations are in the innermost loop body. After
`--convert-scf-to-cf --finalize-memref-to-llvm` they become `llvm.alloca` in a
block that the loop's back edge returns to, and no `llvm.intr.stacksave` /
`stackrestore` pair covers them:

```
mlir-opt repro.mlir --convert-vector-to-scf --expand-strided-metadata \
--lower-affine --convert-scf-to-cf --convert-vector-to-llvm \
--convert-arith-to-llvm --convert-cf-to-llvm --finalize-memref-to-llvm \
--convert-func-to-llvm --reconcile-unrealized-casts
```

```mlir
^bb1(%40: i64): // 2 preds: ^bb0, ^bb32
llvm.cond_br %41, ^bb2, ^bb33
^bb3(%42: i64): // 2 preds: ^bb2, ^bb31
llvm.cond_br %43, ^bb4, ^bb32
^bb5(%44: i64): // 2 preds: ^bb4, ^bb30
llvm.cond_br %45, ^bb6, ^bb31
^bb6: // pred: ^bb5
%47 = llvm.alloca %46 x !llvm.array<1 x array<8 x vector<32xf32>>> : (i64) -> !llvm.ptr
%54 = llvm.alloca %53 x !llvm.array<1 x array<32 x vector<8xf32>>> : (i64) -> !llvm.ptr
%61 = llvm.alloca %60 x !llvm.array<1 x array<8 x vector<8xf32>>> : (i64) -> !llvm.ptr
%68 = llvm.alloca %67 x !llvm.array<1 x array<8 x vector<8xf32>>> : (i64) -> !llvm.ptr
```

`^bb6` is inside all three back edges, and the module contains no
`llvm.intr.stacksave` / `stackrestore` at all — so nothing reclaims these
between iterations.

### Expected output

Either the allocation is hoisted to the nearest *function*-level
`AutomaticAllocationScope`, or the `scf.for` scope is honoured down through the
lowering so the stack is reclaimed per iteration.

### Impact

The stack grows with the total trip count and is never reclaimed within the
function. On a real pipeline (a tiled `linalg.batch_matmul` nest) this is about
2.2 KB per iteration, so an 8 MB stack is exhausted at roughly 3,700
iterations. Measured boundary: 3,600 trips runs, 3,840 trips segfaults.
Doubling the stack to 16 MB doubles the threshold, which is consistent with the
allocation never being freed.

At `-O0` and `-O1` a *single* `batch_matmul` of the shape above is enough to
crash `mlir-runner`.

### Why this has not shown up more often

`-O2` and `-O3` usually promote the allocation out of the loop, so the same IR
runs fine. In our case:

```
-O0 -O1 -O2 -O3
k tile 8 crash crash ok ok
k tile 32 crash crash crash crash
```

So whether the bug is visible depends on the optimisation level and on how much
the loop body allocates — which also means the same IR can pass CI and fail in
a release build with a slightly different shape.

Debugging tools hide it as well: it disappears under `ulimit -s unlimited`
(which also switches the mmap layout), under `setarch -R`, and under ASAN.

### Cause

`VectorToSCF.cpp`:

```cpp
// TODO: Parallelism and threadlocal considerations with a ParallelScope trait.
static Operation *getAutomaticAllocationScope(Operation *op) {
Operation *scope =
op->getParentWithTrait();
assert(scope && "Expected op to be inside automatic allocation scope");
return scope;
}

template
static BufferAllocs allocBuffers(OpBuilder &b, OpTy xferOp) {
...
Operation *scope = getAutomaticAllocationScope(xferOp);
assert(scope->getNumRegions() == 1 &&
"AutomaticAllocationScope with >1 regions");
b.setInsertionPointToStart(&scope->getRegion(0).front());
```

`getParentWithTrait` returns the *nearest* enclosing scope, and `scf.for`
carries `AutomaticAllocationScope`. So inside a loop nest the buffer is
allocated at the start of the innermost loop body rather than at function
entry. There is no check for loop-like ops here.

### Prior art

- D124366 fixed the same class of problem in the vector transfer full/partial
splitting pass: after the trait was added to loop-like constructs, that pass
started placing allocas in the nearest loop instead of the nearest function,
and the fix was to skip loops when looking for the scope. That review left
the question open — whether every alloca-creating site should do this, or
whether there should be a pass that hoists allocas out of loops.
- #69786, `[mlir][sparse] fix stack overflow due to memref.alloca in loops`,
is the same failure mode fixed on the producing side.
- The `memref.alloca_scope` RFC (D97768) anticipated this: it notes that once
`scf` and `alloca` are lowered to `llvm.alloca` plus unstructured control
flow, the allocation is either optimised away by LLVM or, in the worst case,
performs an independent stack allocation on each iteration — and that the
LLVM optimisations are likely but not guaranteed to succeed.
- #62419, `[flang][hlfir] Potential excessive stack usage`, is the same
situation in a different dialect and is still open: an `AllocaOp` inserted
inside a loop, with stacksave/stackrestore or a temporary-creation helper
proposed as the fix.
- #216225 hit the assert in the same `getAutomaticAllocationScope()` from the
other side (no enclosing scope at all). Different symptom, same scope-lookup.

Note also that the documented example for this pattern in `VectorToSCF.h`
places the `alloca` *outside* the loop the pattern itself generates, so
allocating outside a loop appears to be the intended shape; what differs here
is only that the loop is pre-existing rather than generated.

### Workarounds

Both fix it on our pipeline, verified bit-identical output:

- `--buffer-loop-hoisting` right after `--convert-vector-to-scf`
- `--convert-vector-to-scf='full-unroll=true'` (removes the staging buffers;
the descriptor allocas remain)

### Question

Given D124366 and #69786, should `VectorToSCF` skip loop scopes the same way,
or is `--buffer-loop-hoisting` (or an explicit `memref.alloca_scope`) the
intended contract for any pipeline that lowers vector transfers inside a loop?
Happy to send a patch once the direction is settled.

### Version

```
llvm-project @ d7d7992e58674a2e9ad1ae42e4a627b7fff24cef
LLVM version 24.0.0git, Optimized build
```

Reproduced on a clean checkout at that commit with only `mlir-opt` built
(`-DLLVM_ENABLE_PROJECTS=mlir -DCMAKE_BUILD_TYPE=Release`).

Contributor guide

Open the contributing guide

Research direction

Start in VectorToSCF.cpp at getAutomaticAllocationScope and allocBuffers, then compare the loop-scope handling from D124366 and issue #69786. Run the supplied repro.mlir through --convert-vector-to-scf and the listed lowering pipeline. Done means the staging allocations no longer accumulate across loop iterations, with behavior covered by an appropriate regression test.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.