[MLIR][Affine] affine-loop-fusion truncates the producer iteration space to the consumer domain
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
Applying `affine-loop-fusion` to the verifier-valid program below changes the result from `-4` to `0`.
The producer covers `j = [0, 4)`, whereas the consumer covers `j = [0, 3)`. The generated fused nest uses the consumer's smaller `j` domain for both bodies, so the producer iteration at `j = 3` is omitted.
Public `mlir-opt` reproducer: https://godbolt.org/z/8aW11aYT6
## Version
Tested at `llvm/llvm-project` commit `450cf0aafeab9fed6666b112b018d3913a36527c` (`LLVM 24.0.0git`).
The same incorrect target was produced by Release builds with assertions both enabled and disabled.
## Reproducer
```mlir
module {
func.func private @kernel(%ey: memref<4x4xf64>, %hz: memref<4x4xf64>) {
affine.for %t = 0 to 2 {
affine.for %i = 1 to 4 {
affine.for %j = 0 to 4 {
%old = affine.load %ey[%i, %j] : memref<4x4xf64>
%a = affine.load %hz[%i, %j] : memref<4x4xf64>
%b = affine.load %hz[%i - 1, %j] : memref<4x4xf64>
%d = arith.subf %a, %b : f64
%new = arith.subf %old, %d : f64
affine.store %new, %ey[%i, %j] : memref<4x4xf64>
}
}
affine.for %i = 0 to 3 {
affine.for %j = 0 to 3 {
%old = affine.load %hz[%i, %j] : memref<4x4xf64>
%y1 = affine.load %ey[%i + 1, %j] : memref<4x4xf64>
%y0 = affine.load %ey[%i, %j] : memref<4x4xf64>
%d = arith.subf %y1, %y0 : f64
%new = arith.subf %old, %d : f64
affine.store %new, %hz[%i, %j] : memref<4x4xf64>
}
}
}
return
}
func.func @main() -> i64 {
%ey = memref.alloca() : memref<4x4xf64>
%hz = memref.alloca() : memref<4x4xf64>
%zero = arith.constant 0.0 : f64
%two = arith.constant 2.0 : f64
affine.for %i = 0 to 4 {
affine.for %j = 0 to 4 {
affine.store %zero, %ey[%i, %j] : memref<4x4xf64>
affine.store %zero, %hz[%i, %j] : memref<4x4xf64>
}
}
affine.store %two, %hz[1, 3] : memref<4x4xf64>
func.call @kernel(%ey, %hz)
: (memref<4x4xf64>, memref<4x4xf64>) -> ()
%result = affine.load %ey[1, 3] : memref<4x4xf64>
%readable = arith.fptosi %result : f64 to i64
return %readable : i64
}
}
```
## Commands
```sh
LLVM_BUILD=/path/to/llvm-build
LOWER='builtin.module(lower-affine,convert-scf-to-cf,expand-strided-metadata,finalize-memref-to-llvm,convert-func-to-llvm,convert-arith-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)'
SHARED_LIBS="$LLVM_BUILD/lib/libmlir_runner_utils.so,$LLVM_BUILD/lib/libmlir_c_runner_utils.so"
# Baseline.
"$LLVM_BUILD/bin/mlir-opt" input.mlir \
-pass-pipeline="$LOWER" \
-o baseline.mlir
"$LLVM_BUILD/bin/mlir-runner" baseline.mlir \
-e main \
--entry-point-result=i64 \
--shared-libs="$SHARED_LIBS"
# Apply affine-loop-fusion to the same input, then lower and execute it.
"$LLVM_BUILD/bin/mlir-opt" input.mlir \
--affine-loop-fusion \
-o fused.mlir
"$LLVM_BUILD/bin/mlir-opt" fused.mlir \
-pass-pipeline="$LOWER" \
-o fused-lowered.mlir
"$LLVM_BUILD/bin/mlir-runner" fused-lowered.mlir \
-e main \
--entry-point-result=i64 \
--shared-libs="$SHARED_LIBS"
```
## Actual behavior
The baseline and optimized executions produce:
```text
without affine-loop-fusion: -4
with affine-loop-fusion: 0
```
The generated function contains the following loop structure:
```mlir
affine.for %t = 0 to 2 {
affine.for %j = 0 to 3 {
affine.for %i = 1 to 4 {
// producer body
}
affine.for %i = 0 to 3 {
// consumer body
}
}
}
```
The producer body is therefore executed only for `j = [0, 3)`, even though its source iteration space is `j = [0, 4)`. The update of `ey[1, 3]` is lost, and the initialized value `0.0` is returned instead of `-4.0`.
## Expected behavior
`affine-loop-fusion` should preserve the complete producer iteration space, including producer iterations outside the consumer domain, or decline the fusion.
Applying the pass must not change the result of this program.
## Validity
Both the input and generated target pass MLIR verification and lower successfully. The two memrefs are distinct and fully initialized, all memory accesses are in bounds, and the arithmetic and runner entry point are defined. No undefined behavior explains the result difference.
Contributor guide
Research direction
Start with the provided MLIR reproducer and run it through mlir-opt with --affine-loop-fusion, comparing the generated loop structure and results against the baseline commands. Trace the affine-loop-fusion pass behavior around mismatched producer and consumer domains; done means preserving the producer iteration at j = 3 or declining fusion so execution remains -4.
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
- 50/100