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

Use placeholders for the LoadTranspose feature introduced in #2436 + cleanup

Browse Source
This commit is contained in:
Damien Lejeune
2025-07-30 20:17:16 +00:00
parent 96b26b2a3c
commit 49ab76c84b
7 changed files with 99 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp
View File

@@ -252,20 +252,35 @@ struct BlockUniversalGemmArBrCr : public BlockUniversalGemmBase<Problem_, Policy
};
public:
/*
* @FIXME: currently using LoadTranspose as placeholder. The logic needs to be implemented
*/
// C += A * B
template <typename CBlockTensor, typename ARegBlockTensor, typename BRegBlockTensor>
template <typename CBlockTensor,
typename ARegBlockTensor,
typename BRegBlockTensor,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor,
const ARegBlockTensor& a_block_tensor,
const BRegBlockTensor& b_block_tensor)
const BRegBlockTensor& b_block_tensor,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
if(a_load_tr || b_load_tr) {}
block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_tensor);
}
// C = A * B
template <typename ARegBlockTensor, typename BRegBlockTensor>
template <typename ARegBlockTensor,
typename BRegBlockTensor,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor,
const BRegBlockTensor& b_block_tensor)
const BRegBlockTensor& b_block_tensor,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
auto c_block_tensor = Base::MakeCBlockTile();

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

@@ -329,25 +329,40 @@ struct BlockUniversalGemmArBsCr : public BlockUniversalGemmBase<Problem_, Policy
};
public:
template <typename BSmemBlockWindow>
CK_TILE_DEVICE void LocalPrefetchB(const BSmemBlockWindow& b_block_window)
/*
* @FIXME: currently using LoadTranspose as placeholder. The logic needs to be implemented
*/
template <typename BSmemBlockWindow, bool BLoadTranspose = false>
CK_TILE_DEVICE void LocalPrefetchB(const BSmemBlockWindow& b_block_window,
bool_constant<BLoadTranspose> = {})
{
block_gemm_impl_.LocalPrefetchB(b_block_window);
}
// C += A * B
template <typename CBlockTensor, typename ARegBlockTensor, typename BSmemBlockWindow>
template <typename CBlockTensor,
typename ARegBlockTensor,
typename BSmemBlockWindow,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor,
const ARegBlockTensor& a_block_tensor,
const BSmemBlockWindow& b_block_window)
const BSmemBlockWindow& b_block_window,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window);
}
// C = A * B
template <typename ARegBlockTensor, typename BSmemBlockWindow>
template <typename ARegBlockTensor,
typename BSmemBlockWindow,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor,
const BSmemBlockWindow& b_block_window)
const BSmemBlockWindow& b_block_window,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
auto c_block_tensor = Base::MakeCBlockTile();
block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window);

41
include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp
View File

@@ -143,8 +143,10 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase<Problem_, Policy
ALdsTile a_warp_tile_;
BLdsTile b_warp_tile_;
template <typename ASmemBlockWindow>
CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window)
template <typename ASmemBlockWindow, bool ALoadTranspose = false>
CK_TILE_DEVICE void
LocalPrefetchA(const ASmemBlockWindow& a_block_window,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {})
{
if constexpr(std::is_same_v<ADataType, pk_int4_t>)
{
@@ -243,8 +245,10 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase<Problem_, Policy
}
}
template <typename BSmemBlockWindow>
CK_TILE_DEVICE void LocalPrefetchA(const BSmemBlockWindow& b_block_window)
template <typename BSmemBlockWindow, bool ALoadTranspose = false>
CK_TILE_DEVICE void
LocalPrefetchA(const BSmemBlockWindow& b_block_window,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {})
{
LocalPrefetchA<0>(b_block_window);
}
@@ -344,25 +348,40 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase<Problem_, Policy
};
public:
template <typename ASmemBlockWindow>
CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window)
/*
* @FIXME: currently using LoadTranspose as placeholder. The logic needs to be implemented
*/
template <typename ASmemBlockWindow, bool ALoadTranspose = false>
CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window,
bool_constant<ALoadTranspose> a_load_tr = {})
{
block_gemm_impl_.LocalPrefetchA(a_block_window);
block_gemm_impl_.LocalPrefetchA(a_block_window, a_load_tr);
}
// C += A * B
template <typename CBlockTensor, typename ASmemBlockWindow, typename BRegBlockTensor>
template <typename CBlockTensor,
typename ASmemBlockWindow,
typename BRegBlockTensor,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor,
const ASmemBlockWindow& a_block_window,
const BRegBlockTensor& b_block_tensor)
const BRegBlockTensor& b_block_tensor,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
block_gemm_impl_(c_block_tensor, a_block_window, b_block_tensor);
}
// C = A * B
template <typename ASmemBlockWindow, typename BRegBlockTensor>
template <typename ASmemBlockWindow,
typename BRegBlockTensor,
bool ALoadTranspose = false,
bool BLoadTranspose = false>
CK_TILE_DEVICE auto operator()(const ASmemBlockWindow& a_block_window,
const BRegBlockTensor& b_block_tensor)
const BRegBlockTensor& b_block_tensor,
[[maybe_unused]] bool_constant<ALoadTranspose> a_load_tr = {},
[[maybe_unused]] bool_constant<BLoadTranspose> b_load_tr = {})
{
auto c_block_tensor = Base::MakeCBlockTile();
block_gemm_impl_(c_block_tensor, a_block_window, b_block_tensor);

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

@@ -11,9 +11,27 @@
namespace ck_tile {
// A Tile Window: global memory
// B Tile Window: global memory
// C Distributed tensor: register
// ============================================================================
// GEMM Pipeline for Global Memory (A and B) and Register Memory (C)
// ============================================================================
//
// Overview:
// ---------
// - A Tile Window: Global memory
// - B Tile Window: Global memory
// - C Distributed Tensor: Register
//
// This pipeline maximizes memory throughput and computational efficiency,
// supporting both intrawave and interwave scheduling strategies.
//
// Skip LDS Feature:
// -----------------
// - Allows bypassing the Local Data Share (LDS) for A and/or B tiles,
// directly transferring data from Global Memory to Registers when enabled.
// When disabled, data flows through Global Memory -> LDS -> Registers.
// - Supported for both intrawave and interwave pipelines.
// ============================================================================
template <typename Problem>
struct BaseGemmPipelineAgBgCrMem
{
@@ -338,24 +356,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
block_gemm(c_block_tile,
a_block_tiles.get(number<prefetch_idx>{}),
b_lds_gemm_window,
false,
false); // TODO: transpose
is_a_load_tr_v,
is_b_load_tr_v);
}
else if constexpr(SkipALds == false && SkipBLds == true)
{
block_gemm(c_block_tile,
a_lds_gemm_window,
b_block_tiles.get(number<prefetch_idx>{}),
false,
false); // TODO: transpose
is_a_load_tr_v,
is_b_load_tr_v);
}
else
{
block_gemm(c_block_tile,
a_block_tiles.get(number<prefetch_idx>{}),
b_block_tiles.get(number<prefetch_idx>{}),
false,
false); // TODO: transpose
is_a_load_tr_v,
is_b_load_tr_v);
}
};
// -----------------------------------------------------------------------------------------

