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

The parameter const DType ds_ptr could be treated as const void *

Browse Source
This commit is contained in:
Mateusz Ozga
2025-03-22 18:04:36 +00:00
parent 8a75cc80bf
commit 9c9446b3f3
9 changed files with 67 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/03_gemm/gemm_basic.cpp
View File

@@ -22,7 +22,7 @@ template <typename ADataType,
typename DsLayout,
typename CLayout,
typename CDEElementWise = ck_tile::element_wise::PassThrough>
float gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s)
float gemm(const ck_tile::GemmHostArgs</*NumDTensor = 0*/>& args, const ck_tile::stream_config& s)
{
// The kPadM, kPadN, kPadK & kBlockPerCu should also come from the Codegen part.

3
example/ck_tile/03_gemm/gemm_utils.hpp
View File

@@ -219,4 +219,5 @@ auto create_args(int argc, char* argv[])
}
// host API
float gemm_calc(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s);
float gemm_calc(const ck_tile::GemmHostArgs</*NumDTensor = 0*/>& args,
const ck_tile::stream_config& s);

24
example/ck_tile/03_gemm/run_gemm_example.inc
View File

@@ -166,18 +166,18 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf,
int n_warmup,
int n_repeat)
{
ck_tile::GemmHostArgs<> args = {a_m_k_dev_buf.GetDeviceBuffer(),
b_k_n_dev_buf.GetDeviceBuffer(),
{},
c_m_n_dev_buf.GetDeviceBuffer(),
kbatch,
M,
N,
K,
stride_A,
stride_B,
{},
stride_C};
ck_tile::GemmHostArgs</*NumDTensor = 0*/> args = {a_m_k_dev_buf.GetDeviceBuffer(),
b_k_n_dev_buf.GetDeviceBuffer(),
{},
c_m_n_dev_buf.GetDeviceBuffer(),
kbatch,
M,
N,
K,
stride_A,
stride_B,
{},
stride_C};
float ave_time =
gemm<ADataType,

2
example/ck_tile/03_gemm/universal_gemm.cpp
View File

@@ -22,7 +22,7 @@ template <typename ADataType,
typename DsLayout,
typename CLayout,
typename CDEElementWise = ck_tile::element_wise::PassThrough>
float gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s)
float gemm(const ck_tile::GemmHostArgs</*NumDTensor = 0*/>& args, const ck_tile::stream_config& s)
{
using GemmShape = ck_tile::TileGemmShape<

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

@@ -9,7 +9,7 @@
namespace ck_tile {
struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs<>
struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs</*NumDTensor = 0*/>
{
CK_TILE_HOST BatchedGemmHostArgs() = default;
CK_TILE_HOST BatchedGemmHostArgs(const void* a_ptr_,
@@ -56,7 +56,7 @@ struct BatchedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
{
using Base = GemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>;
using GemmKernelArgs = typename ck_tile::GemmKernelArgs<>;
using GemmKernelArgs = typename ck_tile::GemmKernelArgs;
using ADataType = typename Base::ADataType;
using BDataType = typename Base::BDataType;

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

@@ -57,13 +57,11 @@ struct GemmHostArgs
index_t k_batch;
};
// TODO: The parameter const DType ds_ptr could be treated as const void *
template <typename DType = tuple<>>
struct GemmKernelArgs
{
const void* a_ptr;
const void* b_ptr;
const DType ds_ptr;
const void* ds_ptr;
void* c_ptr;
index_t M;
index_t N;
@@ -128,27 +126,21 @@ struct GemmKernel
CK_TILE_HOST static constexpr auto BlockSize() { return dim3(KernelBlockSize); }
CK_TILE_HOST static constexpr GemmKernelArgs<DsGridPointer>
CK_TILE_HOST static constexpr GemmKernelArgs
MakeKernelArgs(const GemmHostArgs<NumDTensor>& hostArgs)
{
DsGridPointer p_ds_grid;
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DDataType_ = remove_cvref_t<std::tuple_element_t<i.value, DsDataType>>;
p_ds_grid(i) = static_cast<const DDataType_*>(hostArgs.ds_ptr[i]);
});
return GemmKernelArgs<DsGridPointer>{hostArgs.a_ptr,
hostArgs.b_ptr,
p_ds_grid,
hostArgs.c_ptr,
hostArgs.M,
hostArgs.N,
hostArgs.K,
hostArgs.stride_A,
hostArgs.stride_B,
hostArgs.stride_Ds.data(),
hostArgs.stride_C,
hostArgs.k_batch};
return GemmKernelArgs{hostArgs.a_ptr,
hostArgs.b_ptr,
static_cast<const void*>(hostArgs.ds_ptr.data()),
hostArgs.c_ptr,
hostArgs.M,
hostArgs.N,
hostArgs.K,
hostArgs.stride_A,
hostArgs.stride_B,
hostArgs.stride_Ds.data(),
hostArgs.stride_C,
hostArgs.k_batch};
}
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
@@ -158,7 +150,7 @@ struct GemmKernel
struct SplitKBatchOffset
{
__device__ SplitKBatchOffset(const GemmKernelArgs<DsGridPointer>& kargs,
__device__ SplitKBatchOffset(const GemmKernelArgs& kargs,
const std::size_t k_id = blockIdx.z)
{
constexpr auto K1 = TilePartitioner::BlockGemmShape::WarpTile::at(number<2>{});
@@ -198,7 +190,7 @@ struct GemmKernel
index_t splitted_k;
};
CK_TILE_HOST static bool IsSupportedArgument(const GemmKernelArgs<DsGridPointer>& kargs)
CK_TILE_HOST static bool IsSupportedArgument(const GemmKernelArgs& kargs)
{
if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 &&
is_any_of<CDataType, fp16_t, bf16_t>::value)
@@ -388,9 +380,9 @@ struct GemmKernel
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr,
const BDataType* b_ptr,
const DsGridPointer ds_ptr,
const DsGridPointer* ds_ptr,
CDataType* c_ptr,
const GemmKernelArgs<DsGridPointer>& kargs,
const GemmKernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset)
{
static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!");
@@ -484,7 +476,7 @@ struct GemmKernel
if constexpr(std::is_same_v<DiLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(
ds_ptr[i],
(*ds_ptr)[i],
make_tuple(kargs.M, kargs.N),
make_tuple(kargs.stride_Ds[i], 1),
number<EpiloguePipeline::GetVectorSizeD(i)>{},
@@ -493,7 +485,7 @@ struct GemmKernel
else
{
return make_naive_tensor_view<address_space_enum::global>(
ds_ptr[i],
(*ds_ptr)[i],
make_tuple(kargs.M, kargs.N),
make_tuple(kargs.stride_Ds[i], 1),
number<EpiloguePipeline::GetVectorSizeD(i)>{},
@@ -689,10 +681,10 @@ struct GemmKernel
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
CK_TILE_DEVICE static void RunGemm(const ADataType* a_ptr,
const BDataType* b_ptr,
const DsGridPointer ds_ptr,
const DsGridPointer* ds_ptr,
CDataType* c_ptr,
void* smem_ptr_0,
const GemmKernelArgs<DsGridPointer>& kargs,
const GemmKernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset,
const index_t block_idx_m,
const index_t block_idx_n)
@@ -747,11 +739,11 @@ struct GemmKernel
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
CK_TILE_DEVICE static void RunGemm2LDS(const ADataType* a_ptr,
const BDataType* b_ptr,
const DsGridPointer ds_ptr,
const DsGridPointer* ds_ptr,
CDataType* c_ptr,
void* __restrict__ smem_ptr_0,
void* __restrict__ smem_ptr_1,
const GemmKernelArgs<DsGridPointer>& kargs,
const GemmKernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset,
const index_t block_idx_m,
const index_t block_idx_n)
@@ -784,7 +776,7 @@ struct GemmKernel
c_block_window, c_block_tile, d_block_window, smem_ptr_0);
}
CK_TILE_DEVICE void operator()(GemmKernelArgs<DsGridPointer> kargs) const
CK_TILE_DEVICE void operator()(GemmKernelArgs kargs) const
{
const auto blockId = __builtin_amdgcn_readfirstlane(blockIdx.x);
const auto [iM, iN] = TilePartitioner{kargs.M, kargs.N}.GetOutputTileIndex(blockId);
@@ -798,6 +790,7 @@ struct GemmKernel
static_cast<const ADataType*>(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset;
const BDataType* b_ptr =
static_cast<const BDataType*>(kargs.b_ptr) + splitk_batch_offset.b_k_split_offset;
const DsGridPointer* ds_ptr = reinterpret_cast<const DsGridPointer*>(kargs.ds_ptr);
CDataType* c_ptr = static_cast<CDataType*>(kargs.c_ptr);
@@ -811,7 +804,7 @@ struct GemmKernel
{
RunGemm2LDS(a_ptr,
b_ptr,
kargs.ds_ptr,
ds_ptr,
c_ptr,
smem_ptr_0,
smem_ptr_1,
@@ -827,7 +820,7 @@ struct GemmKernel
{
RunGemm2LDS<memory_operation_enum::atomic_add>(a_ptr,
b_ptr,
kargs.ds_ptr,
ds_ptr,
c_ptr,
smem_ptr_0,
smem_ptr_1,
@@ -842,15 +835,8 @@ struct GemmKernel
{
if(kargs.k_batch == 1)
{
RunGemm(a_ptr,
b_ptr,
kargs.ds_ptr,
c_ptr,
smem_ptr_0,
kargs,
splitk_batch_offset,
i_m,
i_n);
RunGemm(
a_ptr, b_ptr, ds_ptr, c_ptr, smem_ptr_0, kargs, splitk_batch_offset, i_m, i_n);
}
else
{
@@ -859,7 +845,7 @@ struct GemmKernel
{
RunGemm<memory_operation_enum::atomic_add>(a_ptr,
b_ptr,
kargs.ds_ptr,
ds_ptr,
c_ptr,
smem_ptr_0,
kargs,

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

@@ -13,12 +13,12 @@ namespace ck_tile {
struct GemmTransKernelArg
{
GemmKernelArgs<> group_karg;
GemmKernelArgs group_karg;
ck_tile::index_t block_start;
ck_tile::index_t block_end;
GemmTransKernelArg() = delete;
GemmTransKernelArg(GemmKernelArgs<>&& karg, index_t bl_start, index_t bl_end)
GemmTransKernelArg(GemmKernelArgs&& karg, index_t bl_start, index_t bl_end)
: group_karg{karg}, block_start{bl_start}, block_end{bl_end}
{
}
@@ -106,18 +106,18 @@ struct GroupedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
grid_size += grid_size_grp;
auto karg = GemmKernelArgs<>{type_convert<const ADataType*>(gemm_descs[i].a_ptr),
type_convert<const BDataType*>(gemm_descs[i].b_ptr),
{},
type_convert<CDataType*>(gemm_descs[i].c_ptr),
M,
N,
K,
stride_a,
stride_b,
{},
stride_c,
gemm_descs[i].k_batch};
auto karg = GemmKernelArgs{type_convert<const ADataType*>(gemm_descs[i].a_ptr),
type_convert<const BDataType*>(gemm_descs[i].b_ptr),
{},
type_convert<CDataType*>(gemm_descs[i].c_ptr),
M,
N,
K,
stride_a,
stride_b,
{},
stride_c,
gemm_descs[i].k_batch};
gemm_kernel_args_.emplace_back(std::move(karg), block_start, block_end);
}

5
test/ck_tile/gemm/test_gemm_pipeline_util.hpp
View File

@@ -82,7 +82,8 @@ class TestCkTileGemmPipeline : public ::testing::Test
// TODO: expose tile size through test t-param ?
template <bool PadM, bool PadN, bool PadK>
void invoke_gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s)
void invoke_gemm(const ck_tile::GemmHostArgs</*NumDTensor = 0*/>& args,
const ck_tile::stream_config& s)
{
// TODO: This should be parameterized in tests
constexpr ck_tile::index_t M_Tile = 256;
@@ -423,7 +424,7 @@ class TestCkTileGemmPipeline : public ::testing::Test
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();
ck_tile::GemmHostArgs<> args;
ck_tile::GemmHostArgs</*NumDTensor = 0*/> args;
args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer();
args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer();
args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer();

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

@@ -47,7 +47,7 @@ class TestCkTileGroupedGemm : public ::testing::Test
static const ck_tile::index_t K_Warp_Tile = 8;
};
using grouped_gemm_kargs = ck_tile::GemmHostArgs<>;
using grouped_gemm_kargs = ck_tile::GemmHostArgs</*NumDTensor = 0*/>;
std::size_t get_workspace_size(const std::vector<grouped_gemm_kargs>& gemm_descs)
{
return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg);