feat(ck) [CK] Wavelet gemm pipeline for conv fwd - #10179
Merged
Merged
Conversation
…ces part2 GemmDefault skips dimension padding, causing ~46% wrong values when GEMM M is not a multiple of the tile size (128). All instances in this tuple must use GemmMNKPadding to ensure correctness on arbitrary problem sizes. Co-Authored-By: Claude <noreply@anthropic.com>
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a wave-specialized (“wavelet”) grouped conv forward path in Composable Kernel to reduce VALU pressure by separating address-generation/load work from MFMA math work onto different waves, and wires the new instances into the library plus adds a targeted gtest.
Changes:
- Adds a new wavelet-model XDL+CShuffle V3 grouped conv fwd device implementation and registers new NDHWGC/GKZYXC/NDHWGK FP16 instances.
- Extends the grouped conv3d fwd instance build lists and factory registration to include the new wavelet (and an additional comp “part2”) instance set.
- Adds a new gtest executable covering wavelet conv3d fwd cases and updates test/example CMake accordingly.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| projects/composablekernel/test/grouped_convnd_fwd/test_grouped_convnd_fwd_wavelet.cpp | Adds gtest coverage for wavelet grouped conv3d fwd. |
| projects/composablekernel/test/grouped_convnd_fwd/CMakeLists.txt | Builds/links the new wavelet gtest executable for supported GPU targets. |
| projects/composablekernel/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/ndhwgc/xdl/wavelet/device_grouped_conv3d_fwd_xdl_wavelet_ndhwgc_gkzyxc_ndhwgk_f16_instance.cpp | Defines and registers the concrete wavelet conv3d fwd FP16 instance list. |
| projects/composablekernel/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/ndhwgc/xdl/comp/device_grouped_conv3d_fwd_xdl_ndhwgc_gkzyxc_ndhwgk_f16_comp_part2_instance.cpp | Adds an additional “comp part2” conv3d fwd FP16 instance list. |
| projects/composablekernel/library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd/ndhwgc/CMakeLists.txt | Adds the new wavelet and comp-part2 instance sources to the build. |
| projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward.hpp | Hooks wavelet + comp-part2 instance registration into the instance factory for the relevant type/layout combination. |
| projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_xdl.inc | Declares the new wavelet conv3d fwd instance-adder function. |
| projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_comp_xdl.inc | Declares the new comp-part2 conv3d fwd instance-adder function. |
| projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_comp_instance.hpp | Adds a GEMM specialization constant used for tuning/instance definitions (no functional change observed in diff). |
| projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_wavelet_xdl_instance.hpp | Introduces the wavelet instance tuple definitions used by the new instance registration. |
| projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp | Adds optional DEBUG_LOG diagnostics for validity checks. |
| projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_waveletmodel_cshuffle_v3.hpp | Adds the wavelet-model grouped conv fwd device op implementation (new file). |
| projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp | Moves logging before returning false for MultiABD unsupported path (minor behavior/diagnostic fix). |
| projects/composablekernel/example/09_convnd_fwd/convnd_fwd_wavelet_xdl_fp16.cpp | Adds a wavelet convnd fwd FP16 example using the new device op. |
| projects/composablekernel/example/09_convnd_fwd/CMakeLists.txt | Builds the wavelet example only for a restricted GPU target allowlist. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
bartekxk
approved these changes
Aug 10, 2026
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 10, 2026
feat(ck) [CK] Wavelet gemm pipeline for conv fwd
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
## Motivation
In the current CShuffleV3 conv fwd kernel, the in-kernel conv-to-GEMM
transform generates significant INT32 VALU pressure per MFMA
instruction. On VALU-heavy shapes (e.g., G=1, 3×3, C=256), these index
computation ops compete with MFMA for VALU issue slots, creating a
bottleneck that cannot be resolved by pipeline prefetching alone.
This PR adds a wave-specialized ("wavelet") convolutions forward kernel
that splits workgroup threads into two roles:
- **Load waves**: conv-to-GEMM address computation + global memory loads
+ LDS writes (all VALU/VMEM)
- **Math waves**: LDS reads + MFMA + CShuffle epilogue (no index
computation)
By physically separating the two instruction classes onto different
waves, VALU and MFMA execute on different hardware functional units
without contention.
## Technical Details
**Wave pipeline (modified):**
- `gridwise_gemm_waveletmodel.hpp` — load/math wave pipeline structs
with `sched_group_barrier` scheduling hints to front-load VMEM reads
before address-advance VALU
**Two wave ratios:**
- **(4,4)**: 256 load + 256 math = 512 threads (8 waves). Best on large
shapes.
- **(4,2)**: 256 load + 128 math = 384 threads (6 waves). Best on small
shapes (fewer sync barriers, denser MFMA per math wave).
JIRA ID : ROCM-21620
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 18, 2026
feat(ck) [CK] Wavelet gemm pipeline for conv fwd
## Motivation
In the current CShuffleV3 conv fwd kernel, the in-kernel conv-to-GEMM
transform generates significant INT32 VALU pressure per MFMA
instruction. On VALU-heavy shapes (e.g., G=1, 3×3, C=256), these index
computation ops compete with MFMA for VALU issue slots, creating a
bottleneck that cannot be resolved by pipeline prefetching alone.
This PR adds a wave-specialized ("wavelet") convolutions forward kernel
that splits workgroup threads into two roles:
- **Load waves**: conv-to-GEMM address computation + global memory loads
+ LDS writes (all VALU/VMEM)
- **Math waves**: LDS reads + MFMA + CShuffle epilogue (no index
computation)
By physically separating the two instruction classes onto different
waves, VALU and MFMA execute on different hardware functional units
without contention.
## Technical Details
**Wave pipeline (modified):**
- `gridwise_gemm_waveletmodel.hpp` — load/math wave pipeline structs
with `sched_group_barrier` scheduling hints to front-load VMEM reads
before address-advance VALU
**Two wave ratios:**
- **(4,4)**: 256 load + 256 math = 512 threads (8 waves). Best on large
shapes.
- **(4,2)**: 256 load + 128 math = 384 threads (6 waves). Best on small
shapes (fewer sync barriers, denser MFMA per math wave).
JIRA ID : ROCM-21620
---------
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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.
Motivation
In the current CShuffleV3 conv fwd kernel, the in-kernel conv-to-GEMM transform generates significant INT32 VALU pressure per MFMA instruction. On VALU-heavy shapes (e.g., G=1, 3×3, C=256), these index computation ops compete with MFMA for VALU issue slots, creating a bottleneck that cannot be resolved by pipeline prefetching alone.
This PR adds a wave-specialized ("wavelet") convolutions forward kernel that splits workgroup threads into two roles:
By physically separating the two instruction classes onto different waves, VALU and MFMA execute on different hardware functional units without contention.
Technical Details
Wave pipeline (modified):
gridwise_gemm_waveletmodel.hpp— load/math wave pipeline structs withsched_group_barrierscheduling hints to front-load VMEM reads before address-advance VALUTwo wave ratios:
JIRA ID : ROCM-21620