[Kernel][Perf] Add gfx1151 tile selection for RDNA3 GEMM

[tangzzycc]

Motivation

PR #1028 introduced shape-based tile selection for the RDNA3 GEMM and noted that its gfx1100-specific NUM_CU = 96 thresholds do not transfer directly to gfx11 devices with a different CU count. gfx1151 is an RDNA 3.5 device with 40 physical CUs, and reusing the gfx1100 thresholds leaves repeatable performance gaps for several square GEMM sizes. This PR adds a focused gfx1151 tile-selection profile while preserving the existing gfx1100 behavior.

Technical Details

The existing tile-selection logic is moved to a gfx1100-specific helper without changing its decisions. pick_tile now dispatches by ROCm architecture and uses a gfx1151-specific path on gfx1151 devices.

The gfx1151 path adjusts tile selection only for measured square GEMMs. It uses 32x64x64 for the 256-square case, considers 128x128x32 in CU-relative occupancy regions, selects 256x256x32 only when its grid contains 32--40 workgroups, and retains 64x64x64 for the measured 1024-square long-K case. Shapes outside this scope fall back to the existing gfx1100 heuristic.

The kernel implementation, scheduling options, autotune search space, cache keys, stream handling, and multi-GPU behavior are unchanged. Tests explicitly cover gfx1100 and gfx1151 selection behavior, transition boundaries, held-out shapes, and tile buildability.

This follows the direction discussed in PR #1028: device-specific CU thresholds should not be applied unchanged to other gfx11 parts.

Test Plan

The change was tested on an AMD Radeon 8060S Graphics device (gfx1151, RDNA 3.5, 40 physical CUs) with ROCm 7.14.0, HIP 7.14.60850, and PyTorch 2.12.0+rocm7.14.0.

BF16 and FP16 benchmarks compare the baseline and gfx1151 policies with the same kernel and per-tile options. Measurements use retained JIT launchers, captured graph replay, randomized policy ordering, and two independent order seeds (5713 and 6841). Every timed policy row performs an output correctness check before measurement.

The following validation commands were also run:

python -m compileall kernels/gemm/rdna3_f16_gemm_autotune.py tests/kernels/test_rdna_gemm.py
python -m pytest tests/unit/test_autotune.py -q
python -m pytest tests/kernels/test_rdna_gemm.py -q
git diff --check

Test Result

All 120 timed policy rows passed output correctness checks. Over the 10 shapes whose selected tile changes, BF16 achieved a 1.199x geometric-mean speedup and FP16 achieved a 1.164x geometric-mean speedup.

Shape (M x N x K)	Baseline tile	gfx1151 tile	BF16	FP16
256 x 256 x 256	64x64x64	32x64x64	1.158x	1.048x
768 x 768 x 768	64x64x64	128x128x32	1.325x	1.301x
896 x 896 x 896	128x64x32	128x128x32	1.545x	1.423x
1024 x 1024 x 512	128x64x32	128x128x32	1.104x	1.014x
1024 x 1024 x 1024	128x64x32	64x64x64	1.177x	1.160x
1280 x 1280 x 1280	64x64x64	128x128x32	1.119x	1.127x
1536 x 1536 x 1536	64x64x64	256x256x32	1.115x	1.121x
1792 x 1792 x 1792	64x64x64	128x128x32	1.095x	1.091x
6144 x 6144 x 6144	256x256x32	128x128x32	1.216x	1.210x
8192 x 8192 x 8192	256x256x32	128x128x32	1.204x	1.195x

• python -m compileall: passed
• git diff --check: passed
• tests/unit/test_autotune.py: 42 passed
• tests/kernels/test_rdna_gemm.py: 106 passed, 17 skipped

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Adds gfx1151-specific tile-selection heuristics for RDNA3 GEMM while preserving gfx1100 behavior.

Changes:

• Adds architecture-specific gfx1151 tile selection.
• Extends tile-selection and buildability tests.
• Critical issue remains: architecture detection may select the wrong GPU profile in heterogeneous multi-GPU processes.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
tests/kernels/test_rdna_gemm.py	Adds gfx1100/gfx1151 selection and buildability coverage.
kernels/gemm/rdna3_f16_gemm_autotune.py	Implements architecture-specific dispatch; device-scoped architecture resolution is required.

Suppressed comments (1)

kernels/gemm/rdna3_f16_gemm_autotune.py:145

• Because pick_tile now varies by architecture, the launcher caches need to vary by architecture too. _resolved is keyed only by shape/dtype/rounding/strides, and rdna3_gemm_autotuned consults it before entering the autotuner; after a same-signature call on gfx1100, a call on gfx1151 reuses the first launch function and never evaluates this branch (and _build is likewise not arch-keyed). Include the current device/architecture in the _resolved/_build cache keys, or otherwise prevent cross-architecture reuse.

    arch = str(get_rocm_arch() or "") if arch is None else str(arch)
    if arch.startswith("gfx1151"):

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[coderfeli]

