llvm / llvm/llvm-project

[MLIR] Missed elementwise op fusion when loop bounds depend on producer's outs operands

Open
#223,308 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

`-linalg-fuse-elementwise-ops` fails to fuse an elementwise producer into a reduction consumer when a loop dimension's bound comes from the producer's outs operand.

In the example below, both operations share the iteration domain `(d0, d1)`. The producer's outs operand `%arg1` has a shape `tensor<4x8xf32> (#map1)`, which defines the bound for `d1`. However, running the fusion pass leaves the IR unchanged, resulting in a missed optimization.

### Reproduction Example
```llvm
#map = affine_map<(d0, d1) -> (d0)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
module {
func.func @preserved_producer_out(%arg0: tensor<4xf32>, %arg1: tensor<4x8xf32>, %arg2: tensor<4xf32>) -> tensor<4xf32> {
%0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<4xf32>) outs(%arg1 : tensor<4x8xf32>) {
^bb0(%in: f32, %out: f32):
%2 = arith.addf %in, %out : f32
linalg.yield %2 : f32
} -> tensor<4x8xf32>
%1 = linalg.generic {indexing_maps = [#map1, #map], iterator_types = ["parallel", "reduction"]} ins(%0 : tensor<4x8xf32>) outs(%arg2 : tensor<4xf32>) {
^bb0(%in: f32, %out: f32):
%2 = arith.addf %in, %out : f32
linalg.yield %2 : f32
} -> tensor<4xf32>
return %1 : tensor<4xf32>
}
}
```

### Command
```bash
mlir-opt --linalg-fuse-elementwise-ops
```

### Expected Behavior

The pass should include producer's outs operands during loop bound coverage analysis, allowing `%0` to be fused into `%1`.

Version: 6f54fe6b6fb22f8a335fd19b94049fc769dfea25

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the provided IR with mlir-opt and the -linalg-fuse-elementwise-ops pass, then trace the pass's loop-bound coverage analysis for producer operands. Add a regression case for the shown producer and reduction consumer; done means the producer is fused when its outs operand supplies a loop bound.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.