llvm / llvm/llvm-project

[MLIR][SCF] Miscompilation: `affine-loop-coalescing` miscompiles nested `scf.for` with captured outer iter_arg

Open
#216,289 1 comment 0 reactions 0 assignees View on GitHub
miscompilation mlir
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.