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

Clean up + add TODO comment about A-column-wise/B-row-wise layout

Browse Source
This commit is contained in:
Damien Lejeune
2025-10-01 12:44:32 +00:00
parent b8624c540d
commit c07cec8809
4 changed files with 13 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp
View File

@@ -243,7 +243,6 @@ struct BlockUniversalGemmArBsCr : public BlockUniversalGemmBase<Problem_, Policy
"The CDataType as defined in traits should be the same as correspoinding "
"C block tensor data type!");
// a_warp_tile_.get_thread_buffer() = a_block_tensor.get_thread_buffer();
// hot loop:
static_for<0, KRepeat, 1>{}([&](auto kIter) {
LocalPrefetchB<kIter.value>(b_block_window);

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

@@ -412,16 +412,7 @@ 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())
@@ -436,16 +427,6 @@ 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) {
@@ -500,20 +481,7 @@ 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())
@@ -536,20 +504,6 @@ 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,
@@ -595,19 +549,7 @@ 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())
@@ -625,19 +567,6 @@ 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();
@@ -848,16 +777,7 @@ 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())
@@ -872,16 +792,11 @@ 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;
}
}
// TODO add encoding and support for BRowMajor for SkipBLds, currently:
// 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
// Global prefetch [1, PrefetchStages]
static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) {
@@ -934,20 +849,7 @@ 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())

4
include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp
View File

@@ -68,10 +68,10 @@ struct TileGemmUniversalTraits
static constexpr bool TransposeC = TransposeC_;
/// @brief SkipALds_ Parameterize the data flow for matrix A: if true, global memory ->
/// registers; if false, global memory -> LDS -> registers.
/// VGPR; if false, global memory -> VGPR -> LDS -> VGPR.
static constexpr bool SkipALds = SkipALds_;
/// @brief SkipBLds_ Parameterize the data flow for matrix B: if true, global memory ->
/// registers; if false, global memory -> LDS -> registers.
/// VGPR; if false, global memory -> VGPR -> LDS -> VGPR.
static constexpr bool SkipBLds = SkipBLds_;
static constexpr bool UseStructuredSparsity = UseStructuredSparsity_;
static constexpr bool UsePersistentKernel = UsePersistentKernel_;