llvm / llvm/llvm-project

[mlir][affine] Incorrect results after affine-loop-fusion when updating a memref in place

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.