Cover not divisible m shapes for batch oob ?

[tangzzycc]

Thanks for pointing this out. I checked the history, and non-divisible M is an existing limitation of the gfx11 GEMM rather than something introduced by the gfx1151 tile selection. The kernel has required M % BLOCK_M == 0 since #567. #980 explicitly treated small-M/tail support as separate follow-up work, and #1028 retained the same exact-divisibility constraint.

[vivienfanghuagood]

Nice work — the motivation is well argued and the gfx1100 path really is untouched. I verified that on hardware (W7900, gfx1100, FlyDSL 0.3.2): the PR applies cleanly to main, tests/kernels/test_rdna_gemm.py gives 108 passed / 17 skipped, and across 10 shapes the tile chosen before and after the patch is identical. The torch.cuda.current_stream(A.device) fix is a good independent catch.

Three things I'd like to see addressed before merge:

1. The launcher-cache issue is only half fixed. The second commit added arch to _signature/_resolved, but _build's lru_cache still has no arch in its key, and create_wmma_gemm_module resolves the compile target through the process-level cached get_rocm_arch(). So the same shape+tile reuses one compiled module across architectures. Reproduced on gfx1100:

python

f1 = _build(2048, 2048, 2048, "bf16", "bf16", "rn", *TILE_128x128x32, *strides)
dev.get_rocm_arch = lambda: "gfx1151"
f2 = _build(2048, 2048, 2048, "bf16", "bf16", "rn", *TILE_128x128x32, *strides)
f1 is f2   # True,  _build.cache_info() -> currsize=1

In a heterogeneous multi-GPU process that's a wrong code object, not just a suboptimal tile.

2. _device_arch on the hot path costs ~14%. rdna3_gemm_autotuned now calls get_device_properties(...).gcnArchName on every invocation. Measured on W7900: 5.82 µs per call; a 1024³ GEMM goes 48.44 → 55.09 µs end to end (+6.66 µs, +13.7%). Since arch is fixed for the process lifetime, an lru_cache keyed on device.index would remove this.

3. The M == N gate leaves equivalent shapes behind. Shapes like 4096×512×4096, 512×4096×4096 and 1024×2048×2048 land inside the same wgs_128 occupancy windows the new path keys on, but fall back to the gfx1100 thresholds purely because they aren't square. Your own argument — that gfx1100 CU thresholds don't transfer to a 40-CU part — applies to those just as much. Either widen the predicate or note explicitly in the docstring that coverage is limited to the measured square set.

[tangzzycc]

Thanks for the detailed review and the gfx1100 hardware validation. I have addressed the three points as follows:

1. Architecture-specific launcher caching and code generation

   The device architecture is now included in both the resolved-launcher signature and the _build cache key. It is passed through create_wmma_gemm_module into a new optional, keyword-only @flyc.jit(arch=...) argument.

   The JIT now uses the same frozen backend target for cache identity and GPU module construction. This prevents a launcher compiled for one gfx11 architecture from being reused as a code object for another architecture. Existing @flyc.jit usage remains unchanged when arch is not provided.

2. Device architecture query overhead

   _device_arch is now cached with functools.lru_cache(maxsize=None), keyed by device.index. On gfx1151, the cached lookup measured approximately 0.021 us per call, compared with approximately 0.58 us for querying the PyTorch device properties directly.

3. Rectangular gfx1151 shapes

   I benchmarked the three shapes mentioned in the review, their transposed counterpart, and several neighboring shapes in both BF16 and FP16.

   The workgroup count alone was not sufficient for the strongly skewed 8:1 shapes. With the production tile options, 512x4096x4096 and 4096x512x4096 did not benefit from switching to the 256x256 tile, so they retain the existing 64x64 fallback.

   The 2:1 near-square shapes consistently benefited from the 256x256 tile:

   • 1024x2048x2048: 9.1-10.2% faster in BF16 and 5.7-7.2% faster in FP16.
   • 2048x1024x2048: 5.4-6.4% faster in BF16 and 7.1-7.7% faster in FP16.

   Based on these measurements, the gfx1151 heuristic now extends only the existing 32-40-workgroup 256x256 occupancy band to shapes with an aspect ratio of at most 2. Strongly skewed rectangular shapes continue to use the existing fallback. This avoids applying the square-shape result to a much larger unmeasured space.

Validation:

• tests/kernels/test_rdna_gemm.py: 114 passed, 17 skipped
• Related JIT target/cache/backend tests: 24 passed
• All benchmark correctness checks passed on gfx1151 with ROCm 7.14
• The gfx1100 tile-selection expectations remain unchanged

The follow-up commits are:

• 79aa220c — [Fix] Build RDNA3 GEMM launchers for the selected architecture
• 9794b02f — [Kernel][Perf] Extend gfx1151 tile selection to near-square GEMMs