llvm / llvm/llvm-project

[MLIR] `scf-parallel-loop-fusion` misses fusion with permuted producer-consumer indices

Open
#221,744 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

`scf-parallel-loop-fusion` misses a producer-consumer fusion opportunity between two adjacent `scf.parallel` loops with identical bounds.
The producer writes: `%arg0[i0, i1, i2, i3, i4, i5] = 5.0`. The consumer reads the same buffer with a reversed index permutation: `%v = %arg0[i5, i4, i3, i2, i1, i0]`, and stores it to `%arg1`.

Since both loops iterate over the full same `2x2x2x2x2x2` iteration space, the reversed access is still a one-to-one permutation of the producer iterations. The loops could be fused by mapping each consumer iteration to the corresponding producer iteration. Currently, the loops remain separate, resulting in a missed optimization.

### Input Program

**input.mlir**

```llvm
module {
func.func @permutation_budget(%arg0: memref<2x2x2x2x2x2xf32>, %arg1: memref<2x2x2x2x2x2xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%cst = arith.constant 5.000000e+00 : f32
scf.parallel (%arg2, %arg3, %arg4, %arg5, %arg6, %arg7) = (%c0, %c0, %c0, %c0, %c0, %c0) to (%c2, %c2, %c2, %c2, %c2, %c2) step (%c1, %c1, %c1, %c1, %c1, %c1) {
memref.store %cst, %arg0[%arg2, %arg3, %arg4, %arg5, %arg6, %arg7] : memref<2x2x2x2x2x2xf32>
scf.reduce
}
scf.parallel (%arg2, %arg3, %arg4, %arg5, %arg6, %arg7) = (%c0, %c0, %c0, %c0, %c0, %c0) to (%c2, %c2, %c2, %c2, %c2, %c2) step (%c1, %c1, %c1, %c1, %c1, %c1) {
%0 = memref.load %arg0[%arg7, %arg6, %arg5, %arg4, %arg3, %arg2] : memref<2x2x2x2x2x2xf32>
memref.store %0, %arg1[%arg2, %arg3, %arg4, %arg5, %arg6, %arg7] : memref<2x2x2x2x2x2xf32>
scf.reduce
}
return
}
}
```
### Command
```bash
mlir-opt input.mlir --scf-parallel-loop-fusion
```

Version: b1e0aa39763434060ab6e4dc87af3630ac9a04ac

Contributor guide

Open the contributing guide

Research direction

Run the input.mlir example with mlir-opt --scf-parallel-loop-fusion and inspect the scf-parallel-loop-fusion pass to see how producer and consumer indices are matched. Done means the reversed one-to-one permutation is recognized and the two adjacent scf.parallel loops are fused in the output.

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
Active
Clarity
Mostly clear
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.