From 663b7a703c4346fee57a5fc4b3682dbb469f3e6d Mon Sep 17 00:00:00 2001 From: John Shumway Date: Tue, 5 Aug 2025 15:58:30 +0000 Subject: [PATCH] Clean up namespaces and locations of code. Adds a new ck_tile::runtime for runtime utilities. Also moves functionality into the builder. Kernel is accessed with the ckb::Kernel function that is templated by the builder class. --- experimental/gemm_builder/gemm_builder.h | 40 ++++++++++++++++--- experimental/gemm_builder/gemm_example.cpp | 25 ++++++------ .../gemm_builder/{utils.hpp => run_utils.hpp} | 4 +- 3 files changed, 48 insertions(+), 21 deletions(-) rename experimental/gemm_builder/{utils.hpp => run_utils.hpp} (95%) diff --git a/experimental/gemm_builder/gemm_builder.h b/experimental/gemm_builder/gemm_builder.h index f4c5a0f994..35b4910958 100644 --- a/experimental/gemm_builder/gemm_builder.h +++ b/experimental/gemm_builder/gemm_builder.h @@ -7,6 +7,7 @@ #include "ck_tile/host/kernel_launch.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp" namespace ck_tile::builder { @@ -207,17 +208,44 @@ struct GemmBuilder ck_tile::memory_operation_enum::atomic_add, GemmConfig::NumWaveGroups>>; - using Kernel = ck_tile::GemmKernel; + // TODO: Args should not be a templated class! + template + using KernelArgs = ck_tile::UniversalGemmKernelArgs; + + using KernelClass = ck_tile::GemmKernel; + + // TODO: Args should not be a templated class! + template + static bool Supports(const KernelArgs& args) + { + return KernelClass::IsSupportedArgument(args); + } + + static std::string GetKernelName() { return KernelClass::GetName(); } + + static constexpr dim3 GridDim(index_t M, index_t N, index_t KBatch) + { + return KernelClass::GridSize(M, N, KBatch); + } + + static constexpr dim3 BlockDim() { return KernelClass::BlockSize(); } + + // TODO: Args should not be a templated class! + template + static __device__ void KernelFn(const KernelArgs& args) + { + KernelClass{}(args); + } }; -// Wrap a kernel class's static __device__ operator() method in a __global__ function. +// Wrap a kernel class's static __device__ KerenFn method in a __global__ function. // // Usage: -// launch_kernel<<>>(kernel_args); -template -__global__ void launch_kernel(Args... args) +// Kernel<<>>(kernel_args); +template +__global__ void Kernel(Args... args) { - Kernel{}(args...); + Builder::KernelFn(args...); } } // namespace ck_tile::builder diff --git a/experimental/gemm_builder/gemm_example.cpp b/experimental/gemm_builder/gemm_example.cpp index 57af54c3da..83a8344eaf 100644 --- a/experimental/gemm_builder/gemm_example.cpp +++ b/experimental/gemm_builder/gemm_example.cpp @@ -3,9 +3,10 @@ #include #include "gemm_builder.h" -#include "utils.hpp" +#include "run_utils.hpp" namespace ckb = ck_tile::builder; +namespace ckr = ck_tile::runtime; namespace example { @@ -27,14 +28,13 @@ struct MyGemmLayout }; using Builder = ckb::GemmBuilder; -using Kernel = Builder::Kernel; } // namespace example int main() { // Describe the GEMM kernel. - std::cout << "Kernel name: " << example::Kernel::GetName() << std::endl; + std::cout << "Kernel name: " << example::Builder::GetKernelName() << std::endl; std::cout << "Shape: " << example::Builder::GemmShape::GetName() << std::endl; std::cout << "Problem: " << example::Builder::UniversalGemmProblem::GetName() << std::endl; std::cout << "Pipeline: " << example::Builder::GemmPipeline::GetName() << std::endl; @@ -44,11 +44,11 @@ int main() { const int M = 16, N = 64, K = 128; - auto a_dev = example::AllocDevMem(M * K); - auto b_dev = example::AllocDevMem(K * N); - auto c_dev = example::AllocDevMem(M * N); + auto a_dev = ckr::AllocDevMem(M * K); + auto b_dev = ckr::AllocDevMem(K * N); + auto c_dev = ckr::AllocDevMem(M * N); - auto kernel_args = ck_tile::UniversalGemmKernelArgs{ + auto kernel_args = example::Builder::KernelArgs{ .as_ptr = {a_dev.get()}, // Address of tensor A in device memory. .bs_ptr = {b_dev.get()}, // Address of tensor B in device memory. .ds_ptr = {}, // Unused. @@ -62,13 +62,13 @@ int main() .stride_E = M, // Stride for tensor C_MN (row major). .k_batch = 1 // Batch size is 1 for a single GEMM. }; - if(!example::Kernel::IsSupportedArgument(kernel_args)) + if(!example::Builder::Supports(kernel_args)) { throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); } - dim3 grid_dim = example::Kernel::GridSize(M, N, 1); - dim3 block_dim = example::Kernel::BlockSize(); + dim3 grid_dim = example::Builder::GridDim(M, N, 1); + dim3 block_dim = example::Builder::BlockDim(); std::cout << "Running " << M << " x " << N << " x " << K << " GEMM kernel..." << std::endl; std::cout << "Grid size: " << grid_dim.x << " x " << grid_dim.y << " x " << grid_dim.z @@ -76,10 +76,9 @@ int main() std::cout << "Block size: " << block_dim.x << " x " << block_dim.y << " x " << block_dim.z << std::endl; - ckb::launch_kernel - <<>>(kernel_args); + ckb::Kernel<<>>(kernel_args); - example::CheckHipError(hipDeviceSynchronize()); + ckr::CheckHipError(hipDeviceSynchronize()); std::cout << "GEMM completed successfully!" << std::endl; } catch(const std::exception& e) diff --git a/experimental/gemm_builder/utils.hpp b/experimental/gemm_builder/run_utils.hpp similarity index 95% rename from experimental/gemm_builder/utils.hpp rename to experimental/gemm_builder/run_utils.hpp index 0cf9d56a70..a5786a5303 100644 --- a/experimental/gemm_builder/utils.hpp +++ b/experimental/gemm_builder/run_utils.hpp @@ -4,7 +4,7 @@ #include #include -namespace example { +namespace ck_tile::runtime { inline void CheckHipError(hipError_t err) { @@ -44,4 +44,4 @@ auto AllocDevMem(const size_t n) return d_data; } -} // namespace example +} // namespace ck_tile::runtime