ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-15 03:30:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#8009 (commit 26ab70d)

Browse Source 
[CK Tile] Add WAVELET pipeline for forward grouped
 convolution (#8009)

## Motivation

CK Tile forward grouped convolution trails classic CK on 3x3
convolutions whose
output-channel count is not divisible by 8, where the narrow output
store limits
the compute CShuffle epilogue. This ports the WAVELET pipeline (added
for
backward-weight in #7937) to the forward kernel to close that gap.

## Technical Details

- Kernel (`grouped_convolution_forward_kernel.hpp`): WAVELET
load/math-wave wiring,
mirroring the backward-weight implementation; the non-WAVELET path is
unchanged.
- Generator: implement `parse_native_fwd_instance`, the forward
native-instance parser.
- Registered WAVELET instances: profiler bf16 3 / fp16 5, tests 1 each.

WAVELET requires input channels divisible by 8 (it does not apply to
depthwise).
The bf16/fp16 instance asymmetry is intentional and measured: the VecC=8
tiles
never beat the compute pool in bf16 but win about 20% of divisible-by-8
3x3 shapes
in fp16, so VecC=8 is registered for fp16 only.

## Test Plan

- Correctness (CPU reference) for every registered profiler instance,
across VecC variants.
- Per-shape best-instance performance sweep over the 34 RetinaNet shapes
(bf16) and
a 200-shape cross-model sweep (bf16 and fp16), compared against classic
CK.

## Test Result

- Correctness: PASS for all instances.
- RetinaNet (bf16, vs classic CK): faster on 28 of 34 shapes, geomean
+19.5%; the
not-divisible-by-8 shapes up to 3.7x. One 1x1 stride-2 shape stays ~20%
behind
  classic CK, unrelated to WAVELET.
- Cross-model (200 shapes): WAVELET wins 3x3 not-divisible-by-8 in both
dtypes
(up to 61% over the next-best compute instance); for divisible-by-8 3x3
it wins
  about 20% of shapes in fp16 (3-11%) and none in bf16.

## Submission Checklist

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

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Johannes Graner
2026-06-08 08:57:39 +00:00
committed by assistant-librarian[bot]
parent b7d59e4b5f
commit 0b3c297ee2
6 changed files with 110 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_bf16.conf
View File

@@ -327,3 +327,6 @@ DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,b
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,32,128,64,8,8,32,32,1,2,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),8,Interwave,v2,bf16,bf16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,256,16,256,64,8,8,16,16,1,4,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,16),4,Interwave,v2,bf16,bf16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,256,32,256,64,8,8,32,32,1,2,Seq(8,32,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,32,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,16),8,Interwave,v2,bf16,bf16,false,1>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,64,64,64,2,2,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,128,128,64,2,2,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,4,1,0,0,64,64,64,2,2,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough>

5
experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_fp16.conf
View File

@@ -315,3 +315,8 @@ DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,f
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,32,128,64,8,8,32,32,1,2,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),8,Interwave,v2,fp16,fp16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,256,16,256,64,8,8,16,16,1,4,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,16),4,Interwave,v2,fp16,fp16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,256,32,256,64,8,8,32,32,1,2,Seq(8,32,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,32,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,16),8,Interwave,v2,fp16,fp16,false,1>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,64,64,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,128,128,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,4,1,0,0,64,64,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,64,64,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,8,1,0,0,128,128,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>

1
experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_bf16.conf
View File

@@ -63,3 +63,4 @@ DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,b
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,256,256,16,64,8,8,16,16,4,1,Seq(8,32,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,32,1,8),2,Interwave,v2,bf16,bf16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,32,16,64,8,8,16,16,1,1,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),2,Interwave,v2,bf16,bf16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,bf16,bf16,fp32,bf16,EmptyTuple,bf16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,32,64,64,8,8,32,32,1,1,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),8,Interwave,v2,bf16,bf16,false,1>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,64,64,64,2,2,1,16,16,32,bf16,bf16,WAVELET,Intrawave,0,1,fp32,bf16,EmptyTuple,PassThrough>

1
experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_fp16.conf
View File

@@ -61,3 +61,4 @@ DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,f
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,128,16,64,8,8,16,16,4,1,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),2,Interwave,v2,fp16,fp16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,64,16,16,64,8,8,16,16,1,1,Seq(8,8,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,8,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,4),4,Interwave,v2,fp16,fp16,false,1>
DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<2,NHWGC,GKYXC,EmptyTuple,NHWGK,fp16,fp16,fp32,fp16,EmptyTuple,fp16,PassThrough,PassThrough,PassThrough,OddC,MNKPadding,128,32,128,64,8,8,32,32,1,2,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,Seq(8,16,1),Seq(1,0,2),Seq(1,0,2),2,8,8,false,1,1,Seq(1,16,1,8),8,Interwave,v2,fp16,fp16,false,1>
GroupedConvolutionForwardKernel<2,Default,NHWGC,GKYXC,EmptyTuple,NHWGK,8,8,1,1,0,0,64,64,64,2,2,1,16,16,32,fp16,fp16,WAVELET,Intrawave,0,1,fp32,fp16,EmptyTuple,PassThrough>

