ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-09 00:27:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adding data shuffling when Skipping LDS

Browse Source
This commit is contained in:
Damien Lejeune
2025-09-04 15:10:55 +00:00
parent 4c97724e99
commit 4626ff6e8e
1 changed files with 148 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

153
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp
View File

@@ -412,6 +412,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{}));
a_block_tiles.get(I0{}) = a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -426,6 +436,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func);
}
}
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{}));
b_block_tiles.get(I0{}) = b_shuffle_tmp;
}
}
// Global prefetch [1, PrefetchStages]
static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) {
@@ -480,6 +500,20 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(
a_shuffle_tmp,
a_block_tiles.get(
number<(prefetch_idx + 1) % PrefetchStages>{}));
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -502,6 +536,20 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
b_element_func);
}
}
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(
b_shuffle_tmp,
b_block_tiles.get(
number<(prefetch_idx + 1) % PrefetchStages>{}));
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
b_shuffle_tmp;
}
}
Base::GlobalPrefetch(a_block_tiles.get(number<prefetch_idx>{}),
a_copy_dram_window,
@@ -547,6 +595,19 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(
a_shuffle_tmp,
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}));
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -564,6 +625,19 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
b_element_func);
}
}
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(
b_shuffle_tmp,
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}));
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
b_shuffle_tmp;
}
}
});
block_sync_lds();
@@ -774,6 +848,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{}));
a_block_tiles.get(I0{}) = a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -788,11 +872,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func);
}
}
// TODO add encoding and support for BRowMajor for SkipBLds, current
// MakeShuffledBRegTileDistribution takes into account shuffling encoding which is used
// for [Global -> Vgpr -> Lds -> Vgpr] reading, but for skipping lds we need to have
// different shuffled layout for [Global -> Vgpr] reads, similar for AColMajor
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{}));
b_block_tiles.get(I0{}) = b_shuffle_tmp;
}
}
// Global prefetch [1, PrefetchStages]
static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) {
@@ -845,6 +934,20 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(
a_shuffle_tmp,
a_block_tiles.get(
number<(prefetch_idx + 1) % PrefetchStages>{}));
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -867,6 +970,20 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
b_element_func);
}
}
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(
b_shuffle_tmp,
b_block_tiles.get(
number<(prefetch_idx + 1) % PrefetchStages>{}));
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
b_shuffle_tmp;
}
}
Base::GlobalPrefetch(a_block_tiles.get(number<prefetch_idx>{}),
a_copy_dram_window,
@@ -903,6 +1020,19 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
a_element_func);
}
}
else
{
if constexpr(is_a_col_major && !is_a_load_tr_v())
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
Policy::template MakeShuffledARegTileDistribution<Problem>());
transpose_tile2d(
a_shuffle_tmp,
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}));
a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
a_shuffle_tmp;
}
}
if constexpr(SkipBLds == false)
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
@@ -920,6 +1050,19 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
b_element_func);
}
}
else
{
if constexpr(is_b_row_major && !is_b_load_tr_v())
{
auto b_shuffle_tmp = make_static_distributed_tensor<BDataType>(
Policy::template MakeShuffledBRegTileDistribution<Problem>());
transpose_tile2d(
b_shuffle_tmp,
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}));
b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}) =
b_shuffle_tmp;
}
}
});
block_sync_lds();