llvm / llvm/llvm-project

LinalgSpecializeGenericOpsPass mis-translates linalg.generic with non-identity unary operations in payload

Open
#197,178 1 comment 0 reactions 0 assignees View on GitHub
mlir:linalg
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

The `--linalg-specialize-generic-ops` pass ignores unary, side-effect free, operations in the `linalg.generic` payload.

This is incorrect.

Both the following two `linalg.generic` are specialized to identical `linalg.matmul`, despite the fact the second one negates the return value from the payload.

```
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
module {
func.func @mul(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
%0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) {
^bb0(%in: f32, %in_0: f32, %out: f32):
%1 = arith.mulf %in, %in_0 : f32
%2 = arith.addf %out, %1 : f32
linalg.yield %2 : f32
} -> tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
func.func @mul_with_neg(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
%0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) {
^bb0(%in: f32, %in_0: f32, %out: f32):
%1 = arith.mulf %in, %in_0 : f32
%2 = arith.addf %out, %1 : f32
// NOTE: unary, side-effect free but not identity:
%3 = arith.negf %2 : f32
linalg.yield %3 : f32
} -> tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
}
```

`mlir-opt --linalg-specialize-generic-ops` gives two identical `linalg.matmul`:

```
module {
func.func @mul(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) -> tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
func.func @mul_neg(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<3x3xf32>) -> tensor<3x3xf32> {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<3x3xf32>, tensor<3x3xf32>) outs(%arg2 : tensor<3x3xf32>) -> tensor<3x3xf32>
return %0 : tensor<3x3xf32>
}
}
```

This is on what Compiler Explorer calls "Trunk" but the relevant code seems to be present at least since 2023 (`getSourceSkipUnary()`).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.