50
experimental/grouped_convolution_tile_instances/generate_instances.py
View File

@@ -347,9 +347,53 @@ def parse_native_bwd_weight_instance(args, instance_id, problem_name):
)
def parse_native_fwd_instance(args, instance_id, problem_name):
"""Parse a native CK Tile forward conv instance string."""
raise NotImplementedError("Native forward instance parsing is not yet implemented.")
def parse_native_fwd_instance(args, instance_id, _):
"""Parse a native CK Tile forward conv instance string
(GroupedConvolutionForwardKernel<...>).
Same field layout as backward_weight (fields 0-30) but with no trailing
StreamK fields. Forward has no two-stage path, so two_stage is always False.
"""
spec = args[1]
tile_size = [int(args[12]), int(args[13]), int(args[14])]
warps = [int(args[15]), int(args[16]), int(args[17])]
warp_tile = [int(args[18]), int(args[19]), int(args[20])]
pipeline_name = args[23]
if pipeline_name not in PIPELINE_NAME_TO_VERSION:
raise RuntimeError(
f"Unknown pipeline name '{pipeline_name}' in native instance {instance_id}"
)
pipeline_version = PIPELINE_NAME_TO_VERSION[pipeline_name]
scheduler = args[24]
double_smem_buffer = int(args[25]) != 0
num_wave_groups = int(args[26])
scalar_per_vector = [int(args[6]), int(args[7]), int(args[8])]
num_groups_to_merge = int(args[9])
split_image = int(args[10]) != 0
explicit_gemm = int(args[11]) != 0
return ConvInstanceTemplateParams(
spec,
tile_size,
warps,
warp_tile,
double_smem_buffer,
num_wave_groups,
False, # forward has no two-stage path
pipeline_version,
scheduler,
scalar_per_vector,
num_groups_to_merge,
split_image,
explicit_gemm,
instance_id,
streamk_enabled=False,
streamk_reduction_strategy=None,
streamk_persistent=False,
)
def parse_native_bwd_data_instance(args, instance_id, problem_name):

55
include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp
View File

@@ -572,7 +572,38 @@ struct GroupedConvolutionForwardKernel
using GemmDsLayout = remove_cvref_t<typename EpiloguePipeline_::DsLayout>;
static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor;
static constexpr index_t kBlockSize = Pipeline::BlockSize;
// For wavelet, LaunchBlockSize > BlockSize (extra load-only waves). Use
// LaunchBlockSize for the kernel launch; non-wavelet pipelines fall back to BlockSize.
template <typename T, typename = void>
struct has_launch_block_size : std::false_type
{
};
template <typename T>
struct has_launch_block_size<T, std::void_t<decltype(T::LaunchBlockSize)>> : std::true_type
{
};
static constexpr index_t kBlockSize = []() {
if constexpr(has_launch_block_size<Pipeline>::value)
return Pipeline::LaunchBlockSize;
else
return Pipeline::BlockSize;
}();
// SFINAE helper: detect Pipeline::IsWavelet (load/math wave specialization).
template <typename T, typename = void>
struct has_is_wavelet : std::false_type
{
};
template <typename T>
struct has_is_wavelet<T, std::void_t<decltype(T::IsWavelet)>> : std::true_type
{
};
static constexpr bool kIsWavelet = []() {
if constexpr(has_is_wavelet<Pipeline>::value)
return Pipeline::IsWavelet;
else
return false;
}();
using InDataType = remove_cvref_t<typename Pipeline::ADataType>;
using WeiDataType = remove_cvref_t<typename Pipeline::BDataType>;
@@ -1345,7 +1376,27 @@ struct GroupedConvolutionForwardKernel
Pipeline{}.template operator()(a_block_window, b_block_window, num_loop, smem_ptr_0);
// Run Epilogue Pipeline with k_batch dispatching
if(k_batch == 1)
if constexpr(kIsWavelet)
{
// Wavelet: only math waves hold accumulators and run the epilogue.
// Load waves run a matching barrier sequence to avoid LDS-sync deadlock.
// Forward has no split-K (IsSplitKSupported == false), so only the
// memory_operation_enum::set path is reachable.
if(Pipeline::IsMathWave())
{
auto c_block_window = MakeCBlockWindow<memory_operation_enum::set>(
c_ptr, c_desc, block_idx_m, block_idx_n);
EpiloguePipeline{elfunc}
.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
c_block_window, c_block_tile, ds_block_window, smem_ptr_0);
}
else
{
EpiloguePipeline::RunBarrierStub();
}
}
else if(k_batch == 1)
{
auto c_block_window = MakeCBlockWindow<memory_operation_enum::set>(
c_ptr, c_desc, block_idx_m, block_idx_n);