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

Documentation + cleanup

Browse Source
This commit is contained in:
Damien Lejeune
2025-07-30 21:21:01 +00:00
parent 49ab76c84b
commit 8e878f3914
4 changed files with 33 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp
View File

@@ -37,9 +37,7 @@ struct GemmHostArgs
index_t K_,
index_t stride_A_,
index_t stride_B_,
index_t stride_E_,
bool skip_a_lds_ = false,
bool skip_b_lds_ = false)
index_t stride_E_)
: a_ptr(a_ptr_),
b_ptr(b_ptr_),
e_ptr(e_ptr_),
@@ -49,9 +47,7 @@ struct GemmHostArgs
stride_A(stride_A_),
stride_B(stride_B_),
stride_E(stride_E_),
k_batch(k_batch_),
skip_a_lds(skip_a_lds_),
skip_b_lds(skip_b_lds_)
k_batch(k_batch_)
{
}
@@ -76,8 +72,6 @@ struct GemmHostArgs
};
index_t k_batch;
bool skip_a_lds;
bool skip_b_lds;
};
template <typename TilePartitioner_, typename GemmPipeline_, typename EpiloguePipeline_>

14
include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp
View File

@@ -34,9 +34,7 @@ struct GroupedGemmHostArgs
index_t K_,
index_t stride_A_,
index_t stride_B_,
index_t stride_E_,
bool skip_a_lds_,
bool skip_b_lds_)
index_t stride_E_)
: a_ptr(a_ptr_),
b_ptr(b_ptr_),
e_ptr(e_ptr_),
@@ -46,9 +44,7 @@ struct GroupedGemmHostArgs
stride_A(stride_A_),
stride_B(stride_B_),
stride_E(stride_E_),
k_batch(k_batch_),
skip_a_lds(skip_a_lds_),
skip_b_lds(skip_b_lds_)
k_batch(k_batch_)
{
}
@@ -73,8 +69,6 @@ struct GroupedGemmHostArgs
};
index_t k_batch;
bool skip_a_lds;
bool skip_b_lds;
};
struct GemmTransKernelArg
@@ -233,9 +227,7 @@ struct GroupedGemmKernel
{stride_b},
{/*stride_ds*/},
stride_e,
gemm_descs[i].k_batch,
gemm_descs[i].skip_a_lds,
gemm_descs[i].skip_b_lds};
gemm_descs[i].k_batch};
gemm_kernel_args_.emplace_back(std::move(karg), block_start, block_end);
}

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

@@ -34,6 +34,12 @@ struct TileGemmTraits
static constexpr index_t NumWaveGroups = NumWaveGroups_;
};
/**
* @brief Traits for a universal GEMM pipeline.
*
* @details This struct holds compile-time parameters that configure the behavior of a
* general-purpose tile-based GEMM kernel
*/
template <bool kPadM_,
bool kPadN_,
bool kPadK_,
@@ -60,16 +66,17 @@ struct TileGemmUniversalTraits
using BLayout = BLayout_;
using CLayout = CLayout_;
static constexpr bool TransposeC = TransposeC_;
static constexpr bool SkipALds = SkipALds_;
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.
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.
static constexpr bool SkipBLds = SkipBLds_;
static constexpr bool UseStructuredSparsity = UseStructuredSparsity_;
static constexpr bool UsePersistentKernel = UsePersistentKernel_;
static constexpr index_t NumWaveGroups = NumWaveGroups_;
static constexpr bool Preshuffle = Preshuffle_;
static constexpr bool kSkipALds = SkipALds_;
static constexpr bool kSkipBLds = SkipBLds_;
};
template <bool kPadM_,

45
test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp
View File

@@ -84,23 +84,20 @@ class TestCkTileGroupedGemm : public ::testing::Test
using TilePartitioner = ck_tile::
GemmSpatiallyLocalTilePartitioner<GemmShape, TileParitionerGroupNum, TileParitionerM01>;
using Traits = ck_tile::TileGemmTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
ALayout,
BLayout,
CLayout>;
using GemmUniversalTraits =
ck_tile::TileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
DoubleSmemBuffer,
GroupedGemKernelParam::SkipALds,
GroupedGemKernelParam::SkipBLds,
ALayout,
BLayout,
CLayout,
TransposeC>;
using Traits = ck_tile::TileGemmTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
ALayout,
BLayout,
CLayout>;
using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
DoubleSmemBuffer,
ALayout,
BLayout,
CLayout,
TransposeC>;
using GemmPipelineProblem =
ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
@@ -444,18 +441,8 @@ class TestCkTileGroupedGemm : public ::testing::Test
// TODO add support for kbatch > 1
static constexpr ck_tile::index_t k_batch = 1;
gemm_descs.push_back({p_a,
p_b,
p_c,
k_batch,
M,
N,
K,
stride_As[i],
stride_Bs[i],
stride_Cs[i],
GroupedGemKernelParam::SkipALds,
GroupedGemKernelParam::SkipBLds});
gemm_descs.push_back(
{p_a, p_b, p_c, k_batch, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]});
}
ck_tile::DeviceMem gemm_workspace;