From 40cd09a93d40d18e76f62d7262db5ed8c3367481 Mon Sep 17 00:00:00 2001 From: Yanxing-Shi Date: Thu, 22 May 2025 07:58:10 +0000 Subject: [PATCH] refactor gemm profiler --- tile_engine/ops/gemm/benchmark_gemm.cpp | 120 +--------------------- tile_engine/ops/gemm/benchmark_gemm.hpp | 3 +- tile_engine/ops/gemm/gemm_host_api.hpp | 1 + tile_engine/ops/gemm/gemm_profiler.hpp | 128 +++++++++++++++++++++++- 4 files changed, 129 insertions(+), 123 deletions(-) diff --git a/tile_engine/ops/gemm/benchmark_gemm.cpp b/tile_engine/ops/gemm/benchmark_gemm.cpp index 98f12c12b8..fb56e524d2 100644 --- a/tile_engine/ops/gemm/benchmark_gemm.cpp +++ b/tile_engine/ops/gemm/benchmark_gemm.cpp @@ -5,14 +5,10 @@ #include #include -#include "ck_tile/host.hpp" #include "gemm_profiler.hpp" -#include "gemm_host_api.hpp" #include "benchmark_gemm.hpp" -void benchmark_gemm(const ck_tile::ArgParser& arg_parser, - const std::vector( - ck_tile::GemmHostArgs&, const ck_tile::stream_config&)>>& callables) +void benchmark_gemm(const ck_tile::ArgParser& arg_parser) { GemmProblem gemm_problem{arg_parser.get_int("split_k"), arg_parser.get_int("m"), @@ -42,118 +38,10 @@ void benchmark_gemm(const ck_tile::ArgParser& arg_parser, auto& profiler = GemmProfiler::instance(setting); - const ALayout layout_a = ALayout{}; - const BLayout layout_b = BLayout{}; - const CLayout layout_c = CLayout{}; - - gemm_problem.stride_a_ = ck_tile::get_default_stride( - gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a)); - gemm_problem.stride_b_ = ck_tile::get_default_stride( - gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b)); - gemm_problem.stride_c_ = ck_tile::get_default_stride( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c)); - - ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a))); - ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( - gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b))); - ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); - - if(setting.init_method_ == 0) - { - ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); - ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); - } - else if(setting.init_method_ == 1) - { - ck_tile::FillMonotonicSeq{}(a_m_k); - ck_tile::FillMonotonicSeq{}(b_k_n); - } - else if(setting.init_method_ == 2) - { - ck_tile::FillConstant{static_cast(1)}(a_m_k); - ck_tile::FillConstant{static_cast(1)}(b_k_n); - } - else - { - a_m_k.SetZero(); - b_k_n.SetZero(); - } - - if(gemm_problem.structured_sparsity_) - { - ck_tile::AdjustToStructuredSparsity{}(a_m_k); - } - - ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - - if constexpr(std::is_same_v) - { - // Permute vector pk_i4x4 data for device implementation - ck_tile::HostTensor b_k_n_dev = b_k_n; - // permute_tensor_b(b_k_n_dev); - permute_vectors_i4x4_b(b_k_n_dev); - b_k_n_dev_buf.ToDevice(b_k_n_dev.data()); - } - else - { - b_k_n_dev_buf.ToDevice(b_k_n.data()); - } - - a_m_k_dev_buf.ToDevice(a_m_k.data()); - c_m_n_dev_buf.SetZero(); - c_m_n_dev_result.SetZero(); - - ck_tile::GemmHostArgs gemm_args; - gemm_args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); - gemm_args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); - gemm_args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); - gemm_args.k_batch = gemm_problem.split_k_; - gemm_args.M = gemm_problem.m_; - gemm_args.N = gemm_problem.n_; - gemm_args.K = gemm_problem.k_; - gemm_args.stride_A = gemm_problem.stride_a_; - gemm_args.stride_B = gemm_problem.stride_b_; - gemm_args.stride_C = gemm_problem.stride_c_; - - ck_tile::HostTensor c_m_n_host_result(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); - - if(setting.verify_) - { - gemm_host_reference(setting.verify_, - a_m_k, - b_k_n, - c_m_n_host_result, - a_m_k_dev_buf, - b_k_n_dev_buf, - gemm_problem.m_, - gemm_problem.n_, - gemm_problem.k_, - gemm_problem.stride_a_, - gemm_problem.stride_b_, - gemm_problem.stride_c_); - } - try { - for(auto& callable : callables) - { - profiler.benchmark(gemm_problem, - c_m_n_dev_buf, - c_m_n_host_result, - c_m_n_dev_result, - callable(gemm_args, - ck_tile::stream_config{nullptr, - true, - setting.log_, - setting.n_warmup_, - setting.n_repeat_, - setting.is_gpu_timer_})); - } + auto kernel_func = get_kernel_func_by_trait(arg_parser); + profiler.benchmark(gemm_problem, kernel_func); profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); } catch(const std::exception& e) @@ -169,7 +57,7 @@ int main(int argc, char* argv[]) auto [result, parser] = create_args(argc, argv); if(!result) return EXIT_FAILURE; - benchmark_gemm(parser, get_kernel_func_by_trait(parser)); + benchmark_gemm(parser); return 0; } catch(const std::exception& e) diff --git a/tile_engine/ops/gemm/benchmark_gemm.hpp b/tile_engine/ops/gemm/benchmark_gemm.hpp index c82de1a07c..292d67dad6 100644 --- a/tile_engine/ops/gemm/benchmark_gemm.hpp +++ b/tile_engine/ops/gemm/benchmark_gemm.hpp @@ -8,8 +8,7 @@ #include #include -#include "ck_tile/host.hpp" -#include "gemm_common.hpp" +#include "gemm_host_api.hpp" enum class Metric { diff --git a/tile_engine/ops/gemm/gemm_host_api.hpp b/tile_engine/ops/gemm/gemm_host_api.hpp index c1d5e1fc8a..8cbc3f26f6 100755 --- a/tile_engine/ops/gemm/gemm_host_api.hpp +++ b/tile_engine/ops/gemm/gemm_host_api.hpp @@ -9,6 +9,7 @@ #include "ck_tile/host.hpp" #include "gemm_dispatcher.hpp" +#include "gemm_common.hpp" template struct DataTypeTraits; diff --git a/tile_engine/ops/gemm/gemm_profiler.hpp b/tile_engine/ops/gemm/gemm_profiler.hpp index 5b53e66fae..9170952aa8 100644 --- a/tile_engine/ops/gemm/gemm_profiler.hpp +++ b/tile_engine/ops/gemm/gemm_profiler.hpp @@ -8,6 +8,7 @@ #include #include "ck_tile/host/device_prop.hpp" +#include "ck_tile/ops/gemm.hpp" #include "benchmark_gemm.hpp" class GemmProfiler @@ -19,11 +20,128 @@ class GemmProfiler return instance; } - void benchmark(const GemmProblem& gemm_problem, - ck_tile::DeviceMem& c_m_n_dev_buf, - ck_tile::HostTensor& c_m_n_host_result, - ck_tile::HostTensor& c_m_n_dev_result, - const std::tuple& kernel_run_result) + void benchmark(GemmProblem& gemm_problem, + std::vector( + ck_tile::GemmHostArgs&, const ck_tile::stream_config&)>>& callables) + { + const ALayout layout_a = ALayout{}; + const BLayout layout_b = BLayout{}; + const CLayout layout_c = CLayout{}; + + gemm_problem.stride_a_ = ck_tile::get_default_stride( + gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a)); + gemm_problem.stride_b_ = ck_tile::get_default_stride( + gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b)); + gemm_problem.stride_c_ = ck_tile::get_default_stride( + gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c)); + + ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( + gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a))); + ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( + gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b))); + ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( + gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); + + if(setting_.init_method_ == 0) + { + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); + } + else if(setting_.init_method_ == 1) + { + ck_tile::FillMonotonicSeq{}(a_m_k); + ck_tile::FillMonotonicSeq{}(b_k_n); + } + else if(setting_.init_method_ == 2) + { + ck_tile::FillConstant{static_cast(1)}(a_m_k); + ck_tile::FillConstant{static_cast(1)}(b_k_n); + } + else + { + a_m_k.SetZero(); + b_k_n.SetZero(); + } + + if(gemm_problem.structured_sparsity_) + { + ck_tile::AdjustToStructuredSparsity{}(a_m_k); + } + + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); + + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor b_k_n_dev = b_k_n; + // permute_tensor_b(b_k_n_dev); + permute_vectors_i4x4_b(b_k_n_dev); + b_k_n_dev_buf.ToDevice(b_k_n_dev.data()); + } + else + { + b_k_n_dev_buf.ToDevice(b_k_n.data()); + } + + a_m_k_dev_buf.ToDevice(a_m_k.data()); + c_m_n_dev_buf.SetZero(); + c_m_n_dev_result.SetZero(); + + ck_tile::GemmHostArgs gemm_args; + gemm_args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); + gemm_args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); + gemm_args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); + gemm_args.k_batch = gemm_problem.split_k_; + gemm_args.M = gemm_problem.m_; + gemm_args.N = gemm_problem.n_; + gemm_args.K = gemm_problem.k_; + gemm_args.stride_A = gemm_problem.stride_a_; + gemm_args.stride_B = gemm_problem.stride_b_; + gemm_args.stride_C = gemm_problem.stride_c_; + + ck_tile::HostTensor c_m_n_host_result(ck_tile::host_tensor_descriptor( + gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); + + if(setting_.verify_) + { + gemm_host_reference(setting_.verify_, + a_m_k, + b_k_n, + c_m_n_host_result, + a_m_k_dev_buf, + b_k_n_dev_buf, + gemm_problem.m_, + gemm_problem.n_, + gemm_problem.k_, + gemm_problem.stride_a_, + gemm_problem.stride_b_, + gemm_problem.stride_c_); + } + + for(auto& callable : callables) + { + auto kernel_run_result = callable(gemm_args, + ck_tile::stream_config{nullptr, + true, + setting_.log_, + setting_.n_warmup_, + setting_.n_repeat_, + setting_.is_gpu_timer_}); + process_result(gemm_problem, + c_m_n_dev_buf, + c_m_n_host_result, + c_m_n_dev_result, + kernel_run_result); + } + } + + void process_result(const GemmProblem& gemm_problem, + ck_tile::DeviceMem& c_m_n_dev_buf, + ck_tile::HostTensor& c_m_n_host_result, + ck_tile::HostTensor& c_m_n_dev_result, + const std::tuple& kernel_run_result) { auto [name, avg_time] = kernel_run_result;