diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index 3227efa06d..e520984605 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -22,7 +22,7 @@ template -float gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s) +float gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { // The kPadM, kPadN, kPadK & kBlockPerCu should also come from the Codegen part. diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index eb181b2f19..1170b56aa0 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -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& args, + const ck_tile::stream_config& s); diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 548d326f22..2d5daa30e4 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -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 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 -float gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s) +float gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { using GemmShape = ck_tile::TileGemmShape< diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index 93d43b0c1e..dc258b848f 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -9,7 +9,7 @@ namespace ck_tile { -struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs<> +struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs { CK_TILE_HOST BatchedGemmHostArgs() = default; CK_TILE_HOST BatchedGemmHostArgs(const void* a_ptr_, @@ -56,7 +56,7 @@ struct BatchedGemmKernel : public GemmKernel; - using GemmKernelArgs = typename ck_tile::GemmKernelArgs<>; + using GemmKernelArgs = typename ck_tile::GemmKernelArgs; using ADataType = typename Base::ADataType; using BDataType = typename Base::BDataType; diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index 15adf6c6bb..70b7da98b4 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -57,13 +57,11 @@ struct GemmHostArgs index_t k_batch; }; -// TODO: The parameter const DType ds_ptr could be treated as const void * -template > 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 + CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(const GemmHostArgs& hostArgs) { - DsGridPointer p_ds_grid; - static_for<0, NumDTensor, 1>{}([&](auto i) { - using DDataType_ = remove_cvref_t>; - p_ds_grid(i) = static_cast(hostArgs.ds_ptr[i]); - }); - - return GemmKernelArgs{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(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& 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& kargs) + CK_TILE_HOST static bool IsSupportedArgument(const GemmKernelArgs& kargs) { if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 && is_any_of::value) @@ -388,9 +380,9 @@ struct GemmKernel template 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& 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) { return make_naive_tensor_view( - ds_ptr[i], + (*ds_ptr)[i], make_tuple(kargs.M, kargs.N), make_tuple(kargs.stride_Ds[i], 1), number{}, @@ -493,7 +485,7 @@ struct GemmKernel else { return make_naive_tensor_view( - ds_ptr[i], + (*ds_ptr)[i], make_tuple(kargs.M, kargs.N), make_tuple(kargs.stride_Ds[i], 1), number{}, @@ -689,10 +681,10 @@ struct GemmKernel template 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& 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 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& 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 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(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset; const BDataType* b_ptr = static_cast(kargs.b_ptr) + splitk_batch_offset.b_k_split_offset; + const DsGridPointer* ds_ptr = reinterpret_cast(kargs.ds_ptr); CDataType* c_ptr = static_cast(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(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(a_ptr, b_ptr, - kargs.ds_ptr, + ds_ptr, c_ptr, smem_ptr_0, kargs, diff --git a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp index c26fe37859..b7dc9bdabc 100644 --- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp @@ -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{type_convert(gemm_descs[i].a_ptr), - type_convert(gemm_descs[i].b_ptr), - {}, - type_convert(gemm_descs[i].c_ptr), - M, - N, - K, - stride_a, - stride_b, - {}, - stride_c, - gemm_descs[i].k_batch}; + auto karg = GemmKernelArgs{type_convert(gemm_descs[i].a_ptr), + type_convert(gemm_descs[i].b_ptr), + {}, + type_convert(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); } diff --git a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp index fe92c13888..a1e4bdb574 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp @@ -82,7 +82,8 @@ class TestCkTileGemmPipeline : public ::testing::Test // TODO: expose tile size through test t-param ? template - void invoke_gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s) + void invoke_gemm(const ck_tile::GemmHostArgs& 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 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(); diff --git a/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp b/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp index 6c9a1e1a6d..ce1e48eb8e 100644 --- a/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp +++ b/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp @@ -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; std::size_t get_workspace_size(const std::vector& gemm_descs) { return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg);