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

Clean up namespaces and locations of code.

Browse Source 
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.
This commit is contained in:
John Shumway
2025-08-05 15:58:30 +00:00
parent fda9a28e18
commit 663b7a703c
3 changed files with 48 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
experimental/gemm_builder/gemm_builder.h
View File

@@ -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<TilePartitioner, GemmPipeline, GemmEpilogue>;
// TODO: Args should not be a templated class!
template <index_t NumATensor = 1, index_t NumBTensor = 1, index_t NumDTensor = 0>
using KernelArgs = ck_tile::UniversalGemmKernelArgs<NumATensor, NumBTensor, NumDTensor>;
using KernelClass = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
// TODO: Args should not be a templated class!
template <index_t NumATensor = 1, index_t NumBTensor = 1, index_t NumDTensor = 0>
static bool Supports(const KernelArgs<NumATensor, NumBTensor, NumDTensor>& 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 <index_t NumATensor = 1, index_t NumBTensor = 1, index_t NumDTensor = 0>
static __device__ void KernelFn(const KernelArgs<NumATensor, NumBTensor, NumDTensor>& 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><<<grid_dim, block_dim, 0, hipStreamDefault>>>(kernel_args);
template <typename Kernel, typename... Args>
__global__ void launch_kernel(Args... args)
// Kernel<Builder><<<grid_dim, block_dim, 0, hipStreamDefault>>>(kernel_args);
template <typename Builder, typename... Args>
__global__ void Kernel(Args... args)
{
Kernel{}(args...);
Builder::KernelFn(args...);
}
} // namespace ck_tile::builder

25
experimental/gemm_builder/gemm_example.cpp
View File

@@ -3,9 +3,10 @@
#include <hip/hip_runtime.h>
#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<MyGemmTypes, MyGemmLayout>;
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<ck_tile::bf16_t>(M * K);
auto b_dev = example::AllocDevMem<ck_tile::bf16_t>(K * N);
auto c_dev = example::AllocDevMem<ck_tile::bf16_t>(M * N);
auto a_dev = ckr::AllocDevMem<ck_tile::bf16_t>(M * K);
auto b_dev = ckr::AllocDevMem<ck_tile::bf16_t>(K * N);
auto c_dev = ckr::AllocDevMem<ck_tile::bf16_t>(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<example::Kernel>
<<<grid_dim, block_dim, 0, hipStreamDefault>>>(kernel_args);
ckb::Kernel<example::Builder><<<grid_dim, block_dim, 0, hipStreamDefault>>>(kernel_args);
example::CheckHipError(hipDeviceSynchronize());
ckr::CheckHipError(hipDeviceSynchronize());
std::cout << "GEMM completed successfully!" << std::endl;
}
catch(const std::exception& e)

4
experimental/gemm_builder/utils.hpp → experimental/gemm_builder/run_utils.hpp
View File

@@ -4,7 +4,7 @@
#include <iostream>
#include <stdexcept>
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