iree-org / iree-org/iree

[GPUHeuristics] mi355x GemmSize classification

Open
#23,902 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.