From fca03f72fc7cfd5e4b2a6daeee6ccf034b63a078 Mon Sep 17 00:00:00 2001 From: Qun Lin Date: Thu, 5 Jun 2025 12:58:37 +0800 Subject: [PATCH] support k-split --- example/20_grouped_conv_bwd_weight/common.hpp | 6 +- .../device_grouped_conv_bwd_weight_dl_v4.hpp | 145 ++++++++++++++---- .../run_grouped_conv_bwd_weight_example.inc | 2 +- 3 files changed, 117 insertions(+), 36 deletions(-) diff --git a/example/20_grouped_conv_bwd_weight/common.hpp b/example/20_grouped_conv_bwd_weight/common.hpp index a7cdacd0a3..c56fb32e23 100644 --- a/example/20_grouped_conv_bwd_weight/common.hpp +++ b/example/20_grouped_conv_bwd_weight/common.hpp @@ -74,6 +74,7 @@ struct ExecutionConfig final bool do_verification = true; int init_method = 1; bool time_kernel = false; + int split_k = 1; }; #define DefaultConvParam \ @@ -87,6 +88,7 @@ inline void print_help_msg() std::cerr << "arg1: verification (0=no, 1=yes)\n" << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n" << "arg3: time kernel (0=no, 1=yes)\n" + << "arg4: split k\n" << ck::utils::conv::get_conv_param_parser_helper_msg() << std::endl; } @@ -96,7 +98,7 @@ inline bool parse_cmd_args(int argc, ck::utils::conv::ConvParam& conv_param) { constexpr int num_execution_config_args = - 3; // arguments for do_verification, init_method, time_kernel + 4; // arguments for do_verification, init_method, time_kernel, split_k constexpr int num_conv_param_leading_args = 5; // arguments for num_dim_spatial_, G_, N_, K_, C_ constexpr int threshold_to_catch_partial_args = 1 + num_execution_config_args; @@ -113,6 +115,7 @@ inline bool parse_cmd_args(int argc, config.do_verification = std::stoi(argv[1]); config.init_method = std::stoi(argv[2]); config.time_kernel = std::stoi(argv[3]); + config.split_k = std::stoi(argv[4]); } // catch both ExecutionConfig & ConvParam arguments else if(threshold_to_catch_all_args < argc && ((argc - threshold_to_catch_all_args) % 3 == 0)) @@ -120,6 +123,7 @@ inline bool parse_cmd_args(int argc, config.do_verification = std::stoi(argv[1]); config.init_method = std::stoi(argv[2]); config.time_kernel = std::stoi(argv[3]); + config.split_k = std::stoi(argv[4]); const ck::index_t num_dim_spatial = std::stoi(argv[4]); conv_param = ck::utils::conv::parse_conv_param( diff --git a/example/20_grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_dl_v4.hpp b/example/20_grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_dl_v4.hpp index c6a6eed9d8..44ab4a9b14 100644 --- a/example/20_grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_dl_v4.hpp +++ b/example/20_grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_dl_v4.hpp @@ -33,6 +33,30 @@ __device__ T warp_shuffle_down(const T& v_local, uint32_t lane_delta) return bit_cast(v_remote_tmp); #endif } +template void __device__ global_atomic_add(T* p,T val); + +template <> void __device__ global_atomic_add(float* p, float v) +{ + __builtin_amdgcn_global_atomic_fadd_f32(p, v); +} +template <> void __device__ global_atomic_add(half2_t* p, half2_t v) +{ +#if defined(__gfx942__) || defined(__gfx950__) + __builtin_amdgcn_global_atomic_fadd_v2f16(p, v); +#else + ignore = p; + ignore = v; +#endif +} +template <> void __device__ global_atomic_add(bhalf2_t* p, bhalf2_t v) +{ +#if defined(__gfx942__) || defined(__gfx950__) + __builtin_amdgcn_global_atomic_fadd_v2bf16(p, v); +#else + ignore = p; + ignore = v; +#endif +} template __global__ void @@ -369,18 +393,51 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 } } } + template static void __device__ - write_output(const Argument& arg, index_t g, index_t y, index_t x, WeiDataType acc) + write_output(const Argument& arg, index_t g, index_t y, index_t x, float acc) { const index_t Wei_G_Stride = arg.wei_g_k_c_xs_strides_[0]; const index_t Y_Stride = arg.wei_g_k_c_xs_strides_[3]; const index_t X_Stride = arg.wei_g_k_c_xs_strides_[4]; - if(y < Filter_Y && x < Filter_X) { auto p_wei = arg.p_wei_grid_ + Wei_G_Stride * g + y * Y_Stride + x * X_Stride; - *p_wei = acc; + if(arg.enable_k_split_) + { + if constexpr(is_same::value || + is_same::value) + { + typename vector_type::type pk_acc = {}; + + float acc2 = warp_shuffle_down(acc, 1); + + pk_acc[0] = type_convert(acc); + pk_acc[1] = type_convert(acc2); + if(x % 2 == 0) + { + if(x == Filter_X - 1) + { + pk_acc[1] = pk_acc[0]; + pk_acc[0] = 0; + p_wei--; + } + global_atomic_add( + reinterpret_cast::type*>(p_wei), + pk_acc); + } + } + else + { + global_atomic_add(p_wei, acc); + } + } + else + { + + *p_wei = type_convert(acc); + } } } template @@ -407,12 +464,13 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 static void __device__ Run(Argument arg, char* p_share_in, char* p_share_out) { const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.x); + const index_t k_split_idx = __builtin_amdgcn_readfirstlane(blockIdx.y); const index_t wave_id = __builtin_amdgcn_readfirstlane(threadIdx.x / WaveSize); const index_t tile_id = wave_id / NumWavePerTile; const index_t lane_id = __lane_id(); constexpr index_t ThreadPerBatch = WaveSize / BatchPerWave; - // Debug> xx3; + static_assert(Tile_H % NumWavePerTile == 0); static_assert(TileOut_H % NumWavePerTile == 0); InDataVector tmp_in[math::integer_divide_ceil(TileIn_Pack_H, NumWavePerTile) * @@ -434,6 +492,7 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 num_loop = (n - n_idx) / NBatch - 1; } } + n_idx += n * k_split_idx; // In const index_t hi = arg.in_g_n_c_wis_lengths_[spatial_offset + 0]; @@ -537,6 +596,27 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 const index_t out_x = lane_id % TileOut_Pack_W; const index_t out_y_offset = lane_id / TileOut_Pack_W; + auto share_in_base = share_in; + auto share_out_base = share_out; + // adjust share memory offset for copy + if constexpr(NumWavePerTile > 1) + { + static_assert(RequirePadding == false); + share_in += Copy_Tile_H * TileIn_Align_W * NumVectorPerPixel * wave_id; + share_out += Copy_TileOut_H * TileOut_Align_W * NumVectorPerPixel * wave_id; + } + share_in += (TileIn_Align_W * Pad_H + Pad_W) * NumVectorPerPixel; + + constexpr index_t TileOut_H_batch = math::integer_divide_ceil(Copy_TileOut_H, BatchPerWave); + index_t hout_base = lane_id / ThreadPerBatch * TileOut_H_batch; + if constexpr(NumWavePerTile > 1) + { + hout_base += Copy_TileOut_H * wave_id; + } + index_t x = (lane_id % ThreadPerBatch) % Filter_X; + index_t y = (lane_id % ThreadPerBatch) / Filter_X; + float acc = 0; + // prefetch 0 if(in_x < (Tile_W / InScalarPerVector)) { @@ -551,17 +631,6 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 p_in += NBatch * in_n_stride; p_out += NBatch * out_n_stride; - auto share_in_base = share_in; - auto share_out_base = share_out; - // adjust share memory offset for copy - if constexpr(NumWavePerTile > 1) - { - static_assert(RequirePadding == false); - share_in += Copy_Tile_H * TileIn_Align_W * NumVectorPerPixel * wave_id; - share_out += Copy_TileOut_H * TileOut_Align_W * NumVectorPerPixel * wave_id; - } - - share_in += (TileIn_Align_W * Pad_H + Pad_W) * NumVectorPerPixel; if(in_x < (Tile_W / InScalarPerVector)) { write_data_to_lds( @@ -573,16 +642,6 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 out_x, out_y_offset, tmp_out, share_out); } - constexpr index_t TileOut_H_batch = math::integer_divide_ceil(Copy_TileOut_H, BatchPerWave); - index_t hout_base = lane_id / ThreadPerBatch * TileOut_H_batch; - if constexpr(NumWavePerTile > 1) - { - hout_base += Copy_TileOut_H * wave_id; - } - index_t x = (lane_id % ThreadPerBatch) % Filter_X; - index_t y = (lane_id % ThreadPerBatch) / Filter_X; - float acc = 0; - // if (lane_id == 0) //{ // dump_lds(tmp_in, sizeof(tmp_in)/sizeof(InDataVector), @@ -720,7 +779,7 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 { if(hout_base == 0) { - write_output(arg, g_idx, y, x, acc); + write_output(arg, g_idx, y, x,acc); } } else @@ -735,7 +794,7 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 { acc += bit_cast(p_share_acc[i * WaveSize + lane_id]); } - write_output(arg, g_idx, y, x, type_convert(acc)); + write_output(arg, g_idx, y, x, acc); } } } @@ -749,7 +808,8 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 const std::array& wei_g_k_c_xs_lengths, // weight const std::array& wei_g_k_c_xs_strides, const std::array& out_g_n_k_wos_lengths, // output - const std::array& out_g_n_k_wos_strides) + const std::array& out_g_n_k_wos_strides, + index_t k_batch) : p_in_grid_{p_in_grid}, p_wei_grid_{p_wei_grid}, p_out_grid_{p_out_grid}, @@ -758,8 +818,14 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 wei_g_k_c_xs_lengths_(wei_g_k_c_xs_lengths), wei_g_k_c_xs_strides_(wei_g_k_c_xs_strides), out_g_n_k_wos_lengths_(out_g_n_k_wos_lengths), - out_g_n_k_wos_strides_(out_g_n_k_wos_strides) + out_g_n_k_wos_strides_(out_g_n_k_wos_strides), + enable_k_split_(k_batch > 1) { + if (enable_k_split_) + { + in_g_n_c_wis_lengths_[1] /= k_batch; + out_g_n_k_wos_lengths_[1] /= k_batch; + } } std::size_t GetWorkspaceSizeBytes() const { return 0; } @@ -774,6 +840,7 @@ struct GridwiseGroupedConv2DBwdWeightDlV4 std::array wei_g_k_c_xs_strides_; std::array out_g_n_k_wos_lengths_; std::array out_g_n_k_wos_strides_; + bool enable_k_split_; }; }; @@ -909,11 +976,11 @@ struct DeviceGroupedConvBwdWeightDlV4 : public DeviceGroupedConvBwdWeight (32 * 1024) ? 1 : 2; @@ -932,9 +1000,13 @@ struct DeviceGroupedConvBwdWeightDlV4 : public DeviceGroupedConvBwdWeight; - + const index_t out_buf_size = + arg.wei_g_k_c_xs_lengths_[0] * arg.wei_g_k_c_xs_strides_[0] * sizeof(WeiDataType); + if(arg.k_batch_ > 1) + hipGetErrorString( + hipMemsetAsync(arg.p_wei_grid_, 0, out_buf_size, stream_config.stream_id_)); ave_time += launch_and_time_kernel( - stream_config, kernel, dim3(gdx), dim3(BlockSize), 0, conv_arg); + stream_config, kernel, gdx, dim3(BlockSize), 0, conv_arg); return ave_time; } @@ -959,6 +1031,7 @@ struct DeviceGroupedConvBwdWeightDlV4 : public DeviceGroupedConvBwdWeight 1 - constexpr ck::index_t split_k = 16; + const ck::index_t split_k = config.split_k; const auto in_g_n_c_wis_desc = ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<