Skip to content

[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32) - #7830

Merged
AviralGoelAMD merged 4 commits into
developfrom
users/avirgoel/ck/add-bf16-wmma-16x16x32-gfx1250
Jun 2, 2026
Merged

[CK_Tile][MI450] Add bf16 output wmma instruction (16x16x32)#7830
AviralGoelAMD merged 4 commits into
developfrom
users/avirgoel/ck/add-bf16-wmma-16x16x32-gfx1250

Conversation

@AviralGoelAMD

Copy link
Copy Markdown
Contributor

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

Technical Details

Test Plan

Test Result

Submission Checklist

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 matching WarpGemmAttributeWmmaImpl_* typedef and WarpGemmWmma_bf16_16x16x32_bf16_bf16 alias.
  • 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 against reference_gemm, and a CMake entry gated on gfx1250.

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
AviralGoelAMD marked this pull request as ready for review May 28, 2026 18:32
@AviralGoelAMD
AviralGoelAMD requested a review from a team as a code owner May 28, 2026 18:32

@geyyer geyyer left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

Comment thread projects/composablekernel/test/ck_tile/warp_gemm/CMakeLists.txt
@AviralGoelAMD
AviralGoelAMD merged commit 590fe58 into develop Jun 2, 2026
34 checks passed
@AviralGoelAMD
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants