diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp index 274fec72c1..e695ae00a5 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp @@ -203,10 +203,10 @@ int run(const std::string& in_layout, int argc, char* argv[]) { - // if (num_groups_to_merge == 1) - // { - // return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); - // } + if (num_groups_to_merge == 1) + { + return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); + } // else if (num_groups_to_merge == 2) // { // return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); @@ -232,7 +232,7 @@ int run(const std::string& in_layout, // return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); // } - if (num_groups_to_merge == 8) + else if (num_groups_to_merge == 8) { return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); } diff --git a/include/ck_tile/core/algorithm/static_encoding_pattern.hpp b/include/ck_tile/core/algorithm/static_encoding_pattern.hpp index 29aa2fca5d..913e60d9a3 100644 --- a/include/ck_tile/core/algorithm/static_encoding_pattern.hpp +++ b/include/ck_tile/core/algorithm/static_encoding_pattern.hpp @@ -25,7 +25,7 @@ * (3) number of iterations to cover the entire Y axis. * The raked here represents how data is partitioned across different processing granularity. - * It represents howe we are going to access the data in thread, warp, or blocked in contiguous + * It represents how we are going to access the data in thread, warp, or blocked in contiguous region. * From below, the qualifier for 'raked' is the part of warp/thread hierarchy * in the split of Y tile dimension where the iteration happens, @@ -102,6 +102,11 @@ enum struct tile_distribution_pattern * */ block_raked, + /** + * @brief Sparse rows pattern - when we have very few rows, but potentially many columns. + * + */ + sparse_row }; struct tile_distribution_encoding_pattern @@ -130,6 +135,88 @@ struct tile_distribution_encoding_pattern_2d : public tile_distribution_encoding { }; +// Sparse rows +template +struct tile_distribution_encoding_pattern_2d + : public tile_distribution_encoding_pattern +{ + static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!"); + static_assert(NumWaveGroups == 1, "NumWaveGroups must be 1 for sparse row pattern!"); + + static constexpr index_t warp_size = get_warp_size(); + static constexpr index_t num_warps = BlockSize / warp_size; + + // Calculate optimal vector size + static constexpr index_t LargestVec = max(1, (XPerTile * YPerTile) / (num_warps * warp_size)); + static constexpr index_t X1 = VecSize > LargestVec ? LargestVec : VecSize; + static constexpr index_t X0 = XPerTile / X1; + + // When YPerTile is small, prioritize X dimension distribution + // and handle Y dimension with minimal thread usage. + + // Calculate threads needed for one row. + static constexpr index_t threads_per_row = X0; + + // Calculate how many rows we can process in parallel with available threads + static constexpr index_t max_parallel_rows = min(YPerTile, warp_size / threads_per_row); + + // Y2: Number of rows each warp handles in one iteration + static constexpr index_t Y2 = max_parallel_rows; + + // Y0: Number of warps to use (may be less than total available) + static constexpr index_t warps_needed = (YPerTile + Y2 - 1) / Y2; + static constexpr index_t Y0 = min(warps_needed, num_warps); + + // Y1: Number of iterations needed to cover all rows + static constexpr index_t Y1 = (YPerTile + (Y0 * Y2) - 1) / (Y0 * Y2); + + // Validation + static_assert(Y0 > 0, "Y0 must be greater than 0!"); + static_assert(Y1 > 0, "Y1 must be greater than 0!"); + static_assert(Y2 > 0, "Y2 must be greater than 0!"); + static_assert(X0 > 0, "X0 must be greater than 0!"); + static_assert(X1 > 0, "X1 must be greater than 0!"); + + // Ensure we don't exceed available threads per warp + static_assert(threads_per_row * Y2 <= warp_size, + "Threads per row * rows per warp must not exceed warp size!"); + + // Ensure we cover all elements (may over-cover due to ceiling, but that's OK) + static_assert(Y0 * Y1 * Y2 >= YPerTile, + "Y0 * Y1 * Y2 must cover at least YPerTile rows"); + + CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution() + { + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, // -> , + sequence<1, 2>, + sequence<1, 1>>{}); // -> + } + + CK_TILE_HOST_DEVICE static constexpr auto make_shuffled_2d_static_tile_distribution() + { + return make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<2, 1>>, + tuple, sequence<2, 0>>, // -> , + sequence<1, 2>, + sequence<1, 1>>{}); // -> + } +}; + // Thread raked template : public tile_distribution_encoding_pattern { - // TODO: make pattern where below condition does not need to hold - GGemmMultiDSplitk! - static_assert(YPerTile > 0, "YPerTile must be greater than 0!"); - static_assert(XPerTile > 0, "XPerTile must be greater than 0!"); - static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!"); static constexpr index_t warp_size = get_warp_size(); static constexpr index_t num_warps = BlockSize / get_warp_size(); - static_assert(num_warps > 0, "num_warps must be greater than 0!"); - static_assert(warp_size > 0, "warp_size must be greater than 0!"); - - static constexpr index_t LargestVec = max(1, (XPerTile * YPerTile) / (num_warps * warp_size)); + static constexpr index_t LargestVec = (XPerTile * YPerTile) / (num_warps * warp_size); static constexpr index_t X1 = VecSize > LargestVec ? LargestVec : VecSize; - - static_assert(X1 > 0, "X1 must be greater than 0!"); static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim - static_assert(X0 > 0, "X0 must be greater than 0!"); // # of rows in Y dim accessed by single wavefront in one iteration - static constexpr index_t Y1 = max(1, warp_size / X0); + static constexpr index_t Y1 = warp_size / X0; static_assert(X0 * Y1 == warp_size, "X0 * Y1 must cover whole wavefront!"); - static constexpr index_t Y0 = max(1, num_warps / NumWaveGroups); + static constexpr index_t Y0 = num_warps / NumWaveGroups; // YPerWarp = YPerTile / Y0; // Y2 = YPerWarp / Y1; - static_assert(Y0 > 0, "Y0 must be greater than 0!"); - static_assert(Y1 > 0, "Y1 must be greater than 0!"); - static constexpr index_t Y2 = max(1, YPerTile / (Y1 * Y0)); // # of iters within wavefront + static constexpr index_t Y2 = YPerTile / (Y1 * Y0); // # of iters within wavefront static_assert(X0 * Y1 * Y0 * NumWaveGroups == BlockSize, "X0 * warp_ys * Y0 must cover whole workgroup!"); @@ -241,27 +316,20 @@ struct tile_distribution_encoding_pattern_2d : public tile_distribution_encoding_pattern { - static_assert(YPerTile > 0, "YPerTile must be greater than 0!"); - static_assert(XPerTile > 0, "XPerTile must be greater than 0!"); - static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!"); static constexpr index_t warp_size = get_warp_size(); - static_assert(warp_size > 0, "warp_size must be greater than 0!"); static constexpr index_t num_warps = BlockSize / get_warp_size(); - static_assert(num_warps > 0, "num_warps must be greater than 0!"); - static constexpr index_t LargestVec = max(1, (XPerTile * YPerTile) / (num_warps * warp_size)); + static constexpr index_t LargestVec = (XPerTile * YPerTile) / (num_warps * warp_size); static constexpr index_t X1 = VecSize > LargestVec ? LargestVec : VecSize; - static_assert(X1 > 0, "X1 must be greater than 0!"); static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim static constexpr index_t Y2 = warp_size / X0; // # of rows in Y dim to cover whole wavefront - static_assert(Y2 > 0, "Y2 must be greater than 0!"); static_assert(X0 * Y2 == warp_size, "X0 * Y2 must cover whole wavefront!"); static constexpr index_t Y0 = num_warps; static_assert(X0 * Y2 * Y0 == BlockSize, "X0 * Y2 * Y1 must cover whole workgroup!"); - static constexpr index_t Y1 = max(1, YPerTile / (Y2 * Y0)); // # of iters within wavefront + static constexpr index_t Y1 = YPerTile / (Y2 * Y0); // # of iters within wavefront static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover whole YPerTile"); CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution() @@ -306,7 +374,7 @@ struct tile_distribution_encoding_pattern_2d LargestVec ? LargestVec : VecSize; static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim static constexpr index_t Y2 = warp_size / X0; // # of rows in Y dim to cover whole wavefront @@ -347,6 +415,7 @@ constexpr const char* tile_distribution_pattern_to_string(tile_distribution_patt case tile_distribution_pattern::thread_raked: return "thread_raked"; case tile_distribution_pattern::warp_raked: return "warp_raked"; case tile_distribution_pattern::block_raked: return "block_raked"; + case tile_distribution_pattern::sparse_row: return "sparse_row"; default: return "unknown"; } } diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 84c9befc7a..efe4e599f0 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -338,6 +338,10 @@ struct CShuffleEpilogue sequence<0, 1>, sequence>; + using SFC_dram = space_filling_curve, + sequence<0, 1>, + sequence>; + constexpr index_t num_access = SFC::get_num_of_access(); static_assert(std::is_same_v, @@ -347,7 +351,7 @@ struct CShuffleEpilogue MPerIterationShuffle, NPerIterationShuffle, GetVectorSizeC(), - tile_distribution_pattern::thread_raked, + tile_distribution_pattern::sparse_row, Problem::kNumWaveGroups>; constexpr auto dram_tile_distribution = TileEncodingPattern::make_2d_static_tile_distribution(); @@ -362,22 +366,86 @@ struct CShuffleEpilogue to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; - static_for<0, num_access, 1>{}([&](auto iAccess) { - block_sync_lds(); - constexpr auto idx_y_start = SFC::get_index(iAccess); + // Ensure that we have the expected number of accesses. + if constexpr (NumGroupsToMerge > 1) + { + static_assert(num_access == NumGroupsToMerge * NumGroupsToMerge, + "Number of accesses must be equal to NumGroupsToMerge squared."); - constexpr auto mIter = number{}) / (MPerIterationShuffle)>{}; - constexpr auto nIter = number{}) / (NPerIterationShuffle)>{}; + static_for<0, NumGroupsToMerge, 1>{} + ( + [&](auto group) + { + block_sync_lds(); + constexpr auto iAccess = number{}; + constexpr auto idx_y_start = SFC::get_index(iAccess); + + constexpr auto mIter = number{}) / (MPerIterationShuffle)>{}; + constexpr auto nIter = number{}) / (NPerIterationShuffle)>{}; + + lds_tile.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data( + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence{}, + c_warp_y_lengths)); + + const auto c_warptile_in_tensor_casted = cast_tile(lds_tile); + + store_tile(in_lds_window, c_warptile_in_tensor_casted); + block_sync_lds(); + + auto c_out_tensor = load_tile(make_tile_window(out_lds_window, dram_tile_distribution)); + + const auto ds_tensor = generate_tuple( + [&](auto idx) { return load_tile(d_dram_windows[idx]); }, number{}); + + const auto c_ds_tiles = concat_tuple_of_reference( + tie(c_out_tensor, c_out_tensor), + generate_tie([&](auto idx) -> const auto& { return ds_tensor[idx]; }, + number{})); + + tile_elementwise_inout_unpack(typename Problem::CDElementwise{}, c_ds_tiles); + + if constexpr(MemoryOperation == memory_operation_enum::set) + { + store_tile(out_dram_window, c_out_tensor); + } + else + { + update_tile(out_dram_window, c_out_tensor); + } + + // TODO: This probably doesn't work correctly. + if constexpr(group != NumGroupsToMerge - 1) + { + constexpr auto step = SFC_dram::get_forward_step(group); + + move_tile_window(out_dram_window, {step.at(number<0>{}), step.at(number<1>{})}); + + static_for<0, NumDTensor, 1>{}([&](auto idx) { + move_tile_window(d_dram_windows[idx], + {step.at(number<0>{}), step.at(number<1>{})}); + }); + } + } + ); + } + else + { + static_for<0, num_access, 1>{}([&](auto iAccess) { + block_sync_lds(); + constexpr auto idx_y_start = SFC::get_index(iAccess); + + constexpr auto mIter = number{}) / (MPerIterationShuffle)>{}; + constexpr auto nIter = number{}) / (NPerIterationShuffle)>{}; - // Store only the diagonal blocks when NumGroupsToMerge > 1 - if constexpr (NumGroupsToMerge == 1 || (mIter == nIter)) - { lds_tile.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data( - merge_sequences( - sequence{}, - c_warp_y_index_zeros), - merge_sequences(sequence{}, - c_warp_y_lengths)); + merge_sequences( + sequence{}, + c_warp_y_index_zeros), + merge_sequences(sequence{}, + c_warp_y_lengths)); const auto c_warptile_in_tensor_casted = cast_tile(lds_tile); @@ -405,6 +473,7 @@ struct CShuffleEpilogue update_tile(out_dram_window, c_out_tensor); } + // TODO: This probably doesn't work correctly. if constexpr(iAccess != num_access - 1) { constexpr auto step = SFC::get_forward_step(iAccess); @@ -416,8 +485,8 @@ struct CShuffleEpilogue {step.at(number<0>{}), step.at(number<1>{})}); }); } - } - }); + }); + } } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp index d1de2a3db9..e61e5eff6f 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp @@ -788,46 +788,6 @@ struct GroupedConvolutionBackwardWeightKernel // Run Epilogue Pipeline auto& c_block_window = gemm_tile_windows.at(I3); - // In case we have multiple conv groups per GEMM batch, we need to store only the diagonal elements - // of the c_block_tile. - // constexpr index_t Gm = GroupedConvTraitsType_::NumGroupsToMerge; - // auto c_block_window_per_g = make_tile_window( - // c_block_window, - // c_block_window.get_window_origin() - // ); - // if constexpr(Gm > 1) - // { - // const index_t conv_group_size = kargs.group_stride_c / Gm; - // static_for<0, Gm, 1>{}([&](auto g) - // { - // // TODO: This is pseudocode, need to implement a proper way to slice the tile window - // // to get the submatrix corresponding to the g-th conv group. - // // constexpr auto& c_block_window_per_g = c_block_window.Slice( - // // make_tuple(number{}, - // // number<0>{}), - // // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{}, - // // number{})); - // // constexpr auto& c_block_tile_per_g = c_block_tile.Slice( - // // make_tuple(number{}, - // // number<0>{}), - // // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{}, - // // number{})); - // // constexpr auto & d_block_window_per_g = d_block_window.Slice( - // // make_tuple(number{}, - // // number<0>{}), - // // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{}, - // // number{})); - - // EpiloguePipeline{}.template operator()( - // c_block_window_per_g, c_block_tile_per_g, d_block_window_per_g, smem_ptr_0); - // }); - // } - // else - // { - // EpiloguePipeline{}.template operator()( - // c_block_window, c_block_tile, d_block_window, smem_ptr_0); - // } - // For debugging - results in very slow compilation. // if (blockIdx.x == 0 && threadIdx.x == 0) // {