From 96ac4a44c8cc37ebc91713f24f520058fdf2af46 Mon Sep 17 00:00:00 2001 From: Mateusz Ozga Date: Fri, 21 Mar 2025 15:34:13 +0000 Subject: [PATCH] Initial commit multiple-d gemm --- example/ck_tile/03_gemm/gemm_basic.cpp | 16 +- example/ck_tile/03_gemm/gemm_utils.hpp | 2 +- example/ck_tile/03_gemm/run_gemm_example.inc | 73 ++-- example/ck_tile/03_gemm/universal_gemm.cpp | 17 +- .../ck_tile/16_batched_gemm/batched_gemm.cpp | 16 +- .../ck_tile/16_batched_gemm/batched_gemm.hpp | 1 + .../run_batched_gemm_example.inc | 62 +++- example/ck_tile/17_grouped_gemm/README.md | 2 +- .../ck_tile/17_grouped_gemm/grouped_gemm.cpp | 14 +- .../ck_tile/17_grouped_gemm/grouped_gemm.hpp | 5 +- .../run_grouped_gemm_example.inc | 40 ++- .../ck_tile/18_multi_d_gemm/CMakeLists.txt | 1 + example/ck_tile/18_multi_d_gemm/README.md | 33 ++ .../ck_tile/18_multi_d_gemm/multi_d_gemm.cpp | 313 ++++++++++++++++++ .../ck_tile/18_multi_d_gemm/multi_d_gemm.hpp | 67 ++++ .../run_multi_d_gemm_example.inc | 244 ++++++++++++++ example/ck_tile/18_multi_d_gemm/utils.hpp | 63 ++++ example/ck_tile/CMakeLists.txt | 1 + .../ck_tile/core/tensor/tile_elementwise.hpp | 31 ++ .../ck_tile/host/reference/reference_gemm.hpp | 36 ++ .../unary_element_wise_operation.hpp | 75 +++++ .../ops/epilogue/cshuffle_epilogue.hpp | 56 +++- .../ops/gemm/kernel/batched_gemm_kernel.hpp | 24 +- .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 307 +++++++++++++---- .../ops/gemm/kernel/grouped_gemm_kernel.hpp | 40 ++- test/ck_tile/CMakeLists.txt | 1 + .../batched_gemm/test_batched_gemm_util.hpp | 6 + test/ck_tile/gemm/test_gemm_pipeline_util.hpp | 10 +- .../grouped_gemm/test_grouped_gemm_util.hpp | 22 +- 29 files changed, 1413 insertions(+), 165 deletions(-) create mode 100644 example/ck_tile/18_multi_d_gemm/CMakeLists.txt create mode 100644 example/ck_tile/18_multi_d_gemm/README.md create mode 100644 example/ck_tile/18_multi_d_gemm/multi_d_gemm.cpp create mode 100644 example/ck_tile/18_multi_d_gemm/multi_d_gemm.hpp create mode 100644 example/ck_tile/18_multi_d_gemm/run_multi_d_gemm_example.inc create mode 100644 example/ck_tile/18_multi_d_gemm/utils.hpp diff --git a/example/ck_tile/03_gemm/gemm_basic.cpp b/example/ck_tile/03_gemm/gemm_basic.cpp index 69051423fb..3227efa06d 100644 --- a/example/ck_tile/03_gemm/gemm_basic.cpp +++ b/example/ck_tile/03_gemm/gemm_basic.cpp @@ -14,12 +14,16 @@ template -float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) + typename DsLayout, + typename CLayout, + typename CDEElementWise = ck_tile::element_wise::PassThrough> +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. constexpr bool kPadM = false; @@ -50,15 +54,21 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& using CodegenGemmTraits = ck_tile::TileGemmTraits; + using CodegenPipelineProblem = ck_tile:: GemmPipelineProblem; + using CodegenGemmPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; - using GemmEpilogue = ck_tile::CShuffleEpilogue< + + using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem& 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 6cb40e45d1..548d326f22 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -145,11 +145,14 @@ void permute_vectors_i4x4_b(Tensor& tensor) template + typename DsLayout, + typename CLayout, + typename CDEElementWise = ck_tile::element_wise::PassThrough> float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, ck_tile::DeviceMem& b_k_n_dev_buf, ck_tile::DeviceMem& c_m_n_dev_buf, @@ -163,21 +166,30 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, int n_warmup, int n_repeat) { - 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(); - args.k_batch = kbatch; - args.M = M; - args.N = N; - args.K = K; - args.stride_A = stride_A; - args.stride_B = stride_B; - args.stride_C = 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_calc( - args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + gemm(args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); std::size_t flop = std::size_t(2) * M * N * K; std::size_t num_byte = @@ -296,19 +308,26 @@ int run_gemm_example_with_layouts(int argc, c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); - invoke_gemm( - a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat); + invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat); c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); bool pass = true; diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index eef8d3b60e..9fd7adc281 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -14,12 +14,16 @@ template -float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) + typename DsLayout, + typename CLayout, + typename CDEElementWise = ck_tile::element_wise::PassThrough> +float gemm(const ck_tile::GemmHostArgs<>& args, const ck_tile::stream_config& s) + { using GemmShape = ck_tile::TileGemmShape< ck_tile::sequence, @@ -28,17 +32,19 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& sequence, GemmConfig::PermuteA, GemmConfig::PermuteB>; + using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner; - using Traits = ck_tile::TileGemmTraits; + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmPipelineProblem = ck_tile::GemmPipelineProblem; @@ -78,9 +85,12 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem>; + using Kernel = ck_tile::GemmKernel; auto kargs = Kernel::MakeKernelArgs(args); diff --git a/example/ck_tile/16_batched_gemm/batched_gemm.cpp b/example/ck_tile/16_batched_gemm/batched_gemm.cpp index a0cd18ec74..65d1f5083e 100644 --- a/example/ck_tile/16_batched_gemm/batched_gemm.cpp +++ b/example/ck_tile/16_batched_gemm/batched_gemm.cpp @@ -15,7 +15,16 @@ #include "ck_tile/host.hpp" #include "batched_gemm.hpp" -template +template float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stream_config& s) { #if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) @@ -121,12 +130,16 @@ float batched_gemm(const ck_tile::BatchedGemmHostArgs& args, const ck_tile::stre tail_number_v>; using GemmPipeline = GEMM_PIPELINE; + using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem>; + using Kernel = ck_tile::BatchedGemmKernel; auto kargs = Kernel::MakeKernelArgs(args); diff --git a/example/ck_tile/16_batched_gemm/batched_gemm.hpp b/example/ck_tile/16_batched_gemm/batched_gemm.hpp index 0999c7ad3b..78d915e873 100644 --- a/example/ck_tile/16_batched_gemm/batched_gemm.hpp +++ b/example/ck_tile/16_batched_gemm/batched_gemm.hpp @@ -8,6 +8,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/host/kernel_launch.hpp" #include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #define CK_TILE_PIPELINE_COMPUTE_V3 1 #define CK_TILE_PIPELINE_MEMORY 2 diff --git a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc index 16a31e519a..da1f0aa105 100644 --- a/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc +++ b/example/ck_tile/16_batched_gemm/run_batched_gemm_example.inc @@ -23,7 +23,16 @@ auto calculate_rtol_atol(const ck_tile::index_t K, return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); } -template +template float invoke_batched_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, ck_tile::DeviceMem& b_k_n_dev_buf, ck_tile::DeviceMem& c_m_n_dev_buf, @@ -57,7 +66,16 @@ float invoke_batched_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, args.batch_stride_C = batch_stride_C; args.batch_count = batch_count; - float ave_time = batched_gemm( + float ave_time = batched_gemm( args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); std::string op_name{"Batched Gemm"}; @@ -169,22 +187,30 @@ int run_batched_gemm_example_with_layouts(int argc, c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); - invoke_batched_gemm(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - batch_stride_A, - batch_stride_B, - batch_stride_C, - batch_count, - kbatch, - n_warmup, - n_repeat); + invoke_batched_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + batch_stride_A, + batch_stride_B, + batch_stride_C, + batch_count, + kbatch, + n_warmup, + n_repeat); c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); bool pass = true; diff --git a/example/ck_tile/17_grouped_gemm/README.md b/example/ck_tile/17_grouped_gemm/README.md index d1a0458eda..59396a558b 100644 --- a/example/ck_tile/17_grouped_gemm/README.md +++ b/example/ck_tile/17_grouped_gemm/README.md @@ -1,6 +1,6 @@ # Grouped CShuffle GEMM -This folder contains example for Grouped GEMM using ck_tile tile-programming implementation. Currently, it only supports the basic feature of the CK Tile GEMM, but creates the placeholders for the future support on different GEMM pipeline and different GEMM modules. In the near future, we will gradually migrate all the GEMM features from old CK to CK Tile. +This folder contains example for Grouped GEMM using ck_tile tile-programming implementation. ## build ``` diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp b/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp index 2a9903362d..8a9dd35df5 100644 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp @@ -21,7 +21,16 @@ std::size_t get_workspace_size(const std::vector& gemm_descs return gemm_descs.size() * sizeof(ck_tile::GemmTransKernelArg); } -template +template float grouped_gemm(const std::vector& gemm_descs, const ck_tile::stream_config& s, void* p_workspace_) @@ -132,9 +141,12 @@ float grouped_gemm(const std::vector& gemm_descs, using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem; auto create_args(int argc, char* argv[]) { diff --git a/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc b/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc index f068510d26..1e969bd463 100644 --- a/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc +++ b/example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc @@ -30,7 +30,16 @@ auto calculate_rtol_atol(const ck_tile::index_t K, return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); } -template +template float invoke_gemm(int n_warmup, int n_repeat, int group_count, @@ -40,10 +49,19 @@ float invoke_gemm(int n_warmup, ck_tile::DeviceMem gemm_workspace; gemm_workspace.Realloc(get_workspace_size(args)); - float ave_time = grouped_gemm( - args, - ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}, - gemm_workspace.GetDeviceBuffer()); + float ave_time = + grouped_gemm(args, + ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}, + gemm_workspace.GetDeviceBuffer()); std::string op_name{"Grouped Gemm"}; @@ -172,10 +190,18 @@ int run_grouped_gemm_example_with_layouts(int argc, // TODO Add support for kbatch > 1 in grouped gemm 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]}); + {p_a, p_b, {}, p_c, k_batch, M, N, K, stride_As[i], stride_Bs[i], {}, stride_Cs[i]}); } - invoke_gemm(warmup, repeat, group_count, gemm_descs); + invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(warmup, repeat, group_count, gemm_descs); for(int i = 0; i < group_count; i++) { diff --git a/example/ck_tile/18_multi_d_gemm/CMakeLists.txt b/example/ck_tile/18_multi_d_gemm/CMakeLists.txt new file mode 100644 index 0000000000..e701fc16f6 --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/CMakeLists.txt @@ -0,0 +1 @@ +add_executable(tile_example_multi_d_gemm EXCLUDE_FROM_ALL multi_d_gemm.cpp) diff --git a/example/ck_tile/18_multi_d_gemm/README.md b/example/ck_tile/18_multi_d_gemm/README.md new file mode 100644 index 0000000000..b9d6a4c8b4 --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/README.md @@ -0,0 +1,33 @@ +#Multiple D GEMM + +This folder contains example for Multiple D GEMM using ck_tile tile-programming implementation. + +## build +``` +#in the root of ck_tile +mkdir build && cd build +#you can replace < arch> with the appropriate architecture(for example gfx90a or gfx942) or \ + leave it blank +sh ../script/cmake-ck-dev.sh ../ +#The basic pipeline method on the gemm calculation +make tile_example_multi_d_gemm -j +``` +This will result in an executable `build/bin/tile_example_multi_d_gemm` + +## example +``` +args: + -m M dimensions - (Default: 3840) + -n N dimensions - (Default: 4096) + -k K dimensions - (Default: 4096) +-a_layout Tensor A layout (default:R) +-b_layout Tensor B layout (default:C) +-c_layout Tensor C layout (default:R) +-stride_a Tensor A strides - (Default: 0) +-stride_b Tensor B strides - (Default: 0) +-stride_c Tensor C strides - (Default: 0) +-stride_d Tensor C strides - (Default: 0) +-validate 0. No validation, 1. Validation on GPU. (Default: 1) + -warmup Number of iterations before benchmark the kernel. (Default: 10) + -repeat Number of iterations to benchmark the kernel. (Default: 100) +``` diff --git a/example/ck_tile/18_multi_d_gemm/multi_d_gemm.cpp b/example/ck_tile/18_multi_d_gemm/multi_d_gemm.cpp new file mode 100644 index 0000000000..b545a727f1 --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/multi_d_gemm.cpp @@ -0,0 +1,313 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/host.hpp" +#include "multi_d_gemm.hpp" +#include "utils.hpp" + +template +auto multiple_d_gemm(const multiple_d_gemm_kargs& args, const ck_tile::stream_config& s) -> float +{ +#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) + // Memory friendly for Interwave scheduler + constexpr ck_tile::index_t M_Tile = 128; + constexpr ck_tile::index_t N_Tile = 32; + constexpr ck_tile::index_t K_Tile = 64; + + constexpr ck_tile::index_t M_Warp = 4; + constexpr ck_tile::index_t N_Warp = 1; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = 32; + constexpr ck_tile::index_t N_Warp_Tile = 32; + constexpr ck_tile::index_t K_Warp_Tile = 8; + + constexpr bool DoubleSmemBuffer = false; +#endif +#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3) + // Compute friendly for Intrawave scheduler + constexpr ck_tile::index_t M_Tile = 256; + constexpr ck_tile::index_t N_Tile = 256; + constexpr ck_tile::index_t K_Tile = 64; + + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = 32; + constexpr ck_tile::index_t N_Warp_Tile = 32; + constexpr ck_tile::index_t K_Warp_Tile = 16; + + constexpr bool DoubleSmemBuffer = false; +#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4) + // Compute friendly for Intrawave scheduler + // Using the ping pong reader in the lds level + constexpr ck_tile::index_t M_Tile = 256; + constexpr ck_tile::index_t N_Tile = 256; + constexpr ck_tile::index_t K_Tile = 32; + + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = 32; + constexpr ck_tile::index_t N_Warp_Tile = 32; + constexpr ck_tile::index_t K_Warp_Tile = 16; + + constexpr bool DoubleSmemBuffer = true; +#endif + + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + + constexpr bool TransposeC = false; + + constexpr int kBlockPerCu = 1; + constexpr ck_tile::index_t TileParitionerGroupNum = 8; + constexpr ck_tile::index_t TileParitionerM01 = 4; + + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; + + using Traits = ck_tile::TileGemmTraits; + + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmPipelineProblem = + ck_tile::GemmPipelineProblem; + + using BaseGemmPipeline = UNIVERSAL_GEMM_PIPELINE; + + const ck_tile::index_t k_grain = args.k_batch * K_Tile; + const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + float ave_time{0}; + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr auto scheduler = GEMM_PIPELINE_SCHEDULER; + + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + + using GemmPipeline = GEMM_PIPELINE; + + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem>; + + using Kernel = ck_tile::GemmKernel; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + constexpr dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args:" + << " grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" + << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" + << std::endl; + } + + ave_time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + return ave_time; + }; + + if(has_hot_loop) + { +#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3) + if(tail_num == ck_tile::TailNumber::Full) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else if(tail_num == ck_tile::TailNumber::Odd) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else if(tail_num == ck_tile::TailNumber::Even) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else + { + std::ostringstream err; + err << "For compute pipeline tail number should always be Full, but have \"" << tail_num + << "\" which is not supported! PrefetchStages: " << BaseGemmPipeline::PrefetchStages + << "\n File: " << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } +#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) + // Tail pipeline One to Seven + if(tail_num == ck_tile::TailNumber::One) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else if(tail_num == ck_tile::TailNumber::Full) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + + if constexpr(BaseGemmPipeline::PrefetchStages > 2) + { + if(tail_num == ck_tile::TailNumber::Two) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } + if constexpr(BaseGemmPipeline::PrefetchStages > 3) + { + if(tail_num == ck_tile::TailNumber::Three) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } + if constexpr(BaseGemmPipeline::PrefetchStages > 4) + { + if(tail_num == ck_tile::TailNumber::Four) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } + if constexpr(BaseGemmPipeline::PrefetchStages > 5) + { + if(tail_num == ck_tile::TailNumber::Five) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } + if constexpr(BaseGemmPipeline::PrefetchStages > 6) + { + if(tail_num == ck_tile::TailNumber::Six) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } + if constexpr(BaseGemmPipeline::PrefetchStages > 7) + { + if(tail_num == ck_tile::TailNumber::Seven) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + } +#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4) + if(tail_num == ck_tile::TailNumber::Three) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } +#endif + } + else + { + if(tail_num == ck_tile::TailNumber::Full) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else if(tail_num == ck_tile::TailNumber::Odd) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else if(tail_num == ck_tile::TailNumber::Even) + { + Run(ck_tile::bool_constant{}, + ck_tile::integral_constant{}); + } + else + { + std::ostringstream err; + err << "Num K loop must be larger than number of prefetech stages." + << "\n PrefetchStages: " << BaseGemmPipeline::PrefetchStages + << "\n File: " << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } + } + + return ave_time; +} + +#include "run_multi_d_gemm_example.inc" + +int main(int argc, char* argv[]) { return !run_multiple_d_gemm_example(argc, argv); } diff --git a/example/ck_tile/18_multi_d_gemm/multi_d_gemm.hpp b/example/ck_tile/18_multi_d_gemm/multi_d_gemm.hpp new file mode 100644 index 0000000000..8a3aa981a8 --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/multi_d_gemm.hpp @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" + +#define CK_TILE_PIPELINE_COMPUTE_V3 1 +#define CK_TILE_PIPELINE_MEMORY 2 +#define CK_TILE_PIPELINE_COMPUTE_V4 3 + +#ifndef CK_TILE_PIPELINE_DEFAULT +#define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 +#endif + +#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY) +#define GEMM_PIPELINE ck_tile::GemmPipelineAgBgCrMem +#define UNIVERSAL_GEMM_PIPELINE ck_tile::BaseGemmPipelineAgBgCrMem +#define GEMM_PIPELINE_SCHEDULER ck_tile::GemmPipelineScheduler::Interwave +#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3) +#define GEMM_PIPELINE ck_tile::GemmPipelineAgBgCrCompV3 +#define UNIVERSAL_GEMM_PIPELINE ck_tile::BaseGemmPipelineAgBgCrCompV3 +#define GEMM_PIPELINE_SCHEDULER ck_tile::GemmPipelineScheduler::Intrawave +#elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4) +#define GEMM_PIPELINE ck_tile::GemmPipelineAgBgCrCompV4 +#define UNIVERSAL_GEMM_PIPELINE ck_tile::BaseGemmPipelineAgBgCrCompV4 +#define GEMM_PIPELINE_SCHEDULER ck_tile::GemmPipelineScheduler::Intrawave +#else +#error "unsupported CK_TILE_PIPELINE_DEFAULT value" +#endif + +using ADataType = ck_tile::half_t; +using BDataType = ck_tile::half_t; +using D0DataType = ck_tile::half_t; +using D1DataType = ck_tile::half_t; +using CDataType = ck_tile::half_t; +using DsDataType = ck_tile::tuple; +using AccDataType = float; + +auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "m dimension") + .insert("n", "4096", "n dimension") + .insert("k", "4096", "k dimension") + .insert("a_layout", "R", "A tensor data layout - Row by default") + .insert("b_layout", "C", "B tensor data layout - Col by default") + .insert("c_layout", "R", "C tensor data layout - Row by default") + .insert("stride_a", "0", "Tensor A stride") + .insert("stride_b", "0", "Tensor B stride") + .insert("stride_d", "0", "Tensor Ds stride") + .insert("stride_c", "0", "Tensor C stride") + .insert("v", "1", "0. No validation, 1. Validation on GPU") + .insert("warmup", "50", "number of iterations before benchmark the kernel") + .insert("repeat", "100", "number of iterations to benchmark the kernel"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} + +using multiple_d_gemm_kargs = ck_tile::GemmHostArgs; + +float multiple_d_gemm(const multiple_d_gemm_kargs& kargs, const ck_tile::stream_config& s); diff --git a/example/ck_tile/18_multi_d_gemm/run_multi_d_gemm_example.inc b/example/ck_tile/18_multi_d_gemm/run_multi_d_gemm_example.inc new file mode 100644 index 0000000000..0e61fe7c9d --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/run_multi_d_gemm_example.inc @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once +#include + +template +float invoke_multi_d_gemm(const void* a_m_k_dev_buf, + const void* b_k_n_dev_buf, + const std::array& d_m_n_dev_buf, + void* c_m_n_dev_buf, + ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + ck_tile::index_t StrideA, + ck_tile::index_t StrideB, + const std::array StrideDs, + ck_tile::index_t StrideC, + int n_warmup, + int n_repeat) +{ + multiple_d_gemm_kargs gemm_descs({a_m_k_dev_buf, + b_k_n_dev_buf, + d_m_n_dev_buf, + c_m_n_dev_buf, + /*kbatch */ 1, + M, + N, + K, + StrideA, + StrideB, + StrideDs, + StrideC}); + + float ave_time = multiple_d_gemm( + gemm_descs, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + + std::string op_name{"Multiple-D Gemm"}; + static constexpr ck_tile::index_t NumDTensor = DsDataType::size(); + + std::size_t flop = 0, num_btype = 0; + + flop += std::size_t(2) * M * N * K; + + ck_tile::static_for<0, NumDTensor, 1>{}([&](auto i) { + num_btype += sizeof(ck_tile::remove_cvref_t>) * M * N; + }); + + num_btype += sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Run Multiple-D Gemm kernel with:\n"; + std::cout << "M =" << M << " N =" << N << " K =" << K << "\n"; + std::cout << "StrideA = " << StrideA << " StrideB = " << StrideB << " StrideC = " << StrideC + << "\n"; + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << "\n"; + + return ave_time; +} + +template +int run_multiple_d_gemm_example_with_layouts(int argc, + char* argv[], + const ALayout a_layout = ALayout{}, + const BLayout b_layout = BLayout{}, + const D0Layout d0_layout = D0Layout{}, + const D1Layout d1_layout = D1Layout{}, + const CLayout c_layout = CLayout{}) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + using CDEElementWiseFn = ck_tile::element_wise::ElementWiseAdd; + using DsLayout = ck_tile::tuple; + + ck_tile::index_t M = arg_parser.get_int("m"); + ck_tile::index_t N = arg_parser.get_int("n"); + ck_tile::index_t K = arg_parser.get_int("k"); + + ck_tile::index_t StrideA = arg_parser.get_int("stride_a"); + ck_tile::index_t StrideB = arg_parser.get_int("stride_b"); + ck_tile::index_t StrideD = arg_parser.get_int("stride_d"); + ck_tile::index_t StrideC = arg_parser.get_int("stride_c"); + + ck_tile::index_t StrideD0 = StrideD; + ck_tile::index_t StrideD1 = StrideD; + + const int n_warmup = arg_parser.get_int("warmup"); + const int n_repeat = arg_parser.get_int("repeat"); + + StrideA = f_get_default_stride(M, N, StrideA, a_layout); + StrideB = f_get_default_stride(K, N, StrideB, b_layout); + StrideD0 = f_get_default_stride(M, N, StrideD, d0_layout); + StrideD1 = f_get_default_stride(M, N, StrideD, d1_layout); + StrideC = f_get_default_stride(M, N, StrideC, c_layout); + + ck_tile::HostTensor a_m_k_tesnor(f_host_tensor_descriptor(M, K, StrideA, a_layout)); + ck_tile::HostTensor b_k_n_tensors(f_host_tensor_descriptor(K, N, StrideB, b_layout)); + ck_tile::HostTensor d0_m_n_tensors( + f_host_tensor_descriptor(M, N, StrideD0, d0_layout)); + ck_tile::HostTensor d1_m_n_tensors( + f_host_tensor_descriptor(M, N, StrideD1, d1_layout)); + ck_tile::HostTensor c_m_n_device_result( + f_host_tensor_descriptor(M, N, StrideC, c_layout)); + + ck_tile::FillUniformDistribution{-5.f, 5.f}(a_m_k_tesnor); + ck_tile::FillUniformDistribution{-5.f, 5.f}(b_k_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(d0_m_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(d1_m_n_tensors); + + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k_tesnor.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d0_m_n_dev_buf(d0_m_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d1_m_n_dev_buf(d1_m_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_device_result.get_element_space_size_in_bytes()); + + a_m_k_dev_buf.ToDevice(a_m_k_tesnor.mData.data()); + b_k_n_dev_buf.ToDevice(b_k_n_tensors.mData.data()); + d0_m_n_dev_buf.ToDevice(d0_m_n_tensors.mData.data()); + d1_m_n_dev_buf.ToDevice(d1_m_n_tensors.mData.data()); + + c_m_n_dev_buf.SetZero(); + c_m_n_device_result.SetZero(); + + std::array ds_ptr_buf = {d0_m_n_dev_buf.GetDeviceBuffer(), + d1_m_n_dev_buf.GetDeviceBuffer()}; + + std::array stridesDs = {StrideD0, StrideD1}; + + invoke_multi_d_gemm(a_m_k_dev_buf.GetDeviceBuffer(), + b_k_n_dev_buf.GetDeviceBuffer(), + ds_ptr_buf, + c_m_n_dev_buf.GetDeviceBuffer(), + M, + N, + K, + StrideA, + StrideB, + stridesDs, + StrideC, + n_warmup, + n_repeat); + + c_m_n_dev_buf.FromDevice(c_m_n_device_result.data()); + + ck_tile::HostTensor c_m_n_host_ref( + f_host_tensor_descriptor(M, N, StrideC, c_layout)); + c_m_n_host_ref.SetZero(); + + ck_tile::reference_gemm_multiple_d< + ADataType, + BDataType, + std::tuple, ck_tile::HostTensor>, + AccDataType, + CDataType, + CDEElementWiseFn>( + a_m_k_tesnor, b_k_n_tensors, std::tie(d0_m_n_tensors, d1_m_n_tensors), c_m_n_host_ref); + + bool pass{true}; + if(arg_parser.get_int("v")) + { + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + (void)(max_accumulated_value); + const auto rtol_atol = calculate_rtol_atol(K, 1, max_accumulated_value); + + pass &= ck_tile::check_err(c_m_n_device_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) + << std::endl; + std::cout << "Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + std::cout << "The CPU veification result is: " << (pass ? "correct" : "fail") << std::endl; + } + return pass; +} + +int run_multiple_d_gemm_example(int argc, char* argv[]) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + + const std::string a_layout = arg_parser.get_str("a_layout"); + const std::string b_layout = arg_parser.get_str("b_layout"); + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + + if(a_layout == "R" && b_layout == "C") + { + return run_multiple_d_gemm_example_with_layouts( + argc, argv, Row{}, Col{}, Row{}, Row{}, Row{}); + } + else + { + throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!"); + } +} diff --git a/example/ck_tile/18_multi_d_gemm/utils.hpp b/example/ck_tile/18_multi_d_gemm/utils.hpp new file mode 100644 index 0000000000..9b78f2fe2e --- /dev/null +++ b/example/ck_tile/18_multi_d_gemm/utils.hpp @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +template +constexpr auto +f_host_tensor_descriptor(std::size_t row, std::size_t col, std::size_t stride, TLayout layout) +{ + using namespace ck_tile::literals; + + if constexpr(std::is_same_v) + { + return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride}); + } +} +template +constexpr auto +f_get_default_stride(std::size_t row, std::size_t col, std::size_t stride, TLayout layout) +{ + if(stride == 0) + { + if constexpr(std::is_same_v) + { + return col; + } + else + { + return row; + } + } + else + return stride; +} + +auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeType = + std::conditional_t; + // Calculate thresholds + + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} diff --git a/example/ck_tile/CMakeLists.txt b/example/ck_tile/CMakeLists.txt index 7f4ba2ed35..7f66f1380b 100644 --- a/example/ck_tile/CMakeLists.txt +++ b/example/ck_tile/CMakeLists.txt @@ -17,4 +17,5 @@ add_subdirectory(14_moe_smoothquant) add_subdirectory(15_fused_moe) add_subdirectory(16_batched_gemm) add_subdirectory(17_grouped_gemm) +add_subdirectory(18_multi_d_gemm) add_subdirectory(35_batched_transpose) diff --git a/include/ck_tile/core/tensor/tile_elementwise.hpp b/include/ck_tile/core/tensor/tile_elementwise.hpp index 79018b9ced..0335c223da 100644 --- a/include/ck_tile/core/tensor/tile_elementwise.hpp +++ b/include/ck_tile/core/tensor/tile_elementwise.hpp @@ -59,6 +59,37 @@ CK_TILE_DEVICE auto tile_elementwise_in(const InElementFunc& in_element_func, return out_dstr_tensor; } +/** + * @brief Template function that "unpacks" a tuple and applies an element-wise operation. + * + * @param in_element_func Function to apply element-wise. + * @param t Tuple containing elements to process. + * @return Calls tile_elementwise_inout with unpacked tuple elements. + */ +template +CK_TILE_DEVICE auto tile_elementwise_in_out_unpack_tuple(const InElementFunc& in_element_func, + const Tuple& t, + std::index_sequence) +{ + return tile_elementwise_inout(in_element_func, t[number{}]...); +} + +/** + * @brief Template function that "unpacks" a tuple and applies an element-wise operation. + * + * @param in_element_func Function to apply element-wise. + * @param t Tuple containing elements to process. + * @return Calls the overloaded function, passing an index sequence. + */ +template +CK_TILE_DEVICE auto tile_elementwise_in_out_unpack_tuple(const InElementFunc& in_element_func, + const Tuple& t) +{ + static constexpr auto size = std::tuple_size::value; + return tile_elementwise_in_out_unpack_tuple( + in_element_func, t, std::make_index_sequence{}); +} + template CK_TILE_DEVICE void set_tile(DstrTensors& dstr_tensor, const T& value) { diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index fe5077083c..b2e0ec3fb8 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -71,6 +71,42 @@ CK_TILE_HOST void reference_gemm(const HostTensor& a_m_k, make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); } +template +CK_TILE_HOST void +reference_gemm_multiple_d(const HostTensor& a_m_k, + const HostTensor& b_k_n, + const DsDataType& ds_m_n, + HostTensor& c_m_n, + const ACCElementOp& acc_element_op = {}) +{ + const std::size_t M = a_m_k.get_length(0); + const std::size_t N = b_k_n.get_length(1); + const std::size_t K = a_m_k.get_length(1); + + auto f_mn = [&](auto m, auto n) { + AccDataType v_acc = 0; + + for(std::size_t k = 0; k < K; ++k) + { + ADataType v_a = a_m_k(m, k); + BDataType v_b = b_k_n(k, n); + + v_acc += + ck_tile::type_convert(v_a) * ck_tile::type_convert(v_b); + } + std::apply([&](auto&... di) { ((acc_element_op(v_acc, di(m, n))), ...); }, ds_m_n); + + c_m_n(m, n) = ck_tile::type_convert(v_acc); + }; + + make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); +} + template CK_TILE_DEVICE OutputArray operator()(InputArray const& Input) { return convert(Input); } }; #endif + +/** + * @brief Struct defining element-wise addition operations + */ +struct ElementWiseAdd +{ + /** + * @brief Function call operator for element-wise addition with 3 inputs + * + * @param r Output element (result) + * @param a first input + * @param b second input + * @param c third input + * + * @note [return] Perform element-wise addition and store the result in 'r' + */ + template + CK_TILE_DEVICE auto operator()(ResT& r, const ParamT& a, const ParamT& b, const ParamT& c) const + -> void + { + r = a + b + c; + } + + /** + * @brief Function call operator for element-wise addition with 3 inputs + * + * @param r Output element (result) + * @param a first input + * + * @note [return] Perform element-wise addition and store the result in 'r' + */ + template + CK_TILE_HOST auto operator()(ResT& r, const ParamT& a) const -> void + { + r += a; + } +}; + +/** + * @brief Struct defining element-wise multiplication operations + */ +struct ElementWiseMul +{ + /** + * @brief Function call operator for element-wise multiplication with 3 inputs + * + * @param r Output element (result) + * @param a first input + * @param b second input + * @param c third input + * + * @note [return] Perform element-wise multiplication and store the result in 'r' + */ + template + CK_TILE_DEVICE auto operator()(ResT& r, const ParamT& a, const ParamT& b, const ParamT& c) const + -> void + { + r = a + b + c; + } + + /** + * @brief Function call operator for element-wise addition with 3 inputs + * + * @param r Output element (result) + * @param a first input + * + * @note [return] Perform element-wise addition and store the result in 'r' + */ + template + CK_TILE_HOST auto operator()(ResT& r, const ParamT& a) const -> void + { + r *= a; + } +}; + } // namespace element_wise } // namespace ck_tile diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 0081edcb2e..07242ec636 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -11,9 +11,12 @@ namespace ck_tile { template ; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; + using DsDataType = remove_cvref_t; + using DsLayout = remove_cvref_t; using CLayout = remove_cvref_t; + using ABDELementWise = remove_cvref_t; static constexpr index_t kBlockSize = kBlockSize_; static constexpr index_t kMPerBlock = kM_; static constexpr index_t kNPerBlock = kN_; @@ -39,6 +45,7 @@ struct CShuffleEpilogueProblem static constexpr index_t kNPerXdl = kNPerXdl_; static constexpr index_t kKPerXdl = kKPerXdl_; static constexpr index_t isCTransposed = isCTransposed_; + static constexpr index_t NumDTensor = DsDataType::size(); }; template @@ -49,9 +56,12 @@ struct CShuffleEpilogue using BDataType = remove_cvref_t; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; + using DsDataType = remove_cvref_t; + using DsLayout = remove_cvref_t; using BTypeToUse = std::conditional_t, ODataType, BDataType>; using CLayout = remove_cvref_t; + using ABDELementWise = remove_cvref_t; static constexpr index_t kBlockSize = Problem::kBlockSize; static constexpr index_t kMPerBlock = Problem::kMPerBlock; static constexpr index_t kNPerBlock = Problem::kNPerBlock; @@ -63,6 +73,7 @@ struct CShuffleEpilogue static constexpr index_t isCTransposed = Problem::isCTransposed; static constexpr index_t kMPerIteration = kMPerXdl * kMWave; static constexpr index_t kNPerIteration = kNPerXdl * kNWave; + static constexpr index_t NumDTensor = Problem::NumDTensor; using WG = WarpGemmMfmaDispatcher + CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeD(number index) + { + using DiDataType = remove_cvref_t>; + constexpr index_t MaxVectorStoreSize = 16; + return MaxVectorStoreSize / sizeof(DiDataType); + } + template CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor() { @@ -121,9 +145,12 @@ struct CShuffleEpilogue template - CK_TILE_DEVICE auto - operator()(ODramWindow& out_dram_window, const OAccTile& o_acc_tile, void* p_smem) + CK_TILE_DEVICE auto operator()(ODramWindow& out_dram_window, + const OAccTile& o_acc_tile, + const DDramWindow& ds_dram_window, + void* p_smem) { const index_t iMWarp = get_warp_id() / kNWave; @@ -154,6 +181,14 @@ struct CShuffleEpilogue tile_distribution_pattern::thread_raked>; constexpr auto dram_tile_distribution = TileEncodingPattern::Make2DStaticTileDistribution(); + auto d_dram_small_window = generate_tuple( + [&](auto idx) { return make_tile_window(ds_dram_window[idx], dram_tile_distribution); }, + number{}); + + using elemenet_wise_output_t = + decltype(load_tile(make_tile_window(out_lds_window, dram_tile_distribution))); + elemenet_wise_output_t elemenet_wise_output; + constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; @@ -178,6 +213,17 @@ struct CShuffleEpilogue const auto c_out_tensor = load_tile(make_tile_window(out_lds_window, dram_tile_distribution)); + const auto ds_tensor = + generate_tuple([&](auto idx) { return load_tile(d_dram_small_window[idx]); }, + number{}); + + const auto c_ds_tiles = concat_tuple_of_reference( + tie(elemenet_wise_output, c_out_tensor), + generate_tie( + [&](auto i) -> const auto& { return ds_tensor[i]; }, number{})); + + tile_elementwise_in_out_unpack_tuple(typename Problem::ABDELementWise{}, c_ds_tiles); + if constexpr(out_memory_data_op == memory_operation_enum::set) { store_tile(out_dram_window, c_out_tensor); @@ -189,7 +235,13 @@ struct CShuffleEpilogue if constexpr(iAccess != num_access - 1) { constexpr auto step = SFC::get_forward_step(iAccess); + move_tile_window(out_dram_window, {step.at(number<0>{}), step.at(number<1>{})}); + + static_for<0, NumDTensor, 1>{}([&](auto idx) { + move_tile_window(d_dram_small_window[idx], + {step.at(number<0>{}), step.at(number<1>{})}); + }); } }); } 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 dfb6bfae58..93d43b0c1e 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_, @@ -26,8 +26,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_), + : GemmHostArgs(a_ptr_, + b_ptr_, + {}, + c_ptr_, + k_batch_, + M_, + N_, + K_, + stride_A_, + stride_B_, + {}, + stride_C_), batch_stride_A(batch_stride_A_), batch_stride_B(batch_stride_B_), batch_stride_C(batch_stride_C_), @@ -46,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; @@ -94,12 +104,14 @@ struct BatchedGemmKernel : public GemmKernelRunGemm(a_ptr, b_ptr, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n); + this->RunGemm(a_ptr, b_ptr, {}, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n); } else { this->template RunGemm( - a_ptr, b_ptr, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n); + a_ptr, b_ptr, {}, c_ptr, smem_ptr, kargs, splitk_batch_offset, i_m, i_n); } } }; diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index 9435855d0a..15adf6c6bb 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -12,28 +12,13 @@ namespace ck_tile { -struct GemmProblem -{ - CK_TILE_HOST GemmProblem() = default; - CK_TILE_HOST GemmProblem( - index_t M_, index_t N_, index_t K_, index_t stride_A_, index_t stride_B_, index_t stride_C_) - : M(M_), N(N_), K(K_), stride_A(stride_A_), stride_B(stride_B_), stride_C(stride_C_) - { - } - - index_t M; - index_t N; - index_t K; - index_t stride_A; - index_t stride_B; - index_t stride_C; -}; - -struct GemmHostArgs : public GemmProblem +template +struct GemmHostArgs { CK_TILE_HOST GemmHostArgs() = default; CK_TILE_HOST GemmHostArgs(const void* a_ptr_, const void* b_ptr_, + const std::array& ds_ptr_, void* c_ptr_, index_t k_batch_, index_t M_, @@ -41,31 +26,51 @@ struct GemmHostArgs : public GemmProblem index_t K_, index_t stride_A_, index_t stride_B_, + const std::array& stride_Ds_, index_t stride_C_) - : GemmProblem(M_, N_, K_, stride_A_, stride_B_, stride_C_), - a_ptr(a_ptr_), + : a_ptr(a_ptr_), b_ptr(b_ptr_), + ds_ptr(ds_ptr_), c_ptr(c_ptr_), + M(M_), + N(N_), + K(K_), + stride_A(stride_A_), + stride_B(stride_B_), + stride_Ds(stride_Ds_), + stride_C(stride_C_), k_batch(k_batch_) { } const void* a_ptr; const void* b_ptr; - void* c_ptr; - index_t k_batch; -}; - -struct GemmKernelArgs -{ - const void* a_ptr; - const void* b_ptr; + const std::array ds_ptr; void* c_ptr; index_t M; index_t N; index_t K; index_t stride_A; index_t stride_B; + const std::array stride_Ds; + index_t stride_C; + 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; + void* c_ptr; + index_t M; + index_t N; + index_t K; + index_t stride_A; + index_t stride_B; + const index_t* stride_Ds; index_t stride_C; index_t k_batch; }; @@ -73,12 +78,14 @@ struct GemmKernelArgs template struct GemmKernel { - using TilePartitioner = remove_cvref_t; - using GemmPipeline = remove_cvref_t; - using EpiloguePipeline = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using TilePartitioner = remove_cvref_t; + using GemmPipeline = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + using ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + using DsLayout = remove_cvref_t; + using DsDataType = remove_cvref_t; static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize; using ADataType = remove_cvref_t; @@ -86,9 +93,13 @@ struct GemmKernel // Below type is actually accumulation data type - the output of block GEMM. using CDataType = remove_cvref_t; + using Empty_Tuple = ck_tile::tuple<>; + static constexpr index_t NumDTensor = DsDataType::size(); + static constexpr auto I0 = number<0>(); static constexpr auto I1 = number<1>(); static constexpr auto I2 = number<2>(); + static constexpr auto I3 = number<3>{}; [[nodiscard]] CK_TILE_HOST static const std::string GetName() { @@ -97,6 +108,19 @@ struct GemmKernel // clang-format on } + CK_TILE_HOST static constexpr auto MakeDsGridPointer() + { + return generate_tuple( + [&](auto i) { + using DDataType = remove_cvref_t>; + + return static_cast(nullptr); + }, + number{}); + } + + using DsGridPointer = decltype(MakeDsGridPointer()); + CK_TILE_HOST static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) { return dim3(TilePartitioner::GridSize(M, N), 1, KBatch); @@ -104,18 +128,27 @@ struct GemmKernel CK_TILE_HOST static constexpr auto BlockSize() { return dim3(KernelBlockSize); } - CK_TILE_HOST static constexpr GemmKernelArgs MakeKernelArgs(const GemmHostArgs& hostArgs) + CK_TILE_HOST static constexpr GemmKernelArgs + MakeKernelArgs(const GemmHostArgs& hostArgs) { - return GemmKernelArgs{hostArgs.a_ptr, - hostArgs.b_ptr, - hostArgs.c_ptr, - hostArgs.M, - hostArgs.N, - hostArgs.K, - hostArgs.stride_A, - hostArgs.stride_B, - hostArgs.stride_C, - hostArgs.k_batch}; + 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}; } CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() @@ -125,7 +158,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>{}); @@ -165,7 +198,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) @@ -264,6 +297,51 @@ struct GemmKernel } } + bool DTesnorIsValid = {true}; + static_for<0, NumDTensor, 1>{}([&](auto index) { + using DiLayout = remove_cvref_t>; + if constexpr(std::is_same_v) + { + if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support N that is not a multiple of NPerBlock without padding!"); + } + DTesnorIsValid = false; + } + if(kargs.N % EpiloguePipeline::GetVectorSizeD(index) != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("N is not a multiple of vector load size for D tensor!"); + } + DTesnorIsValid = false; + } + } + else + { + if(kargs.M % TilePartitioner::MPerBlock != 0 && GemmPipeline::kPadM == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support M that is not a multiple of MPerBlock without padding!"); + } + DTesnorIsValid = false; + } + if(kargs.M % EpiloguePipeline::GetVectorSizeD(index) != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("M is not a multiple of vector load size for D tensor!"); + } + DTesnorIsValid = false; + } + } + }); + if constexpr(std::is_same_v) { if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) @@ -304,14 +382,15 @@ struct GemmKernel return false; } } - return true; + return DTesnorIsValid && true; } template CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr, const BDataType* b_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!"); @@ -399,6 +478,29 @@ struct GemmKernel } }(); + // TODO: enable vector write for D in ColMajor + const auto& d_tensor_view = [&](auto i) { + using DiLayout = remove_cvref_t>; + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + ds_ptr[i], + make_tuple(kargs.M, kargs.N), + make_tuple(kargs.stride_Ds[i], 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + ds_ptr[i], + make_tuple(kargs.M, kargs.N), + make_tuple(kargs.stride_Ds[i], 1), + number{}, + number<1>{}); + } + }; + // TODO: enable vector write for C in ColMajor const auto& c_tensor_view = [&]() { if constexpr(std::is_same_v) @@ -421,7 +523,10 @@ struct GemmKernel } }(); - return make_tuple(a_tensor_view, b_tensor_view, c_tensor_view); + return make_tuple(a_tensor_view, + b_tensor_view, + generate_tuple(d_tensor_view, number{}), + c_tensor_view); } template @@ -463,9 +568,28 @@ struct GemmKernel } }(); + // TODO vector write in for D in ColMajor + const auto& d_pad_view = [&](auto i) { + const auto& d_tensor_view = views.at(I2); + if constexpr(std::is_same_v) + { + return pad_tensor_view(d_tensor_view[i], + make_tuple(number{}, + number{}), + sequence{}); + } + else + { + return pad_tensor_view(d_tensor_view[i], + make_tuple(number{}, + number{}), + sequence{}); + } + }; + // TODO vector write in for C in ColMajor const auto& c_pad_view = [&]() { - const auto& c_tensor_view = views.at(I2); + const auto& c_tensor_view = views.at(I3); if constexpr(std::is_same_v) { return pad_tensor_view(c_tensor_view, @@ -482,7 +606,8 @@ struct GemmKernel } }(); - return make_tuple(a_pad_view, b_pad_view, c_pad_view); + return make_tuple( + a_pad_view, b_pad_view, generate_tuple(d_pad_view, number{}), c_pad_view); } template @@ -491,7 +616,8 @@ struct GemmKernel { const auto& a_pad_view = views.at(I0); const auto& b_pad_view = views.at(I1); - const auto& c_pad_view = views.at(I2); + const auto& d_pad_view = views.at(I2); + const auto& c_pad_view = views.at(I3); const auto& a_block_window = [&]() { if constexpr(std::is_same_v) @@ -527,12 +653,22 @@ struct GemmKernel } }(); + const auto d_block_window = [&](auto i) { + return make_tile_window(d_pad_view[i], + make_tuple(number{}, + number{}), + {i_m, i_n}); + }; + auto c_block_window = make_tile_window( c_pad_view, make_tuple(number{}, number{}), {i_m, i_n}); - return make_tuple(a_block_window, b_block_window, c_block_window); + return make_tuple(a_block_window, + b_block_window, + generate_tuple(d_block_window, number{}), + c_block_window); } /** @@ -540,6 +676,7 @@ struct GemmKernel * * @param a_ptr input A pointer * @param b_ptr input B pointer + * @param ds_ptr input Ds pointer * @param c_ptr output C pointer * @param smem_ptr_0 The start memory pointer of the shared memory block. * @param kargs GEMM kernel arguments @@ -552,16 +689,17 @@ struct GemmKernel template CK_TILE_DEVICE static void RunGemm(const ADataType* a_ptr, const BDataType* b_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) { // Create Gemm tensor views, pad views and tile windows - const auto& gemm_tensor_views_tuple = - MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, kargs, splitk_batch_offset); + const auto& gemm_tensor_views_tuple = MakeGemmTensorViews( + a_ptr, b_ptr, ds_ptr, c_ptr, kargs, splitk_batch_offset); const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple); auto gemm_tile_windows = MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n); @@ -572,16 +710,20 @@ struct GemmKernel // Run GEMM cooperatively by whole workgroup. const auto& a_block_window = gemm_tile_windows.at(I0); const auto& b_block_window = gemm_tile_windows.at(I1); + const auto& d_block_window = gemm_tile_windows.at(I2); const auto& c_block_tile = GemmPipeline{}.template operator()( a_block_window, b_block_window, num_loop, smem_ptr_0); // Run Epilogue Pipeline - auto& c_block_window = gemm_tile_windows.at(I2); + auto& c_block_window = gemm_tile_windows.at(I3); EpiloguePipeline{} - .template operator()( - c_block_window, c_block_tile, smem_ptr_0); + .template operator()( + c_block_window, c_block_tile, d_block_window, smem_ptr_0); } /** @@ -591,6 +733,7 @@ struct GemmKernel * * @param a_ptr input A pointer * @param b_ptr input B pointer + * @param ds_ptr input Ds pointer * @param c_ptr output C pointer * @param smem_ptr_0 The starting pointer of 1st shared memory block. * @param smem_ptr_1 The starting pointer of 2nd shared memory block. @@ -604,17 +747,18 @@ struct GemmKernel template CK_TILE_DEVICE static void RunGemm2LDS(const ADataType* a_ptr, const BDataType* b_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) { // Create Gemm tensor views, pad views and tile windows - const auto& gemm_tensor_views_tuple = - MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, kargs, splitk_batch_offset); + const auto& gemm_tensor_views_tuple = MakeGemmTensorViews( + a_ptr, b_ptr, ds_ptr, c_ptr, kargs, splitk_batch_offset); const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple); auto gemm_tile_windows = MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n); @@ -624,19 +768,23 @@ struct GemmKernel // Run GEMM cooperatively by whole workgroup. const auto& a_block_window = gemm_tile_windows.at(I0); const auto& b_block_window = gemm_tile_windows.at(I1); + const auto& d_block_window = gemm_tile_windows.at(I2); const auto& c_block_tile = GemmPipeline{}.template operator()( a_block_window, b_block_window, num_loop, smem_ptr_0, smem_ptr_1); // Run Epilogue Pipeline - auto& c_block_window = gemm_tile_windows.at(I2); + auto& c_block_window = gemm_tile_windows.at(I3); EpiloguePipeline{} - .template operator()( - c_block_window, c_block_tile, smem_ptr_0); + .template operator()( + 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); @@ -644,11 +792,13 @@ struct GemmKernel const index_t i_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock); const SplitKBatchOffset splitk_batch_offset(kargs); + // options const ADataType* a_ptr = 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; + CDataType* c_ptr = static_cast(kargs.c_ptr); // allocate LDS @@ -661,6 +811,7 @@ struct GemmKernel { RunGemm2LDS(a_ptr, b_ptr, + kargs.ds_ptr, c_ptr, smem_ptr_0, smem_ptr_1, @@ -676,6 +827,7 @@ struct GemmKernel { RunGemm2LDS(a_ptr, b_ptr, + kargs.ds_ptr, c_ptr, smem_ptr_0, smem_ptr_1, @@ -690,15 +842,30 @@ struct GemmKernel { if(kargs.k_batch == 1) { - RunGemm(a_ptr, b_ptr, c_ptr, smem_ptr_0, kargs, splitk_batch_offset, i_m, i_n); + RunGemm(a_ptr, + b_ptr, + kargs.ds_ptr, + c_ptr, + smem_ptr_0, + kargs, + splitk_batch_offset, + i_m, + i_n); } else { if constexpr(!(EpiloguePipeline::GetVectorSizeC() % 2 != 0 && is_any_of::value)) { - RunGemm( - a_ptr, b_ptr, c_ptr, smem_ptr_0, kargs, splitk_batch_offset, i_m, i_n); + RunGemm(a_ptr, + b_ptr, + kargs.ds_ptr, + c_ptr, + smem_ptr_0, + kargs, + splitk_batch_offset, + i_m, + i_n); } } } 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 5577cb083a..c26fe37859 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() = default; - GemmTransKernelArg(GemmKernelArgs&& karg, index_t bl_start, index_t bl_end) + GemmTransKernelArg() = delete; + GemmTransKernelArg(GemmKernelArgs<>&& karg, index_t bl_start, index_t bl_end) : group_karg{karg}, block_start{bl_start}, block_end{bl_end} { } @@ -55,15 +55,16 @@ struct GroupedGemmKernel : public GemmKernel& gemm_descs) - -> std::size_t + __host__ static auto + GetWorkSpaceSize(const std::vector>& gemm_descs) -> std::size_t { return gemm_descs.size() * sizeof(GemmTransKernelArg); } __host__ static constexpr auto BlockSize() -> dim3 { return dim3(KernelBlockSize); } - __host__ static constexpr auto GridSize(const std::vector& gemm_descs) + __host__ static constexpr auto + GridSize(const std::vector>& gemm_descs) { index_t grid_size = 0; for(const auto& it_desc : gemm_descs) @@ -74,7 +75,8 @@ struct GroupedGemmKernel : public GemmKernel& gemm_descs) + CK_TILE_HOST static auto + MakeKargs(const std::vector>& gemm_descs) -> std::vector { std::vector gemm_kernel_args_; @@ -104,16 +106,18 @@ struct GroupedGemmKernel : public GemmKernel(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); } @@ -144,7 +148,7 @@ struct GroupedGemmKernel : public GemmKernelRunGemm( - a_ptr, b_ptr, c_ptr, smem_ptr, kargs.group_karg, splitk_batch_offset, i_m, i_n); + a_ptr, b_ptr, {}, c_ptr, smem_ptr, kargs.group_karg, splitk_batch_offset, i_m, i_n); } CK_TILE_DEVICE void operator()(const void CK_CONSTANT_ADDRESS_SPACE* gemm_descs_const, diff --git a/test/ck_tile/CMakeLists.txt b/test/ck_tile/CMakeLists.txt index 8f9d7ac89b..9a75b0ab6b 100644 --- a/test/ck_tile/CMakeLists.txt +++ b/test/ck_tile/CMakeLists.txt @@ -2,4 +2,5 @@ add_subdirectory(image_to_column) add_subdirectory(gemm) add_subdirectory(batched_gemm) add_subdirectory(grouped_gemm) +add_subdirectory(multiple_d_gemm) add_subdirectory(data_type) diff --git a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp index 0af3ef3b34..e79e45dee8 100644 --- a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp +++ b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp @@ -11,6 +11,7 @@ #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" #include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" template class TestCkTileBatchedGemm : public ::testing::Test @@ -23,6 +24,8 @@ class TestCkTileBatchedGemm : public ::testing::Test using BDataType = std::tuple_element_t<4, Tuple>; using AccDataType = std::tuple_element_t<5, Tuple>; using CDataType = std::tuple_element_t<6, Tuple>; + using DsLayout = ck_tile::tuple<>; + using DsDataType = ck_tile::tuple<>; template void invoke_batched_gemm(const ck_tile::BatchedGemmHostArgs& args, @@ -99,9 +102,12 @@ class TestCkTileBatchedGemm : public ::testing::Test using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem; static constexpr auto Scheduler = std::tuple_element_t<7, Tuple>::value; static constexpr auto PipelineType = std::tuple_element_t<8, Tuple>::value; + + using DsLayout = ck_tile::tuple<>; + using DsDataType = ck_tile::tuple<>; // 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; @@ -157,9 +160,12 @@ class TestCkTileGemmPipeline : public ::testing::Test using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem 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 b125d19762..6c9a1e1a6d 100644 --- a/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp +++ b/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp @@ -11,6 +11,7 @@ #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" template class TestCkTileGroupedGemm : public ::testing::Test @@ -23,6 +24,8 @@ class TestCkTileGroupedGemm : public ::testing::Test using BDataType = std::tuple_element_t<4, Tuple>; using AccDataType = std::tuple_element_t<5, Tuple>; using CDataType = std::tuple_element_t<6, Tuple>; + using DsLayout = ck_tile::tuple<>; + using DsDataType = ck_tile::tuple<>; struct GroupedGemKernelParam { @@ -44,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); @@ -120,9 +123,12 @@ class TestCkTileGroupedGemm : public ::testing::Test using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem 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]}); } ck_tile::DeviceMem gemm_workspace;