56
include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
View File

@@ -613,62 +613,6 @@ struct UniversalGemmBasePolicy
return smem_size_a + smem_size_b;
}
// template <typename Problem_, typename Policy_>
// struct GemmTraits
// {
// using Problem = remove_cvref_t<Problem_>;
// using Policy = remove_cvref_t<Policy_>;
// using ADataType = remove_cvref_t<typename Problem::ADataType>;
// using BDataType = remove_cvref_t<typename Problem::BDataType>;
// using ComputeDataType = remove_cvref_t<typename Problem::ComputeDataType>;
// using CDataType = remove_cvref_t<typename Problem::CDataType>;
// using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
// static constexpr index_t kBlockSize = Problem::kBlockSize;
// static constexpr auto Scheduler = Problem::Scheduler;
// static constexpr index_t MPerBlock = BlockGemmShape::kM;
// static constexpr index_t NPerBlock = BlockGemmShape::kN;
// static constexpr index_t KPerBlock = BlockGemmShape::kK;
// static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
// using WarpGemm = remove_cvref_t<decltype(config.template at<0>())>;
// static constexpr index_t MWarp = config.template at<1>();
// static constexpr index_t NWarp = config.template at<2>();
// using I0 = number<0>;
// using I1 = number<1>;
// static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}),
// "Error! WarpGemm's MWarp is not consistent with BlockGemmShape!");
// static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}),
// "Error! WarpGemm's NWarp is not consistent with BlockGemmShape!");
// static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}),
// "Error! WarpGemm's M is not consistent with BlockGemmShape!");
// static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}),
// "Error! WarpGemm's N is not consistent with BlockGemmShape!");
// static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM);
// static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN);
// static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK;
// static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock,
// "Error! Warps should cover all Block tile!");
// static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock,
// "Error! Warps should cover all Block tile!");
// static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM;
// static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN;
// static constexpr index_t KPerBlockPerIter = WarpGemm::kK;
// static constexpr index_t InterWaveSchedulingMacClusters = 1;
// static constexpr index_t KPack = WarpGemm::kKPerThread;
// static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread;
// };
};
// Forward declarations, to avoid circular dependencies

2
test/ck_tile/gemm/test_gemm_pipeline_smoke_util.hpp
View File

@@ -101,7 +101,7 @@ struct GemmConfigMemoryIntrawave : public GemmConfigBase
static constexpr bool DoubleSmemBuffer = false;
static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY;
static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave
static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
};
template <typename PrecType>

16
test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc
View File

@@ -1,5 +1,5 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
@@ -87,18 +87,4 @@ TYPED_TEST(TEST_SUITE_NAME, LargeMatrix)
this->Run(M, N, K);
}
// TYPED_TEST(TEST_SUITE_NAME, NotSupportedArgument)
// {
// constexpr int M = 512;
// constexpr int N = 1025;
// constexpr int K = 513;
// constexpr bool PadM = false;
// constexpr bool PadN = false;
// constexpr bool PadK = false;
// EXPECT_THROW((this->template Run<PadM, PadN, PadK>(M, N, K)), std::runtime_error);
// }
#endif