[Dialect][Kernel][Perf] Add gfx120x.wmma atom and port the RDNA4 GEMM to the layout API - #943
Open
vlluvia wants to merge 1 commit into
Open
[Dialect][Kernel][Perf] Add gfx120x.wmma atom and port the RDNA4 GEMM to the layout API#943vlluvia wants to merge 1 commit into
vlluvia wants to merge 1 commit into
Conversation
sjfeng1999
reviewed
Jul 31, 2026
| // gfx1250 is a different target: it has 16x16x32 and the mods/reuse operand | ||
| // form, so it keeps its own `gfx1250.wmma` atom. | ||
|
|
||
| def FlyROCDL_MmaOpGFX12_WMMA : FlyROCDL_MmaOp<"MmaOpGFX12_WMMA", "gfx12.wmma", []> { |
Collaborator
There was a problem hiding this comment.
If RDNA4 includes only GFX1200 and GFX1201. GFX120X_WMMA would be a better name because it distinguishes cleanly from GFX1250, The same applies elsewhere.
Author
There was a problem hiding this comment.
Renamed to GFX120X throughout — type, mnemonic, directory, test files and diagnostics. Also narrowed the dispatch guard in fx.rocdl.WMMA() from gfx12 to gfx120 for the same reason, so it no longer overlaps gfx1250.
sjfeng1999
reviewed
Jul 31, 2026
Comment on lines
+268
to
+270
| arg_a_2d = fx.Tensor(fx.make_view(fx.get_iter(arg_a), fx.make_layout((M, K), (K, 1)))) | ||
| arg_bt_2d = fx.Tensor(fx.make_view(fx.get_iter(arg_bt), fx.make_layout((N, K), (K, 1)))) | ||
| arg_c_2d = fx.Tensor(fx.make_view(fx.get_iter(arg_c), fx.make_layout((M, N), (N, 1)))) |
Collaborator
There was a problem hiding this comment.
The fx.Tensor wrapper is unnecessary actually here .
sjfeng1999
requested changes
Jul 31, 2026
sjfeng1999
left a comment
sjfeng1999
left a comment
Collaborator
There was a problem hiding this comment.
Everything else looks great.
RDNA4 (gfx1200 / gfx1201) has no usable WMMA atom today: fx.rocdl.WMMA()
routed every gfx12* target to MmaOpGFX1250_WMMAType, whose verifier only
accepts K=32 on the f16/bf16 paths, while RDNA4 only has the 16x16x16
shapes. The RDNA4 GEMM therefore had to spell the v8 fragment layout out as
hand-written lane arithmetic instead of deriving it from an atom.
Add a dedicated gfx120x.wmma atom. RDNA4 sits between the two existing WMMA
atoms: it keeps the gfx11 16x16x16 instruction shapes and the 3-operand
ROCDL intrinsic form, but uses the gfx1250 v8 register ABI, so the fragment
layouts are shared with the gfx1250 helpers and only codegen is new. The name
is gfx120x rather than gfx12 because gfx1250 is also a gfx12* target but keeps
its own atom; fx.rocdl.WMMA() matches RDNA4 on the correspondingly narrowed
gfx120 prefix, which is disjoint from gfx1250.
With the atom in place, rewrite kernels/gemm/rdna_f16_gemm.py against
make_tiled_copy_{A,B,C} and the tiled_mma thread slices (348 -> 281 lines).
Selecting an LDS partition needs a compile-time constant stage index, so the
k-loop now processes two tiles per iteration with the odd/even remainder
peeled.
Output is bit-identical to the previous kernel (torch.equal) at 2048^3,
4096^3 and 8192^3 for both bf16 and f16, and it is 5-10% faster (geomean
1.08x on gfx1201). The win is not from doing less work: per k-tile the ISA
still has 32 v_wmma, 16 ds_load_b128, 8 ds_store_b128 and 8
buffer_load_b128. Making the stage index constant folds the LDS addresses
into ds_load immediate offsets, which cuts per-k-tile address VALU from 22
to 3 ops and frees enough registers to keep 24 fragment quads live instead
of 6. The old loop had to drain LDS (s_wait_dscnt 0x0) 9 times per k-tile to
recycle those registers; the new one drains 1.5 times and issues runs of up
to 16 WMMA with no wait in between. VGPR use rises 189 -> 247, which costs
nothing here because the 40KB LDS footprint already caps the CU at one
workgroup.
Signed-off-by: vlluvia <925716329@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
sjfeng1999
approved these changes
Jul 31, 2026
vlluvia
changed the title
vivienfanghuagood
approved these changes
Jul 31, 2026
Collaborator
|
good job! |
coderfeli
approved these changes
Jul 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RDNA4 (gfx1200 / gfx1201) has no usable WMMA atom today: fx.rocdl.WMMA() routed every gfx12* target to MmaOpGFX1250_WMMAType, whose verifier only accepts K=32 on the f16/bf16 paths, while RDNA4 only has the 16x16x16 shapes. The RDNA4 GEMM therefore had to spell the v8 fragment layout out as hand-written lane arithmetic instead of deriving it from an atom.
Add a dedicated gfx12.wmma atom. RDNA4 sits between the two existing WMMA atoms: it keeps the gfx11 16x16x16 instruction shapes and the 3-operand ROCDL intrinsic form, but uses the gfx1250 v8 register ABI, so the fragment layouts are shared with the gfx1250 helpers and only codegen is new. Fix the dispatch order in fx.rocdl.WMMA() accordingly -- gfx1250 must be matched before the gfx12 prefix would otherwise claim it.
With the atom in place, rewrite kernels/gemm/rdna_f16_gemm.py against make_tiled_copy_{A,B,C} and the tiled_mma thread slices (348 -> 281 lines). Selecting an LDS partition needs a compile-time constant stage index, so the k-loop now processes two tiles per iteration with the odd/even remainder peeled.
Output is bit-identical to the previous kernel (torch.equal) at 2048^3, 4096^3 and 8192^3 for both bf16 and f16, and it is 5-10% faster (geomean 1.08x on gfx1201). The win is not from doing less work: per k-tile the ISA still has 32 v_wmma, 16 ds_load_b128, 8 ds_store_b128 and 8 buffer_load_b128. Making the stage index constant folds the LDS addresses into ds_load immediate offsets, which cuts per-k-tile address VALU from 22 to 3 ops and frees enough registers to keep 24 fragment quads live instead of 6. The old loop had to drain LDS (s_wait_dscnt 0x0) 9 times per k-tile to recycle those registers; the new one drains 1.5 times and issues runs of up to 16 WMMA with no wait in between. VGPR use rises 189 -> 247, which costs nothing here because the 40KB LDS footprint already caps the CU at one workgroup.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist