[mlir][linalg] bitcast is not preserved during matmul specialisation from generic
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
```
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul_bitcast_int_to_float(%A: tensor<16x8xi32>,
%B: tensor<8x32xi32>,
%Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.generic
{indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]}
ins(%A, %B : tensor<16x8xi32>, tensor<8x32xi32>) outs(%Out : tensor<16x32xf32>) {
^bb0(%in: i32, %in_0: i32, %out: f32):
%1 = arith.bitcast %in : i32 to f32
%2 = arith.bitcast %in_0 : i32 to f32
%3 = arith.mulf %1, %2 : f32
%4 = arith.addf %out, %3 : f32
linalg.yield %4 : f32
} -> tensor<16x32xf32>
return %0 : tensor<16x32xf32>
}
```
specialising the above linalg.generic like this
```
./mlir-opt --linalg-specialize-generic-ops c1.mlir -o c2.mlir
```
leads to
```
module {
func.func @op_matmul_bitcast_int_to_float(%arg0: tensor<16x8xi32>, %arg1: tensor<8x32xi32>, %arg2: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<16x8xi32>, tensor<8x32xi32>) outs(%arg2 : tensor<16x32xf32>) -> tensor<16x32xf32>
return %0 : tensor<16x32xf32>
}
}
```
and converting the above to linalg.generic again
```
./mlir-opt --linalg-generalize-named-ops c2.mlir -o c3.mlir
```
leads to the following
```
#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 @op_matmul_bitcast_int_to_float(%arg0: tensor<16x8xi32>, %arg1: tensor<8x32xi32>, %arg2: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<16x8xi32>, tensor<8x32xi32>) outs(%arg2 : tensor<16x32xf32>) {
^bb0(%in: i32, %in_0: i32, %out: f32):
%1 = arith.sitofp %in : i32 to f32
%2 = arith.sitofp %in_0 : i32 to f32
%3 = arith.mulf %1, %2 : f32
%4 = arith.addf %out, %3 : f32
linalg.yield %4 : f32
} -> tensor<16x32xf32>
return %0 : tensor<16x32xf32>
}
}
```
Notice that the `arith.bitcast` op was not preserved during the roundtrip and we instead see `arith.sitofop`. Ideally, we should not be loosing any information here.
Contributor guide
Research direction
Start by running the provided reproducer with --linalg-specialize-generic-ops and --linalg-generalize-named-ops, then inspect the implementations of those passes and the handling of arith.bitcast. Done means specializing the linalg.generic and generalizing it again preserves arith.bitcast rather than producing arith.sitofp.
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
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100