[MLIR][SCF] Miscompilation: `affine-loop-coalescing` miscompiles nested `scf.for` with captured outer iter_arg
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
Coalescing this loop nest effectively identifies the inner loop's `iter_arg` with the outer loop's `iter_arg`. This is sound when the outer `iter_arg` is used only to initialize the inner loop-carried value. However, the inner loop body may also capture and read the outer `iter_arg` directly. In that case, the two values have different semantics before coalescing: the outer `iter_arg` remains fixed throughout one execution of the inner loop, whereas the inner `iter_arg` changes after every inner iteration.
In the reproducer below, the inner loop starts with `%b = %a` and executes `%b = %b + %a` four times, while `%a` remains fixed. Thus, each outer iteration maps `a` to `5a`, and four outer iterations produce `625`. After `--affine-loop-coalescing`, the program prints `65536 = 2^16`. The coalesced loop effectively computes `a + a` on every merged iteration, so the loop-carried value doubles sixteen times.
## Reproducer
### small.mlir
```mlir
func.func @f(%init: i64) -> i64 {
%c0 = arith.constant 0 : i32
%c4 = arith.constant 4 : i32
%c1 = arith.constant 1 : i32
%r = scf.for %i = %c0 to %c4 step %c1 iter_args(%a = %init) -> (i64) : i32 {
%s = scf.for %j = %c0 to %c4 step %c1 iter_args(%b = %a) -> (i64) : i32 {
// %a is the OUTER loop's iter_arg: fixed for the whole inner loop.
%t = arith.addi %b, %a : i64
scf.yield %t : i64
}
scf.yield %s : i64
}
return %r : i64
}
func.func @main() {
%one = arith.constant 1 : i64
%v = func.call @f(%one) : (i64) -> i64
vector.print %v : i64
return
}
```
### To reproduce:
```bash
$ mlir-opt small.mlir \
--convert-scf-to-cf --convert-vector-to-llvm --convert-arith-to-llvm \
--convert-func-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \
| mlir-runner -e main --entry-point-result=void \
--shared-libs=$LLVM_BUILD/lib/libmlir_c_runner_utils.so
625
$ mlir-opt small.mlir --affine-loop-coalescing \
--convert-scf-to-cf --convert-vector-to-llvm --convert-arith-to-llvm \
--convert-func-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \
| mlir-runner -e main --entry-point-result=void \
--shared-libs=$LLVM_BUILD/lib/libmlir_c_runner_utils.so
65536
```
### Actual Output:
```
65536
```
### Expected Output:
```
625
```
Contributor guide
Research direction
Start with the small.mlir reproducer and run the two mlir-opt pipelines, comparing behavior with and without --affine-loop-coalescing. Trace the affine-loop-coalescing pass handling of captured outer scf.for iter_args; the fix is complete when the transformed pipeline preserves the expected output of 625.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100