[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32) - #7830
Merged
AviralGoelAMD merged 4 commits intoJun 2, 2026
Conversation
Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250, enabling bf16-input bf16-output WMMA at the warp GEMM level. - Add WmmaTraits specialization for <gfx125_t, bf16, bf16, bf16, 16,16,32> - Add WarpGemmAttributeWmmaImpl typedef and WarpGemmWmma alias - Add Dispatcher entry for bf16->bf16 16x16x32 - Add warp_gemm test with reference GEMM validation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds CK Tile warp-level WMMA support for the bf16 x bf16 -> bf16 16x16x32 instruction on gfx1250, plumbing it from the low-level trait through the dispatcher and adding a runtime gtest.
Changes:
- Specializes
WmmaTraits<gfx125_t, bf16_t, bf16_t, bf16_t, 16,16,32>to invoke__builtin_amdgcn_wmma_bf16_16x16x32_bf16, plus the matchingWarpGemmAttributeWmmaImpl_*typedef andWarpGemmWmma_bf16_16x16x32_bf16_bf16alias. - Routes
Dispatcher<bf16_t,bf16_t,bf16_t,16,16,32,...>to the new warp-gemm alias under__gfx125__. - Adds a gtest (
test_wmma_bf16_16x16x32_gfx1250.cpp) parameterized over bf16→bf16 and bf16→float that compares the WMMA dispatcher output againstreference_gemm, and a CMake entry gated ongfx1250.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp | New WmmaTraits specialization wrapping the bf16-output builtin under __gfx125__. |
| include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp | New WarpGemmAttributeWmmaImpl_bf16_16x16x32_bf16_bf16 typedef binding the trait. |
| include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp | New WarpGemmWmma_bf16_16x16x32_bf16_bf16 alias built on WarpGemmImpl. |
| include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp | gfx125 dispatcher route for bf16→bf16 16x16x32. |
| test/ck_tile/warp_gemm/test_wmma_bf16_16x16x32_gfx1250.cpp | Typed gtest validating dispatcher vs reference_gemm for bf16→bf16 and bf16→float. |
| test/ck_tile/warp_gemm/CMakeLists.txt | Registers the new test under GPU_TARGETS MATCHES "gfx1250". |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
AviralGoelAMD
marked this pull request as ready for review
May 28, 2026 18:32
JiaLuo-CAN
reviewed
Jun 2, 2026
AviralGoelAMD
deleted the
users/avirgoel/ck/add-bf16-wmma-16x16x32-gfx1250
branch
June 2, 2026 13:52
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Jun 2, 2026
[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32) (#7830) Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250, enabling bf16-input bf16-output WMMA at the warp GEMM level. - Add WmmaTraits specialization for <gfx125_t, bf16, bf16, bf16, 16,16,32> - Add WarpGemmAttributeWmmaImpl typedef and WarpGemmWmma alias - Add Dispatcher entry for bf16->bf16 16x16x32 - Add warp_gemm test with reference GEMM validation ## Motivation <!-- Explain the purpose of this PR and the goals it aims to achieve. --> ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 18, 2026
[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32) Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250, enabling bf16-input bf16-output WMMA at the warp GEMM level. - Add WmmaTraits specialization for <gfx125_t, bf16, bf16, bf16, 16,16,32> - Add WarpGemmAttributeWmmaImpl typedef and WarpGemmWmma alias - Add Dispatcher entry for bf16->bf16 16x16x32 - Add warp_gemm test with reference GEMM validation ## Motivation <!-- Explain the purpose of this PR and the goals it aims to achieve. --> ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
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
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.
Wire __builtin_amdgcn_wmma_bf16_16x16x32_bf16 into CK Tile for gfx1250, enabling bf16-input bf16-output WMMA at the warp GEMM level.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist