From 29e3112b9b31afaa420511d81fb4e902da86cdda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= Date: Mon, 22 Sep 2025 15:38:02 +0000 Subject: [PATCH] Epilogue fixes. --- .../ops/epilogue/cshuffle_epilogue.hpp | 103 ++++++++-------- .../transform_conv_bwd_weight_to_gemm.hpp | 116 +++++++++--------- 2 files changed, 116 insertions(+), 103 deletions(-) diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 0dab7fc22d..a0589332f9 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -257,7 +257,14 @@ struct CShuffleEpilogue CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { - return MPerIterationShuffle * NPerIterationShuffle * sizeof(ODataType); + if constexpr(NumGroupsToMerge > 1) + { + return kMPerBlock * kNPerBlock * sizeof(ODataType); + } + else + { + return MPerIterationShuffle * NPerIterationShuffle * sizeof(ODataType); + } } template @@ -267,17 +274,17 @@ struct CShuffleEpilogue void* p_smem) { - if constexpr (NumGroupsToMerge > 1) - { + //if constexpr (NumGroupsToMerge > 1) + //{ // When NumGroupsToMerge > 1, we want to write out only the diagonal blocks. // Hence, we configure the shuffle such that it iterates one merge group block at a time. return merged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem); - } - else - { + //} + //else + //{ // When NumGroupsToMerge == 1, we want to write out all the blocks. - return unmerged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem); - } + // return unmerged_op(out_dram_window, o_acc_tile, ds_dram_windows, p_smem); + //} } template @@ -289,7 +296,11 @@ struct CShuffleEpilogue constexpr auto LdsTileDistr = make_static_tile_distribution(MakeLdsDistributionEncode()); auto lds_tile = make_static_distributed_tensor(LdsTileDistr); - constexpr auto lds_block_desc = MakeLdsBlockDescriptor(); + //constexpr auto lds_block_desc = MakeLdsBlockDescriptor(); + constexpr auto lds_block_desc = make_naive_tensor_descriptor( + make_tuple(number{}, number{}), + make_tuple(number{}, number<1>{})); // Row-major layout + auto o_lds_block = make_tensor_view( static_cast(p_smem), lds_block_desc); @@ -314,21 +325,6 @@ struct CShuffleEpilogue static_assert(std::is_same_v, "Currently, the CShuffle Epilogue only supports the Row Major Output layout"); - using TileEncodingPattern = tile_distribution_encoding_pattern_2d; - constexpr auto dram_tile_distribution = - TileEncodingPattern::make_2d_static_tile_distribution(); - - auto d_dram_windows = generate_tuple( - [&](auto idx) { - return make_tile_window(ds_dram_windows[idx], dram_tile_distribution); - }, - number{}); - constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; @@ -362,27 +358,46 @@ struct CShuffleEpilogue }); block_sync_lds(); - // Copy diagonal block from LDS to global memory. + // Set-up LDS to global memory copy. + // We copy only the diagonal blocks from LDS to global memory. + // Hence, we must configure the tile distrinbution and SFCs to work + // on the group size blocks. constexpr index_t MPerGroup = kMPerBlock / NumGroupsToMerge; constexpr index_t NPerGroup = kNPerBlock / NumGroupsToMerge; - auto out_lds_window = make_tile_window( - o_lds_block, - make_tuple(number{}, number{}), - {0, 0}); - using SFC_lds = space_filling_curve, - sequence<0, 1>, - sequence>; + // TODO: Within the subtile, we have column-major data layout. + constexpr auto dram_tile_encoding = tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<2>>, + tuple, sequence<0>>, + sequence<1, 2>, + sequence<1, 1>>{}; - using SFC_dram = space_filling_curve, - sequence<0, 1>, - sequence>; + constexpr auto dram_tile_distribution = make_static_tile_distribution(dram_tile_encoding); + + auto d_dram_windows = generate_tuple( + [&](auto idx) { + return make_tile_window(ds_dram_windows[idx], dram_tile_distribution); + }, + number{}); static_for<0, NumGroupsToMerge, 1>{} ( [&](auto group) { - auto c_out_tensor = load_tile(make_tile_window(out_lds_window, dram_tile_distribution)); + // Create LDS window at the correct diagonal position for this group + constexpr auto lds_start_m = group * number{}; + constexpr auto lds_start_n = group * number{}; + + auto current_lds_window = make_tile_window( + o_lds_block, + make_tuple(number{}, number{}), + {lds_start_m, lds_start_n}, + dram_tile_distribution); + + block_sync_lds(); + auto c_out_tensor = load_tile(current_lds_window); + //make_tile_window(current_lds_window, dram_tile_distribution)); const auto ds_tensor = generate_tuple( [&](auto idx) { return load_tile(d_dram_windows[idx]); }, number{}); @@ -405,21 +420,13 @@ struct CShuffleEpilogue if constexpr(group != NumGroupsToMerge - 1) { + constexpr auto step_m = number{}; + // Move the global memory window - constexpr auto step_dram = SFC_dram::get_forward_step(group); - - move_tile_window(out_dram_window, {step_dram.at(number<0>{}), step_dram.at(number<1>{})}); - + move_tile_window(out_dram_window, {step_m, 0}); static_for<0, NumDTensor, 1>{}([&](auto idx) { - move_tile_window(d_dram_windows[idx], - {step_dram.at(number<0>{}), step_dram.at(number<1>{})}); + move_tile_window(d_dram_windows[idx], {step_m, 0}); }); - - // Move the LDS window - constexpr auto iAccess = number{}; - constexpr auto next_iAccess = number<(group+1) * NumGroupsToMerge + (group+1)>{}; - constexpr auto step_lds = SFC_lds::get_step_between(iAccess, next_iAccess); - move_tile_window(out_lds_window, {step_lds.at(number<0>{}), step_lds.at(number<1>{})}); } } ); diff --git a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp index 0bd817f73f..db512997d6 100644 --- a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp +++ b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp @@ -418,17 +418,18 @@ struct TransformConvBwdWeightToGemm { // NWGK const index_t NDoHoWoStride = G_ * K_; - const index_t GStride = K_; - constexpr auto KStride = I1; if constexpr (NumGroupsToMerge > 1) { + const index_t KStride = G_; + constexpr auto GStride = I1; return make_naive_tensor_descriptor( - make_tuple(K_, NumGroupsToMerge, N_ * Wo_), - make_tuple(KStride, GStride, NDoHoWoStride)); + make_tuple(NumGroupsToMerge, K_, N_ * Wo_), + make_tuple(GStride, KStride, NDoHoWoStride)); } else { + constexpr auto KStride = I1; return make_naive_tensor_descriptor( make_tuple(K_, N_ * Wo_), make_tuple(KStride, NDoHoWoStride)); @@ -441,17 +442,18 @@ struct TransformConvBwdWeightToGemm // NWGC const index_t NStride = Wi_ * G_ * C_; const index_t WiStride = G_ * C_; - const index_t GStride = C_; - constexpr auto CStride = I1; if constexpr (NumGroupsToMerge > 1) { + const index_t CStride = G_; + constexpr auto GStride = I1; return make_naive_tensor_descriptor( - make_tuple(N_, Wi_, NumGroupsToMerge, C_), + make_tuple(N_, Wi_, C_, NumGroupsToMerge), make_tuple(NStride, WiStride, GStride, CStride)); } else { + constexpr auto CStride = I1; return make_naive_tensor_descriptor( make_tuple(N_, Wi_, C_), make_tuple(NStride, WiStride, CStride)); @@ -488,17 +490,18 @@ struct TransformConvBwdWeightToGemm { // NHWGK const index_t NDoHoWoStride = G_ * K_; - const index_t GStride = K_; - constexpr auto KStride = I1; if constexpr (NumGroupsToMerge > 1) { + const index_t KStride = G_; + constexpr auto GStride = I1; return make_naive_tensor_descriptor( - make_tuple(K_, NumGroupsToMerge, N_ * Ho_ * Wo_), - make_tuple(KStride, GStride, NDoHoWoStride)); + make_tuple(NumGroupsToMerge, K_, N_ * Ho_ * Wo_), + make_tuple(GStride, KStride, NDoHoWoStride)); } else { + constexpr auto KStride = I1; return make_naive_tensor_descriptor( make_tuple(K_, N_ * Ho_ * Wo_), make_tuple(KStride, NDoHoWoStride)); @@ -512,17 +515,18 @@ struct TransformConvBwdWeightToGemm const index_t NStride = Hi_ * Wi_ * G_ * C_; const index_t HiStride = Wi_ * G_ * C_; const index_t WiStride = G_ * C_; - const index_t GStride = C_; - constexpr auto CStride = I1; - + if constexpr (NumGroupsToMerge > 1) { + const index_t CStride = G_; + constexpr auto GStride = I1; return make_naive_tensor_descriptor( - make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_), - make_tuple(NStride, HiStride, WiStride, GStride, CStride)); + make_tuple(N_, Hi_, Wi_, C_, NumGroupsToMerge), + make_tuple(NStride, HiStride, WiStride, CStride, GStride)); } else { + constexpr auto CStride = I1; return make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, C_), make_tuple(NStride, HiStride, WiStride, CStride)); @@ -559,17 +563,18 @@ struct TransformConvBwdWeightToGemm { // NDHWGK const index_t NDoHoWoStride = G_ * K_; - const index_t GStride = K_; - constexpr auto KStride = I1; if constexpr (NumGroupsToMerge > 1) { + const index_t KStride = G_; + constexpr auto GStride = I1; return make_naive_tensor_descriptor( - make_tuple(K_, NumGroupsToMerge, N_ * Do_ * Ho_ * Wo_), - make_tuple(KStride, GStride, NDoHoWoStride)); + make_tuple(NumGroupsToMerge, K_, N_ * Do_ * Ho_ * Wo_), + make_tuple(GStride, KStride, NDoHoWoStride)); } else { + constexpr auto KStride = I1; return make_naive_tensor_descriptor( make_tuple(K_, N_ * Do_ * Ho_ * Wo_), make_tuple(KStride, NDoHoWoStride)); @@ -583,17 +588,18 @@ struct TransformConvBwdWeightToGemm const index_t DiStride = Hi_ * Wi_ * G_ * C_; const index_t HiStride = Wi_ * G_ * C_; const index_t WiStride = G_ * C_; - const index_t GStride = C_; - constexpr auto CStride = I1; - + if constexpr (NumGroupsToMerge > 1) { + constexpr auto GStride = I1; + const index_t CStride = G_; return make_naive_tensor_descriptor( - make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge, C_), - make_tuple(NStride, DiStride, HiStride, WiStride, GStride, CStride)); + make_tuple(N_, Di_, Hi_, Wi_, C_, NumGroupsToMerge), + make_tuple(NStride, DiStride, HiStride, WiStride, CStride, GStride)); } else { + constexpr auto CStride = I1; return make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, C_), make_tuple(NStride, DiStride, HiStride, WiStride, CStride)); @@ -637,12 +643,12 @@ struct TransformConvBwdWeightToGemm { // Output tensor transformation // [0, 1, 2] -> [0, 1] - // [K, Gm, (N*Wo)] -> [(K*Gm), (N*Wo)] + // [Gm, K, (N*Wo)] -> [(Gm*K), (N*Wo)] const auto out_gemm_m_gemm_k_grid_desc = transform_tensor_descriptor( out_grid_desc, make_tuple( - make_merge_transform(make_tuple(K_, NumGroupsToMerge)), + make_merge_transform(make_tuple(NumGroupsToMerge, K_)), make_pass_through_transform(N_ * Wo_)), make_tuple(sequence<0, 1>{}, sequence<2>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -705,38 +711,38 @@ struct TransformConvBwdWeightToGemm // Input tensor transformation, part 1. // [N, Wi, Gm, C] -> [N, (Wi + InLeftPadW + InRightPadW), Gm, C] = [N, Wip, Gm, C] - const auto in_n_wip_gm_c_grid_desc = transform_tensor_descriptor( + const auto in_n_wip_c_gm_grid_desc = transform_tensor_descriptor( in_grid_desc, make_tuple( make_pass_through_transform(N_), make_pad_transform(Wi_, InLeftPadW_, InRightPadW_), - make_pass_through_transform(NumGroupsToMerge), - make_pass_through_transform(C_)), + make_pass_through_transform(C_), + make_pass_through_transform(NumGroupsToMerge)), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{})); // Input tensor transformation, part 2. - // [N, Wip, Gm, C] -> [N, X, Wo, Gm, C] + // [N, Wip, C, Gm] -> [N, X, Wo, C, Gm] const auto in_n_x_wo_gm_c_grid_desc = transform_tensor_descriptor( - in_n_wip_gm_c_grid_desc, + in_n_wip_c_gm_grid_desc, make_tuple( make_pass_through_transform(N_), make_embed_transform( make_tuple(X_, Wo_), make_tuple(ConvDilationW_, ConvStrideW_)), - make_pass_through_transform(NumGroupsToMerge), - make_pass_through_transform(C_)), + make_pass_through_transform(C_), + make_pass_through_transform(NumGroupsToMerge)), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}), make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3>{}, sequence<4>{})); // Input tensor transformation, part 3. - // [0, 1, 2, 3, 4] -> [0, 1] - // [N, X, Wo, Gm, C] -> [(X*Gm*C), (N*Wo)] + // [0, 1, 2, 3, 4] -> [0, 1] + // [N, X, Wo, C, Gm] -> [(Gm*X*C), (N*Wo)] const auto in_gemm_n_gemm_k_grid_desc = transform_tensor_descriptor( in_n_x_wo_gm_c_grid_desc, make_tuple( - make_merge_transform(make_tuple(X_, NumGroupsToMerge, C_)), + make_merge_transform(make_tuple(X_, C_, NumGroupsToMerge)), make_merge_transform(make_tuple(N_, Wo_))), make_tuple(sequence<1, 3, 4>{}, sequence<0, 2>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -788,12 +794,12 @@ struct TransformConvBwdWeightToGemm { // Output tensor transformation // [0, 1, 2] -> [0, 1] - // [K, Gm, (N*Ho*Wo)] -> [(K*Gm), (N*Ho*Wo)] + // [Gm, K, (N*Ho*Wo)] -> [(K*Gm), (N*Ho*Wo)] const auto out_gemm_m_gemm_k_grid_desc = transform_tensor_descriptor( out_grid_desc, make_tuple( - make_merge_transform(make_tuple(K_, NumGroupsToMerge)), + make_merge_transform(make_tuple(NumGroupsToMerge, K_)), make_pass_through_transform(N_ * Ho_ * Wo_)), make_tuple(sequence<0, 1>{}, sequence<2>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -855,22 +861,22 @@ struct TransformConvBwdWeightToGemm make_tuple(sequence<0>{}, sequence<1>{})); // Input tensor transformation, part 1. - // [N, Hi, Wi, Gm, C] -> [N, Hip, Wip, Gm, C] - const auto in_n_hip_wip_gm_c_grid_desc = transform_tensor_descriptor( + // [N, Hi, Wi, C, Gm] -> [N, Hip, Wip, C, Gm] + const auto in_n_hip_wip_c_gm_grid_desc = transform_tensor_descriptor( in_grid_desc, make_tuple( make_pass_through_transform(N_), - make_pad_transform(Wi_, InLeftPadH_, InRightPadH_), + make_pad_transform(Hi_, InLeftPadH_, InRightPadH_), make_pad_transform(Wi_, InLeftPadW_, InRightPadW_), - make_pass_through_transform(NumGroupsToMerge), - make_pass_through_transform(C_)), + make_pass_through_transform(C_), + make_pass_through_transform(NumGroupsToMerge)), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{}), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{})); // Input tensor transformation, part 2. - // [N, Hip, Wip, Gm, C] -> [N, (Y, Wo), (X, Wo), Gm, C] - const auto in_n_y_ho_x_wo_gm_c_grid_desc = transform_tensor_descriptor( - in_n_hip_wip_gm_c_grid_desc, + // [N, Hip, Wip, C, Gm] -> [N, (Y, Wo), (X, Wo), C, Gm] + const auto in_n_y_ho_x_wo_c_gm_grid_desc = transform_tensor_descriptor( + in_n_hip_wip_c_gm_grid_desc, make_tuple( make_pass_through_transform(N_), make_embed_transform( @@ -879,19 +885,19 @@ struct TransformConvBwdWeightToGemm make_embed_transform( make_tuple(X_, Wo_), make_tuple(ConvDilationW_, ConvStrideW_)), - make_pass_through_transform(NumGroupsToMerge), - make_pass_through_transform(C_)), + make_pass_through_transform(C_), + make_pass_through_transform(NumGroupsToMerge)), make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}, sequence<4>{}), make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3, 4>{}, sequence<5>{}, sequence<6>{})); // Input tensor transformation, part 3. - // [0, 1, 2, 3, 4 5 6] -> [0, 1] - // [N, Y, Ho, X, Wo, Gm, C] -> [(Y*X*Gm*C), (N*Ho*Wo)] + // [0, 1, 2, 3, 4 5 6] -> [0, 1] + // [N, Y, Ho, X, Wo, C, Gm] -> [(Gm*Y*X*C), (N*Ho*Wo)] const auto in_gemm_n_gemm_k_grid_desc = transform_tensor_descriptor( - in_n_y_ho_x_wo_gm_c_grid_desc, + in_n_y_ho_x_wo_c_gm_grid_desc, make_tuple( - make_merge_transform(make_tuple(Y_, X_, NumGroupsToMerge, C_)), + make_merge_transform(make_tuple(Y_, X_, C_, NumGroupsToMerge)), make_merge_transform(make_tuple(N_, Ho_, Wo_))), make_tuple(sequence<1, 3, 5, 6>{}, sequence<0, 2, 4>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -942,12 +948,12 @@ struct TransformConvBwdWeightToGemm { // Output tensor transformation // [0, 1, 2] -> [0, 1] - // [K, Gm, (N*Do*Ho*Wo)] -> [(K*Gm), (N*Do*Ho*Wo)] + // [Gm, K, (N*Do*Ho*Wo)] -> [(K*Gm), (N*Do*Ho*Wo)] const auto out_gemm_m_gemm_k_grid_desc = transform_tensor_descriptor( out_grid_desc, make_tuple( - make_merge_transform(make_tuple(K_, NumGroupsToMerge)), + make_merge_transform(make_tuple(NumGroupsToMerge, K_)), make_pass_through_transform(N_ * Do_ * Ho_ * Wo_)), make_tuple(sequence<0, 1>{}, sequence<2>{}), make_tuple(sequence<0>{}, sequence<1>{}));