From 0b3c297ee2ddf761dc855914cfc253676bdaa679 Mon Sep 17 00:00:00 2001 From: Johannes Graner <67631091+johannes-graner@users.noreply.github.com> Date: Mon, 8 Jun 2026 08:57:39 +0000 Subject: [PATCH] [rocm-libraries] ROCm/rocm-libraries#8009 (commit 26ab70d) [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) --- .../configs/forward/profiler/nhwgc_bf16.conf | 3 + .../configs/forward/profiler/nhwgc_fp16.conf | 5 ++ .../configs/forward/tests/nhwgc_bf16.conf | 1 + .../configs/forward/tests/nhwgc_fp16.conf | 1 + .../generate_instances.py | 50 ++++++++++++++++- .../grouped_convolution_forward_kernel.hpp | 55 ++++++++++++++++++- 6 files changed, 110 insertions(+), 5 deletions(-) diff --git a/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_bf16.conf b/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_bf16.conf index 66e325510a..aa8f030343 100644 --- a/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_bf16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_bf16.conf @@ -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> diff --git a/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_fp16.conf b/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_fp16.conf index 8448c32117..7ad3448446 100644 --- a/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_fp16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/forward/profiler/nhwgc_fp16.conf @@ -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> diff --git a/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_bf16.conf b/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_bf16.conf index ea2a00e86c..cd17ff7600 100644 --- a/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_bf16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_bf16.conf @@ -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> diff --git a/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_fp16.conf b/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_fp16.conf index cebad8227e..8eb4a4ef7d 100644 --- a/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_fp16.conf +++ b/experimental/grouped_convolution_tile_instances/configs/forward/tests/nhwgc_fp16.conf @@ -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> diff --git a/experimental/grouped_convolution_tile_instances/generate_instances.py b/experimental/grouped_convolution_tile_instances/generate_instances.py index e8898a854e..ba72c58403 100755 --- a/experimental/grouped_convolution_tile_instances/generate_instances.py +++ b/experimental/grouped_convolution_tile_instances/generate_instances.py @@ -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): diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp index af548e5e6f..48979b09a2 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp @@ -572,7 +572,38 @@ struct GroupedConvolutionForwardKernel using GemmDsLayout = remove_cvref_t; 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 + struct has_launch_block_size : std::false_type + { + }; + template + struct has_launch_block_size> : std::true_type + { + }; + static constexpr index_t kBlockSize = []() { + if constexpr(has_launch_block_size::value) + return Pipeline::LaunchBlockSize; + else + return Pipeline::BlockSize; + }(); + + // SFINAE helper: detect Pipeline::IsWavelet (load/math wave specialization). + template + struct has_is_wavelet : std::false_type + { + }; + template + struct has_is_wavelet> : std::true_type + { + }; + static constexpr bool kIsWavelet = []() { + if constexpr(has_is_wavelet::value) + return Pipeline::IsWavelet; + else + return false; + }(); using InDataType = remove_cvref_t; using WeiDataType = remove_cvref_t; @@ -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( + c_ptr, c_desc, block_idx_m, block_idx_n); + + EpiloguePipeline{elfunc} + .template operator()( + 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( c_ptr, c_desc, block_idx_m, block_idx_n);