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

Expose skip lds to batch/group gemm kernels api

Browse Source
This commit is contained in:
Damien Lejeune
2025-07-29 07:01:32 +00:00
parent 00b153b99c
commit 0d40254873
4 changed files with 51 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp
View File

@@ -28,8 +28,18 @@ struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs
ck_tile::index_t batch_stride_B_,
ck_tile::index_t batch_stride_C_,
ck_tile::index_t batch_count_)
: GemmHostArgs(
a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_, skip_a_lds_, skip_b_lds_),
: GemmHostArgs(a_ptr_,
b_ptr_,
c_ptr_,
k_batch_,
M_,
N_,
K_,
stride_A_,
stride_B_,
stride_C_,
skip_a_lds_,
skip_b_lds_),
batch_stride_A(batch_stride_A_),
batch_stride_B(batch_stride_B_),
batch_stride_C(batch_stride_C_),
@@ -103,7 +113,9 @@ struct BatchedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
hostArgs.stride_A,
hostArgs.stride_B,
hostArgs.stride_C,
hostArgs.k_batch},
hostArgs.k_batch,
hostArgs.skip_a_lds,
hostArgs.skip_b_lds},
hostArgs.batch_stride_A,
hostArgs.batch_stride_B,
hostArgs.batch_stride_C,

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

@@ -113,7 +113,9 @@ struct GroupedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
stride_a,
stride_b,
stride_c,
gemm_descs[i].k_batch};
gemm_descs[i].k_batch,
gemm_descs[i].skip_a_lds,
gemm_descs[i].skip_b_lds};
gemm_kernel_args_.emplace_back(std::move(karg), block_start, block_end);
}

5
test/ck_tile/batched_gemm/test_batched_gemm_util.hpp
View File

@@ -42,6 +42,9 @@ class TestCkTileBatchedGemm : public ::testing::Test
constexpr bool DoubleSmemBuffer = false;
constexpr bool kSkipALds = false;
constexpr bool kSkipBLds = false;
constexpr bool kPadM = false;
constexpr bool kPadN = false;
constexpr bool kPadK = false;
@@ -64,6 +67,8 @@ class TestCkTileBatchedGemm : public ::testing::Test
kPadN,
kPadK,
DoubleSmemBuffer,
kSkipALds,
kSkipBLds,
ALayout,
BLayout,
CLayout,

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

@@ -30,6 +30,9 @@ class TestCkTileGroupedGemm : public ::testing::Test
static const bool kPadN = false;
static const bool kPadK = false;
static const bool SkipALds = false;
static const bool SkipBLds = false;
static const int kBlockPerCu = 1;
static const ck_tile::index_t M_Tile = 128;
static const ck_tile::index_t N_Tile = 128;
@@ -74,20 +77,23 @@ class TestCkTileGroupedGemm : public ::testing::Test
using TilePartitioner = ck_tile::
GemmSpatiallyLocalTilePartitioner<GemmShape, TileParitionerGroupNum, TileParitionerM01>;
using Traits = ck_tile::TileGemmTraits<GroupedGemKernelParam::kPadM,
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 GemmUniversalTraits =
ck_tile::TileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
GroupedGemKernelParam::kPadN,
GroupedGemKernelParam::kPadK,
DoubleSmemBuffer,
GroupedGemKernelParam::SkipALds,
GroupedGemKernelParam::SkipBLds,
ALayout,
BLayout,
CLayout,
TransposeC>;
using GemmPipelineProblem =
ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
@@ -317,8 +323,19 @@ 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]});
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});
}
ck_tile::DeviceMem gemm_workspace;