iree-org / iree-org/iree

[SPIRV] `linalg.matmul_transpose_b` fails to vectorize

Open
#14,701 0 comments 0 reactions 0 assignees View on GitHub
bug 🐞
Dominant language
C++
Stars
3.9k
Forks
1k
Avg merge
4d 16h
Merged PRs (30d)
47

Description

### What happened?

When compiling the following example IR:
```
module {
func.func @transpose_matmul(%arg0: tensor<1x4096xf32>, %arg1: tensor<32000x4096xf32>) -> tensor<1x32000xf32> {
%cst = arith.constant 0.000000e+00 : f32
%0 = tensor.empty() : tensor<1x32000xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1x32000xf32>) -> tensor<1x32000xf32>
%2 = linalg.matmul_transpose_b ins(%arg0, %arg1 : tensor<1x4096xf32>, tensor<32000x4096xf32>) outs(%1 : tensor<1x32000xf32>) -> tensor<1x32000xf32>
return %2 : tensor<1x32000xf32>
}
}
```
the `linalg.matmul_transpose_b` op fails to vectorize, and the op remains after the `LinalgVectorizationPattern`, and you get this error:
```
transpose_matmul/sources/module_transpose_matmul_dispatch_0.mlir:19:8: error: 'linalg.matmul_transpose_b' op should not remain after vectorization
%7 = linalg.matmul_transpose_b ins(%3, %4 : tensor<1x4096xf32>, tensor<32000x4096xf32>) outs(%6 : tensor<1x32000xf32>) -> tensor<1x32000xf32>
^
```

### Steps to reproduce your issue

1. Compiled on https://github.com/openxla/iree/commit/dce0e41a8c4d499b01f8f068aef47052da7126a1
2. Compile command:
```
iree-compile matmul_transpose_b.mlir \
--iree-input-type=none \
--iree-vm-bytecode-module-output-format=flatbuffer-binary \
--iree-hal-target-backends=vulkan \
--mlir-print-op-on-diagnostic=false \
--iree-llvmcpu-target-cpu-features=host \
--iree-stream-resource-index-bits=64 \
--iree-vm-target-index-bits=64 \
--iree-vm-bytecode-module-strip-source-map=true \
--iree-util-zero-fill-elided-attrs \
--iree-vm-target-truncate-unsupported-floats \
--iree-codegen-check-ir-before-llvm-conversion=false \
--iree-vulkan-target-triple=rdna3-unknown-linux \
--iree-opt-const-expr-hoisting=false \
-o transpose_matmul/transpose_matmul.vmfb \
```

### What component(s) does this issue relate to?

_No response_

### Version information

_No response_

### Additional context

This is happening because the `linalg.matmul_transpose_b` op is going down the `SPIRVBaseVectorize` pipeline instead of the `SPIRVMatmulPromoteVectorize` pipeline, causing it to fail to vectorize.
This IR:
```
module {
func.func @transpose_matmul(%arg0: tensor<2x4096xf32>, %arg1: tensor<32000x4096xf32>) -> tensor<2x32000xf32> {
%cst = arith.constant 0.000000e+00 : f32
%0 = tensor.empty() : tensor<2x32000xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<2x32000xf32>) -> tensor<2x32000xf32>
%2 = linalg.matmul_transpose_b ins(%arg0, %arg1 : tensor<2x4096xf32>, tensor<32000x4096xf32>) outs(%1 : tensor<2x32000xf32>) -> tensor<2x32000xf32>
return %2 : tensor<2x32000xf32>
}
}
```
where the unit dim in `%arg0` is changed from 1 to 2 compiles e2e with the `SPIRVMatmulPromoteVectorize` pipeline. The reason the first IR fails to select the `SPIRVMatmulPromoteVectorize` pipeline seems to be because of the unit dim on the shape of `%arg0`, making the op effectively a matvec transpose, which is not supported by the current `KernelConfig`. Support for cases like this will probably need to be added in order for this to lower properly.

The config setting fails here https://github.com/openxla/iree/blob/dce0e41a8c4d499b01f8f068aef47052da7126a1/compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp#L501 where the subgroup size equals the total number of threads.
Debug output:
```
//===-------------------------------------------===//
trying to matmul tensorcore config...
(B, M, N, K) indices = (-1, 0, 1, 2)
input matmul shape (B, M, N, K) = (-1, 1, 32000, 4096)
trying to deduce config as matmul...
(B, M, N, K) indices = (-1, 0, 1, 2)
best thread tile size (M, N, K) = (8, 4, 16)
best workgroup size (X, Y) = (256, 1)
workgroup tile size before promotion = (1, 256, 0)
reduction tile size before promotion = (0, 0, 16)
workgroup size before promotion = (64, 1, 1)
subgroup size = 64
initial total thread = 64
```

Contributor guide

Open the contributing guide

Research direction

Start in compiler/src/iree/compiler/Codegen/SPIRV/KernelConfig.cpp around line 501 and trace why the unit-dimension matmul_transpose_b case is not selected for SPIRVMatmulPromoteVectorize. Use the supplied MLIR and iree-compile command to reproduce it; done means the 1x4096 transpose matmul selects a supported pipeline and no linalg.matmul_transpose_b remains after vectorization.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.