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

LDS to global memory copy.

Browse Source
This commit is contained in:
Ville Pietilä
2025-09-18 14:59:32 +00:00
parent 0e09504057
commit 6bcdb0947e
4 changed files with 186 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp
View File

@@ -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<InPrecType, 1>(in_layout, wei_layout, out_layout, argc, argv);
// }
if (num_groups_to_merge == 1)
{
return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 1>(in_layout, wei_layout, out_layout, argc, argv);
}
// else if (num_groups_to_merge == 2)
// {
// return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 2>(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<InPrecType, 64>(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<InPrecType, 8>(in_layout, wei_layout, out_layout, argc, argv);
}

123
include/ck_tile/core/algorithm/static_encoding_pattern.hpp
View File

@@ -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 <index_t BlockSize,
index_t YPerTile,
index_t XPerTile,
index_t VecSize,
index_t NumWaveGroups>
struct tile_distribution_encoding_pattern_2d<BlockSize,
YPerTile,
XPerTile,
VecSize,
tile_distribution_pattern::sparse_row,
NumWaveGroups>
: 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<sequence<1>,
tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
tuple<sequence<1>, sequence<1, 2>>,
tuple<sequence<0>, sequence<2, 0>>, // -> <Y0>, <Y2, X0>
sequence<1, 2>,
sequence<1, 1>>{}); // -> <Y1, X1>
}
CK_TILE_HOST_DEVICE static constexpr auto make_shuffled_2d_static_tile_distribution()
{
return make_static_tile_distribution(
tile_distribution_encoding<sequence<1>,
tuple<sequence<X0, X1>, sequence<Y0, Y1, Y2>>,
tuple<sequence<2>, sequence<2, 1>>,
tuple<sequence<0>, sequence<2, 0>>, // -> <Y0>, <Y2, X0>
sequence<1, 2>,
sequence<1, 1>>{}); // -> <X1, Y1>
}
};
// Thread raked
template <index_t BlockSize,
index_t YPerTile,
@@ -144,34 +231,22 @@ struct tile_distribution_encoding_pattern_2d<BlockSize,
NumWaveGroups>
: 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<BlockSize,
NumWaveGroups>
: 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<BlockSize,
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 constexpr index_t LargestVec = (XPerTile * YPerTile) / (num_warps * warp_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; // # 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";
}
}

101
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -338,6 +338,10 @@ struct CShuffleEpilogue
sequence<0, 1>,
sequence<MPerIterationShuffle, NPerIterationShuffle>>;
using SFC_dram = space_filling_curve<sequence<kMPerBlock, kNPerBlock / NumGroupsToMerge>,
sequence<0, 1>,
sequence<MPerIterationShuffle, NPerIterationShuffle>>;
constexpr index_t num_access = SFC::get_num_of_access();
static_assert(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>,
@@ -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<CWarpDstr::NDimY, 0>{};
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<idx_y_start.at(number<0>{}) / (MPerIterationShuffle)>{};
constexpr auto nIter = number<idx_y_start.at(number<1>{}) / (NPerIterationShuffle)>{};
static_for<0, NumGroupsToMerge, 1>{}
(
[&](auto group)
{
block_sync_lds();
constexpr auto iAccess = number<group * NumGroupsToMerge + group>{};
constexpr auto idx_y_start = SFC::get_index(iAccess);
constexpr auto mIter = number<idx_y_start.at(number<0>{}) / (MPerIterationShuffle)>{};
constexpr auto nIter = number<idx_y_start.at(number<1>{}) / (NPerIterationShuffle)>{};
lds_tile.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data(
merge_sequences(
sequence<mIter * NumMXdlPerWavePerShuffle, nIter * NumNXdlPerWavePerShuffle>{},
c_warp_y_index_zeros),
merge_sequences(sequence<NumMXdlPerWavePerShuffle, NumNXdlPerWavePerShuffle>{},
c_warp_y_lengths));
const auto c_warptile_in_tensor_casted = cast_tile<ODataType>(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<NumDTensor>{});
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<NumDTensor>{}));
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<idx_y_start.at(number<0>{}) / (MPerIterationShuffle)>{};
constexpr auto nIter = number<idx_y_start.at(number<1>{}) / (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<mIter * NumMXdlPerWavePerShuffle, nIter * NumNXdlPerWavePerShuffle>{},
c_warp_y_index_zeros),
merge_sequences(sequence<NumMXdlPerWavePerShuffle, NumNXdlPerWavePerShuffle>{},
c_warp_y_lengths));
merge_sequences(
sequence<mIter * NumMXdlPerWavePerShuffle, nIter * NumNXdlPerWavePerShuffle>{},
c_warp_y_index_zeros),
merge_sequences(sequence<NumMXdlPerWavePerShuffle, NumNXdlPerWavePerShuffle>{},
c_warp_y_lengths));
const auto c_warptile_in_tensor_casted = cast_tile<ODataType>(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

40
include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp
View File

@@ -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<g.value * TilePartitioner::MPerBlock / Gm>{},
// // number<0>{}),
// // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{},
// // number<TilePartitioner::NPerBlock>{}));
// // constexpr auto& c_block_tile_per_g = c_block_tile.Slice(
// // make_tuple(number<g.value * TilePartitioner::MPerBlock / Gm>{},
// // number<0>{}),
// // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{},
// // number<TilePartitioner::NPerBlock>{}));
// // constexpr auto & d_block_window_per_g = d_block_window.Slice(
// // make_tuple(number<g.value * TilePartitioner::MPerBlock / Gm>{},
// // number<0>{}),
// // make_tuple(number<(g.value + 1) * TilePartitioner::MPerBlock / Gm>{},
// // number<TilePartitioner::NPerBlock>{}));
// EpiloguePipeline{}.template operator()<decltype(c_block_window_per_g), decltype(c_block_tile_per_g)>(
// c_block_window_per_g, c_block_tile_per_g, d_block_window_per_g, smem_ptr_0);
// });
// }
// else
// {
// EpiloguePipeline{}.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
// 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)
// {