[mlir][affine] Incorrect results after affine-loop-fusion when updating a memref in place
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
I think `-affine-loop-fusion` might be incorrectly changing the semantics of the program below which has a producer loop that updates `%p` in place, then copies the value to `%d` in a second loop.
Tested on `main` (`6a33b69d8bae`).
### Reproducer
```mlir
func.func @rmw_producer(%p: memref<4xf32>, %d: memref<2xf32>, %c: f32) {
affine.for %i = 0 to 4 {
%0 = affine.load %p[%i] : memref<4xf32>
%1 = arith.mulf %0, %c : f32
affine.store %1, %p[%i] : memref<4xf32>
}
affine.for %i = 0 to 2 {
%0 = affine.load %p[%i] : memref<4xf32>
affine.store %0, %d[%i] : memref<2xf32>
}
return
}
```
```
mlir-opt -affine-loop-fusion repro.mlir
```
Compiler Explorer: https://godbolt.org/z/fzPT8hn6G
### Output
```mlir
module {
func.func @rmw_producer(%arg0: memref<4xf32>, %arg1: memref<2xf32>, %arg2: f32) {
%alloc = memref.alloc() : memref<1xf32>
affine.for %arg3 = 0 to 4 {
%0 = affine.load %arg0[%arg3] : memref<4xf32>
%1 = arith.mulf %0, %arg2 : f32
affine.store %1, %arg0[%arg3] : memref<4xf32>
}
affine.for %arg3 = 0 to 2 {
%0 = affine.load %arg0[%arg3] : memref<4xf32>
%1 = arith.mulf %0, %arg2 : f32
affine.store %1, %alloc[0] : memref<1xf32>
%2 = affine.load %alloc[0] : memref<1xf32>
affine.store %2, %arg1[%arg3] : memref<2xf32>
}
return
}
}
```
If `p = [1, 1, 1, 1]` and `c = 2.0`, then the input program scales `p` once and then copies a prefix to `d`, which results in `p = [2, 2, 2, 2]` and `d = [2, 2]`.
But in the fused output, the multiply is applied a second time so it results in `d = [4, 4]`.
I wasn't sure whether this is intended since `@same_memref_load_store` in `mlir/test/Dialect/Affine/loop-fusion-4.mlir` tests a similar pattern.
Contributor guide
Research direction
Run mlir-opt -affine-loop-fusion on the reproducer and compare its behavior with mlir/test/Dialect/Affine/loop-fusion-4.mlir, especially the @same_memref_load_store case. Trace affine-loop-fusion's handling of in-place memref updates and verify that fusion preserves the original result: p is scaled once and d receives [2, 2] for the given inputs.
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
- Clearly specified
- Newbie friendliness
- 55/100