From 1ee7564ae55953e68a6706b2cb0b063e047314d5 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 23 Oct 2025 14:57:21 -0500 Subject: [PATCH] implement gemm universal with a visitor --- .../03_gemm/universal_gemm_invoker.hpp | 75 +++++++- .../42_gemm_dispatcher/gemm_dispatcher.cpp | 177 +++++++++++++++++- 2 files changed, 242 insertions(+), 10 deletions(-) diff --git a/example/ck_tile/03_gemm/universal_gemm_invoker.hpp b/example/ck_tile/03_gemm/universal_gemm_invoker.hpp index d0fd69b1e2..ce8e37f46d 100644 --- a/example/ck_tile/03_gemm/universal_gemm_invoker.hpp +++ b/example/ck_tile/03_gemm/universal_gemm_invoker.hpp @@ -69,9 +69,9 @@ struct UniversalInvoker 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_, - const auto memory_operation_) { + const auto kernel_launch_visitor = [&](const auto has_hot_loop_, + const auto tail_number_, + const auto memory_operation_) { constexpr bool has_hot_loop_v = has_hot_loop_.value; constexpr auto tail_number_v = tail_number_.value; constexpr auto scheduler = GemmConfig::Scheduler; @@ -182,17 +182,76 @@ struct UniversalInvoker return ave_time; }; - const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) { - if(args.k_batch == 1) + using KTrue = ck_tile::integral_constant; + using KFalse = ck_tile::integral_constant; + + using BoolVariant = std::variant; + + auto make_bool_variant = [](bool b) -> BoolVariant { + if(b) { - return Run(has_hot_loop_, tail_number_, MemoryOpSet{}); + return KTrue{}; } else { - return Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{}); + return KFalse{}; } }; - return ave_time = BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num); + using KTailNumOdd = + ck_tile::integral_constant; + using KTailNumEven = + ck_tile::integral_constant; + using KTailNumFull = + ck_tile::integral_constant; + using TailNumVariant = std::variant; + + auto make_tailnum_variant = [](ck_tile::TailNumber tail_number) -> TailNumVariant { + switch(tail_number) + { + case ck_tile::TailNumber::Full: return KTailNumFull{}; + case ck_tile::TailNumber::Odd: return KTailNumOdd{}; + case ck_tile::TailNumber::Even: return KTailNumEven{}; + case ck_tile::TailNumber::Empty: + case ck_tile::TailNumber::One: + case ck_tile::TailNumber::Two: + case ck_tile::TailNumber::Three: + case ck_tile::TailNumber::Four: + case ck_tile::TailNumber::Five: + case ck_tile::TailNumber::Six: + case ck_tile::TailNumber::Seven: + default: throw std::runtime_error("Case not handled"); + } + }; + + using KSet = ck_tile::integral_constant; + using KAtomic = ck_tile::integral_constant; + + using MemoryOperationVariant = std::variant; + + auto make_memory_op_variant = + [](ck_tile::memory_operation_enum ev) -> MemoryOperationVariant { + if(ev == ck_tile::memory_operation_enum::set) + { + return KSet{}; + } + else + { + return KAtomic{}; + } + }; + + ave_time = std::visit( + kernel_launch_visitor, + make_bool_variant(has_hot_loop), + make_tailnum_variant(tail_num), + make_memory_op_variant(args.k_batch == 1 ? ck_tile::memory_operation_enum::set + : ck_tile::memory_operation_enum::atomic_add) + + ); + + return ave_time; } }; diff --git a/example/ck_tile/42_gemm_dispatcher/gemm_dispatcher.cpp b/example/ck_tile/42_gemm_dispatcher/gemm_dispatcher.cpp index 91befdfcc9..64486d9311 100644 --- a/example/ck_tile/42_gemm_dispatcher/gemm_dispatcher.cpp +++ b/example/ck_tile/42_gemm_dispatcher/gemm_dispatcher.cpp @@ -1,9 +1,182 @@ +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm.hpp" +// template +// float gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) +// { +// using GemmShape = ck_tile::TileGemmShape< +// ck_tile::sequence, +// ck_tile::sequence, +// ck_tile:: +// 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; + +// using BaseGemmPipeline = typename PipelineTypeTraits< +// GemmConfig::Pipeline>::template UniversalGemmPipeline; + +// // const ck_tile::index_t k_grain = args.k_batch * GemmConfig::K_Tile; +// // const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * +// GemmConfig::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); + +// constexpr auto scheduler = GemmConfig::Scheduler; + +// using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + +// using GemmPipeline = typename PipelineTypeTraits< +// GemmConfig::Pipeline>::template GemmPipeline; + +// using GemmEpilogue = ck_tile::CShuffleEpilogue< +// ck_tile::CShuffleEpilogueProblem>; + +// using Kernel = ck_tile::GemmKernel; +// auto kargs = Kernel::MakeKernelArgs(args); + +// const dim3 grids = Persistent ? Kernel::MaxOccupancyGridSize(s) +// : Kernel::GridSize(args.M, args.N, args.k_batch); +// const dim3 blocks = Kernel::BlockSize(); + +// if(!Kernel::IsSupportedArgument(kargs)) +// { +// throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); +// } + +// float ave_time = ck_tile::launch_kernel_time_mask( +// s, +// ck_tile::make_kernel(Kernel{}, grids, blocks, 0, +// kargs)); + +// return ave_time +// } int main(int argc, char** argv) { - (void) argc; (void) argv; - return 0; + (void)argc; + (void)argv; + + using KTrue = ck_tile::integral_constant; + using KFalse = ck_tile::integral_constant; + + using BoolVariant = std::variant; + + auto make_bool_variant = [](bool b) -> BoolVariant { + if(b) + { + return KTrue{}; + } + else + { + return KFalse{}; + } + }; + + std::visit([](auto&& bv) { std::cout << bv.value << std::endl; }, make_bool_variant(true)); + std::visit([](auto&& bv) { std::cout << bv.value << std::endl; }, make_bool_variant(false)); + + using KTailNumOdd = ck_tile::integral_constant; + using KTailNumEven = ck_tile::integral_constant; + // using KTailNumFull = ck_tile::integral_constant; + + using TailNumVariant = std::variant; + auto make_tailnum_variant = [](ck_tile::index_t tail_num) -> TailNumVariant { + if(tail_num % 2 == 1) + { + return KTailNumOdd{}; + } + else + { + return KTailNumEven{}; + } + }; + + std::visit( + [](auto&& bv, auto&& tnv) { + std::cout << "bool val: " << bv.value << ", tn val: " << tnv.value << std::endl; + }, + make_bool_variant(false), + make_tailnum_variant(5)); + + using KSet = ck_tile::integral_constant; + using KAtomic = ck_tile::integral_constant; + + using MemoryOperationVariant = std::variant; + // std::cout << make_bool_variant(true).value; + // std::cout << make_bool_variant(false).value; + + return 0; }