[GPUHeuristics] mi355x GemmSize classification
- Dominant language
- C++
- Stars
- 3.9k
- Forks
- 1k
- Avg merge
- 4d 16h
- Merged PRs (30d)
- 47
Description
## Description
For 2048x2048x2048 BF16 matmul on mi355x, using `--iree-rocm-target=mi355x` produces a slower config than `--iree-rocm-target=gfx950`, which is unexpected as mi355x is a more specific target and should deliver equal or better performance.
The 2048x2048x2048 matmul has compute intensity (CI) = 1365.3. The mi355x target has chip-specific perf/bandwidth data (`perfTflops=2500, memBW=8.0 Tbps`), giving a `largeGemmCutoff = 5 * (2500/8) = 1562.5`. Since CI < 1562.5, it is classified as **MediumGemm**. The gfx950 target has no chip details and uses defaults (`largeGemmCutoff = 1000`), so CI > 1000 classifies it as **LargeGemm**.
MediumGemm uses `bestKTileCountPerSubgroup=4` while LargeGemm uses `bestKTileCountPerSubgroup=2`, resulting in different heuristic configs.
**Input mlir:**
```mlir
func.func @mm_2048(%lhs: tensor<2048x2048xbf16>, %rhs: tensor<2048x2048xbf16>) -> tensor<2048x2048xf32> {
%cst = arith.constant 0.0 : f32
%init = tensor.empty() : tensor<2048x2048xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<2048x2048xf32>) -> tensor<2048x2048xf32>
%result = linalg.matmul ins(%lhs, %rhs : tensor<2048x2048xbf16>, tensor<2048x2048xbf16>) outs(%fill : tensor<2048x2048xf32>) -> tensor<2048x2048xf32>
return %result : tensor<2048x2048xf32>
}
```
**Steps:**
```bash
iree-compile mm_2048.mlir --iree-hal-target-device=hip --iree-rocm-target=mi355x \
--iree-llvmgpu-use-direct-load -o mm_2048_mi355x.vmfb
iree-compile mm_2048.mlir --iree-hal-target-device=hip --iree-rocm-target=gfx950 \
--iree-llvmgpu-use-direct-load -o mm_2048_gfx950.vmfb
iree-benchmark-module --module= --device=hip \
--function=mm_2048 --input=2048x2048xbf16 --input=2048x2048xbf16
```
**heuristic-gfx950 (LargeGemm, tile_k=1):**
`lowering_config = #iree_gpu.lowering_config<{..., reduction = [0, 0, 1], subgroup = [2, 2, 0], workgroup = [128, 128, 0]}>`
Time: **86 us**
**heuristic-mi355x (MediumGemm, tile_k=4):**
`lowering_config = #iree_gpu.lowering_config<{..., reduction = [0, 0, 4], subgroup = [2, 2, 0], workgroup = [64, 128, 0]}>`
Time: **117 us** (1.36x slower)
Contributor guide
Research direction
Start with the GPU heuristics involved in iree-compile and reproduce the two commands using the supplied mm_2048.mlir input. Compare the mi355x and gfx950 GemmSize classifications, lowering configs, and benchmark results; done means the mi355x configuration no longer regresses against gfx950 for this case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100