From 63d664d59991d353421e24e6b18ae604b06b44e2 Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Tue, 29 Sep 2026 20:25:30 +0800 Subject: [PATCH 01/11] feat(grouped_gemm): add GroupedGemmMultipleDXdlCShuffleTileLoop --- example_hcu/15_grouped_gemm/CMakeLists.txt | 14 + .../grouped_gemm_multiple_d_xdl_fp16.cpp | 81 ++ .../run_grouped_gemm_multiple_d_example.inc | 342 ++++++ .../gpu/device/device_grouped_gemm.hpp | 124 +- .../device/device_grouped_gemm_tile_loop.hpp | 107 ++ ...gemm_multiple_d_xdl_cshuffle_tile_loop.hpp | 1017 +++++++++++++++++ .../gpu/grid/block_to_ctile_map.hpp | 46 + .../cpu/reference_gemm_multiple_d.hpp | 166 +++ .../gpu/grouped_gemm_tile_loop.hpp | 109 ++ .../gpu/grouped_gemm_tile_loop/CMakeLists.txt | 11 + ...ile_loop_f16_f16_f16_mk_kn_mn_instance.cpp | 70 ++ ...ile_loop_f16_f16_f16_mk_nk_mn_instance.cpp | 70 ++ 12 files changed, 2156 insertions(+), 1 deletion(-) create mode 100644 example_hcu/15_grouped_gemm/CMakeLists.txt create mode 100644 example_hcu/15_grouped_gemm/grouped_gemm_multiple_d_xdl_fp16.cpp create mode 100644 example_hcu/15_grouped_gemm/run_grouped_gemm_multiple_d_example.inc create mode 100644 include/ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp create mode 100644 include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp create mode 100644 library/include/ck/library/reference_tensor_operation/cpu/reference_gemm_multiple_d.hpp create mode 100644 library/include/ck/library/tensor_operation_instance/gpu/grouped_gemm_tile_loop.hpp create mode 100644 library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/CMakeLists.txt create mode 100644 library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp create mode 100644 library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp diff --git a/example_hcu/15_grouped_gemm/CMakeLists.txt b/example_hcu/15_grouped_gemm/CMakeLists.txt new file mode 100644 index 0000000..1a9203a --- /dev/null +++ b/example_hcu/15_grouped_gemm/CMakeLists.txt @@ -0,0 +1,14 @@ +# Copyright (c) 2026 Hygon Info Technologies Ltd. +# SPDX-License-Identifier: MIT +# Modified by Hygon Information Technology Co., Ltd. +# hcu grouped gemm tile loop examples +add_example_executable(example_hcu_grouped_gemm_multiple_d_xdl_fp16 + grouped_gemm_multiple_d_xdl_fp16.cpp) + +add_test(NAME example_hcu_grouped_gemm_multiple_d_xdl_fp16_smoke + COMMAND $ + 1 1 0 64,127,255 252,252,252 4608,4608,4608 4608,4608,4608 4608,4608,4608 252,252,252 +) +set_tests_properties(example_hcu_grouped_gemm_multiple_d_xdl_fp16_smoke + PROPERTIES LABELS "hcu;grouped_gemm;multi_d;smoke" +) diff --git a/example_hcu/15_grouped_gemm/grouped_gemm_multiple_d_xdl_fp16.cpp b/example_hcu/15_grouped_gemm/grouped_gemm_multiple_d_xdl_fp16.cpp new file mode 100644 index 0000000..573f649 --- /dev/null +++ b/example_hcu/15_grouped_gemm/grouped_gemm_multiple_d_xdl_fp16.cpp @@ -0,0 +1,81 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp" +#include "ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#include +#include + +#include "ck/library/utility/check_err.hpp" +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/literals.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_gemm_multiple_d.hpp" + +#if __clang_major__ >= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +#endif + +template +using S = ck::Sequence; + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; +using AddAdd = ck::tensor_operation::element_wise::AddAdd; + +using ADataType = F16; +using BDataType = F16; +using AccDataType = F32; +using CShuffleDataType = F32; +using DDataType = F16; +using DsDataType = ck::Tuple; +using EDataType = F16; + +using ALayout = Row; +using BLayout = Col; +using DLayout = Row; +using DsLayout = ck::Tuple; +using ELayout = Row; + +using AElementOp = PassThrough; +using BElementOp = PassThrough; +using CDEElementOp = AddAdd; + +static constexpr auto GemmMNKPadding = ck::tensor_operation::device::GemmSpecialization::MNKPadding; +static constexpr int NumDs = 2; + +using DeviceGemmInstance = + ck::tensor_operation::device::DeviceGroupedGemmMultipleDXdlCShuffleTileLoop + // clang-format off +//######| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| +//######| | | | | Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| +//######| | | | | | | | | | | Operation| Operation| Operation| | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| +//######| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + < ALayout, BLayout, DsLayout, ELayout, ADataType, BDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 8, 16, 16, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<4,4,4>>; +// clang-format on + +#include "run_grouped_gemm_multiple_d_example.inc" + +int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); } + +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/example_hcu/15_grouped_gemm/run_grouped_gemm_multiple_d_example.inc b/example_hcu/15_grouped_gemm/run_grouped_gemm_multiple_d_example.inc new file mode 100644 index 0000000..0a1483d --- /dev/null +++ b/example_hcu/15_grouped_gemm/run_grouped_gemm_multiple_d_example.inc @@ -0,0 +1,342 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +struct ProblemSize final +{ + std::vector Ms; + std::vector Ns; + std::vector Ks; + + std::vector stride_As; + std::vector stride_Bs; + std::vector> stride_Ds; + std::vector stride_Cs; + + ck::index_t group_count; +}; + +struct ExecutionConfig final +{ + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; +}; + +bool run_grouped_gemm(const ProblemSize& problem_size, const ExecutionConfig& config) +{ + auto group_count = problem_size.group_count; + + using KernelArguments = ck::tensor_operation::device::GroupedGemmKernelArgument; + using GemmDesc = ck::tensor_operation::device::GemmDesc; + + // GEMM shape + std::vector gemm_descs; + std::vector ggemm_kargs; + std::vector p_Cs; + std::vector p_As; + std::vector p_Bs; + std::vector> p_Ds = {}; + + gemm_descs.reserve(group_count); + ggemm_kargs.reserve(group_count); + p_As.reserve(group_count); + p_Bs.reserve(group_count); + p_Ds.reserve(group_count); + + auto f_host_tensor_descriptor = + [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { + using namespace ck::literals; + + if(std::is_same::value) + { + return HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return HostTensorDescriptor({row, col}, {1_uz, stride}); + } + }; + + std::vector> a_tensors; + std::vector> b_tensors; + std::vector, NumDs>> d_tensors; + std::vector> c_host_tensors; + std::vector> c_device_result_tensors; + + a_tensors.reserve(group_count); + b_tensors.reserve(group_count); + d_tensors.reserve(group_count); + c_host_tensors.reserve(group_count); + c_device_result_tensors.reserve(group_count); + + using DeviceMemPtr = std::unique_ptr; + + std::vector a_tensors_device, b_tensors_device, c_tensors_device; + std::vector> d_tensors_device; + + a_tensors_device.reserve(group_count); + b_tensors_device.reserve(group_count); + c_tensors_device.reserve(group_count); + d_tensors_device.resize(group_count); // reserve and update vector size + + std::size_t flop = 0, num_btype = 0; + + for(int i = 0; i < group_count; i++) + { + a_tensors.push_back(Tensor(f_host_tensor_descriptor( + problem_size.Ms[i], problem_size.Ks[i], problem_size.stride_As[i], ALayout{}))); + b_tensors.push_back(Tensor(f_host_tensor_descriptor( + problem_size.Ks[i], problem_size.Ns[i], problem_size.stride_Bs[i], BLayout{}))); + + auto d0_tensor = Tensor(f_host_tensor_descriptor( + problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], DLayout{})); + auto d1_tensor = Tensor(f_host_tensor_descriptor( + problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], DLayout{})); + + std::array, NumDs> d_tens = {d0_tensor, d1_tensor}; + d_tensors.push_back(d_tens); + c_host_tensors.push_back(Tensor(f_host_tensor_descriptor( + problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], ELayout{}))); + c_device_result_tensors.push_back(Tensor(f_host_tensor_descriptor( + problem_size.Ms[i], problem_size.Ns[i], problem_size.stride_Cs[i], ELayout{}))); + std::cout << "gemm[" << i << "] a_m_k: " << a_tensors[i].mDesc + << " b_k_n: " << b_tensors[i].mDesc + << " c_m_n: " << c_device_result_tensors[i].mDesc << std::endl; + + flop += std::size_t(2) * problem_size.Ms[i] * problem_size.Ks[i] * problem_size.Ns[i]; + num_btype += sizeof(ADataType) * a_tensors[i].GetElementSize() + + sizeof(BDataType) * b_tensors[i].GetElementSize() + + sizeof(DDataType) * d_tensors[i][0].GetElementSize() * NumDs + + sizeof(EDataType) * c_device_result_tensors[i].GetElementSize(); + + switch(config.init_method) + { + case 0: break; + case 1: + a_tensors[i].GenerateTensorValue(GeneratorTensor_2{-5, 5}); + b_tensors[i].GenerateTensorValue(GeneratorTensor_2{-5, 5}); + for(int j = 0; j < NumDs; ++j) + { + d_tensors[i][j].GenerateTensorValue(GeneratorTensor_2{-5, 5}); + } + break; + case 2: + a_tensors[i].GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b_tensors[i].GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + for(int j = 0; j < NumDs; ++j) + { + d_tensors[i][j].GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + } + break; + default: + a_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<0>{}); + b_tensors[i].GenerateTensorValue(GeneratorTensor_Sequential<1>{}); + for(int j = 0; j < NumDs; ++j) + { + d_tensors[i][j].GenerateTensorValue(GeneratorTensor_Sequential<0>{}); + } + } + } + + for(int i = 0; i < group_count; i++) + { + a_tensors_device.emplace_back( + std::make_unique(a_tensors[i].GetElementSpaceSize() * sizeof(ADataType))); + b_tensors_device.emplace_back( + std::make_unique(b_tensors[i].GetElementSpaceSize() * sizeof(BDataType))); + c_tensors_device.emplace_back(std::make_unique( + c_device_result_tensors[i].GetElementSpaceSize() * sizeof(EDataType))); + + for(int j = 0; j < NumDs; ++j) + { + d_tensors_device[i].emplace_back(std::make_unique( + d_tensors[i][j].GetElementSpaceSize() * sizeof(DDataType))); + } + + a_tensors_device[i]->ToDevice(a_tensors[i].mData.data()); + b_tensors_device[i]->ToDevice(b_tensors[i].mData.data()); + for(int j = 0; j < NumDs; ++j) + { + d_tensors_device[i][j]->ToDevice(d_tensors[i][j].mData.data()); + } + c_tensors_device[i]->SetZero(); + + p_As.push_back(a_tensors_device[i]->GetDeviceBuffer()); + p_Bs.push_back(b_tensors_device[i]->GetDeviceBuffer()); + p_Ds.push_back( + {d_tensors_device[i][0]->GetDeviceBuffer(), d_tensors_device[i][1]->GetDeviceBuffer()}); + p_Cs.push_back(c_tensors_device[i]->GetDeviceBuffer()); + + // The device op does not have to know M problem size at lunch time. + gemm_descs.push_back({0, + problem_size.Ns[i], + problem_size.Ks[i], + problem_size.stride_As[i], + problem_size.stride_Bs[i], + problem_size.stride_Cs[i], + {problem_size.stride_Cs[i], problem_size.stride_Cs[i]}}); + ggemm_kargs.push_back( + {a_tensors_device[i]->GetDeviceBuffer(), + b_tensors_device[i]->GetDeviceBuffer(), + {d_tensors_device[i][0]->GetDeviceBuffer(), d_tensors_device[i][1]->GetDeviceBuffer()}, + c_tensors_device[i]->GetDeviceBuffer(), + problem_size.Ms[i], + problem_size.Ns[i], + problem_size.Ks[i], + problem_size.stride_As[i], + problem_size.stride_Bs[i], + {problem_size.stride_Cs[i], problem_size.stride_Cs[i]}, + problem_size.stride_Cs[i]}); + } + auto a_element_op = AElementOp{}; + auto b_element_op = BElementOp{}; + auto cde_element_op = CDEElementOp{}; + + auto gemm = DeviceGemmInstance{}; + auto invoker = gemm.MakeInvoker(); + + // do GEMM + auto argument = gemm.MakeArgument( + p_As, p_Bs, p_Ds, p_Cs, gemm_descs, a_element_op, b_element_op, cde_element_op); + if(!gemm.IsSupportedArgument(argument)) + { + throw std::runtime_error( + "wrong! device_gemm with the specified compilation parameters does " + "not support this GEMM problem"); + } + + DeviceMem gemm_arg_dev_mem(gemm.GetDeviceKernelArgSize(&argument)); + hip_check_error(hipMemcpy(gemm_arg_dev_mem.GetDeviceBuffer(), + ggemm_kargs.data(), + gemm.GetDeviceKernelArgSize(&argument), + hipMemcpyHostToDevice)); + gemm.SetDeviceKernelArgs(argument, gemm_arg_dev_mem.GetDeviceBuffer()); + + invoker.Run(argument, StreamConfig{nullptr, false, 1}); + + bool pass = true; + if(config.do_verification) + { + using ReferenceGemmInstance = + ck::tensor_operation::host::ReferenceGemmMultipleD; + + for(std::size_t i = 0; i < gemm_descs.size(); i++) + { + auto karg = ggemm_kargs[i]; + auto dev_res_tensor = + Tensor(f_host_tensor_descriptor(karg.M, karg.N, karg.StrideE, ELayout{})); + c_tensors_device[i]->FromDevice(c_device_result_tensors[i].mData.data()); + auto ref_gemm = ReferenceGemmInstance{}; + auto ref_invoker = ref_gemm.MakeInvoker(); + + auto ref_argument = ref_gemm.MakeArgument(a_tensors[i], + b_tensors[i], + d_tensors[i], + c_host_tensors[i], + a_element_op, + b_element_op, + cde_element_op); + + ref_invoker.Run(ref_argument); + pass &= ck::utils::check_err(c_device_result_tensors[i], c_host_tensors[i]); + } + + std::cout << "Verification: " << (pass ? "SUCCESS" : "FAILURE") << "!" << std::endl; + } + + if(config.time_kernel) + { + float ave_time = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel}); + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec + << " GB/s, " << gemm.GetTypeString() << std::endl; + } + + return pass; +} + +std::vector argToIntArray(char* input) +{ + std::vector out; + std::istringstream in(input); + std::string item; + + while(std::getline(in, item, ',')) + { + out.push_back(std::stoi(item)); + } + return out; +} + +bool run_grouped_gemm_example(int argc, char* argv[]) +{ + ProblemSize problem_size; + ExecutionConfig config; + + if(argc < 10) + { + std::vector Ms{64, 127, 255, 129, 260, 190, 77}; + problem_size.group_count = Ms.size(); + + for(int i = 0; i < problem_size.group_count; i++) + { + problem_size.Ms.push_back(Ms[i]); + problem_size.Ns.push_back(252); + problem_size.Ks.push_back(4608); + + problem_size.stride_As.push_back(problem_size.Ks[i]); + problem_size.stride_Bs.push_back(problem_size.Ks[i]); + problem_size.stride_Cs.push_back(problem_size.Ns[i]); + + problem_size.stride_Ds.push_back({}); + for(int j = 0; j < NumDs; ++j) + { + problem_size.stride_Ds[i].push_back(problem_size.Ns[i]); + } + } + + std::cout + << "Usage:\n" + << "arg1: verification (0=no, 1=yes)\n" + << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n" + << "arg3: time kernel (0=n0, 1=yes)\n" + << "arg4 to 9: Ms, Ns, Ks, StrideAs, StrideBs, StrideCs (e.g., 256,256 128,128 64,64 " + "64,64 64,64 128,128)\n" + << "... setting default values." << std::endl; + } + else + { + config.do_verification = std::stoi(argv[1]); + config.init_method = std::stoi(argv[2]); + config.time_kernel = std::stoi(argv[3]); + + problem_size.Ms = argToIntArray(argv[4]); + problem_size.Ns = argToIntArray(argv[5]); + problem_size.Ks = argToIntArray(argv[6]); + + problem_size.stride_As = argToIntArray(argv[7]); + problem_size.stride_Bs = argToIntArray(argv[8]); + problem_size.stride_Cs = argToIntArray(argv[9]); + + for(int j = 0; j < NumDs; ++j) + { + problem_size.stride_Ds.push_back(problem_size.stride_Cs); + } + + problem_size.group_count = problem_size.Ms.size(); + } + + return run_grouped_gemm(problem_size, config); +} diff --git a/include/ck/tensor_operation/gpu/device/device_grouped_gemm.hpp b/include/ck/tensor_operation/gpu/device/device_grouped_gemm.hpp index 0c4aa68..3ee28f8 100644 --- a/include/ck/tensor_operation/gpu/device/device_grouped_gemm.hpp +++ b/include/ck/tensor_operation/gpu/device/device_grouped_gemm.hpp @@ -1,17 +1,79 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT // Modified by Hygon Information Technology Co., Ltd. - #pragma once + +#include #include +#include +#include #include #include "device_base.hpp" +#include "ck/utility/ignore.hpp" namespace ck { namespace tensor_operation { namespace device { +/// +/// @brief Structure representing single GEMM problem arguments. +/// +/// The pointer to the vector of those structures is passed to the GroupedGEMM entry +/// point kernel. +/// +/// @tparam NumDTensor The number of D input tensors. +/// +template +struct GroupedGemmKernelArgument +{ + __host__ __device__ GroupedGemmKernelArgument(const void* p_a_grid_, + const void* p_b_grid_, + std::array p_ds_grid_, + void* p_e_grid_, + index_t M_, + index_t N_, + index_t K_, + index_t StrideA_, + index_t StrideB_, + std::array StrideDs_, + index_t StrideE_) + : p_a_grid{p_a_grid_}, + p_b_grid{p_b_grid_}, + p_ds_grid{p_ds_grid_}, + p_e_grid{p_e_grid_}, + M{M_}, + N{N_}, + K{K_}, + StrideA{StrideA_}, + StrideB{StrideB_}, + StrideDs{StrideDs_}, + StrideE{StrideE_} + { + } + const void* p_a_grid; + const void* p_b_grid; + std::array p_ds_grid; + void* p_e_grid; + index_t M; + index_t N; + index_t K; + index_t StrideA; + index_t StrideB; + std::array StrideDs; + index_t StrideE; + + void Print() const + { + std::stringstream str; + for(auto sd : StrideDs) + str << sd << ","; + + std::cout << "arg {" << "M:" << M << ", " << "N:" << N << ", " << "K:" << K << ", " + << "SA:" << StrideA << ", " << "SB:" << StrideB << ", " << "SE:" << StrideE + << ", " << "SDs: {" << str.str() << "}" << "}" << std::endl; + } +}; struct GemmDesc { ck::index_t M_, N_, K_; @@ -48,6 +110,66 @@ struct DeviceGroupedGemm : public BaseOperator CElementwiseOperation c_element_op) = 0; virtual std::unique_ptr MakeInvokerPointer() = 0; + + //--------------------------------------------------------------------------------------------- + /// @brief Sets the device kernel arguments pointer and may copy data to device. + /// + /// TODO: Add which kernels are using this (TileLoop * FixedNK ??) + /// + /// @param p_arg The pointer to the Argument we're going to update. + /// @param[in] p_dev_kernel_args The pointer to the device memory which will contain kernel + /// arguments. + /// @param[in] p_host_kernel_args The pointer to the host memory which contains kernel + /// arguments that should be copied to device memory. + /// + virtual void SetDeviceKernelArgs(BaseArgument* p_arg, + void* p_dev_kernel_args, + const void* p_host_kernel_args) const + { + ignore = p_arg; + ignore = p_dev_kernel_args; + ignore = p_host_kernel_args; + + std::ostringstream err; + err << "This function is not implemented by the kernel: " << this->GetTypeString() + << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } + + //---------------------------------------------------------------------------------------------- + /// @brief Sets the device kernel arguments pointer and may copy data to device. + /// + /// @param p_arg The pointer to the Argument we're going to update. + /// @param[in] p_dev_kernel_args The pointer to the device memory which contains kernel + /// arguments. + /// + virtual void SetDeviceKernelArgs(BaseArgument* p_arg, void* p_dev_kernel_args) const + { + ignore = p_arg; + ignore = p_dev_kernel_args; + + std::ostringstream err; + err << "This function is not implemented by the kernel: " << this->GetTypeString() + << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } + + //---------------------------------------------------------------------------------------------- + /// @brief Gets the device kernel argument size. + /// + /// @param[in] p_arg The pointer to the Device op Argument. + /// + /// @return The device kernel argument size. + /// + virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const + { + ignore = p_arg; + + std::ostringstream err; + err << "This function is not implemented by the kernel: " << this->GetTypeString() + << __FILE__ << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } }; } // namespace device diff --git a/include/ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp b/include/ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp new file mode 100644 index 0000000..3c18e69 --- /dev/null +++ b/include/ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp @@ -0,0 +1,107 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include "ck/ck.hpp" +#include "ck/utility/common_header.hpp" +#include "ck/host_utility/device_prop.hpp" +#include "ck/host_utility/stream_utility.hpp" + +#include "device_grouped_gemm.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { + +/// @brief Grouped GEMM kernel using output Tile Looping algorithm +/// +/// @par This kernel does not require any knowledge about input data sizes (GEMM M/N/K) +/// It requires only the number of groups to launch. Other information like +/// data pointers and GEMM sizes, packed into gemm kernel args may be all dynamic +/// (known only at kernel run-time). +/// +/// @note This kernel does not support SplitK. + +template +struct DeviceGroupedGemmTileLoop : public DeviceGroupedGemm +{ +}; + +template +struct TileLoopKernelConfig +{ + // The oversubscription factor for the number of blocks that can simultaneously reside on + // GPU. + static constexpr int BLOCK_SUBSCRIPTION_FACTOR = 1; + // static constexpr int BLOCK_WAVES = BlockSize / get_warp_size(); + static constexpr int CU_SIMDS = 4; + // Assume we want to have at most 2 waves per SIMD + // static constexpr int CU_BLOCKS = math::integer_divide_floor(2 * CU_SIMDS, BLOCK_WAVES); + static int GetCuBlocks() + { + int BLOCK_WAVES = BlockSize / get_warp_size(); + return ck::math::integer_divide_floor(2 * CU_SIMDS, BLOCK_WAVES); + } + + template + static int CalculateMaxOccupancyGridSize(const KernelFunction& kernel, + const StreamConfig& stream_config) + { + // Calculate max number of workgroups that can simultaneously reside on the CU. + int occ_num_blocks = GetKernelOccupancy(kernel); + int cu_count = getAvailableComputeUnitCount(stream_config); + + if(stream_config.log_level_ > 0) + { + std::cout << "MaxActiveBlocksPerCU: " << occ_num_blocks + << ", available CUs count: " << cu_count << ", occup. grid size: " + << ck::math::min(occ_num_blocks, GetCuBlocks()) * cu_count << std::endl; + } + + return cu_count * ck::math::min(occ_num_blocks, GetCuBlocks()); + } + + template + static int GetKernelOccupancy(const KernelFunction& kernel) + { + int occupancy = 0; + hip_check_error( + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel, BlockSize, 0)); + return std::max(occupancy, 1); + } + + static int GetComputeUnitCount() + { + hipDeviceProp_t dev_prop; + hipDevice_t dev; + hip_check_error(hipGetDevice(&dev)); + hip_check_error(hipGetDeviceProperties(&dev_prop, dev)); + return dev_prop.multiProcessorCount; + } +}; + +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp new file mode 100644 index 0000000..a32a595 --- /dev/null +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp @@ -0,0 +1,1017 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include +#include +#include +#include + +#include "ck/host_utility/device_prop.hpp" +#include "ck/host_utility/kernel_launch.hpp" +#include "ck/host_utility/hip_check_error.hpp" +#include "ck/host_utility/stream_utility.hpp" +#include "ck/utility/common_header.hpp" +#include "ck/utility/loop_scheduler.hpp" +#include "ck/tensor_description/tensor_descriptor.hpp" +#include "ck/tensor_description/tensor_descriptor_helper.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include +#include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp" // stare wywalic +#include "ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_selector.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { + +// Dummy kernel to use as a fallback in the kernel selection logic +// Is not used in practice, but only used in case of misconfigured parameters +template +__global__ void kernel_dummy(const void CK_CONSTANT_ADDRESS_SPACE*, + const index_t, + const AElementwiseOperation, + const BElementwiseOperation, + const CDEElementwiseOperation) +{ +} +/// +/// @brief Entry point kernel for device-wide Grouped GEMM operation. +/// +/// @param[in] gemm_descs_const The pointer to the array of GEMM descriptor structures. +/// @param[in] group_count The number of together processed GEMMs. +/// +/// @tparam GridwiseGemm The specific GridwiseGEMM algorithm implementation. +/// @tparam GemmDesc The structure holding all necessary descriptors and +/// other data needed for grouped gemm calculation and work +/// distribution. +/// @tparam LocalBlock2ETileMap The structure providing mapping between workgroup ids, +/// the data tiles to process and the output tiles. +/// +template +__global__ void +#if CK_USE_LAUNCH_BOUNDS +__launch_bounds__(GridwiseGemm::MaxBlockSize, CK_MIN_BLOCK_PER_CU) +#endif + kernel_grouped_gemm_multiple_d_xdl(const void CK_CONSTANT_ADDRESS_SPACE* gemm_descs_const, + const index_t group_count, + const AElementwiseOperation a_element_op, + const BElementwiseOperation b_element_op, + const CDEElementwiseOperation cde_element_op) +{ +#if defined(__gfx9__) || defined(__gfx11__) || defined(__gfx12__) || defined(__gfx926__) || \ + defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ + defined(__gfx946__) + if constexpr(GridwiseGemm::template IsValidCompilationParameter<>()) + { + constexpr index_t shared_size = + GridwiseGemm::GetSharedMemoryNumberOfByte(get_device_arch()); + __shared__ uint8_t p_shared[shared_size]; + __shared__ uint8_t p_shared1[shared_size]; + + const auto gemm_desc_ptr = reinterpret_cast( + cast_pointer_to_generic_address_space(gemm_descs_const)); + + constexpr auto NumDTensor = DsDataType::Size(); + index_t tile_id = get_block_1d_id(); + index_t tile_offset = 0; + index_t group_id = -1; + index_t group_offset = 0; + index_t grid_size_grp = 0; + + index_t gemm_tile_id_start = 0; + index_t gemm_tile_id_end = 0; + + index_t M = 0, N = 0, K = 0; + + auto b2c_tile_map = OffsettedBlockToCTileMap(LocalBlock2ETileMap(1, 1), 1, 1); + + do + { + // Find corresponding GEMM group for our tile + while(!(tile_id >= gemm_tile_id_start && tile_id < gemm_tile_id_end) && + group_id < group_count) + { + group_offset += grid_size_grp; + group_id++; + + if(group_id >= group_count) + return; + + M = gemm_desc_ptr[group_id].M; + N = gemm_desc_ptr[group_id].N; + K = gemm_desc_ptr[group_id].K; + + if(M == 0 || N == 0 || K == 0) + { + grid_size_grp = 0; + continue; + } + + b2c_tile_map = OffsettedBlockToCTileMap( + LocalBlock2ETileMap(M, N, 4), group_offset, tile_offset); + grid_size_grp = b2c_tile_map.CalculateGridSize(M, N); + + gemm_tile_id_start = group_offset; + gemm_tile_id_end = group_offset + grid_size_grp; + } + + using DsGridPointer = decltype(GridwiseGemm::MakeDsGridPointer()); + DsGridPointer p_ds_grid; + + static_for<0, NumDTensor, 1>{}([&](auto i) { + using DDataType = remove_cvref_t>; + p_ds_grid(i) = static_cast(gemm_desc_ptr[group_id].p_ds_grid[i]); + }); + + static constexpr index_t kbatch = 1; + static constexpr index_t k_grain = kbatch * KPerBlock; + index_t K_split = (K + k_grain - 1) / k_grain * KPerBlock; + + const bool has_main_k_block_loop = GridwiseGemm::CalculateHasMainKBlockLoop(K_split); + + // Update tile offset if we have moved within group + b2c_tile_map.UpdateTileOffset(tile_offset); + + using Problem = typename GridwiseGemm::Problem; + auto problem = Problem(gemm_desc_ptr[group_id].M, + gemm_desc_ptr[group_id].N, + gemm_desc_ptr[group_id].K, + gemm_desc_ptr[group_id].StrideA, + gemm_desc_ptr[group_id].StrideB, + gemm_desc_ptr[group_id].StrideDs, + gemm_desc_ptr[group_id].StrideE, + kbatch); + + if(has_main_k_block_loop) + { + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1 || + BlkGemmPipelineVer == BlockGemmPipelineVersion::v3) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Full) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Three) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Four) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Five) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Seven) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + } + // Tail number could be Odd or Even + else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v4) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) + { + GridwiseGemm::template Run_2Lds( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + static_cast(p_shared1), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + else + { + GridwiseGemm::template Run_2Lds( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + static_cast(p_shared1), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + } + else + { + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1) + { + GridwiseGemm::template Run( + static_cast(gemm_desc_ptr[group_id].p_a_grid), + static_cast(gemm_desc_ptr[group_id].p_b_grid), + p_ds_grid, + static_cast(gemm_desc_ptr[group_id].p_e_grid), + static_cast(p_shared), + problem, + a_element_op, + b_element_op, + cde_element_op, + b2c_tile_map); + } + } + + tile_id += get_grid_size(); + tile_offset += get_grid_size(); + block_sync_lds(); + + } while(group_id < group_count); + } +#else + ignore = gemm_descs_const; + ignore = group_count; + ignore = a_element_op; + ignore = b_element_op; + ignore = cde_element_op; +#endif // end of if (defined(__gfx9__)) +} + +template + +struct DeviceGroupedGemmMultipleDXdlCShuffleTileLoop + : public DeviceGroupedGemmTileLoop +{ + using DeviceOp = DeviceGroupedGemmMultipleDXdlCShuffleTileLoop; + static constexpr auto WarpTileConfig64 = GetWarpTileConfig(); + static constexpr auto WarpTileConfig32 = GetWarpTileConfig(); + static constexpr auto NXdlPerWave64 = WarpTileConfig64.At(3); + static constexpr auto NXdlPerWave32 = WarpTileConfig32.At(3); + static constexpr index_t NumDTensor = DsDataType::Size(); + + template + using GridwiseGemmBase = GridwiseGemmMultiD_xdl_cshuffle_v3< + ALayout, + BLayout, + DsLayout, + ELayout, + ADataType, + BDataType, + AccDataType, + CShuffleDataType, + DsDataType, + EDataType, + AElementwiseOperation, + BElementwiseOperation, + CDEElementwiseOperation, + GemmSpec, + BlockSize, + MPerBlock, + NPerBlock, + KPerBlock, + AK1, + BK1, + WarpTileConfig::At(0), + WarpTileConfig::At(1), + WarpTileConfig::At(2), + WarpTileConfig::At(3), + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + false, // AThreadTransferSrcResetCoordinateAfterRun, + ABlockLdsExtraM, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + false, // BThreadTransferSrcResetCoordinateAfterRun, + BBlockLdsExtraN, + WarpTileConfig::At(4), + WarpTileConfig::At(5), + CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + CDEShuffleBlockTransferScalarPerVectors, + BlkGemmPipeSched, + BlkGemmPipelineVer, + ComputeTypeA, + ComputeTypeB>; + using GridwiseGemm64 = GridwiseGemmBase; + using GridwiseGemm32 = GridwiseGemmBase; + + using KernelConfig = TileLoopKernelConfig; + using KernelArguments = GroupedGemmKernelArgument; + using Block2ETileMap = BlockToCTileMap_Grouped_M00_N0_M01Adapt<8, MPerBlock, NPerBlock>; + using OffsettedLocalBlock2ETileMap = OffsettedBlockToCTileMap2; + + // Argument + struct Argument : public BaseArgument + { + Argument(std::vector& /* p_As */, + std::vector& /* p_Bs */, + std::vector>& /* p_Ds */, + std::vector& /* p_Es */, + const std::vector& gemm_descs, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op, + int occupancy_num_blocks, + int gpu_cu_count) + : group_count_{static_cast(gemm_descs.size())}, + occupancy_num_blocks_{occupancy_num_blocks}, + gpu_cu_count_{gpu_cu_count}, + gemm_descs_{gemm_descs}, + a_element_op_{a_element_op}, + b_element_op_{b_element_op}, + cde_element_op_{cde_element_op}, + tile_count_{0} + { + for(const auto& desc : gemm_descs) + { + const auto M = desc.M_; + const auto N = desc.N_; + const auto b2c_tile_map = Block2ETileMap(M, N); + tile_count_ += b2c_tile_map.CalculateGridSize(M, N); + } + } + + index_t group_count_; + const void* p_dev_gemm_args_; + int occupancy_num_blocks_; + int gpu_cu_count_; + const std::vector& gemm_descs_; + AElementwiseOperation a_element_op_; + BElementwiseOperation b_element_op_; + CDEElementwiseOperation cde_element_op_; + index_t tile_count_; + }; + + // Invoker + struct Invoker : public BaseInvoker + { + /// + /// @brief Launch Grouped Gemm kernel. + /// + /// @note This function overload is using user provided device buffer for kernel + /// arguments. + /// + /// @param[in] arg The structure containing kernel arguments (in host + /// memory). + /// @param[in] dev_gemm_args The pointer to device memory with kernel arguments. + /// @param[in] stream_config The device stream configuration. + /// + /// @return The average kernel execution time (if time measurement is enabled.) + /// + template + float Run(const Argument& arg, + const void* dev_gemm_args, + const StreamConfig& stream_config = StreamConfig{}) + { + if(dev_gemm_args == nullptr) + { + std::ostringstream err; + err << "The gemm arguments device buffer is not allocated!" << " In " << __FILE__ + << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } + + float ave_time = 0; + ave_time = DispatchKernel(arg, dev_gemm_args, stream_config); + + return ave_time; + } + + /// + /// @brief Launch Grouped Gemm kernel. + /// + /// @note This function overload is using device buffers (for kernel arguments and + /// for kernel auxiliary workspace) provided with an argument. The user should + /// call @see GetDeviceKernelArgSize, and @see SetDeviceKernelArgs, on arg + /// parameter to properly allocate those buffers. + /// + /// @param[in] arg The structure containing kernel arguments (in host memory). + /// @param[in] stream_config The device stream configuration. + /// + /// @return The average kernel execution time (if time measurement is enabled.) + /// + template + float RunImp(const Argument& arg, const StreamConfig& stream_config = StreamConfig{}) + { + if(arg.p_dev_gemm_args_ == nullptr) + { + std::ostringstream err; + err << "The gemm arguments device buffer is not allocated!" << " In " << __FILE__ + << ":" << __LINE__ << ", in function: " << __func__; + throw std::runtime_error(err.str()); + } + + return Run(arg, arg.p_dev_gemm_args_, stream_config); + } + + INVOKER_RUN_IMPL + + float Run(const BaseArgument* p_arg, + const StreamConfig& stream_config = StreamConfig{}) override + { + return Run(*dynamic_cast(p_arg), stream_config); + } + + private: + template + float DispatchKernel(const Argument& arg, + const void* dev_gemm_args, + const StreamConfig& stream_config) const + { + const auto kernel = GetKernelFunction(); + return LaunchKernel(kernel, arg, dev_gemm_args, stream_config); + } + + template + float LaunchKernel(const KernelFunction& kernel, + const Argument& arg, + const void* dev_gemm_args, + const StreamConfig& stream_config) const + { + int grid_size = KernelConfig::CalculateMaxOccupancyGridSize(kernel, stream_config); + + if(stream_config.log_level_ > 0) + { + std::cout << "grid_size: " << grid_size << " tile_count: " << arg.tile_count_ + << std::endl; + } + + // run multiple kernels + + return launch_and_time_kernel(stream_config, + kernel, + dim3(grid_size), + dim3(BlockSize), + 0, + cast_pointer_to_constant_address_space(dev_gemm_args), + arg.group_count_, + arg.a_element_op_, + arg.b_element_op_, + arg.cde_element_op_); + } + }; + + static constexpr bool IsValidCompilationParameter() + { + // TODO: properly implement this check + return true; + } + + static bool IsSupportedArgument(const Argument& arg) + { + if(!ck::is_xdl_wmma_supported()) + { + return false; + } + bool supported = true; + + constexpr index_t k_batch = 1; + bool isWave64 = get_warp_size() == 64; + for(index_t i = 0; i < arg.group_count_; ++i) + { + std::array placeholder_p_ds_grid{}; + std::array stride_Ds; + std::copy_n(arg.gemm_descs_[i].stride_Ds_.begin(), NumDTensor, stride_Ds.begin()); + if((arg.gemm_descs_[i].K_ % AK1 != 0 || arg.gemm_descs_[i].K_ % BK1 != 0) && + !(GemmSpec == GemmSpecialization::MKPadding || + GemmSpec == GemmSpecialization::NKPadding || + GemmSpec == GemmSpecialization::MNKPadding || + GemmSpec == GemmSpecialization::KPadding)) + { + return false; + } + if(isWave64) + { + if constexpr(NXdlPerWave64 > 0) + { + using GridArg = typename GridwiseGemm64::Argument; + GridArg gridwise_arg(nullptr, // p_a_grid, + nullptr, // p_b_grid, + placeholder_p_ds_grid, // p_ds_grid, + nullptr, // p_e_grid , + arg.gemm_descs_[i].M_, + arg.gemm_descs_[i].N_, + arg.gemm_descs_[i].K_, + arg.gemm_descs_[i].stride_A_, + arg.gemm_descs_[i].stride_B_, + stride_Ds, + arg.gemm_descs_[i].stride_C_, + k_batch, + arg.a_element_op_, + arg.b_element_op_, + arg.cde_element_op_); + + supported = supported && GridwiseGemm64::CheckValidity(gridwise_arg); + } + else + { + supported = false; + } + } + else + { + if constexpr(NXdlPerWave32 > 0) + { + using GridArg = typename GridwiseGemm32::Argument; + GridArg gridwise_arg(nullptr, // p_a_grid, + nullptr, // p_b_grid, + placeholder_p_ds_grid, // p_ds_grid, + nullptr, // p_e_grid , + arg.gemm_descs_[i].M_, + arg.gemm_descs_[i].N_, + arg.gemm_descs_[i].K_, + arg.gemm_descs_[i].stride_A_, + arg.gemm_descs_[i].stride_B_, + stride_Ds, + arg.gemm_descs_[i].stride_C_, + k_batch, + arg.a_element_op_, + arg.b_element_op_, + arg.cde_element_op_); + + supported = supported && GridwiseGemm32::CheckValidity(gridwise_arg); + } + else + { + supported = false; + } + } + } + + return supported; + } + + bool IsSupportedArgument(const BaseArgument* p_arg) override + { + return IsSupportedArgument(*dynamic_cast(p_arg)); + } + + template + static auto GetKernelFunction() + { + const auto kernel = kernel_grouped_gemm_multiple_d_xdl; + return kernel; + } + + static auto GetKernelFunction() + { + if(get_warp_size() == 64) + { + if constexpr(NXdlPerWave64 > 0) + { + const auto kernel = GetKernelFunction(); + return kernel; + } + } + else + { + if constexpr(NXdlPerWave32 > 0) + { + const auto kernel = GetKernelFunction(); + return kernel; + } + } + + // This is here to handle the case where MXdlPerWave/NxdPerWave is too small + // This is caught by IsSupportedArgument(), but as GetKernelFunction is sometimes called + // before we need a fallback kernel to return here. + return kernel_dummy; + } + + static int GetKernelOccupancy() + { + const auto kernel = GetKernelFunction(); + return KernelConfig::GetKernelOccupancy(kernel); + } + + static auto MakeArgument(std::vector& p_As, + std::vector& p_Bs, + std::vector>& p_Ds, + std::vector& p_Es, + std::vector& gemm_descs, + AElementwiseOperation a_elementwise_op, + BElementwiseOperation b_elementwise_op, + CDEElementwiseOperation cde_elementwise_op) + { + int occupancy = GetKernelOccupancy(); + int num_cu = KernelConfig::GetComputeUnitCount(); + + return Argument{p_As, + p_Bs, + p_Ds, + p_Es, + gemm_descs, + a_elementwise_op, + b_elementwise_op, + cde_elementwise_op, + occupancy, + num_cu}; + } + + std::unique_ptr + MakeArgumentPointer(std::vector& p_As, + std::vector& p_Bs, + std::vector>& p_Ds, + std::vector& p_Es, + std::vector& gemm_descs, + AElementwiseOperation a_elementwise_op, + BElementwiseOperation b_elementwise_op, + CDEElementwiseOperation cde_elementwise_op) override + { + int occupancy = GetKernelOccupancy(); + int num_cu = KernelConfig::GetComputeUnitCount(); + + return std::make_unique(p_As, + p_Bs, + p_Ds, + p_Es, + gemm_descs, + a_elementwise_op, + b_elementwise_op, + cde_elementwise_op, + occupancy, + num_cu); + } + + static auto MakeInvoker() { return Invoker{}; } + + std::unique_ptr MakeInvokerPointer() override + { + return std::make_unique(Invoker{}); + } + + std::string GetTypeString() const override + { + auto str = std::ostringstream(); + + std::map BlkGemmPipelineSchedulerToString{ + {BlockGemmPipelineScheduler::Intrawave, "Intrawave"}, + {BlockGemmPipelineScheduler::Interwave, "Interwave"}}; + + std::map BlkGemmPipelineVersionToString{ + {BlockGemmPipelineVersion::v1, "v1"}, + {BlockGemmPipelineVersion::v2, "v2"}, + {BlockGemmPipelineVersion::v3, "v3"}, + {BlockGemmPipelineVersion::v4, "v4"}, + {BlockGemmPipelineVersion::v5, "v5"}}; + + // clang-format off + str << "DeviceGroupedGemmMultipleDXdlCShuffleTileLoop" + << "<" + << std::string(ALayout::name)[0] << "," + << std::string(BLayout::name)[0] << "," + << std::string(ELayout::name)[0] << "," + << BlockSize << ", " + << MPerBlock << ", " + << NPerBlock << ", " + << KPerBlock << ", " + << AK1 << ", " + << BK1 << ", " + << MPerXDL << ", " + << NPerXDL << ", " + << MXdlPerWave << ", " + << NXdlPerWave << ", " + << ABlockTransferSrcScalarPerVector << ", " + << BBlockTransferSrcScalarPerVector << ", " + << CShuffleMXdlPerWavePerShuffle << ", " + << CShuffleNXdlPerWavePerShuffle << ", " + << getGemmSpecializationString(GemmSpec) << ", " + << "BlkGemmPipelineScheduler: " + << BlkGemmPipelineSchedulerToString[BlkGemmPipeSched] << ", " + << "BlkGemmPipelineVersion: " + << BlkGemmPipelineVersionToString[BlkGemmPipelineVer] + << ">"; + // clang-format on + + return str.str(); + } + + void SetDeviceKernelArgs(Argument& arg, + void* p_dev_kernel_args, + const void* p_host_kernel_args) const + { + arg.p_dev_gemm_args_ = p_dev_kernel_args; + hip_check_error(hipMemcpyAsync(p_dev_kernel_args, + p_host_kernel_args, + GetDeviceKernelArgSize(&arg), + hipMemcpyHostToDevice)); + } + + virtual void SetDeviceKernelArgs(BaseArgument* p_arg, + void* p_dev_kernel_args, + const void* p_host_kernel_args) const override + { + return SetDeviceKernelArgs( + *dynamic_cast(p_arg), p_dev_kernel_args, p_host_kernel_args); + } + + void SetDeviceKernelArgs(Argument& arg, void* p_dev_kernel_args) const + { + arg.p_dev_gemm_args_ = p_dev_kernel_args; + } + + virtual void SetDeviceKernelArgs(BaseArgument* p_arg, void* p_dev_kernel_args) const override + { + return SetDeviceKernelArgs(*dynamic_cast(p_arg), p_dev_kernel_args); + } + + size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const override + { + return dynamic_cast(p_arg)->group_count_ * sizeof(KernelArguments); + } +}; + +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp index 0c04c2d..9201fb8 100644 --- a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp +++ b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp @@ -813,6 +813,52 @@ struct OffsettedBlockToCTileMap index_t block_start_; }; +// second version with 2 offsets +template +struct OffsettedBlockToCTileMap2 +{ + using underlying_type = UnderlyingBlockToCTileMap; + + __host__ __device__ OffsettedBlockToCTileMap2(UnderlyingBlockToCTileMap block_to_ctile_map, + index_t group_offset, + index_t tile_offset) + : block_to_ctile_map_{block_to_ctile_map}, + group_offset_{group_offset}, + tile_offset_{tile_offset} + { + } + + template + __host__ __device__ constexpr auto CalculateBottomIndex(const TopIdx& idx_top) const + { + return block_to_ctile_map_.CalculateBottomIndex( + make_multi_index(idx_top[Number<0>{}] + tile_offset_ - group_offset_)); + } + + template + __host__ __device__ bool ValidCTileIndex(const CTileIdx& c_tile_idx, + const CTileDim& c_tile_dim) const + { + return block_to_ctile_map_.ValidCTileIndex(c_tile_idx, c_tile_dim); + } + + template + __host__ constexpr bool CheckValidity(const CGridDesc_M_N& c_grid_desc_m_n) const + { + return block_to_ctile_map_.CheckValidity(c_grid_desc_m_n); + } + + __host__ __device__ constexpr index_t CalculateGridSize(index_t M, index_t N) const + { + return block_to_ctile_map_.CalculateGridSize(M, N); + } + + __device__ void UpdateTileOffset(index_t offset) { tile_offset_ = offset; } + UnderlyingBlockToCTileMap block_to_ctile_map_; + index_t group_offset_; + index_t tile_offset_; +}; + /** * @brief Simple tile mapping which creates 3D grid of block of threads. * diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_gemm_multiple_d.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_gemm_multiple_d.hpp new file mode 100644 index 0000000..bc45823 --- /dev/null +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_gemm_multiple_d.hpp @@ -0,0 +1,166 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include +#include + +#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/device/device_base.hpp" +#include "ck/library/utility/host_tensor.hpp" + +#if __clang_major__ >= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +#endif +namespace ck { +namespace tensor_operation { +namespace host { + +// assumption: every D matrix has the same layout and the same datatype +template +struct ReferenceGemmMultipleD : public device::BaseOperator +{ + using DDataType = remove_cvref_t>; + // Argument + struct Argument : public device::BaseArgument + { + Argument(const Tensor& a_m_k, + const Tensor& b_k_n, + const std::array, DsDataType::Size()>& ds_m_n, + Tensor& c_m_n, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) + : a_m_k_{a_m_k}, + b_k_n_{b_k_n}, + ds_m_n_{ds_m_n}, + c_m_n_{c_m_n}, + a_element_op_{a_element_op}, + b_element_op_{b_element_op}, + cde_element_op_{cde_element_op} + { + } + + const Tensor& a_m_k_; + const Tensor& b_k_n_; + const std::array, DsDataType::Size()>& ds_m_n_; + Tensor& c_m_n_; + + AElementwiseOperation a_element_op_; + BElementwiseOperation b_element_op_; + CDEElementwiseOperation cde_element_op_; + }; + + // Invoker + struct Invoker : public device::BaseInvoker + { + using Argument = ReferenceGemmMultipleD::Argument; + + float Run(const Argument& arg) + { + auto f_mk_kn_mn = [&](auto m, auto n) { + const int K = arg.a_m_k_.mDesc.GetLengths()[1]; + + AccDataType v_acc = 0; + ComputeTypeA v_a = 0; + ComputeTypeB v_b = 0; + + for(int k = 0; k < K; ++k) + { + arg.a_element_op_(v_a, arg.a_m_k_(m, k)); + arg.b_element_op_(v_b, arg.b_k_n_(k, n)); + + v_acc += + ck::type_convert(v_a) * ck::type_convert(v_b); + } + + CDataType v_c = 0; + + if constexpr(DsDataType::Size() == 0) + { + arg.cde_element_op_(v_c, v_acc); + } + else if constexpr(DsDataType::Size() == 1) + { + arg.cde_element_op_(v_c, v_acc, arg.ds_m_n_[0](m, n)); + } + else if constexpr(DsDataType::Size() == 2) + { + arg.cde_element_op_(v_c, v_acc, arg.ds_m_n_[0](m, n), arg.ds_m_n_[1](m, n)); + } + + arg.c_m_n_(m, n) = v_c; + }; + + make_ParallelTensorFunctor( + f_mk_kn_mn, arg.c_m_n_.mDesc.GetLengths()[0], arg.c_m_n_.mDesc.GetLengths()[1])( + std::thread::hardware_concurrency()); + + return 0; + } + + float Run(const device::BaseArgument* p_arg, + const StreamConfig& /* stream_config */ = StreamConfig{}) override + { + return Run(*dynamic_cast(p_arg)); + } + }; + + static constexpr bool IsValidCompilationParameter() + { + // TODO: properly implement this check + return true; + } + + bool IsSupportedArgument(const device::BaseArgument*) override { return true; } + + static auto MakeArgument(const Tensor& a_m_k, + const Tensor& b_k_n, + const std::array, DsDataType::Size()>& ds_m_n, + Tensor& c_m_n, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) + { + return Argument{a_m_k, b_k_n, ds_m_n, c_m_n, a_element_op, b_element_op, cde_element_op}; + } + + static auto MakeInvoker() { return Invoker{}; } + + virtual std::unique_ptr MakeInvokerPointer() + { + return std::make_unique(Invoker{}); + } + + std::string GetTypeString() const override + { + auto str = std::stringstream(); + + // clang-format off + str << "ReferenceGemmMultipleD" + << std::endl; + // clang-format on + + return str.str(); + } +}; + +} // namespace host +} // namespace tensor_operation +} // namespace ck + +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_gemm_tile_loop.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_gemm_tile_loop.hpp new file mode 100644 index 0000000..0d53fa7 --- /dev/null +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_gemm_tile_loop.hpp @@ -0,0 +1,109 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/device_grouped_gemm_tile_loop.hpp" + +#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +#ifdef CK_ENABLE_FP16 +// fp16_output +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instances( + std::vector>>& instances); + +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instances( + std::vector>>& instances); +#endif + +template +struct DeviceOperationInstanceFactory< + ck::tensor_operation::device::DeviceGroupedGemmTileLoop> +{ + using DeviceOp = DeviceGroupedGemmTileLoop; + + static auto GetInstances() + { + std::vector> op_ptrs; + +#ifdef CK_ENABLE_FP16 + // fp16_output + if constexpr(is_same_v && is_same_v && + is_same_v) + { + if constexpr(is_same_v && is_same_v && + is_same_v) + { + add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instances(op_ptrs); + } + if constexpr(is_same_v && is_same_v && + is_same_v) + { + add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instances(op_ptrs); + } + } +#endif + return op_ptrs; + } +}; + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/CMakeLists.txt b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/CMakeLists.txt new file mode 100644 index 0000000..074dca8 --- /dev/null +++ b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/CMakeLists.txt @@ -0,0 +1,11 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# Modified by Hygon Information Technology Co., Ltd. + +# GroupedGEMM tile loop instances for HCU (gfx936 family, 16x16 XDL tiles only). +set(GROUPED_GEMM_TILE_LOOP_HCU_INSTANCES + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp + ) + +add_instance_library(device_grouped_gemm_tile_loop_instance ${GROUPED_GEMM_TILE_LOOP_HCU_INSTANCES}) diff --git a/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp new file mode 100644 index 0000000..2b559cf --- /dev/null +++ b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp @@ -0,0 +1,70 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. +// HCU (gfx936 family) instance set for GroupedGEMM tile loop. +// HCU MMAC only supports 16x16 XDL tiles; this config is derived from the +// validated example configuration (MPerXDL/NPerXDL = 16). + +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +template +using S = ck::Sequence; + +using DsDataType = ck::Tuple<>; + +using DsLayout = ck::Tuple<>; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +static constexpr auto GemmMNKPadding = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +using device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_irregular_tile_instances = std::tuple< + // clang-format off + //###########################################| A| B| Ds| E| AData| BData| AccData| CShuffle| DsData| EData| A| B| C| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| + //###########################################| Layout| Layout| Layout| Layout| Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| + //###########################################| | | | | | | | | | | Operation| Operation| Operation| | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| + //###########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 8, 16, 16, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 8>, S<4>> + // clang-format on + >; + +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instances( + std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_irregular_tile_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp new file mode 100644 index 0000000..01e29f5 --- /dev/null +++ b/library/src/tensor_operation_instance_hcu/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp @@ -0,0 +1,70 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. +// HCU (gfx936 family) instance set for GroupedGEMM tile loop. +// HCU MMAC only supports 16x16 XDL tiles; this config is the exact tile +// geometry validated by example/15_grouped_gemm/grouped_gemm_multiple_d_xdl_fp16. + +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +template +using S = ck::Sequence; + +using DsDataType = ck::Tuple<>; + +using DsLayout = ck::Tuple<>; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +static constexpr auto GemmMNKPadding = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +using device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_irregular_tile_instances = std::tuple< + // clang-format off + //###########################################| A| B| Ds| E| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| + //###########################################| Layout| Layout| Layout| Layout| Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| + //###########################################| | | | | | | | | | | Operation| Operation| Operation| | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| + //###########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 8, 16, 16, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<4>> + // clang-format on + >; + +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instances( + std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_irregular_tile_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck From 8fc01b7a05c7f7d23af13383174f03e17c1ceadd Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Wed, 30 Sep 2026 11:46:39 +0800 Subject: [PATCH 02/11] feat(gemm_ab_scale): add GemmMultiD_ABScale --- example_hcu/65_gemm_ab_scale/CMakeLists.txt | 13 + .../gemm_ab_scale_xdl_bf16.cpp | 336 +++++ ...gemm_pipeline_xdlops_ab_scale_selector.hpp | 113 ++ ...kwise_gemm_pipeline_xdlops_v1_ab_scale.hpp | 880 +++++++++++ ...kwise_gemm_pipeline_xdlops_v2_ab_scale.hpp | 728 ++++++++++ ...kwise_gemm_pipeline_xdlops_v3_ab_scale.hpp | 652 +++++++++ .../device_gemm_multiple_d_ab_scale.hpp | 460 ++++++ ...mm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp | 624 ++++++++ ..._gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp | 1289 +++++++++++++++++ include/ck/utility/index_expression.hpp | 169 +++ .../ck/utility/thread_buf_to_vec_loader.hpp | 163 +++ .../gpu/gemm_ab_scale.hpp | 108 ++ .../gpu/grouped_gemm_tile_loop/CMakeLists.txt | 13 + ...ile_loop_f16_f16_f16_mk_kn_mn_instance.cpp | 76 + ...ile_loop_f16_f16_f16_mk_nk_mn_instance.cpp | 78 + .../gpu/gemm_ab_scale/CMakeLists.txt | 12 + ...e_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp | 79 + 17 files changed, 5793 insertions(+) create mode 100644 example_hcu/65_gemm_ab_scale/CMakeLists.txt create mode 100644 example_hcu/65_gemm_ab_scale/gemm_ab_scale_xdl_bf16.cpp create mode 100644 include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp create mode 100644 include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp create mode 100644 include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp create mode 100644 include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp create mode 100644 include/ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp create mode 100644 include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp create mode 100644 include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp create mode 100644 include/ck/utility/index_expression.hpp create mode 100644 include/ck/utility/thread_buf_to_vec_loader.hpp create mode 100644 library/include/ck/library/tensor_operation_instance/gpu/gemm_ab_scale.hpp create mode 100644 library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/CMakeLists.txt create mode 100644 library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp create mode 100644 library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp create mode 100644 library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/CMakeLists.txt create mode 100644 library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp diff --git a/example_hcu/65_gemm_ab_scale/CMakeLists.txt b/example_hcu/65_gemm_ab_scale/CMakeLists.txt new file mode 100644 index 0000000..27ec15f --- /dev/null +++ b/example_hcu/65_gemm_ab_scale/CMakeLists.txt @@ -0,0 +1,13 @@ +# Copyright (c) 2026 Hygon Info Technologies Ltd. +# SPDX-License-Identifier: MIT +# Modified by Hygon Information Technology Co., Ltd. +# hcu gemm ab_scale (per-block A/B dequantization scales) example +add_example_executable(example_hcu_gemm_ab_scale_xdl_bf16 + gemm_ab_scale_xdl_bf16.cpp) + +add_test(NAME example_hcu_gemm_ab_scale_xdl_bf16_smoke + COMMAND $ 1 1 0 256 256 512 0 +) +set_tests_properties(example_hcu_gemm_ab_scale_xdl_bf16_smoke + PROPERTIES LABELS "hcu;gemm;ab_scale;smoke" +) diff --git a/example_hcu/65_gemm_ab_scale/gemm_ab_scale_xdl_bf16.cpp b/example_hcu/65_gemm_ab_scale/gemm_ab_scale_xdl_bf16.cpp new file mode 100644 index 0000000..d511aec --- /dev/null +++ b/example_hcu/65_gemm_ab_scale/gemm_ab_scale_xdl_bf16.cpp @@ -0,0 +1,336 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. +// HCU (gfx936) adaptation: BF16 inputs with FP32 per-block dequantization +// scales (gfx936 has no dense FP8 MMAC, so the upstream fp8 ab_scale example +// is validated here with BF16 inputs instead). + +#include +#include +#include +#include + +#include "ck/ck.hpp" + +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp" + +#include "ck/library/utility/device_memory.hpp" +#include "ck/library/utility/host_tensor.hpp" +#include "ck/library/utility/host_tensor_generator.hpp" +#include "ck/library/utility/literals.hpp" +#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp" +#include "ck/library/utility/check_err.hpp" + +#include "ck/utility/blkgemmpipe_scheduler.hpp" + +using ::ck::DeviceMem; + +template +using S = ck::Sequence; + +using BF16 = ck::bhalf_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +using A0DataType = BF16; +using A1DataType = F32; +using B0DataType = BF16; +using B1DataType = F32; +using AccDataType = F32; +using CShuffleDataType = F32; +using DsDataType = ck::Tuple<>; +using EDataType = BF16; + +using A0Layout = Row; +using B0Layout = Col; +using D0Layout = Row; +using D1Layout = Col; +using DsLayout = ck::Tuple<>; +using ELayout = Row; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +using AElementOp = PassThrough; +using BElementOp = PassThrough; +using CDEElementOp = PassThrough; + +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; + +static constexpr ck::index_t Scale_Block_M = 1; +static constexpr ck::index_t Scale_Block_N = 128; +static constexpr ck::index_t Scale_Block_K = 32; + +using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3 + // clang-format off + , S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, + S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, + 1, 1, S<1, 32, 1, 8>, S<4>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16>; +// clang-format on + +int main(int argc, char* argv[]) +{ + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; + bool flush_cache = false; + + // GEMM shape + ck::index_t M = 256; + ck::index_t N = 256; + ck::index_t K = 512; + + ck::index_t StrideA = K; + ck::index_t StrideB = K; + ck::index_t StrideE = N; + + ck::index_t KBatch = 1; + + if(argc == 1) + { + // use default case + } + else if(argc == 4) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + } + else if(argc == 8 || argc == 9) + { + do_verification = std::stoi(argv[1]); + init_method = std::stoi(argv[2]); + time_kernel = std::stoi(argv[3]); + + M = std::stoi(argv[4]); + N = std::stoi(argv[5]); + K = std::stoi(argv[6]); + + flush_cache = std::stoi(argv[7]); + + if(argc == 9) + { + KBatch = std::stoi(argv[8]); + } + + StrideA = K; + StrideB = K; + StrideE = N; + } + else + { + printf("arg1: verification (0=no, 1=yes)\n"); + printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"); + printf("arg3: time kernel (0=no, 1=yes)\n"); + printf("arg4 to 6: M, N, K\n"); + printf("arg7: flush both I$ and L2$ (0=no, 1=yes)\n"); + printf("arg8: KBatch (default: 1)\n"); + exit(0); + } + + ck::index_t Scale_Stride_AM = (K + Scale_Block_K - 1) / Scale_Block_K; + ck::index_t Scale_Stride_BN = (K + Scale_Block_K - 1) / Scale_Block_K; + + auto f_host_tensor_descriptor = + [](std::size_t row, std::size_t col, std::size_t stride, auto layout) { + using namespace ck::literals; + + if(std::is_same::value) + { + return HostTensorDescriptor({row, col}, {stride, 1_uz}); + } + else + { + return HostTensorDescriptor({row, col}, {1_uz, stride}); + } + }; + + Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{})); + Tensor a1_m_k(f_host_tensor_descriptor((M + Scale_Block_M - 1) / Scale_Block_M, + (K + Scale_Block_K - 1) / Scale_Block_K, + Scale_Stride_AM, + A0Layout{})); + Tensor b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{})); + Tensor b1_k_n(f_host_tensor_descriptor((K + Scale_Block_K - 1) / Scale_Block_K, + (N + Scale_Block_N - 1) / Scale_Block_N, + Scale_Stride_BN, + B0Layout{})); + Tensor e_m_n_host_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{})); + Tensor e_m_n_device_result(f_host_tensor_descriptor(M, N, StrideE, ELayout{})); + + std::cout << "a0_m_k: " << a0_m_k.mDesc << std::endl; + std::cout << "a1_m_k: " << a1_m_k.mDesc << std::endl; + std::cout << "b0_k_n: " << b0_k_n.mDesc << std::endl; + std::cout << "b1_k_n: " << b1_k_n.mDesc << std::endl; + std::cout << "e_m_n: " << e_m_n_host_result.mDesc << std::endl; + + switch(init_method) + { + case 0: break; + case 1: + a0_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 2}); + b0_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 2}); + a1_m_k.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + b1_k_n.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + break; + case 2: + a0_m_k.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + b0_k_n.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + a1_m_k.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + b1_k_n.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + break; + case 3: + // diagnostic: constant unit scales -> device must match plain GEMM + a0_m_k.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + b0_k_n.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + a1_m_k.GenerateTensorValue(GeneratorTensor_3{1.0, 1.0}); + b1_k_n.GenerateTensorValue(GeneratorTensor_3{1.0, 1.0}); + break; + default: + a0_m_k.GenerateTensorValue(GeneratorTensor_2{-2, 2}); + b0_k_n.GenerateTensorValue(GeneratorTensor_2{-2, 2}); + a1_m_k.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + b1_k_n.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); + } + + DeviceMem a0_device_buf(sizeof(A0DataType) * a0_m_k.mDesc.GetElementSpaceSize()); + DeviceMem a1_device_buf(sizeof(A1DataType) * a1_m_k.mDesc.GetElementSpaceSize()); + DeviceMem b0_device_buf(sizeof(B0DataType) * b0_k_n.mDesc.GetElementSpaceSize()); + DeviceMem b1_device_buf(sizeof(B1DataType) * b1_k_n.mDesc.GetElementSpaceSize()); + DeviceMem e_device_buf(sizeof(EDataType) * e_m_n_device_result.mDesc.GetElementSpaceSize()); + + a0_device_buf.ToDevice(a0_m_k.mData.data()); + a1_device_buf.ToDevice(a1_m_k.mData.data()); + b0_device_buf.ToDevice(b0_k_n.mData.data()); + b1_device_buf.ToDevice(b1_k_n.mData.data()); + + auto a_element_op = AElementOp{}; + auto b_element_op = BElementOp{}; + auto cde_element_op = CDEElementOp{}; + + constexpr ck::index_t NumDTensor = DsDataType::Size(); + + // do GEMM + auto device_op = DeviceOpInstance{}; + auto invoker = device_op.MakeInvoker(); + auto argument = device_op.MakeArgument(a0_device_buf.GetDeviceBuffer(), + b0_device_buf.GetDeviceBuffer(), + std::array{}, + e_device_buf.GetDeviceBuffer(), + M, + N, + K, + StrideA, + StrideB, + std::array{}, + StrideE, + a1_device_buf.GetDeviceBuffer(), + b1_device_buf.GetDeviceBuffer(), + a_element_op, + b_element_op, + cde_element_op); + argument.KBatch = KBatch; + + if(!device_op.IsSupportedArgument(argument)) + { + throw std::runtime_error( + "wrong! device_gemm with the specified compilation parameters does " + "not support this GEMM problem"); + } + + std::size_t flop = std::size_t(2) * M * N * K; + std::size_t num_btype = + sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(EDataType) * M * N; + + float ave_time = .0; + + if(flush_cache) + { + int rotating_buf = (512 * 1024 * 1024 + num_btype - 1) / num_btype; + + ave_time = invoker.Run(argument, + StreamConfig{nullptr, time_kernel, 0, 50, 100, true, rotating_buf}); + } + else + { + ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel, 0, 50, 100}); + } + + float tflops = static_cast(flop) / 1.E9 / ave_time; + + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s" + << std::endl; + + if(do_verification) + { + Tensor c_m_n({M, N}); + Tensor a_m_k({M, K}); + Tensor b_k_n({K, N}); + + for(int m = 0; m < M; m++) + { + for(int k = 0; k < K; k++) + { + a_m_k(m, k) = ck::type_convert(a0_m_k(m, k)) * + a1_m_k(m / Scale_Block_M, k / Scale_Block_K); + } + } + + for(int n = 0; n < N; n++) + { + for(int k = 0; k < K; k++) + { + b_k_n(k, n) = ck::type_convert(b0_k_n(k, n)) * + b1_k_n(k / Scale_Block_K, n / Scale_Block_N); + } + } + + using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm; + auto ref_gemm = ReferenceGemmInstance{}; + auto ref_invoker = ref_gemm.MakeInvoker(); + + auto ref_argument = + ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n, PassThrough{}, PassThrough{}, PassThrough{}); + + ref_invoker.Run(ref_argument); + + for(int m = 0; m < M; ++m) + { + for(int n = 0; n < N; ++n) + { + e_m_n_host_result(m, n) = ck::type_convert(c_m_n(m, n)); + } + } + + e_device_buf.FromDevice(e_m_n_device_result.mData.data()); + + return ck::utils::check_err( + e_m_n_device_result, e_m_n_host_result, "Error: Incorrect results!", 1e-2, 1e-2) + ? 0 + : 1; + } + + return 0; +} diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp new file mode 100644 index 0000000..ced35e5 --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp @@ -0,0 +1,113 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp" +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp" +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp" + +namespace ck { + +template +constexpr auto BlockGemmABScalePipeline_Selector() +{ + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1) + { + return BlockwiseGemmXdlops_pipeline_v1_ab_scale{}; + } + else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2) + { + return BlockwiseGemmXdlops_pipeline_v2_ab_scale{}; + } + else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3) + { + return BlockwiseGemmXdlops_pipeline_v3_ab_scale{}; + } + else + { + std::cerr << "BlockGemmPipeline configuration is not available" << std::endl; + } +} + +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp new file mode 100644 index 0000000..0155763 --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp @@ -0,0 +1,880 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/utility/thread_buf_to_vec_loader.hpp" +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp" + +namespace ck { + +// Compute optimized pipeline +// GlobalPrefetchStages: 2 +// LocalPreFillStages: 1 +// LocalPreFetchStages: 1 +// LocalSharedMemoryBuffer: 1 + +template +struct BlockwiseGemmXdlops_pipeline_v1_ab_scale +{ +}; + +template +struct BlockwiseGemmXdlops_pipeline_v1_ab_scale + : BlockwiseGemmXdlops_pipeline_base + +{ + using Base = BlockwiseGemmXdlops_pipeline_base; + using Base::A_K1; + using Base::B_K1; + using Base::I0; + using Base::I1; + using Base::KRepeat; + using Base::xdlops_gemm; + using typename Base::HotLoopInstList; + + using Base::CalculateCThreadOriginDataIndex; + using Base::CalculateCThreadOriginDataIndex8D; + using Base::GetCBlockDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::GetCThreadBuffer; + using Base::GetCThreadDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::GetWaveIdx; + using Base::MakeCGridDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + + using Base::a_block_desc_m0_m1_m2_k; + using Base::b_block_desc_n0_n1_n2_k; + + using ComputeDataTypeBuf = typename Base::ComputeDataTypeBuf; + + static constexpr index_t AMmaKStride = xdlops_gemm.K0PerXdlops * KPack; + static constexpr index_t BMmaKStride = xdlops_gemm.K0PerXdlops * KPack; + + static constexpr index_t PrefetchStages = 2; + static constexpr index_t PrefillStages = 1; + static constexpr index_t GlobalBufferNum = 1; + + // Force mfma not cross the scaleblock + __device__ static auto CalculateAThreadOriginDataIndex() + { + const auto wave_idx = GetWaveIdx(); + + const auto waveId_m = wave_idx[I0]; + + const auto xdlops_a_idx = xdlops_gemm.CalculateAThreadOriginDataIndex(); + + return make_tuple(0, waveId_m, xdlops_a_idx[I1], KPack * xdlops_a_idx[I0]); + } + + __device__ static auto CalculateBThreadOriginDataIndex() + { + const auto wave_idx = GetWaveIdx(); + + const auto waveId_n = wave_idx[I1]; + + const auto xdlops_b_idx = xdlops_gemm.CalculateBThreadOriginDataIndex(); + + return make_tuple(0, waveId_n, xdlops_b_idx[I1], KPack * xdlops_b_idx[I0]); + } + + __host__ static constexpr bool BlockHasHotloop(index_t num_loop) + { + return num_loop > PrefetchStages; + } + + __host__ static constexpr TailNumber BlockLoopTailNum(index_t num_loop) + { + return num_loop == 1 ? TailNumber::Odd : TailNumber::Full; + } + + __device__ static constexpr auto HotLoopScheduler() + { + // A/B split schedule + // compiler is likely to use ds_read2 when instruction width smaller than 16bytes + constexpr auto num_ds_read_inst_a = + HotLoopInstList::A_LDS_Read_Width * sizeof(ADataType) == 16 + ? HotLoopInstList::A_LDS_Read_Inst_Num + : HotLoopInstList::A_LDS_Read_Inst_Num / 2; + constexpr auto num_ds_read_inst_b = + HotLoopInstList::B_LDS_Read_Width * sizeof(BDataType) == 16 + ? HotLoopInstList::B_LDS_Read_Inst_Num + : HotLoopInstList::B_LDS_Read_Inst_Num / 2; + + constexpr auto num_ds_write_inst_a = HotLoopInstList::A_LDS_Write_Inst_Num; + constexpr auto num_ds_write_inst_b = HotLoopInstList::B_LDS_Write_Inst_Num; + + constexpr auto num_buffer_load_inst_a = HotLoopInstList::A_Buffer_Load_Inst_Num; + constexpr auto num_buffer_load_inst_b = HotLoopInstList::B_Buffer_Load_Inst_Num; + + constexpr auto num_mfma_inst = HotLoopInstList::C_MFMA_Inst_Num; + + constexpr auto mfma_cycle = HotLoopInstList::C_MFMA_Inst_Cycle; + constexpr auto ds_read_a_issue_cycle = + HotLoopInstList::A_LDS_Read_Width * sizeof(ADataType) == 16 ? 8 : 4; + constexpr auto ds_read_b_issue_cycle = + HotLoopInstList::B_LDS_Read_Width * sizeof(BDataType) == 16 ? 8 : 4; + constexpr auto ds_read_a_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_a_issue_cycle - 1) / (2 * ds_read_a_issue_cycle); + constexpr auto ds_read_b_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_b_issue_cycle - 1) / (2 * ds_read_b_issue_cycle); + + constexpr auto num_dsread_a_mfma = + (num_ds_read_inst_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; + constexpr auto num_dsread_b_mfma = + (num_ds_read_inst_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; + + // stage 1 + // Separate this part? + // constexpr auto num_mfma_per_ds_read = sizeof(ComputeDataTypeBuf) / sizeof(ADataType) > + // sizeof(ComputeDataTypeBuf) / + // sizeof(BDataType) + // ? sizeof(ComputeDataTypeBuf) / + // sizeof(ADataType) : sizeof(ComputeDataTypeBuf) + // / sizeof(BDataType); + constexpr auto num_mfma_stage1 = num_mfma_inst - (num_dsread_a_mfma + num_dsread_b_mfma); + constexpr auto num_mfma_per_issue = + num_mfma_stage1 / (num_buffer_load_inst_a + num_buffer_load_inst_b); + constexpr auto num_dswrite_per_issue_a = num_ds_write_inst_a / num_buffer_load_inst_a; + constexpr auto num_dswrite_per_issue_b = num_ds_write_inst_b / num_buffer_load_inst_b; + + static_for<0, num_buffer_load_inst_a, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_a, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_a, 0); // MFMA + }); + static_for<0, num_buffer_load_inst_b, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_b, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_b, 0); // MFMA + }); + + // stage 2 + static_for<0, num_dsread_a_mfma, 1>{}([&](auto i) { + if constexpr((num_ds_read_inst_a - (i + 1) * ds_read_a_mfma_rate) >= + ds_read_a_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier(0x100, + num_ds_read_inst_a - (num_dsread_a_mfma - 1) * + ds_read_a_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + + static_for<0, num_dsread_b_mfma, 1>{}([&](auto i) { + if constexpr((num_ds_read_inst_b - (i + 1) * ds_read_b_mfma_rate) >= + ds_read_b_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier(0x100, + num_ds_read_inst_b - (num_dsread_b_mfma - 1) * + ds_read_b_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + } + + template + __device__ void Run( + // ABlockCopy + const AGridDesc& a_grid_desc, + const ABlockDesc& a_block_desc, + ABlockTransfer& a_blockwise_copy, + const AGridBuffer& a_grid_buf, + ABlockBuffer& a_block_buf, + const ABlockTransferStep& a_block_copy_step, + // BBlockCopy + const BGridDesc& b_grid_desc, + const BBlockDesc& b_block_desc, + BBlockTransfer& b_blockwise_copy, + const BGridBuffer& b_grid_buf, + BBlockBuffer& b_block_buf, + const BBlockTransferStep& b_block_copy_step, + // CThread + const CScaleThreadDesc& c_scale_thread_desc, + CThreadBuffer& c_thread_buf, + // AScaleThreadCopy + const AScaleGridDesc& a_scale_grid_desc, + const AScaleThreadDesc& a_scale_thread_desc, + AScaleThreadTransfer& a_scale_thread_copy, + const AScaleGridBuffer& a_scale_grid_buf, + const AScaleThreadTransferStep& a_scale_thread_copy_step, + // BScaleThreadCopy + const BScaleGridDesc& b_scale_grid_desc, + const BScaleThreadDesc& b_scale_thread_desc, + BScaleThreadTransfer& b_scale_thread_copy, + const BScaleGridBuffer& b_scale_grid_buf, + const BScaleThreadTransferStep& b_scale_thread_copy_step, + // num_loop + index_t num_loop) const + { + __builtin_amdgcn_sched_barrier(0); + // assume kperblock = scaleblockk + auto a_thread_buf = make_static_buffer( + a_thread_desc_.GetElementSpaceSize()); + auto b_thread_buf = make_static_buffer( + b_thread_desc_.GetElementSpaceSize()); + auto a_scale_thread_buf = make_static_buffer( + a_scale_thread_desc.GetElementSpaceSize()); + auto b_scale_thread_buf = make_static_buffer( + b_scale_thread_desc.GetElementSpaceSize()); + auto c_scale_thread_buf = make_static_buffer( + c_scale_thread_desc.GetElementSpaceSize()); + + // Global prefetch 1 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + constexpr auto num_scale_k_block = CScaleThreadDesc{}.GetLength(Number<0>{}); + constexpr auto num_scale_m_block = CScaleThreadDesc{}.GetLength(Number<1>{}); + constexpr auto num_scale_n_block = CScaleThreadDesc{}.GetLength(Number<2>{}); + + static_ford>{}( + [&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto k0 = Number{}]>{}; + constexpr index_t c_offset = + CScaleThreadDesc{}.CalculateOffset(make_tuple(k0, m0, n0)); + constexpr index_t a_offset = AScaleThreadDesc{}.CalculateOffset(make_tuple(m0, k0)); + constexpr index_t b_offset = BScaleThreadDesc{}.CalculateOffset(make_tuple(n0, k0)); + + c_scale_thread_buf(Number{}) = + a_scale_thread_buf[Number{}] * b_scale_thread_buf[Number{}]; + }); + + // Local prefill 1 + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); + + // Global prefetch 2 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + // Initialize C + c_thread_buf.Clear(); + + StaticBufferTupleOfVector + c_thread_buf_per_scale; + + // Local prefetch 1 + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k0, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k0, I0), + b_thread_buf); + }); + }); + + __builtin_amdgcn_sched_barrier(0); + + // main body + if constexpr(HasMainLoop) + { + index_t i = 0; + do + { + block_sync_lds(); + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); + + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_ford>{}([&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto kscale0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat / num_scale_k_block, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + constexpr auto var = kscale0 * KRepeat / num_scale_k_block + k0; + auto loadA = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, + a_thread_buf}; + auto loadB = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, + b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + constexpr index_t cscale_offset = CScaleThreadDesc{}.CalculateOffset( + make_tuple(kscale0, m0, n0 * num_scale_n_block / NRepeat)); + + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[Number{}]); + }); + }); + + static_ford>{}( + [&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto k0 = Number{}]>{}; + constexpr index_t c_offset = + CScaleThreadDesc{}.CalculateOffset(make_tuple(k0, m0, n0)); + constexpr index_t a_offset = + AScaleThreadDesc{}.CalculateOffset(make_tuple(m0, k0)); + constexpr index_t b_offset = + BScaleThreadDesc{}.CalculateOffset(make_tuple(n0, k0)); + + c_scale_thread_buf(Number{}) = + a_scale_thread_buf[Number{}] * + b_scale_thread_buf[Number{}]; + }); + + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + i += 1; + } while(i < (num_loop - 2)); + } + + // tail + if constexpr(TailNum == TailNumber::Full) + { + block_sync_lds(); + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); + + static_ford>{}([&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto kscale0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat / num_scale_k_block, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + constexpr auto var = kscale0 * KRepeat / num_scale_k_block + k0; + auto loadA = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + constexpr index_t cscale_offset = CScaleThreadDesc{}.CalculateOffset( + make_tuple(kscale0, m0, n0 * num_scale_n_block / NRepeat)); + + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[Number{}]); + }); + }); + + static_ford>{}([&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto k0 = Number{}]>{}; + constexpr index_t c_offset = + CScaleThreadDesc{}.CalculateOffset(make_tuple(k0, m0, n0)); + constexpr index_t a_offset = AScaleThreadDesc{}.CalculateOffset(make_tuple(m0, k0)); + constexpr index_t b_offset = BScaleThreadDesc{}.CalculateOffset(make_tuple(n0, k0)); + + c_scale_thread_buf(Number{}) = + a_scale_thread_buf[Number{}] * b_scale_thread_buf[Number{}]; + }); + + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + + static_ford>{}([&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto kscale0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat / num_scale_k_block, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + constexpr auto var = kscale0 * KRepeat / num_scale_k_block + k0; + auto loadA = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + constexpr index_t cscale_offset = CScaleThreadDesc{}.CalculateOffset( + make_tuple(kscale0, m0, n0 * num_scale_n_block / NRepeat)); + + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[Number{}]); + }); + }); + __builtin_amdgcn_sched_barrier(0); + } + else if constexpr(TailNum == TailNumber::Odd) + { + static_ford>{}([&](auto mnk) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + constexpr auto kscale0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat / num_scale_k_block, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + constexpr auto var = kscale0 * KRepeat / num_scale_k_block + k0; + auto loadA = + thread_buf_to_vec_loader, + Number, + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + Number, + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + constexpr index_t cscale_offset = CScaleThreadDesc{}.CalculateOffset( + make_tuple(kscale0, m0, n0 * num_scale_n_block / NRepeat)); + + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[Number{}]); + }); + }); + __builtin_amdgcn_sched_barrier(0); + } + } + + protected: + using Base::a_thread_desc_; + using Base::b_thread_desc_; + using Base::c_thread_desc_; + using AThreadCopy = ThreadwiseTensorSliceTransfer_v4, + Sequence<0, 1, 2, 3>, + 3, + A_K1, + A_K1>; + + using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, + Sequence<0, 1, 2, 3>, + 3, + B_K1, + B_K1>; + + AThreadCopy a_thread_copy_{CalculateAThreadOriginDataIndex()}; + BThreadCopy b_thread_copy_{CalculateBThreadOriginDataIndex()}; +}; + +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp new file mode 100644 index 0000000..2b1d888 --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp @@ -0,0 +1,728 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/utility/thread_buf_to_vec_loader.hpp" +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp" + +namespace ck { + +// Maximum Global Memory throughput pipeline with >=32KB data in fly +// GlobalPrefetchStages: >=2 +// LocalPreFillStages: 1 +// LocalPreFetchStages: 0 +// LocalSharedMemoryBuffer: 1 + +template +struct BlockwiseGemmXdlops_pipeline_v2_ab_scale +{ +}; + +template +struct BlockwiseGemmXdlops_pipeline_v2_ab_scale + : BlockwiseGemmXdlops_pipeline_base + +{ + using Base = BlockwiseGemmXdlops_pipeline_base; + using Base::I0; + using Base::KRepeat; + using Base::xdlops_gemm; + + using Base::CalculateCThreadOriginDataIndex; + using Base::CalculateCThreadOriginDataIndex8D; + using Base::GetCBlockDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::GetCThreadBuffer; + using Base::GetCThreadDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::MakeCGridDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + + using Base::a_block_desc_m0_m1_m2_k; + using Base::b_block_desc_n0_n1_n2_k; + + using Base::AMmaKStride; + using Base::BMmaKStride; + using Base::WaveSize; + + using ComputeDataTypeBuf = typename Base::ComputeDataTypeBuf; + + static constexpr index_t WgpPerCU = + (4 * WaveSize / BlockSize) >= 1 ? 4 * WaveSize / BlockSize : 1; + static constexpr index_t FullMemBandPrefetchStages = math::integer_divide_ceil( + 32768 / WgpPerCU, + (MPerBlock * sizeof(ADataType) + NPerBlock * sizeof(BDataType)) * KPerBlock); + static constexpr index_t PrefetchStages = + FullMemBandPrefetchStages >= 2 + ? FullMemBandPrefetchStages <= 8 ? FullMemBandPrefetchStages : 8 + : 2; + + static constexpr index_t PrefillStages = 1; + static constexpr index_t GlobalBufferNum = PrefetchStages; + + __host__ static constexpr bool BlockHasHotloop(index_t num_loop) + { + return num_loop > PrefetchStages; + } + + __host__ static constexpr TailNumber BlockLoopTailNum(index_t num_loop) + { + if(num_loop % PrefetchStages == 1) + { + return TailNumber::One; + } + else if(num_loop % PrefetchStages == 2) + { + return TailNumber::Two; + } + else if(num_loop % PrefetchStages == 3) + { + return TailNumber::Three; + } + else if(num_loop % PrefetchStages == 4) + { + return TailNumber::Four; + } + else if(num_loop % PrefetchStages == 5) + { + return TailNumber::Five; + } + else if(num_loop % PrefetchStages == 6) + { + return TailNumber::Six; + } + else if(num_loop % PrefetchStages == 7) + { + return TailNumber::Seven; + } + else + { + return TailNumber::Full; + } + } + + template + __device__ void Run( + // ABlockCopy + const AGridDesc& a_grid_desc, + const ABlockDesc& a_block_desc, + ABlockTransfer& a_blockwise_copy, + const AGridBuffer& a_grid_buf, + ABlockBuffer& a_block_buf, + const ABlockTransferStep& a_block_copy_step, + // BBlockCopy + const BGridDesc& b_grid_desc, + const BBlockDesc& b_block_desc, + BBlockTransfer& b_blockwise_copy, + const BGridBuffer& b_grid_buf, + BBlockBuffer& b_block_buf, + const BBlockTransferStep& b_block_copy_step, + // CThread + CThreadBuffer& c_thread_buf, + // AScaleThreadCopy + const AScaleGridDesc& a_scale_grid_desc, + const AScaleThreadDesc& a_scale_thread_desc, + AScaleThreadTransfer& a_scale_thread_copy, + const AScaleGridBuffer& a_scale_grid_buf, + const AScaleThreadTransferStep& a_scale_thread_copy_step, + // BScaleThreadCopy + const BScaleGridDesc& b_scale_grid_desc, + const BScaleThreadDesc& b_scale_thread_desc, + BScaleThreadTransfer& b_scale_thread_copy, + const BScaleGridBuffer& b_scale_grid_buf, + const BScaleThreadTransferStep& b_scale_thread_copy_step, + // num_loop + index_t num_loop, + index_t num_loop_per_scale) const + { + // assume kperblock = scaleblockk + ignore = num_loop_per_scale; + auto a_thread_buf = make_static_buffer( + a_thread_desc_.GetElementSpaceSize()); + auto b_thread_buf = make_static_buffer( + b_thread_desc_.GetElementSpaceSize()); + auto a_scale_thread_buf = make_static_buffer( + a_scale_thread_desc.GetElementSpaceSize()); + auto b_scale_thread_buf = make_static_buffer( + b_scale_thread_desc.GetElementSpaceSize()); + + // Global prefetch 1 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, I0); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, I0); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<0>{})); + }); + + if(num_loop_per_scale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + // Local prefill 1 + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf, I0); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, I0); + + // Initialize C + c_thread_buf.Clear(); + + // Global prefetch [2, PrefetchStages] + static_for<1, PrefetchStages, 1>{}([&](auto iprefetch) { + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, iprefetch); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, iprefetch); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + }); + + auto c_thread_buf_per_scale = remove_cvref_t(); + + // main body + if constexpr(HasMainLoop) + { + index_t i = 0; + do + { + static_for<0, PrefetchStages, 1>{}([&](auto iprefetch) { + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + c_thread_buf_per_scale.Clear(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{ + a_thread_vec, a_thread_buf}; + auto loadB = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{ + b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale[Number{}] * + type_convert(a_scale_thread_buf[m0]) * + type_convert(b_scale_thread_buf[I0]); + }); + }); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<0>{})); + }); + + if(num_loop_per_scale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, + b_scale_thread_copy_step); + + block_sync_lds(); + a_blockwise_copy.RunWrite( + a_block_desc, a_block_buf, Number<(iprefetch + 1) % PrefetchStages>{}); + b_blockwise_copy.RunWrite( + b_block_desc, b_block_buf, Number<(iprefetch + 1) % PrefetchStages>{}); + + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf, iprefetch); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, iprefetch); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + }); + + i += PrefetchStages; + } while(i < (num_loop - PrefetchStages)); + } + + // tail + auto LoopTailFunc = [&](auto tail_num) { + static_for<1, tail_num, 1>{}([&](auto iprefetch) { + block_sync_lds(); + static_ford>{}([&](auto km) { + constexpr auto k = Number{}]>{}; + constexpr auto m0 = Number{}]>{}; + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + c_thread_buf_per_scale.Clear(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, + a_thread_buf}; + auto loadB = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, + b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale[Number{}] * + type_convert(a_scale_thread_buf[m0]) * + type_convert(b_scale_thread_buf[I0]); + }); + }); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<0>{})); + }); + + if(num_loop_per_scale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + block_sync_lds(); + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf, iprefetch); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, iprefetch); + }); + + block_sync_lds(); + static_ford>{}([&](auto km) { + constexpr auto k = Number{}]>{}; + constexpr auto m0 = Number{}]>{}; + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + c_thread_buf_per_scale.Clear(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = + thread_buf_to_vec_loader, + decltype(k0), + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + decltype(k0), + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>(a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale[Number{}] * + type_convert(a_scale_thread_buf[m0]) * + type_convert(b_scale_thread_buf[I0]); + }); + }); + }; + + if constexpr(TailNum == TailNumber::One) + { + block_sync_lds(); + static_ford>{}([&](auto km) { + constexpr auto k = Number{}]>{}; + constexpr auto m0 = Number{}]>{}; + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + c_thread_buf_per_scale.Clear(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = + thread_buf_to_vec_loader, + decltype(k0), + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + decltype(k0), + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>(a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(I0)); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale[Number{}] * + type_convert(a_scale_thread_buf[m0]) * + type_convert(b_scale_thread_buf[I0]); + }); + }); + } + else if constexpr(TailNum == TailNumber::Two) + { + LoopTailFunc(Number<2>{}); + } + else if constexpr(TailNum == TailNumber::Three) + { + LoopTailFunc(Number<3>{}); + } + else if constexpr(TailNum == TailNumber::Four) + { + LoopTailFunc(Number<4>{}); + } + else if constexpr(TailNum == TailNumber::Five) + { + LoopTailFunc(Number<5>{}); + } + else if constexpr(TailNum == TailNumber::Six) + { + LoopTailFunc(Number<6>{}); + } + else if constexpr(TailNum == TailNumber::Seven) + { + LoopTailFunc(Number<7>{}); + } + else if constexpr(TailNum == TailNumber::Full) + { + LoopTailFunc(Number{}); + } + } + + protected: + using Base::a_thread_copy_; + using Base::a_thread_desc_; + using Base::b_thread_copy_; + using Base::b_thread_desc_; + using Base::c_thread_desc_; +}; + +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp new file mode 100644 index 0000000..173d6d2 --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp @@ -0,0 +1,652 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/utility/thread_buf_to_vec_loader.hpp" +#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp" + +namespace ck { + +// Compute optimized pipeline +// GlobalPrefetchStages: 2 +// LocalPreFillStages: 1 +// LocalPreFetchStages: 1 +// LocalSharedMemoryBuffer: 1 + +template +struct BlockwiseGemmXdlops_pipeline_v3_ab_scale +{ +}; + +template +struct BlockwiseGemmXdlops_pipeline_v3_ab_scale + : BlockwiseGemmXdlops_pipeline_base + +{ + using Base = BlockwiseGemmXdlops_pipeline_base; + using Base::I0; + using Base::KRepeat; + using Base::xdlops_gemm; + using typename Base::HotLoopInstList; + + using Base::CalculateCThreadOriginDataIndex; + using Base::CalculateCThreadOriginDataIndex8D; + using Base::GetCBlockDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::GetCThreadBuffer; + using Base::GetCThreadDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_N3_N4; + using Base::MakeCGridDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2; + using Base::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2; + + using Base::a_block_desc_m0_m1_m2_k; + using Base::b_block_desc_n0_n1_n2_k; + + using Base::AMmaKStride; + using Base::BMmaKStride; + + using ComputeDataTypeBuf = typename Base::ComputeDataTypeBuf; + + static constexpr index_t PrefetchStages = 2; + static constexpr index_t PrefillStages = 1; + static constexpr index_t GlobalBufferNum = 1; + + __host__ static constexpr bool BlockHasHotloop(index_t num_loop) + { + return num_loop > PrefetchStages; + } + + __host__ static constexpr TailNumber BlockLoopTailNum(index_t num_loop) + { + ignore = num_loop; + return TailNumber::Full; + } + + __device__ static constexpr auto HotLoopScheduler() + { + // A/B split schedule + // compiler is likely to use ds_read2 when instruction width smaller than 16bytes + constexpr auto num_ds_read_inst_a = + HotLoopInstList::A_LDS_Read_Width * sizeof(ADataType) == 16 + ? HotLoopInstList::A_LDS_Read_Inst_Num + : HotLoopInstList::A_LDS_Read_Inst_Num / 2; + constexpr auto num_ds_read_inst_b = + HotLoopInstList::B_LDS_Read_Width * sizeof(BDataType) == 16 + ? HotLoopInstList::B_LDS_Read_Inst_Num + : HotLoopInstList::B_LDS_Read_Inst_Num / 2; + + constexpr auto num_ds_write_inst_a = HotLoopInstList::A_LDS_Write_Inst_Num; + constexpr auto num_ds_write_inst_b = HotLoopInstList::B_LDS_Write_Inst_Num; + + constexpr auto num_buffer_load_inst_a = HotLoopInstList::A_Buffer_Load_Inst_Num; + constexpr auto num_buffer_load_inst_b = HotLoopInstList::B_Buffer_Load_Inst_Num; + + constexpr auto num_mfma_inst = HotLoopInstList::C_MFMA_Inst_Num; + + constexpr auto mfma_cycle = HotLoopInstList::C_MFMA_Inst_Cycle; + constexpr auto ds_read_a_issue_cycle = + HotLoopInstList::A_LDS_Read_Width * sizeof(ADataType) == 16 ? 8 : 4; + constexpr auto ds_read_b_issue_cycle = + HotLoopInstList::B_LDS_Read_Width * sizeof(BDataType) == 16 ? 8 : 4; + constexpr auto ds_read_a_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_a_issue_cycle - 1) / (2 * ds_read_a_issue_cycle); + constexpr auto ds_read_b_mfma_rate = + (mfma_cycle - 4 + 2 * ds_read_b_issue_cycle - 1) / (2 * ds_read_b_issue_cycle); + + constexpr auto num_dsread_a_mfma = + (num_ds_read_inst_a + ds_read_a_mfma_rate - 1) / ds_read_a_mfma_rate; + constexpr auto num_dsread_b_mfma = + (num_ds_read_inst_b + ds_read_b_mfma_rate - 1) / ds_read_b_mfma_rate; + + // stage 1 + // Separate this part? + // constexpr auto num_mfma_per_ds_read = sizeof(ComputeDataTypeBuf) / sizeof(ADataType) > + // sizeof(ComputeDataTypeBuf) / sizeof(BDataType) + // ? sizeof(ComputeDataTypeBuf) / sizeof(ADataType) + // : sizeof(ComputeDataTypeBuf) / sizeof(BDataType); + constexpr auto num_mfma_stage1 = num_mfma_inst - (num_dsread_a_mfma + num_dsread_b_mfma); + constexpr auto num_mfma_per_issue = + num_mfma_stage1 / (num_buffer_load_inst_a + num_buffer_load_inst_b); + constexpr auto num_dswrite_per_issue_a = num_ds_write_inst_a / num_buffer_load_inst_a; + constexpr auto num_dswrite_per_issue_b = num_ds_write_inst_b / num_buffer_load_inst_b; + + static_for<0, num_buffer_load_inst_a, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_a, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_a, 0); // MFMA + }); + static_for<0, num_buffer_load_inst_b, 1>{}([&](auto i) { + ignore = i; + static_for<0, num_dswrite_per_issue_b, 1>{}([&](auto idswrite) { + ignore = idswrite; + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + __builtin_amdgcn_sched_group_barrier( + 0x008, num_mfma_per_issue - num_dswrite_per_issue_b, 0); // MFMA + }); + + // stage 2 + static_for<0, num_dsread_a_mfma, 1>{}([&](auto i) { + if constexpr((num_ds_read_inst_a - (i + 1) * ds_read_a_mfma_rate) >= + ds_read_a_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_a_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier(0x100, + num_ds_read_inst_a - (num_dsread_a_mfma - 1) * + ds_read_a_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + + static_for<0, num_dsread_b_mfma, 1>{}([&](auto i) { + if constexpr((num_ds_read_inst_b - (i + 1) * ds_read_b_mfma_rate) >= + ds_read_b_mfma_rate) + { + __builtin_amdgcn_sched_group_barrier(0x100, ds_read_b_mfma_rate, 0); // DS read + } + else + { + __builtin_amdgcn_sched_group_barrier(0x100, + num_ds_read_inst_b - (num_dsread_b_mfma - 1) * + ds_read_b_mfma_rate, + 0); // DS read + } + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + }); + } + + template + __device__ void Run( + // ABlockCopy + const AGridDesc& a_grid_desc, + const ABlockDesc& a_block_desc, + ABlockTransfer& a_blockwise_copy, + const AGridBuffer& a_grid_buf, + ABlockBuffer& a_block_buf, + const ABlockTransferStep& a_block_copy_step, + // BBlockCopy + const BGridDesc& b_grid_desc, + const BBlockDesc& b_block_desc, + BBlockTransfer& b_blockwise_copy, + const BGridBuffer& b_grid_buf, + BBlockBuffer& b_block_buf, + const BBlockTransferStep& b_block_copy_step, + // CThread + const CScaleThreadDesc& c_scale_thread_desc, + CThreadBuffer& c_thread_buf, + // AScaleThreadCopy + const AScaleGridDesc& a_scale_grid_desc, + const AScaleThreadDesc& a_scale_thread_desc, + AScaleThreadTransfer& a_scale_thread_copy, + const AScaleGridBuffer& a_scale_grid_buf, + const AScaleThreadTransferStep& a_scale_thread_copy_step, + // BScaleThreadCopy + const BScaleGridDesc& b_scale_grid_desc, + const BScaleThreadDesc& b_scale_thread_desc, + BScaleThreadTransfer& b_scale_thread_copy, + const BScaleGridBuffer& b_scale_grid_buf, + const BScaleThreadTransferStep& b_scale_thread_copy_step, + // num_loop + index_t num_loop) const + { + __builtin_amdgcn_sched_barrier(0); + static_assert(CScaleThreadDesc{}.GetLength(Number<0>{}) == 1, + "Pipeline v3 only support scaleblocksliceK=1"); + static_assert(CScaleThreadDesc{}.GetLength(Number<2>{}) == 1, + "Pipeline v3 only support scaleblocksliceN=1"); + // assume kperblock = scaleblockk + auto a_thread_buf = make_static_buffer( + a_thread_desc_.GetElementSpaceSize()); + auto b_thread_buf = make_static_buffer( + b_thread_desc_.GetElementSpaceSize()); + auto a_scale_thread_buf = make_static_buffer( + a_scale_thread_desc.GetElementSpaceSize()); + auto b_scale_thread_buf = make_static_buffer( + b_scale_thread_desc.GetElementSpaceSize()); + auto c_scale_thread_buf = make_static_buffer( + c_scale_thread_desc.GetElementSpaceSize()); + + // Global prefetch 1 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + c_scale_thread_buf(m0) = a_scale_thread_buf[m0] * b_scale_thread_buf[I0]; + }); + + // Local prefill 1 + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); + + // Global prefetch 2 + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc, + a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + // Initialize C + c_thread_buf.Clear(); + + StaticBufferTupleOfVector + c_thread_buf_per_scale; + + // Local prefetch 1 + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k0) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k0, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k0, I0), + b_thread_buf); + }); + }); + + __builtin_amdgcn_sched_barrier(0); + + // main body + if constexpr(HasMainLoop) + { + index_t i = 0; + do + { + block_sync_lds(); + a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); + + a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf); + b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); + + a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step); + + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, + a_thread_buf}; + auto loadB = thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, + b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[m0]); + }); + }); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + c_scale_thread_buf(m0) = a_scale_thread_buf[m0] * b_scale_thread_buf[I0]; + }); + + block_sync_lds(); + static_for<0, KRepeat, 1>{}([&](auto k) { + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, + make_tuple(m0, I0, I0, Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(m0, I0, k, I0), + a_thread_buf); + }); + static_for<0, NRepeat, 1>{}([&](auto n0) { + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(n0, I0, k, I0), + b_thread_buf); + }); + }); + + HotLoopScheduler(); + __builtin_amdgcn_sched_barrier(0); + + static_for<0, MRepeat, 1>{}([&](auto m0) { + a_scale_thread_copy.Run(a_scale_grid_desc, + a_scale_grid_buf, + a_scale_thread_desc, + make_tuple(m0, I0), + a_scale_thread_buf); + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<0>{})); + }); + + if constexpr(NumKBlockPerScale == 1) + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{})); + } + else + { + a_scale_thread_copy.MoveSrcSliceWindow( + a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{})); + } + + b_scale_thread_copy.Run(b_scale_grid_desc, + b_scale_grid_buf, + b_scale_thread_desc, + make_tuple(I0, I0), + b_scale_thread_buf); + + b_scale_thread_copy.MoveSrcSliceWindow(b_scale_grid_desc, b_scale_thread_copy_step); + + i += 1; + } while(i < (num_loop - 1)); + } + + // tail + if constexpr(TailNum == TailNumber::Full) + { + static_ford>{}([&](auto mn) { + constexpr auto m0 = Number{}]>{}; + constexpr auto n0 = Number{}]>{}; + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()(Number{}) = 0; + }); + static_for<0, KRepeat, 1>{}([&](auto k0) { + vector_type a_thread_vec; + vector_type b_thread_vec; + + auto loadA = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{a_thread_vec, a_thread_buf}; + auto loadB = + thread_buf_to_vec_loader, + Number<0>, + Number, + index_expression::Ik>{b_thread_vec, b_thread_buf}; + + static_for<0, KPack, 1>{}(MakeFunctorInvoker(loadA, loadB)); + + using mfma_input_type = + typename vector_type::type; + + xdlops_gemm.template Run<>( + a_thread_vec.template AsType(), + b_thread_vec.template AsType(), + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{})); + }); + static_for<0, xdlops_gemm.GetRegSizePerXdlops(), 1>{}([&](auto t) { + constexpr index_t c_offset = + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t)); + c_thread_buf(Number{}) += + c_thread_buf_per_scale.GetVectorTypeReference(Number<0>{}) + .template AsType()[Number{}] * + type_convert(c_scale_thread_buf[m0]); + }); + }); + __builtin_amdgcn_sched_barrier(0); + } + } + + protected: + using Base::a_thread_copy_; + using Base::a_thread_desc_; + using Base::b_thread_copy_; + using Base::b_thread_desc_; + using Base::c_thread_desc_; +}; + +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp b/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp new file mode 100644 index 0000000..8eef34c --- /dev/null +++ b/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp @@ -0,0 +1,460 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include + +#include "ck/tensor_operation/gpu/device/device_base.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { + +// GEMM: +// input : A[M, K], B[K, N], +// input : D0[M, N], D1[M, N], ... +// output : E[M, N] +// C = a_op(A) * b_op(B) +// E = cde_op(C, D0, D1, ...) +// Assume: +// D0, D1, ... and E have the same layout +template +struct DeviceGemmMultipleD_ABScale : public BaseOperator +{ + static constexpr index_t NumDTensor = DsDataType::Size(); + + virtual std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) = 0; + + virtual std::unique_ptr MakeInvokerPointer() = 0; + + virtual void SetKBatch(BaseArgument* arg, int KBatch) const = 0; +}; + +template +struct DeviceGemmMultipleD_BlockScale_BPreshuffle : public BaseOperator +{ + static constexpr index_t NumDTensor = DsDataType::Size(); + + virtual std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) = 0; + + virtual std::unique_ptr MakeInvokerPointer() = 0; + + virtual int GetPreShuffleParameters() = 0; +}; + +template +struct DeviceGemmMultipleD_BlockScale_BPreshuffleSplitK : public BaseOperator +{ + static constexpr index_t NumDTensor = DsDataType::Size(); + + virtual std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op, + index_t KBatch) = 0; + + virtual std::unique_ptr MakeInvokerPointer() = 0; + + virtual int GetPreShuffleParameters() = 0; +}; + +/// @brief Wrapper for backward compatibility that allows to use instances of +/// DeviceGemmMultipleD_BlockScale_BPreshuffleSplitK in contexts where +// DeviceGemmMultipleD_BlockScale_BPreshuffle is expected. +/// +/// @note The main area where it can be used is DeviceOperationInstanceFactory::GetInstances(). +/// The only difference between API of DeviceGemmMultipleD_BlockScale_BPreshuffle and +// DeviceGemmMultipleD_BlockScale_BPreshuffleSplitK is +/// that DeviceGemmMultipleD_BlockScale_BPreshuffleSplitK::MakeArgumentPointer requires +// an additional parameter KBatch which is explicitly passed as 1 by this wrapper. +template +struct DeviceGemmMultipleD_BlockScale_BPreshuffleWrapper + : public DeviceGemmMultipleD_BlockScale_BPreshuffle +{ + using DeviceOp = DeviceGemmMultipleD_BlockScale_BPreshuffleSplitK; + + static constexpr index_t NumDTensor = DsDataType::Size(); + +#ifndef __HIPCC_RTC__ + + explicit DeviceGemmMultipleD_BlockScale_BPreshuffleWrapper(std::unique_ptr p_op) + : p_op_(std::move(p_op)) + { + } + + bool IsSupportedArgument(const BaseArgument* p_arg) override + { + return p_op_->IsSupportedArgument(p_arg); + } + std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) override + { + return p_op_->MakeArgumentPointer(p_a, + p_b, + p_ds, + p_e, + M, + N, + K, + StrideA, + StrideB, + StrideDs, + StrideE, + p_a_scale, + p_b_scale, + a_element_op, + b_element_op, + cde_element_op, + 1); // KBatch + } + + std::unique_ptr MakeInvokerPointer() override + { + return p_op_->MakeInvokerPointer(); + } + + int GetPreShuffleParameters() override { return p_op_->GetPreShuffleParameters(); } + + std::string GetTypeString() const override { return p_op_->GetTypeString(); } + + private: + std::unique_ptr p_op_; + +#endif // __HIPCC_RTC__ +}; + +// GEMM: +// input : A[M, K], B[K, N], +// input : D0[M, N], D1[M, N], ... +// output : E[M, N] +// C = a_op(A) * b_op(B) +// E = cde_op(C, D0, D1, ...) +// Assume: +// D0, D1, ... and E have the same layout +template +struct DeviceGemmMultipleD_ABScaleSplitK : public BaseOperator +{ + static constexpr index_t NumDTensor = DsDataType::Size(); + + virtual std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op, + index_t KBatch) = 0; + + virtual std::unique_ptr MakeInvokerPointer() = 0; + + virtual void SetKBatch(BaseArgument* arg, int KBatch) const = 0; +}; + +/// @brief Wrapper for backward compatibility that allows to use instances of +/// DeviceGemmMultipleD_ABScaleSplitK in contexts where DeviceGemmMultipleD_ABScale is +/// expected. +/// +/// @note The main area where it can be used is DeviceOperationInstanceFactory::GetInstances(). +/// The only difference between API of DeviceGemmMultipleD_ABScale and +/// DeviceGemmMultipleD_ABScaleSplitK is that +/// DeviceGemmMultipleD_ABScaleSplitK::MakeArgumentPointer requires a additional parameter +/// KBatch which is explicitly passed as 1 by this wrapper. +template +struct DeviceGemmMultipleD_ABScaleSplitKWrapper + : public DeviceGemmMultipleD_ABScale +{ + + using DeviceOp = DeviceGemmMultipleD_ABScaleSplitK; + + static constexpr index_t NumDTensor = DsDataType::Size(); + +#ifndef __HIPCC_RTC__ + + explicit DeviceGemmMultipleD_ABScaleSplitKWrapper(std::unique_ptr p_op) + : p_op_(std::move(p_op)) + { + } + + bool IsSupportedArgument(const BaseArgument* p_arg) override + { + return p_op_->IsSupportedArgument(p_arg); + } + std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_e, + const ck::index_t M, + const ck::index_t N, + const ck::index_t K, + const ck::index_t StrideA, + const ck::index_t StrideB, + const std::array StrideDs, + const ck::index_t StrideE, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CDEElementwiseOperation cde_element_op) override + { + return p_op_->MakeArgumentPointer(p_a, + p_b, + p_ds, + p_e, + M, + N, + K, + StrideA, + StrideB, + StrideDs, + StrideE, + p_a_scale, + p_b_scale, + a_element_op, + b_element_op, + cde_element_op, + 1); // KBatch + } + + void SetKBatch(BaseArgument* arg, int KBatch) const override { p_op_->SetKBatch(arg, KBatch); } + + std::unique_ptr MakeInvokerPointer() override + { + return p_op_->MakeInvokerPointer(); + } + + std::string GetTypeString() const override { return p_op_->GetTypeString(); } + + private: + std::unique_ptr p_op_; + +#endif // __HIPCC_RTC__ +}; + +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp new file mode 100644 index 0000000..b2fea54 --- /dev/null +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp @@ -0,0 +1,624 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include +#include + +#include "ck/utility/common_header.hpp" +#include "ck/tensor_description/tensor_descriptor.hpp" +#include "ck/tensor_description/tensor_descriptor_helper.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp" +#include "ck/host_utility/device_prop.hpp" +#include "ck/host_utility/kernel_launch.hpp" +#include "ck/host_utility/flush_cache.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { + +template +struct DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3 + : public DeviceGemmMultipleD_ABScale +{ + static constexpr auto WarpTileConfig64 = GetWarpTileConfig(); + static constexpr auto WarpTileConfig32 = GetWarpTileConfig(); + static constexpr auto NXdlPerWave64 = WarpTileConfig64.At(3); + static constexpr auto NXdlPerWave32 = WarpTileConfig32.At(3); + static constexpr index_t NumDTensor = DsDataType::Size(); + + // GridwiseGemm + template + using GridwiseGemmBase = GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3< + ALayout, + BLayout, + DsLayout, + CLayout, + ADataType, + BDataType, + GemmAccDataType, + CShuffleDataType, + DsDataType, + CDataType, + AElementwiseOperation, + BElementwiseOperation, + CElementwiseOperation, + GemmSpec, + BlockSize, + ScaleBlockM, + ScaleBlockN, + ScaleBlockK, + MPerBlock, + NPerBlock, + KPerBlock, + AK1, + BK1, + WarpTileConfig::At(0), + WarpTileConfig::At(1), + WarpTileConfig::At(2), + WarpTileConfig::At(3), + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + false, + ABlockLdsExtraM, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + false, + BBlockLdsExtraN, + WarpTileConfig::At(4), + WarpTileConfig::At(5), + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + CDEShuffleBlockTransferScalarPerVectors, + BlkGemmPipeSched, + BlkGemmPipelineVer, + ComputeTypeA, + ComputeTypeB, + LDSTypeA, + LDSTypeB>; + using GridwiseGemm64 = GridwiseGemmBase; + using GridwiseGemm32 = GridwiseGemmBase; + + using Argument = typename GridwiseGemm64::Argument; + + // Invoker + struct Invoker : public BaseInvoker + { + template + float RunImp(const typename GridwiseGemm::Argument& arg, + const StreamConfig& stream_config = StreamConfig{}) + { + if(stream_config.log_level_ > 0) + { + arg.Print(); + } + + if(!GridwiseGemm::CheckValidity(arg)) + { + throw std::runtime_error("wrong! GridwiseGemm has invalid setting"); + } + + index_t gdx, gdy, gdz; + std::tie(gdx, gdy, gdz) = GridwiseGemm::CalculateGridSize(arg.M, arg.N, arg.KBatch); + + float ave_time = 0; + + index_t k_grain = arg.KBatch * KPerBlock; + index_t K_split = (arg.K + k_grain - 1) / k_grain * KPerBlock; + + const bool has_main_k_block_loop = GridwiseGemm::CalculateHasMainKBlockLoop(K_split); + + const auto Run = [&](const auto& kernel) { + if(stream_config.flush_cache) + { + auto arg_ = arg; + + const auto a_grid_desc_ak0_m_ak1 = GridwiseGemm::MakeAGridDescriptor_AK0_M_AK1( + arg_.M, arg_.MPadded, arg_.K, arg_.KPadded, arg_.StrideA, arg_.AK0); + const auto b_grid_desc_bk0_n_bk1 = GridwiseGemm::MakeBGridDescriptor_BK0_N_BK1( + arg_.K, arg_.KPadded, arg_.N, arg_.NPadded, arg_.StrideB, arg_.BK0); + + auto size_a_buffer = + a_grid_desc_ak0_m_ak1.GetElementSpaceSize() * sizeof(ADataType); + auto size_b_buffer = + b_grid_desc_bk0_n_bk1.GetElementSpaceSize() * sizeof(BDataType); + + ck::utility::RotatingMemWrapper rotating_mem( + arg_, stream_config.rotating_count, size_a_buffer, size_b_buffer); + rotating_mem.Print(); + + auto run_flush_cache = [&]() { + // flush icache + ck::utility::flush_icache(); + // rotating mem + rotating_mem.Next(); + // clear c mem + if(arg_.KBatch > 1 && !arg_.skip_zero_init) + hipGetErrorString(hipMemsetAsync(arg_.p_c_grid, + 0, + arg_.M * arg_.N * sizeof(CDataType), + stream_config.stream_id_)); + }; + + ave_time = ck::utility::launch_and_time_kernel_with_preprocess( + stream_config, + run_flush_cache, + kernel, + dim3(gdx, gdy, gdz), + dim3(BlockSize), + 0, + arg_); + } + else + { + if(arg.KBatch > 1 && !arg.skip_zero_init) + hipGetErrorString(hipMemsetAsync(arg.p_c_grid, + 0, + arg.M * arg.N * sizeof(CDataType), + stream_config.stream_id_)); + + ave_time = launch_and_time_kernel( + stream_config, kernel, dim3(gdx, gdy, gdz), dim3(BlockSize), 0, arg); + } + }; + + constexpr index_t minimum_occupancy = [&]() { + if constexpr(is_same_v && + is_same_v) + { + // FIXME: many instances have many spills with occupancy > 1, a better solution + // needed to get best performance + return 1; + } + else + { + return (BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave && + MPerBlock * NPerBlock / BlockSize > 64) + ? 1 + : 2; + } + }(); + + if(has_main_k_block_loop) + { + // Tail number always full + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1 || + BlkGemmPipelineVer == BlockGemmPipelineVersion::v3) + { + if(arg.KBatch > 1) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + else + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + } + } + else + { + // Tail number always 1 + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1) + { + if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Full) + { + if(arg.KBatch > 1) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + else + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + } + else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) + { + if(arg.KBatch > 1) + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + else + { + const auto kernel = + kernel_gemm_xdl_cshuffle_v3; + Run(kernel); + } + } + } + } + return ave_time; + } + + float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{}) + { + if(get_warp_size() == 64) + { + if constexpr(NXdlPerWave64 > 0) + { + return RunImp(arg, stream_config); + } + } + else + { + if constexpr(NXdlPerWave32 > 0) + { + return RunImp( + reinterpret_cast(arg), stream_config); + } + } + return 0; + } + + // polymorphic + float Run(const BaseArgument* p_arg, + const StreamConfig& stream_config = StreamConfig{}) override + { + return Run(*dynamic_cast(p_arg), stream_config); + } + }; + + void SetKBatch(BaseArgument* base_arg, int KBatch) const override + { + auto& arg = *dynamic_cast(base_arg); + arg.KBatch = KBatch; + if(get_warp_size() == 64) + { + if constexpr(NXdlPerWave64 > 0) + { + arg.KRead = GridwiseGemm64::CalculateKRead(arg.K, KBatch); + arg.KPadded = GridwiseGemm64::CalculateKPadded(arg.K, KBatch); + arg.AK0 = GridwiseGemm64::CalculateAK0Padded(arg.K, KBatch); + arg.BK0 = GridwiseGemm64::CalculateBK0Padded(arg.K, KBatch); + } + } + else + { + if constexpr(NXdlPerWave32 > 0) + { + arg.KRead = GridwiseGemm32::CalculateKRead(arg.K, KBatch); + arg.KPadded = GridwiseGemm32::CalculateKPadded(arg.K, KBatch); + arg.AK0 = GridwiseGemm32::CalculateAK0Padded(arg.K, KBatch); + arg.BK0 = GridwiseGemm32::CalculateBK0Padded(arg.K, KBatch); + } + } + } + + static constexpr bool IsValidCompilationParameter() + { + // TODO: properly implement this check + return true; + } + + static bool IsSupportedArgument(const Argument& arg) + { + // with splitk the implementation doesn't work + // when KRead % ScaleBlockK != 0, independently of K padding + if(arg.KBatch > 1 && arg.KRead % ScaleBlockK != 0) + { + return false; + } + + if(!ck::is_xdl_wmma_supported()) + { + return false; + } + + // if(ScaleBlockM % MPerBlock != 0 || ScaleBlockN % NPerBlock != 0 || ScaleBlockK != + // KPerBlock) + // { + // return false; + // } + + if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding || + GemmSpec == GemmSpecialization::NKPadding || + GemmSpec == GemmSpecialization::MNKPadding || + GemmSpec == GemmSpecialization::KPadding)) + { + return false; + } + + if(get_warp_size() == 64) + { + if constexpr(NXdlPerWave64 > 0) + { + return GridwiseGemm64::CheckValidity(arg); + } + } + else + { + if constexpr(NXdlPerWave32 > 0) + { + return GridwiseGemm32::CheckValidity( + reinterpret_cast(arg)); + } + } + return false; + } + + // polymorphic + bool IsSupportedArgument(const BaseArgument* p_arg) override + { + return IsSupportedArgument(*dynamic_cast(p_arg)); + } + + static auto MakeArgument(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_c, + const index_t M, + const index_t N, + const index_t K, + const index_t StrideA, + const index_t StrideB, + const std::array StrideDs, + const index_t StrideC, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CElementwiseOperation c_element_op) + { + index_t StrideScaleA = ck::is_same_v + ? math::integer_divide_ceil(K, ScaleBlockK) + : math::integer_divide_ceil(M, ScaleBlockM); + + index_t StrideScaleB = ck::is_same_v + ? math::integer_divide_ceil(K, ScaleBlockK) + : math::integer_divide_ceil(N, ScaleBlockN); + + return Argument{static_cast(p_a), + static_cast(p_b), + p_ds, + static_cast(p_c), + M, + N, + K, + StrideA, + StrideB, + StrideDs, + StrideC, + StrideScaleA, + StrideScaleB, + static_cast(p_a_scale), + static_cast(p_b_scale), + 1, + a_element_op, + b_element_op, + c_element_op}; + } + + static auto MakeInvoker() { return Invoker{}; } + + // polymorphic + std::unique_ptr + MakeArgumentPointer(const void* p_a, + const void* p_b, + std::array p_ds, + void* p_c, + const index_t M, + const index_t N, + const index_t K, + const index_t StrideA, + const index_t StrideB, + const std::array StrideDs, + const index_t StrideC, + const void* p_a_scale, + const void* p_b_scale, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CElementwiseOperation c_element_op) override + { + index_t StrideScaleA = ck::is_same_v + ? math::integer_divide_ceil(K, ScaleBlockK) + : math::integer_divide_ceil(M, ScaleBlockM); + + index_t StrideScaleB = ck::is_same_v + ? math::integer_divide_ceil(K, ScaleBlockK) + : math::integer_divide_ceil(N, ScaleBlockN); + + return std::make_unique(static_cast(p_a), + static_cast(p_b), + p_ds, + static_cast(p_c), + M, + N, + K, + StrideA, + StrideB, + StrideDs, + StrideC, + StrideScaleA, + StrideScaleB, + static_cast(p_a_scale), + static_cast(p_b_scale), + 1, + a_element_op, + b_element_op, + c_element_op); + } + + // polymorphic + std::unique_ptr MakeInvokerPointer() override + { + return std::make_unique(Invoker{}); + } + + // polymorphic + std::string GetTypeString() const override + { + auto str = std::stringstream(); + + std::map BlkGemmPipelineSchedulerToString{ + {BlockGemmPipelineScheduler::Intrawave, "Intrawave"}, + {BlockGemmPipelineScheduler::Interwave, "Interwave"}}; + + std::map BlkGemmPipelineVersionToString{ + {BlockGemmPipelineVersion::v1, "v1"}, + {BlockGemmPipelineVersion::v2, "v2"}, + {BlockGemmPipelineVersion::v3, "v3"}}; + + // clang-format off + str << "DeviceGemmXdlUniversal" + << "<" + << getGemmSpecializationString(GemmSpec) << ", " + << std::string(ALayout::name)[0] + << std::string(BLayout::name)[0] + << std::string(CLayout::name)[0] + << ">" + << " BlkSize: " + << BlockSize << ", " + << "BlkTile: " + << MPerBlock<<"x"<= 23 +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" +#endif +namespace ck { + +// Currently we do not have a elegant way to put single lds buffer & double lds buffer pipe in same +// kernel function Blockers: +// 1. Two separted declaration of __shared__ pointer is the key to make sure data access operate on +// two lds chunks. +// 2. Occupied __shared__ won't release until whole shader end, a.k.a AB and C may not use same lds +// buffer when we declare __shared__ inside blkgemmpipe +template +__global__ void +#if CK_USE_LAUNCH_BOUNDS +__launch_bounds__(GridwiseGemm::MaxBlockSize, MinimumOccupancy) +#endif + kernel_gemm_xdl_cshuffle_v3(typename GridwiseGemm::Argument karg) +{ +#if defined(__gfx9__) || defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || \ + defined(__gfx938__) || defined(__gfx11__) || defined(__gfx12__) + if constexpr(GridwiseGemm::template IsValidCompilationParameter()) + { + __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte(get_device_arch())]; + + auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg); + + GridwiseGemm::template Run( + karg.p_a_grid + splitk_batch_offset.a_k_split_offset, + karg.p_b_grid + splitk_batch_offset.b_k_split_offset, + karg.p_ds_grid, + karg.p_c_grid, + karg.p_a_scale_grid + splitk_batch_offset.scale_a_k_split_offset, + karg.p_b_scale_grid + splitk_batch_offset.scale_b_k_split_offset, + p_shared, + karg, + karg.a_element_op, + karg.b_element_op, + karg.c_element_op); + } +#else + ignore = karg; +#endif // end of if (defined(__gfx908__) || defined(__gfx90a__)) +} + +template +struct GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3 + : public GridwiseGemm_xdl_cshuffle_base< + ALayout, + BLayout, + CLayout, + LDSTypeA, + LDSTypeB, + AccDataType, + CShuffleDataType, + DsDataType, + CDataType, + AElementwiseOperation, + BElementwiseOperation, + BlockSize, + MPerBlock, + NPerBlock, + KPerBlock, + AK1Value, + BK1Value, + MPerXDL, + NPerXDL, + MXdlPerWave, + NXdlPerWave, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + AThreadTransferSrcResetCoordinateAfterRun, + ABlockLdsExtraM, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + BThreadTransferSrcResetCoordinateAfterRun, + BBlockLdsExtraN, + CShuffleMXdlPerWavePerShuffle, + CShuffleNXdlPerWavePerShuffle, + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + CDEShuffleBlockTransferScalarPerVectors, + ComputeTypeA, + ComputeTypeB, + false> // ForceNaiveLayout +{ + using Base = GridwiseGemm_xdl_cshuffle_base< + ALayout, + BLayout, + CLayout, + LDSTypeA, + LDSTypeB, + AccDataType, + CShuffleDataType, + DsDataType, + CDataType, + AElementwiseOperation, + BElementwiseOperation, + BlockSize, + MPerBlock, + NPerBlock, + KPerBlock, + AK1Value, + BK1Value, + MPerXDL, + NPerXDL, + MXdlPerWave, + NXdlPerWave, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + AThreadTransferSrcResetCoordinateAfterRun, + ABlockLdsExtraM, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + BThreadTransferSrcResetCoordinateAfterRun, + BBlockLdsExtraN, + CShuffleMXdlPerWavePerShuffle, + CShuffleNXdlPerWavePerShuffle, + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + CDEShuffleBlockTransferScalarPerVectors, + ComputeTypeA, + ComputeTypeB, + false>; // ForceNaiveLayout + + using Base::AK0Number; + using Base::AK1Number; + using Base::BK0Number; + using Base::BK1Number; + using Base::I0; + using Base::I1; + using Base::I2; + using ThisThreadBlock = typename Base::ThisThreadBlock; + using Base::CShuffleBlockTransferScalarPerVector_NPerBlock; + using Base::GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1; + using Base::GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1; + using Base::NumDTensor; + + using AScaleType = float; + using BScaleType = float; + + static constexpr auto MakeDsGridPointer() + { + return generate_tuple( + [&](auto i) { + using DDataType = remove_cvref_t>; + + return static_cast(nullptr); + }, + Number{}); + } + + using DsGridPointer = decltype(MakeDsGridPointer()); + + static constexpr auto lcm_AK1_BK1 = math::lcm(AK1Number, BK1Number); + static constexpr index_t KPack = + math::max(lcm_AK1_BK1, + MfmaSelector::selected_mfma.k_per_blk); + + __host__ static auto CalculateGridSize(index_t M, index_t N, index_t KBatch) + { + return std::make_tuple(Block2CTileMap::CalculateGridSize(M, N), 1, KBatch); + } + + __host__ static auto CalculateMPadded(index_t M) + { + return math::integer_least_multiple(M, MPerBlock); + } + + __host__ static auto CalculateNPadded(index_t N) + { + return math::integer_least_multiple(N, NPerBlock); + } + + __host__ static auto CalculateKPadded(index_t K) + { + return math::integer_divide_ceil(K, KPerBlock) * KPerBlock; + } + + __host__ static auto CalculateAK0Padded(index_t K, index_t K_Batch = 1) + { + auto K_t = K_Batch * KPerBlock; + return (K + K_t - 1) / K_t * (KPerBlock / AK1Value); + } + + __host__ static auto CalculateBK0Padded(index_t K, index_t K_Batch = 1) + { + auto K_t = K_Batch * KPerBlock; + return (K + K_t - 1) / K_t * (KPerBlock / BK1Value); + } + + __host__ static auto CalculateKPadded(index_t K, index_t K_Batch = 1) + { + auto K_t = K_Batch * KPerBlock; + return (K + K_t - 1) / K_t * KPerBlock; + } + + __host__ static auto CalculateKRead(index_t K, index_t K_Batch = 1) + { + constexpr auto KReadVec = math::lcm(AK1Number, BK1Number); + auto K_t = K_Batch * KReadVec; + return (K + K_t - 1) / K_t * KReadVec; + } + + __host__ static auto CalculateMBlock(index_t M) + { + return math::integer_divide_ceil(M, MPerBlock); + } + + __host__ static auto CalculateNBlock(index_t N) + { + return math::integer_divide_ceil(N, NPerBlock); + } + + template + __host__ __device__ static constexpr auto MakeGemmMmaTileDescriptor(const TileDesc_K0_MN_K1&) + { + constexpr index_t K0 = TileDesc_K0_MN_K1{}.GetLength(Number<0>{}); + constexpr index_t K1 = TileDesc_K0_MN_K1{}.GetLength(Number<2>{}); + + return transform_tensor_descriptor( + TileDesc_K0_MN_K1{}, + make_tuple(make_merge_transform_v3_division_mod(make_tuple(Number{}, Number{})), + make_unmerge_transform(make_tuple( + Number{}, Number{}, Number{}))), + make_tuple(Sequence<0, 2>{}, Sequence<1>{}), + make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{})); + } + + __host__ __device__ static auto MakeAGridDescriptor_AK0_M_AK1( + index_t M, index_t MPad, index_t K, index_t KPad, index_t StrideA, index_t AK0) + { + const auto a_grid_desc_mraw_kraw = [&]() { + if constexpr(is_same_v) + { + return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(StrideA, I1)); + } + else if constexpr(is_same_v) + { + return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(I1, StrideA)); + } + }(); + + using GemmSpecialization = tensor_operation::device::GemmSpecialization; + + if constexpr(GemmSpec == GemmSpecialization::MKPadding || + GemmSpec == GemmSpecialization::MNKPadding) + { + // pad both M and K + const auto a_grid_desc_m_k = + transform_tensor_descriptor(a_grid_desc_mraw_kraw, + make_tuple(make_right_pad_transform(M, MPad - M), + make_right_pad_transform(K, KPad - K)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); + + const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor( + a_grid_desc_m_k, + make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)), + make_pass_through_transform(MPad)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return a_grid_desc_ak0_m_ak1; + } + else if constexpr(GemmSpec == GemmSpecialization::MPadding || + GemmSpec == GemmSpecialization::MNPadding) + { + // pad M, but not K + const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor( + a_grid_desc_mraw_kraw, + make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)), + make_right_pad_transform(M, MPad - M)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return a_grid_desc_ak0_m_ak1; + } + else if constexpr(GemmSpec == GemmSpecialization::KPadding || + GemmSpec == GemmSpecialization::NKPadding) + { + // pad K, but not M + const auto a_grid_desc_m_k = transform_tensor_descriptor( + a_grid_desc_mraw_kraw, + make_tuple(make_pass_through_transform(M), make_right_pad_transform(K, KPad - K)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); + + const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor( + a_grid_desc_m_k, + make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)), + make_pass_through_transform(M)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return a_grid_desc_ak0_m_ak1; + } + else + { + // not pad M or K + const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor( + a_grid_desc_mraw_kraw, + make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)), + make_pass_through_transform(M)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return a_grid_desc_ak0_m_ak1; + } + } + + __host__ __device__ static auto MakeBGridDescriptor_BK0_N_BK1( + index_t K, index_t KPad, index_t N, index_t NPad, index_t StrideB, index_t BK0) + { + const auto b_grid_desc_nraw_kraw = [&]() { + if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(N, K), make_tuple(I1, StrideB)); + } + else if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(N, K), make_tuple(StrideB, I1)); + } + }(); + + using GemmSpecialization = tensor_operation::device::GemmSpecialization; + + if constexpr(GemmSpec == GemmSpecialization::NKPadding || + GemmSpec == GemmSpecialization::MNKPadding) + { + // pad both N and K + const auto b_grid_desc_n_k = + transform_tensor_descriptor(b_grid_desc_nraw_kraw, + make_tuple(make_right_pad_transform(N, NPad - N), + make_right_pad_transform(K, KPad - K)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); + + const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor( + b_grid_desc_n_k, + make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)), + make_pass_through_transform(NPad)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return b_grid_desc_bk0_n_bk1; + } + else if constexpr(GemmSpec == GemmSpecialization::NPadding || + GemmSpec == GemmSpecialization::MNPadding) + { + // pad N, but not K + const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor( + b_grid_desc_nraw_kraw, + make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)), + make_right_pad_transform(N, NPad - N)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return b_grid_desc_bk0_n_bk1; + } + else if constexpr(GemmSpec == GemmSpecialization::KPadding || + GemmSpec == GemmSpecialization::MKPadding) + { + // pad K, but not N + const auto b_grid_desc_n_k = transform_tensor_descriptor( + b_grid_desc_nraw_kraw, + make_tuple(make_pass_through_transform(N), make_right_pad_transform(K, KPad - K)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); + + const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor( + b_grid_desc_n_k, + make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)), + make_pass_through_transform(N)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return b_grid_desc_bk0_n_bk1; + } + else + { + // not pad N or K + const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor( + b_grid_desc_nraw_kraw, + make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)), + make_pass_through_transform(N)), + make_tuple(Sequence<1>{}, Sequence<0>{}), + make_tuple(Sequence<0, 2>{}, Sequence<1>{})); + + return b_grid_desc_bk0_n_bk1; + } + } + + __host__ __device__ static constexpr auto + MakeAScaleGridDesciptor_M_K(index_t M, index_t K, index_t StrideScaleA) + { + const auto BM = math::integer_divide_ceil(M, ScaleBlockM); + const auto BK = math::integer_divide_ceil(K, ScaleBlockK); + if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(BM, BK), make_tuple(StrideScaleA, I1)); + } + else if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(BM, BK), make_tuple(I1, StrideScaleA)); + } + } + + __host__ __device__ static constexpr auto + MakeBScaleGridDesciptor_N_K(index_t N, index_t K, index_t StrideScaleB) + { + const auto BN = math::integer_divide_ceil(N, ScaleBlockN); + const auto BK = math::integer_divide_ceil(K, ScaleBlockK); + if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(BN, BK), make_tuple(StrideScaleB, I1)); + } + else if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(BN, BK), make_tuple(I1, StrideScaleB)); + } + } + + template + __host__ __device__ static constexpr auto + MakeAMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&) + { + constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXDL); + + return MakeGemmMmaTileDescriptor(ABlockDesc_AK0_M_AK1{}); + } + + template + __host__ __device__ static constexpr auto + MakeBMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&) + { + constexpr index_t NWaves = + NXdlPerWave * NPerXDL == 0 ? 1 : NPerBlock / (NXdlPerWave * NPerXDL); + + return MakeGemmMmaTileDescriptor(BBlockDesc_BK0_N_BK1{}); + } + + template + __host__ __device__ static auto + MakeCGridDescriptor_M_N(index_t M, index_t MPad, index_t N, index_t NPad, index_t StrideC) + { + const auto c_grid_desc_mraw_nraw = [&]() { + if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(StrideC, I1)); + } + else if constexpr(is_same::value) + { + return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(I1, StrideC)); + } + }(); + + // pad M and N + return transform_tensor_descriptor(c_grid_desc_mraw_nraw, + make_tuple(make_right_pad_transform(M, MPad - M), + make_right_pad_transform(N, NPad - N)), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0>{}, Sequence<1>{})); + } + + __host__ __device__ static auto MakeDsGridDescriptor_M_N( + index_t M, index_t MPad, index_t N, index_t NPad, std::array StrideDs) + { + return generate_tuple( + [&](auto i) { + using DLayout = remove_cvref_t>; + return MakeCGridDescriptor_M_N(M, MPad, N, NPad, StrideDs[i]); + }, + Number{}); + } + + template + __device__ static constexpr auto MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + const DsGridDesc& ds_grid_desc_m_n, index_t MBlock, index_t NBlock) + { + return generate_tuple( + [&](auto i) { + return MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + ds_grid_desc_m_n[i], MBlock, NBlock); + }, + Number{}); + } + + using DsGridDesc_M_N = remove_cvref_t; + + struct Problem + { + __host__ Problem(index_t M_, + index_t N_, + index_t K_, + index_t StrideA_, + index_t StrideB_, + std::array StrideDs_, + index_t StrideC_, + index_t StrideScaleA_, + index_t StrideScaleB_, + index_t KBatch_) + : M{M_}, + N{N_}, + K{K_}, + StrideA{StrideA_}, + StrideB{StrideB_}, + StrideDs{StrideDs_}, + StrideC{StrideC_}, + StrideScaleA{StrideScaleA_}, + StrideScaleB{StrideScaleB_}, + KBatch{KBatch_}, + MPadded{CalculateMPadded(M_)}, + NPadded{CalculateNPadded(N_)}, + KRead{CalculateKRead(K_, KBatch_)}, + KPadded{CalculateKPadded(K_, KBatch_)}, + AK0{CalculateAK0Padded(K_, KBatch_)}, + BK0{CalculateBK0Padded(K_, KBatch_)}, + MBlock{CalculateMBlock(M_)}, + NBlock{CalculateNBlock(N_)} + { + } + + __host__ void Print() const + { + std::cout << "problem {" << "M:" << M << ", " << "N:" << N << ", " << "K:" << K << ", " + << "SA:" << StrideA << ", " << "SB:" << StrideB << ", " << "SC:" << StrideC + << ", " << "MP:" << MPadded << ", " << "NP:" << NPadded << ", " + << "KRead:" << KRead << ", " << "KP:" << KPadded << ", " << "AK0:" << AK0 + << ", " << "BK0:" << BK0 << ", " << "MBlock: " << MBlock << ", " + << "NBlock: " << NBlock << "}" << std::endl; + } + + index_t M; + index_t N; + index_t K; + index_t StrideA; + index_t StrideB; + std::array StrideDs; + index_t StrideC; + index_t StrideScaleA; + index_t StrideScaleB; + index_t KBatch; + index_t MPadded; + index_t NPadded; + index_t KRead; + index_t KPadded; + index_t AK0; + index_t BK0; + index_t MBlock; + index_t NBlock; + // When true, the caller guarantees p_c_grid is already zeroed before + // launch, so the device invoker skips its own hipMemsetAsync for the + // KBatch > 1 split-K atomic-accumulation path. Defaults to false, which + // preserves the original zero-init behavior. + bool skip_zero_init = false; + }; + + // Argument + struct Argument : public tensor_operation::device::BaseArgument, public Problem + { + __host__ Argument(const ADataType* p_a_grid_, + const BDataType* p_b_grid_, + std::array p_ds_grid_, + CDataType* p_c_grid_, + index_t M_, + index_t N_, + index_t K_, + index_t StrideA_, + index_t StrideB_, + std::array StrideDs_, + index_t StrideC_, + index_t StrideScaleA_, + index_t StrideScaleB_, + const AScaleType* p_a_scale_grid_, + const BScaleType* p_b_scale_grid_, + index_t k_batch_, + AElementwiseOperation a_element_op_, + BElementwiseOperation b_element_op_, + CElementwiseOperation c_element_op_) + : Problem{M_, + N_, + K_, + StrideA_, + StrideB_, + StrideDs_, + StrideC_, + StrideScaleA_, + StrideScaleB_, + k_batch_}, + p_a_grid{p_a_grid_}, + p_b_grid{p_b_grid_}, + p_ds_grid{}, + p_c_grid{p_c_grid_}, + p_a_scale_grid{p_a_scale_grid_}, + p_b_scale_grid{p_b_scale_grid_}, + a_element_op{a_element_op_}, + b_element_op{b_element_op_}, + c_element_op{c_element_op_} + { + + // populate pointer, desc for Ds + static_for<0, NumDTensor, 1>{}([&](auto i) { + using DDataType_ = remove_cvref_t>; + + // D pointer + p_ds_grid(i) = static_cast(p_ds_grid_[i]); + }); + } + + const ADataType* p_a_grid; + const BDataType* p_b_grid; + DsGridPointer p_ds_grid; + CDataType* p_c_grid; + + const AScaleType* p_a_scale_grid; + const BScaleType* p_b_scale_grid; + + const AElementwiseOperation a_element_op; + const BElementwiseOperation b_element_op; + const CElementwiseOperation c_element_op; + }; + + struct SplitKBatchOffset + { + __device__ SplitKBatchOffset(Argument& karg) + { + if constexpr(is_same_v) + { + a_k_split_offset = blockIdx.z * karg.KRead; + } + else if constexpr(is_same_v) + { + a_k_split_offset = blockIdx.z * karg.KRead * karg.M; + } + + if constexpr(is_same_v) + { + b_k_split_offset = blockIdx.z * karg.KRead * karg.N; + } + else if constexpr(is_same_v) + { + b_k_split_offset = blockIdx.z * karg.KRead; + } + + // Calculate A scale offset + if constexpr(is_same_v) + { + scale_a_k_split_offset = blockIdx.z * (karg.KRead / ScaleBlockK); + } + else if constexpr(is_same_v) + { + scale_a_k_split_offset = + blockIdx.z * (karg.KRead / ScaleBlockK) * karg.StrideScaleA; + } + + // Calculate B scale offset + if constexpr(is_same_v) + { + scale_b_k_split_offset = + blockIdx.z * (karg.KRead / ScaleBlockK) * karg.StrideScaleB; + } + else if constexpr(is_same_v) + { + scale_b_k_split_offset = blockIdx.z * (karg.KRead / ScaleBlockK); + } + + if(blockIdx.z < static_cast(karg.KBatch - 1)) + { + karg.K = karg.KRead; + } + else + { + karg.K = karg.K - karg.KRead * (karg.KBatch - 1); + } + } + + index_t a_k_split_offset; + index_t b_k_split_offset; + index_t scale_a_k_split_offset; // A scale matrix offset + index_t scale_b_k_split_offset; // B scale matrix offset + }; + + using BlockwiseGemmPipe = remove_cvref_t< + decltype(BlockGemmABScalePipeline_Selector< + BlkGemmPipelineVer, + BlkGemmPipeSched, + BlockSize, + LDSTypeA, + LDSTypeB, + ComputeTypeA, + AccDataType, + decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch())), + decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch())), + decltype(MakeAMmaTileDescriptor_M0_M1_M2_K( + GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch()))), + decltype(MakeBMmaTileDescriptor_N0_N1_N2_K( + GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch()))), + ABlockTransferSrcScalarPerVector, + BBlockTransferSrcScalarPerVector, + MPerBlock, + NPerBlock, + KPerBlock, + MPerXDL, + NPerXDL, + MXdlPerWave, + NXdlPerWave, + KPack>())>; + + IS_VALID_COMPILATION_PARAMETER_IMPL(CDataType) + + // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01} + __host__ static constexpr bool CheckValidity(const Argument& karg) + { + static_assert((MPerBlock % (MPerXDL * MXdlPerWave) == 0) && + (NPerBlock % (NXdlPerWave * NPerXDL)) == 0, + "Invalid tuning param!"); + + if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + !(is_same::value)) + { + if(!(karg.M % MPerBlock == 0)) + { +#if DEBUG_LOG + std::cout << "Arg M value is not a multiple of MPerBlock! M: " << karg.M << " " + << __FILE__ << ":" << __LINE__ << ", in function: " << __func__ + << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + + if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) && + (is_same::value)) + { + if(!(karg.N % NPerBlock == 0)) + { +#if DEBUG_LOG + std::cout << "Arg N value is not a multiple of NPerBlock! N: " << karg.N << " " + << __FILE__ << ":" << __LINE__ << ", in function: " << __func__ + << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + + if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::KPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding || + GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding)) + { + + auto K_t = karg.KBatch * KPerBlock; + if(!(karg.K % K_t == 0)) + { +#if DEBUG_LOG + std::cout << "Arg K value is not a multiple of K_Batch * K0PerBlock * K1! K: " + << karg.K << " " << __FILE__ << ":" << __LINE__ + << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + else + { + constexpr auto KReadVec = math::lcm(AK1Number, BK1Number); + auto K_t = karg.KBatch * KReadVec; + auto KReadPadSplited = math::integer_divide_ceil(karg.K, K_t) * KReadVec; + if((KReadPadSplited * (karg.KBatch - 1)) >= karg.K) + { + return false; + } + } + + if constexpr(is_same::value) + { + if(karg.K % ABlockTransferSrcScalarPerVector != 0) + { +#if DEBUG_LOG + std::cout << "Arg K (" << karg.K + << ") value is not a multiple of ABlockTransferSrcScalarPerVector (" + << ABlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":" + << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + else + { + if(karg.M % ABlockTransferSrcScalarPerVector != 0) + { +#if DEBUG_LOG + std::cout << "Arg M (" << karg.M + << ") value is not a multiple of ABlockTransferSrcScalarPerVector (" + << ABlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":" + << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + + if constexpr(is_same::value) + { + if(karg.N % BBlockTransferSrcScalarPerVector != 0) + { +#if DEBUG_LOG + std::cout << "Arg N (" << karg.N + << ") value is not a multiple of BBlockTransferSrcScalarPerVector (" + << BBlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":" + << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + else + { + if(karg.K % BBlockTransferSrcScalarPerVector != 0) + { +#if DEBUG_LOG + std::cout << "Arg K (" << karg.K + << ") value is not a multiple of BBlockTransferSrcScalarPerVector (" + << BBlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":" + << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + + if constexpr(is_same::value) + { + if(karg.N % CShuffleBlockTransferScalarPerVector_NPerBlock != 0) + { +#if DEBUG_LOG + std::cout << "Arg N (" << karg.N + << ") value is not a multiple of " + "CShuffleBlockTransferScalarPerVector_NPerBlock (" + << CShuffleBlockTransferScalarPerVector_NPerBlock << " )! " << __FILE__ + << ":" << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + else + { + if(karg.M % CShuffleBlockTransferScalarPerVector_NPerBlock != 0) + { +#if DEBUG_LOG + std::cout << "Arg M (" << karg.M + << ") value is not a multiple of " + "CShuffleBlockTransferScalarPerVector_NPerBlock (" + << CShuffleBlockTransferScalarPerVector_NPerBlock << " )! " << __FILE__ + << ":" << __LINE__ << ", in function: " << __func__ << std::endl; + +#endif // DEBUG_LOG + return false; + } + } + + // check gridwise gemm pipeline + const auto num_k_loop = karg.AK0 / (KPerBlock / AK1Value); + + if constexpr(BlkGemmPipelineVer != BlockGemmPipelineVersion::v1) + { + if(num_k_loop <= BlockwiseGemmPipe::PrefetchStages) + { + return false; + } + } + + // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc) + return true; + } + + __host__ static constexpr bool CalculateHasMainKBlockLoop(index_t K) + { + const index_t num_loop = K / KPerBlock; + + return BlockwiseGemmPipe::BlockHasHotloop(num_loop); + } + + __host__ static constexpr TailNumber CalculateKBlockLoopTailNum(index_t K) + { + const index_t num_loop = K / KPerBlock; + + return BlockwiseGemmPipe::BlockLoopTailNum(num_loop); + } + + template + __device__ static constexpr auto MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + const CGridDesc& c_grid_desc_m_n, index_t MBlock, index_t NBlock) + { + const auto c_grid_desc_mblock_mperblock_nblock_nperblock = transform_tensor_descriptor( + c_grid_desc_m_n, + make_tuple(make_unmerge_transform(make_tuple(MBlock, Number{})), + make_unmerge_transform(make_tuple(NBlock, Number{}))), + make_tuple(Sequence<0>{}, Sequence<1>{}), + make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{})); + + return c_grid_desc_mblock_mperblock_nblock_nperblock; + } + + // return block_id to C matrix tile idx (m0, n0) mapping + // if arch = gfx942 + using Block2CTileMap = BlockToCTileMap_Grouped_M00_N0_M01Adapt<8, MPerBlock, NPerBlock>; + // using Block2CTileMap = BlockToCTileMap_3DGrid_KSplit; + + template + __device__ static void Run(const ADataType* p_a_grid, + const BDataType* p_b_grid, + DsGridPointer& p_ds_grid, + CDataType* p_c_grid, + const AScaleType* p_a_scale_grid, + const BScaleType* p_b_scale_grid, + void* p_shared, + const Problem& problem, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CElementwiseOperation c_element_op) + { + const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1( + problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0); + const auto b_grid_desc_bk0_n_bk1 = MakeBGridDescriptor_BK0_N_BK1( + problem.K, problem.KPadded, problem.N, problem.NPadded, problem.StrideB, problem.BK0); + const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N( + problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC); + + const auto a_scale_grid_desc_am_ak = + MakeAScaleGridDesciptor_M_K(problem.M, problem.K, problem.StrideScaleA); + const auto b_scale_grid_desc_bn_ak = + MakeBScaleGridDesciptor_N_K(problem.N, problem.K, problem.StrideScaleB); + + const auto c_grid_desc_mblock_mperblock_nblock_nperblock = + MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + c_grid_desc_m_n, problem.MBlock, problem.NBlock); + + const auto a_grid_buf = make_dynamic_buffer( + p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize()); + const auto b_grid_buf = make_dynamic_buffer( + p_b_grid, b_grid_desc_bk0_n_bk1.GetElementSpaceSize()); + + const auto a_scale_grid_buf = make_dynamic_buffer( + p_a_scale_grid, a_scale_grid_desc_am_ak.GetElementSpaceSize()); + + const auto b_scale_grid_buf = make_dynamic_buffer( + p_b_scale_grid, b_scale_grid_desc_bn_ak.GetElementSpaceSize()); + + // divide block work by [M, N] + const auto block_2_ctile_map = Block2CTileMap{problem.M, problem.N, 4}; + + const auto block_work_idx = + block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id())); + + if(!block_2_ctile_map.ValidCTileIndex( + block_work_idx, + make_tuple(c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I0), + c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I2)))) + { + return; + } + + const index_t block_m_id = __builtin_amdgcn_readfirstlane(block_work_idx[I0]); + const index_t block_n_id = __builtin_amdgcn_readfirstlane(block_work_idx[I1]); + + // HACK: this force m/n_block_data_idx_on_grid into SGPR + const index_t m_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(block_m_id * MPerBlock); + + const index_t n_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(block_n_id * NPerBlock); + + // lds max alignment + constexpr auto max_lds_align = math::lcm(AK1Number, BK1Number); + + // A matrix in LDS memory, dst of blockwise copy + constexpr auto a_block_desc_ak0_m_ak1 = + GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch()); + + // B matrix in LDS memory, dst of blockwise copy + constexpr auto b_block_desc_bk0_n_bk1 = + GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch()); + + // A matrix blockwise copy + auto a_blockwise_copy = + ThreadGroupTensorSliceTransfer_v4r1, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + LDSTypeA, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + ABlockTransferSrcVectorDim, + 2, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + 1, + 1, + AThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + a_grid_desc_ak0_m_ak1, + make_multi_index(0, m_block_data_idx_on_grid, 0), + a_element_op, + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + + // B matrix blockwise copy + auto b_blockwise_copy = + ThreadGroupTensorSliceTransfer_v4r1, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BDataType, + LDSTypeB, + decltype(b_grid_desc_bk0_n_bk1), + decltype(b_block_desc_bk0_n_bk1), + BBlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + BBlockTransferSrcVectorDim, + 2, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + 1, + 1, + BThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + b_grid_desc_bk0_n_bk1, + make_multi_index(0, n_block_data_idx_on_grid, 0), + b_element_op, + b_block_desc_bk0_n_bk1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + + // LDS allocation for A and B: be careful of alignment + constexpr auto a_block_space_size_aligned = math::integer_least_multiple( + a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align); + + // Cast after lds + auto a_block_buf = make_dynamic_buffer( + static_cast(p_shared), a_block_desc_ak0_m_ak1.GetElementSpaceSize()); + + auto b_block_buf = make_dynamic_buffer( + static_cast(p_shared) + + a_block_space_size_aligned * sizeof(LDSTypeA) / sizeof(LDSTypeB), + b_block_desc_bk0_n_bk1.GetElementSpaceSize()); + + constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0); + constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock / BK1Number, 0, 0); + + // Blockwise GEMM pipeline + static_assert(std::is_default_constructible_v); + auto blockwise_gemm_pipeline = BlockwiseGemmPipe{}; + auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer(); + + const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane( + (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) / + KPerBlock); + + constexpr index_t ScaleSliceSizeM = MXdlPerWave; + constexpr index_t ScaleSliceSizeN = math::integer_divide_ceil(NPerBlock, ScaleBlockN); + constexpr index_t ScaleSliceSizeK = math::integer_divide_ceil(KPerBlock, ScaleBlockK); + + // ScaleSliceSizeK is last dimension in A/B scale for vector memory access + // ScaleSliceSizeK is first dimension in C scale for packed math + constexpr auto a_scale_thread_desc = make_naive_tensor_descriptor_packed( + make_tuple(Number{}, Number{})); + + constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXDL); + constexpr index_t NWaves = NPerBlock / (NXdlPerWave * NPerXDL); + constexpr index_t WaveSize = BlockSize / (MWaves * NWaves); + auto a_thread_offset = get_thread_local_1d_id() % MPerXDL + + (get_thread_local_1d_id() / WaveSize) / NWaves * MPerXDL; + + constexpr auto b_scale_thread_desc = make_naive_tensor_descriptor_packed( + make_tuple(Number{}, Number{})); + + constexpr auto c_scale_thread_desc = make_naive_tensor_descriptor_packed(make_tuple( + Number{}, Number{}, Number{})); + + auto a_scale_thread_copy = + ThreadwiseTensorSliceTransfer_v2, + Sequence<0, 1>, + 1, + ScaleSliceSizeK, + 1, + false>( + a_scale_grid_desc_am_ak, + make_multi_index(block_m_id * MPerBlock / ScaleBlockM + a_thread_offset, 0)); + + auto b_scale_thread_copy = + ThreadwiseTensorSliceTransfer_v2, + Sequence<0, 1>, + 1, + ScaleSliceSizeK, + 1, + false>( + b_scale_grid_desc_bn_ak, make_multi_index(block_n_id * NPerBlock / ScaleBlockN, 0)); + + // constexpr auto a_scale_thread_slice_copy_step = make_multi_index(0, 1); + constexpr auto a_scale_thread_slice_copy_step = + make_tuple(make_multi_index(MWaves * MPerXDL, 0), + make_multi_index(-MPerBlock, 0), + make_multi_index(-MPerBlock, ScaleSliceSizeK)); + constexpr auto b_scale_thread_slice_copy_step = make_multi_index(0, ScaleSliceSizeK); + + constexpr auto NumKBlockPerScale = math::integer_divide_ceil(ScaleBlockK, KPerBlock); + + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_buf, + a_block_slice_copy_step, + b_grid_desc_bk0_n_bk1, + b_block_desc_bk0_n_bk1, + b_blockwise_copy, + b_grid_buf, + b_block_buf, + b_block_slice_copy_step, + + c_scale_thread_desc, + c_thread_buf, + + a_scale_grid_desc_am_ak, + a_scale_thread_desc, + a_scale_thread_copy, + a_scale_grid_buf, + a_scale_thread_slice_copy_step, + + b_scale_grid_desc_bn_ak, + b_scale_thread_desc, + b_scale_thread_copy, + b_scale_grid_buf, + b_scale_thread_slice_copy_step, + + num_k_block_main_loop); + + // shuffle C and write out + const auto ds_grid_desc_m_n = MakeDsGridDescriptor_M_N( + problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideDs); + + const auto ds_grid_desc_mblock_mperblock_nblock_nperblock = + MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + ds_grid_desc_m_n, problem.MBlock, problem.NBlock); + Base::template RunMultiDEpilogue( + blockwise_gemm_pipeline, + ds_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_desc_mblock_mperblock_nblock_nperblock, + c_thread_buf, + block_m_id, + block_n_id, + p_shared, + p_ds_grid, + p_c_grid, + c_element_op); + } +}; + +} // namespace ck +#if __clang_major__ >= 23 +#pragma clang diagnostic pop +#endif diff --git a/include/ck/utility/index_expression.hpp b/include/ck/utility/index_expression.hpp new file mode 100644 index 0000000..3de0b83 --- /dev/null +++ b/include/ck/utility/index_expression.hpp @@ -0,0 +1,169 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/utility/common_header.hpp" + +/** + * @file index_expression.hpp + * @brief Compile-time index expression evaluation utilities + * @details Provides types and operations for building and evaluating index expressions at compile + * time. Supports arithmetic operations (Add, Mult, Div, Mod) on compile-time constants and index + * variables. Expressions are built using template types and evaluated via the eval() function + * family. + */ + +namespace ck::index_expression { + +/** + * @brief Represents an index variable in a compile-time expression + * @details Used as a placeholder for index values that are determined at evaluation time. + * When an expression containing Ik is evaluated with eval(), the Ik type evaluates to the + * provided ik value. + */ +struct Ik +{ +}; + +/** + * @brief Binary addition operation for compile-time index expressions + * @tparam L Left operand type (must be evaluable to index_t) + * @tparam R Right operand type (must be evaluable to index_t) + */ +template +struct Add +{ +}; + +/** + * @brief Binary multiplication operation for compile-time index expressions + * @tparam L Left operand type + * @tparam R Right operand type + */ +template +struct Mult +{ +}; + +/** + * @brief Binary division operation for compile-time index expressions + * @tparam L Left operand type + * @tparam R Right operand type + */ +template +struct Div +{ +}; + +/** + * @brief Binary modulo operation for compile-time index expressions + * @tparam L Left operand type + * @tparam R Right operand type + * @note Both operands must evaluate to integer types + */ +template +struct Mod +{ +}; + +template +/** + * @brief Evaluates a literal Number value + * @tparam ik The index variable value (unused for literal expressions) + * @tparam v The literal constant value + * @return The constant value v + */ +constexpr index_t eval(Number) +{ + return v; +} + +/** + * @brief Evaluates an index variable to its provided value + * @tparam ik The value to substitute for the index variable + * @return The index value ik + */ +template +constexpr index_t eval(Ik) +{ + return ik; +} + +/** + * @brief Evaluates an addition expression + * @tparam ik The index variable value + * @tparam L Type of left operand + * @tparam R Type of right operand + * @return Sum of eval(L) + eval(R) + */ +template +constexpr index_t eval(Add) +{ + return eval(L{}) + eval(R{}); +} + +/** + * @brief Evaluates a multiplication expression + * @tparam ik The index variable value + * @tparam L Type of left operand + * @tparam R Type of right operand + * @return Product of eval(L) * eval(R) + */ +template +constexpr index_t eval(Mult) +{ + return eval(L{}) * eval(R{}); +} + +/** + * @brief Evaluates a division expression + * @tparam ik The index variable value + * @tparam L Type of left operand + * @tparam R Type of right operand (must evaluate to non-zero) + * @return Quotient of eval(L) / eval(R) + */ +template +constexpr index_t eval(Div) +{ + constexpr index_t d = eval(R{}); + static_assert(d != 0, + "ck::index_expression::Div: division by zero in compile-time index expression"); + return eval(L{}) / d; +} + +/** + * @brief Evaluates a modulo expression + * @tparam ik The index variable value + * @tparam L Type of left operand + * @tparam R Type of right operand + * @return Remainder of eval(L) % eval(R) + * @details Both operands must evaluate to integer index types + */ +template +constexpr index_t eval(Mod) +{ + using lhs_t = decltype(eval(L{})); + using rhs_t = decltype(eval(R{})); + static_assert(std::is_integral_v && std::is_integral_v, + "ck::index_expression::Mod: modulo requires integer operands"); + return eval(L{}) % eval(R{}); +} + +/** + * @brief Compile-time evaluation helper for index expressions + * @tparam Expr The expression to evaluate + * @tparam ik The index variable value + * @details Helper variable template that evaluates the given expression at compile time + * for the provided index value. + * @example + * using Expr = Add>; + * static_assert(eval_v == 8); // Evaluates to 3 + 5 = 8 + */ +template +inline constexpr index_t eval_v = eval(Expr{}); + +} // namespace ck::index_expression diff --git a/include/ck/utility/thread_buf_to_vec_loader.hpp b/include/ck/utility/thread_buf_to_vec_loader.hpp new file mode 100644 index 0000000..b842b97 --- /dev/null +++ b/include/ck/utility/thread_buf_to_vec_loader.hpp @@ -0,0 +1,163 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + + +#pragma once + +#include "ck/utility/common_header.hpp" +#include "ck/utility/data_type.hpp" +#include "ck/utility/index_expression.hpp" + +namespace ck { + +/** + * @brief Invokes multiple functors based on an index parameter + * @tparam Funcs Parameter pack of functor types + * @details Stores a tuple of functors and provides an operator() that invokes all of them + * with the same index parameter. Uses static_for to iterate through the functors. + */ +template +struct FunctorInvoker +{ + ck::Tuple funcs; + + __host__ __device__ constexpr FunctorInvoker(Funcs... fs) : funcs(ck::forward(fs)...) {} + + /** + * @brief Invokes all functors with the given index + * @tparam I The index to pass to each functor + * @param i Number wrapper containing the index value + */ + template + __host__ __device__ constexpr void operator()(ck::Number i) const + { + invoke(i, std::index_sequence_for{}); + } + + private: + template + __host__ __device__ constexpr void invoke(ck::Number i, std::index_sequence) const + { + (funcs[ck::Number(Is)>{}](i), ...); + } +}; + +// required for CTAD to work with __host__ __device__ qualifiers +template +__host__ __device__ constexpr auto MakeFunctorInvoker(Fs&&... fs) +{ + return FunctorInvoker{ck::forward(fs)...}; +} + +/** + * @brief Helper struct for evaluating compile-time index expressions + * @tparam T The expression type to evaluate + * @tparam ik The index variable value + * @details Provides a value member that evaluates the index expression T using + * the index_expression::eval_v + */ +template +struct IndexEval; + +template +struct IndexEval : IndexEval +{ +}; + +template +struct IndexEval, ik> +{ + static constexpr index_t value = v; +}; + +template +struct IndexEval +{ + static constexpr index_t value = ik; +}; + +template +struct IndexEval, ik> +{ + static constexpr index_t value = IndexEval::value + IndexEval::value; +}; + +template +struct IndexEval, ik> +{ + static constexpr index_t value = IndexEval::value * IndexEval::value; +}; + +template +struct IndexEval, ik> +{ + static constexpr index_t divisor = IndexEval::value; + static_assert(divisor != 0, + "ck::index_expression::Div: division by zero in compile-time index expression"); + static constexpr index_t value = IndexEval::value / divisor; +}; + +template +struct IndexEval, ik> +{ + static constexpr index_t divisor = IndexEval::value; + static_assert(divisor != 0, + "ck::index_expression::Mod: modulo by zero in compile-time index expression"); + static constexpr index_t value = IndexEval::value % divisor; +}; + +/** + * @brief Loads thread elements from buffer to vector using compile-time index expressions + * @tparam ThreadVec The vector type to load into + * @tparam ThreadBuf The buffer type to load from + * @tparam ThreadDesc The descriptor for thread memory layout + * @tparam ComputeType The computation type for the result + * @tparam IdxExpr Parameter pack of compile-time index expressions + * @details Uses index expressions to compute offsets in ThreadBuf and loads the values + * into the ThreadVec. The operator() accepts a compile-time index and evaluates all + * index expressions for that particular index value. + * + * Example: + * @code + * // Load from buffer using index expressions Ik (the loop index) and Number<5> + * using Loader = thread_buf_to_vec_loader>; + * Loader loader{thread_vec, thread_buf}; + * loader(Number<3>{}); // Loads at offset computed by evaluating expressions with ik=3 + * @endcode + */ +template +struct thread_buf_to_vec_loader +{ + ThreadVec& thread_vec; + ThreadBuf& thread_buf; + + __host__ __device__ constexpr thread_buf_to_vec_loader(ThreadVec& tv, ThreadBuf& tb) + : thread_vec(tv), thread_buf(tb) + { + } + + /** + * @brief Loads a single element from buffer to vector for the given index + * @tparam ik The index value for which to evaluate the index expressions + */ + template + __host__ __device__ constexpr void operator()(Number) const + { + // TODO c++20: ThreadDesc could be an auto parameter, but clang doesn't support auto + // non-type template parameters yet + constexpr auto thread_desc = ThreadDesc{}; + constexpr auto idx_tuple = ck::make_tuple(Number::value>{}...); + constexpr auto offset = thread_desc.CalculateOffset(idx_tuple); + + auto& target = thread_vec.template AsType()(Number{}); + target = thread_buf[Number{}]; + } +}; +} // namespace ck diff --git a/library/include/ck/library/tensor_operation_instance/gpu/gemm_ab_scale.hpp b/library/include/ck/library/tensor_operation_instance/gpu/gemm_ab_scale.hpp new file mode 100644 index 0000000..014b583 --- /dev/null +++ b/library/include/ck/library/tensor_operation_instance/gpu/gemm_ab_scale.hpp @@ -0,0 +1,108 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp" + +#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +#ifdef CK_ENABLE_BF16 +void add_device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances( + std::vector>>& instances); +#endif + +template +struct DeviceOperationInstanceFactory< + ck::tensor_operation::device::DeviceGemmMultipleD_ABScale> +{ + using DeviceOp = DeviceGemmMultipleD_ABScale; + + static auto GetInstances() + { + std::vector> op_ptrs; + +#ifdef CK_ENABLE_BF16 + if constexpr(is_same_v && is_same_v && + is_same_v && is_same_v && + is_same_v) + { + if constexpr(is_same_v && is_same_v && + is_same_v) + { + add_device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances(op_ptrs); + } + } +#endif + return op_ptrs; + } +}; + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/CMakeLists.txt new file mode 100644 index 0000000..da5b4b3 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/CMakeLists.txt @@ -0,0 +1,13 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# Modified by Hygon Information Technology Co., Ltd. + +# ONLY XDL_KERNELS (ported from ck-develop; WMMA tile loop instances not ported) +set(GROUPED_GEMM_TILE_LOOP_INSTANCES) + +list(APPEND GROUPED_GEMM_TILE_LOOP_INSTANCES + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp + ) + +add_instance_library(device_grouped_gemm_tile_loop_instance ${GROUPED_GEMM_TILE_LOOP_INSTANCES}) diff --git a/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp new file mode 100644 index 0000000..c2f43cf --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instance.cpp @@ -0,0 +1,76 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +template +using S = ck::Sequence; + +using DsDataType = ck::Tuple<>; + +using DsLayout = ck::Tuple<>; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +static constexpr auto GemmMNKPadding = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +using device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_irregular_tile_instances = std::tuple< + // clang-format off + //###########################################| A| B| Ds| E| AData| BData| AccData| CShuffle| DsData| EData| A| B| C| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| + //###########################################| Layout| Layout| Layout| Layout| Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| + //###########################################| | | | | | | | | | | Operation| Operation| Operation| | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| + //###########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 256, 128, 32, 8, 8, 32, 32, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 128, 32, 8, 8, 32, 32, 2, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 2, 32, 32, 2, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<16,16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, 0, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 2, 32, 32, 1, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, 0, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 64, 128, 32, 8, 8, 32, 32, 1, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 128, 64, 32, 8, 2, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, 0, 1, 1, S<1, 32, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 128, 64, 32, 8, 8, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 64, 128, 32, 8, 2, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, 0, 1, 1, S<1, 16, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Row, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 64, 128, 32, 8, 8, 32, 32, 2, 2, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1, 1, 1, S<1, 16, 1, 8>, S<8>> + // clang-format on + >; + +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_instances( + std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_kn_mn_irregular_tile_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp new file mode 100644 index 0000000..08655ea --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_gemm_tile_loop/device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instance.cpp @@ -0,0 +1,78 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Modified by Hygon Information Technology Co., Ltd. + +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_xdl_cshuffle_tile_loop.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using F16 = ck::half_t; +using F32 = float; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +template +using S = ck::Sequence; + +using DsDataType = ck::Tuple<>; +using DsLayout = ck::Tuple<>; + +using PassThrough = ck::tensor_operation::element_wise::PassThrough; + +static constexpr auto GemmMNKPadding = ck::tensor_operation::device::GemmSpecialization::MNKPadding; + +using device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_irregular_tile_instances = std::tuple< + // clang-format off + //###########################################| A| B| Ds| E| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| NumGemmK| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| + //###########################################| Layout| Layout| Layout| Layout| Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch| Size| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| + //###########################################| | | | | | | | | | | Operation| Operation| Operation| | Stage| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| + //###########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 256, 64, 8, 8, 32, 32, 2, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 128, 64, 64, 8, 8, 32, 32, 2, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 256, 64, 128, 64, 8, 8, 32, 32, 1, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 128, 128, 64, 8, 8, 32, 32, 4, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 128, 64, 64, 8, 8, 32, 32, 2, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 64, 128, 64, 8, 8, 32, 32, 2, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 128, 32, 64, 8, 8, 32, 32, 2, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 32, 128, 64, 8, 8, 32, 32, 1, 2, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 128, 32, 256, 64, 8, 8, 32, 32, 1, 4, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 64, 64, 64, 64, 8, 8, 32, 32, 2, 2, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 64, 64, 32, 64, 8, 8, 32, 32, 2, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, S<8>>, + DeviceGroupedGemmMultipleDXdlCShuffleTileLoop< Row, Col, DsLayout, Row, F16, F16, F32, F32, DsDataType, F16, PassThrough, PassThrough, PassThrough, GemmMNKPadding, 1, 64, 32, 64, 64, 8, 8, 32, 32, 1, 2, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 4>, S<8>> + // clang-format on + >; + +void add_device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_instances( + std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_gemm_xdl_tile_loop_f16_f16_f16_mk_nk_mn_irregular_tile_instances{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/CMakeLists.txt b/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/CMakeLists.txt new file mode 100644 index 0000000..2cc1996 --- /dev/null +++ b/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/CMakeLists.txt @@ -0,0 +1,12 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Hygon Information Technology Co., Ltd. +# Modified by Hygon Information Technology Co., Ltd. + +# GEMM ab_scale (per-block A/B dequantization scales) instances for HCU +# (gfx936 family, 16x16 XDL tiles only). +set(GEMM_AB_SCALE_HCU_INSTANCES + device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp + ) + +add_instance_library(device_gemm_ab_scale_instance ${GEMM_AB_SCALE_HCU_INSTANCES}) diff --git a/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp b/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp new file mode 100644 index 0000000..6e4b8a0 --- /dev/null +++ b/library/src/tensor_operation_instance_hcu/gpu/gemm_ab_scale/device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instance.cpp @@ -0,0 +1,79 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// Modified by Hygon Information Technology Co., Ltd. +// HCU (gfx936 family) instance set for GEMM with per-block A/B dequantization +// scales. HCU MMAC only supports 16x16 XDL tiles; the tuning below is derived +// from the validated HCU bf16 16x16 GEMM configuration +// (MPerBlock=NPerBlock=128, KPerBlock=32, AK1=BK1=8). + +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp" + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +using BF16 = bhalf_t; +using F32 = float; + +using Row = tensor_layout::gemm::RowMajor; +using Col = tensor_layout::gemm::ColumnMajor; + +template +using S = Sequence; + +using PassThrough = element_wise::PassThrough; + +template +using device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances = std::tuple< + // clang-format off + //############################| ALayout| BLayout| DsLayout| ELayout| AData| AScale| BData| BScale| DsData| EData| AccData| Cshuffle| A| B| C| GEMM| Block| Scale| Scale| Scale| MPer| NPer| KPer| AK1| BK1|MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| Block-wiseGemm| Block-wiseGemm| Compute| + //############################| | | | | Type| Type| Type| Type| Type| Type| Type| Type| Elementwise| Elementwise| Elementwise|Specialization| Size| Block| Block| Block| Block| Block| Block| | | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MXdlPerWave_MWaveMPerXdl| ScalarPerVector| Pipeline| Pipeline| TypeA| + //############################| | | | | | | | | | | | | Operation| Operation| Operation| | | M| N| K| | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| Scheduler| Verision| | + //############################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + + // Compute friendly + DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3< Row, Col, Tuple<>, Row, BF16, F32, BF16, F32, Tuple<>, BF16, F32, F32, PassThrough, PassThrough, PassThrough, GemmSpec, 256, 1, 128, 32, 128, 128, 32, 8, 8, 16, 16, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<4>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v1, BF16> + // clang-format on + >; + +void add_device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances( + std::vector, + Row, + BF16, + F32, + BF16, + F32, + Tuple<>, + BF16, + 1, + 128, + 32, + PassThrough, + PassThrough, + PassThrough>>>& instances) +{ + add_device_operation_instances( + instances, + device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances< + GemmSpecialization::Default>{}); + add_device_operation_instances( + instances, + device_gemm_ab_scale_xdl_bf16_bf16_bf16_mk_nk_mn_instances< + GemmSpecialization::MNKPadding>{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck From a870f1b190999a47cc1bfc82c0afd6dbc878a658 Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Sat, 10 Oct 2026 11:36:36 +0800 Subject: [PATCH 03/11] feat(fused_moe): add fused_moegemm_pipeline_flatmm_uk --- .../ck_tile/17_fused_moe/CMakeLists.txt | 25 + .../17_fused_moe/fused_moe_flatmm_uk_fp16.cpp | 842 +++++++++++++++ .../instances/fused_moegemm_api_internal.hpp | 4 +- .../kernel/fused_moegemm_uk_kernel.hpp | 142 +++ .../fused_moegemm_flatmm_uk_engine.hpp | 966 ++++++++++++++++++ .../fused_moegemm_pipeline_flatmm_uk.hpp | 230 +++++ 6 files changed, 2208 insertions(+), 1 deletion(-) create mode 100644 example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp create mode 100644 include/ck_tile/ops/fused_moe/kernel/fused_moegemm_uk_kernel.hpp create mode 100644 include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_flatmm_uk_engine.hpp create mode 100644 include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp diff --git a/example_hcu/ck_tile/17_fused_moe/CMakeLists.txt b/example_hcu/ck_tile/17_fused_moe/CMakeLists.txt index c709cfa..b37d5d5 100644 --- a/example_hcu/ck_tile/17_fused_moe/CMakeLists.txt +++ b/example_hcu/ck_tile/17_fused_moe/CMakeLists.txt @@ -1,6 +1,31 @@ # Copyright (c) 2026 Hygon Information Technology Co., Ltd. # SPDX: SPDX-License-Identifier: MIT + +add_example_executable(example_hcu_ck_tile_fused_moe_flatmm_uk + fused_moe_flatmm_uk_fp16.cpp) + +add_test(NAME example_hcu_ck_tile_fused_moe_flatmm_uk_smoke + COMMAND $ 64 4 2 256 256 1 2 2 0) + +add_test(NAME example_hcu_ck_tile_fused_moe_flatmm_uk_partial_interm + COMMAND $ 64 4 2 256 1232 1 2 2 0) + +add_test(NAME example_hcu_ck_tile_fused_moe_flatmm_uk_mock_id + COMMAND $ 64 4 2 256 256 1 2 2 0 11939 2) + +add_test(NAME example_hcu_ck_tile_fused_moe_flatmm_uk_tiny_hidden + COMMAND $ 8192 128 64 1 1 1 2 2 0) +add_test(NAME example_hcu_ck_tile_fused_moe_flatmm_uk_ragged_hidden + COMMAND $ 64 4 2 3 256 1 2 2 0) + +set_tests_properties(example_hcu_ck_tile_fused_moe_flatmm_uk_smoke + example_hcu_ck_tile_fused_moe_flatmm_uk_partial_interm + example_hcu_ck_tile_fused_moe_flatmm_uk_mock_id + example_hcu_ck_tile_fused_moe_flatmm_uk_tiny_hidden + example_hcu_ck_tile_fused_moe_flatmm_uk_ragged_hidden + PROPERTIES LABELS "hcu;ck_tile;fused_moe;smoke") + set(TILE_EXAPMLE_FUSED_MOE "tile_example_fused_moe") # not using add_example_executable() to add this target, since we don't want this to have # to be included in "make all/install/check" diff --git a/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp b/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp new file mode 100644 index 0000000..013dc39 --- /dev/null +++ b/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp @@ -0,0 +1,842 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Info Technologies Ltd. +// Modified by Hygon Information Technology Co., Ltd. +// +// Example/validation harness for the ported FusedMoeGemmPipeline_FlatmmUk +// (HCU MMAC engine). Runs a 2-GEMM fused MoE (gate+up, silu*mul bridge, +// down-proj with topk-weighted fp32 atomic accumulation) and checks against +// a fp32 CPU reference. +// +// Weight layout (natural, no host pre-shuffle), same as fork's +// FusedMoeGemmKernel: W1/G[e] = [2*interm, hidden] row-major (gate|up), +// W2/D[e] = [hidden, interm] row-major. sorted_token_ids holds raw token +// ids (default, enc=1); padding slots use the num_tokens sentinel, weight 0. +// enc=2 packs ids the way the repo sorting kernel does under +// CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID (token | (topk_idx << 24)). +// +// args: tokens experts topk hidden interm [v=1] [warmup=5] [repeat=20] +// [init=0|1] [seed] [enc=1|0|2] [shuf=0] +// +// hidden is auto-padded up to the gmem load-vector alignment (16B = 8 fp16 +// here): the engine's A/G/U loads mask out-of-range accesses per whole +// vector, so an unpadded hidden would read across row ends. The pad region +// is zero on both the device buffers and the reference, and only logical +// columns are validated. intermediate_size has no such constraint (the D +// loads are guarded element-wise). + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/host/stream_config.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/elementwise.hpp" +#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_problem.hpp" +#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_traits.hpp" +#include "ck_tile/ops/fused_moe/kernel/fused_moegemm_shape.hpp" +#include "ck_tile/ops/fused_moe/kernel/fused_moegemm_tile_partitioner.hpp" +#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp" +#include "ck_tile/ops/fused_moe/kernel/fused_moegemm_uk_kernel.hpp" + +// ---------------------------------------------------------------- utils +static uint64_t lcg_state; +static void lcg_seed(uint64_t s) { lcg_state = s ? s : 1; } +static float lcg_float() // [-0.5, 0.5) +{ + lcg_state = lcg_state * 6364136223846793005ULL + 1442695040888963407ULL; + return static_cast(static_cast(lcg_state >> 41) & 0x7fffff) / + static_cast(0x800000) - + 0.5f; +} + +// Host fp16 conversions are done with these inlined integer-only helpers +// instead of _Float16 casts: the DTK aicc -c (object) offload-bundle path +// miscompiles the cast call sites in this TU (garbage arguments reach +// __truncsfhf2/__extendhfsf2: device buffers get wrong bits and the CPU +// reference silently reads zeros, so validation degrades to a vacuous PASS; +// see the "sanity:" line below). The bit math is immune, matches the device +// v_cvt rounding (round-to-nearest-even), and mirrors the software fallback +// in ck_tile/core/numeric/half.hpp. Device code is unaffected. +static inline uint16_t f32_to_f16_bits(float x) +{ + uint32_t f; + std::memcpy(&f, &x, 4); + const uint32_t sign = (f >> 16) & 0x8000u; + const int e = static_cast((f >> 23) & 0xffu); + uint32_t m = f & 0x7fffffu; + if(e == 255) + return static_cast(sign | 0x7c00u | (m ? 1u : 0u)); // inf/nan + int ee = e - 127 + 15; + if(ee >= 31) + return static_cast(sign | 0x7c00u); // overflow -> inf + if(ee <= 0) + { + if(ee < -10) + return static_cast(sign); // underflow -> zero + m |= 0x800000u; + const int shift = 14 - ee; + const uint32_t half = 1u << (shift - 1); + const uint32_t rem = m & ((1u << shift) - 1); + m >>= shift; + if(rem > half || (rem == half && (m & 1u))) + m++; // RNE + if(m & 0x400u) + return static_cast(sign | (1u << 10)); // carried to min normal + return static_cast(sign | m); + } + const uint32_t half = 1u << 12; + const uint32_t rem = m & 0x1fffu; + m >>= 13; + if(rem > half || (rem == half && (m & 1u))) + { + m++; + if(m == 0x400u) + { + m = 0; + ee++; + if(ee >= 31) + return static_cast(sign | 0x7c00u); + } + } + return static_cast(sign | (static_cast(ee) << 10) | m); +} + +static inline float f16_bits_to_f32(uint16_t h) +{ + const uint32_t sign = static_cast(h & 0x8000u) << 16; + const int e = (h >> 10) & 0x1f; + uint32_t m = h & 0x3ffu; + uint32_t f; + if(e == 0) + { + if(m == 0) + f = sign; + else + { + int ee = -14; + while(!(m & 0x400u)) + { + m <<= 1; + ee--; + } + m &= 0x3ffu; + f = sign | (static_cast(ee + 127) << 23) | (m << 13); + } + } + else if(e == 31) + f = sign | 0x7f800000u | (m << 13); + else + f = sign | (static_cast(e - 15 + 127) << 23) | (m << 13); + float out; + std::memcpy(&out, &f, 4); + return out; +} + +static inline ck_tile::fp16_t h_fp16(float x) +{ + const uint16_t bits = f32_to_f16_bits(x); + ck_tile::fp16_t h; + std::memcpy(&h, &bits, 2); + return h; +} + +static inline float h_f32(ck_tile::fp16_t h) +{ + uint16_t bits; + std::memcpy(&bits, &h, 2); + return f16_bits_to_f32(bits); +} + +#define HIP_CHECK(cmd) \ + do \ + { \ + hipError_t e = (cmd); \ + if(e != hipSuccess) \ + { \ + printf("HIP error %s at %s:%d\n", hipGetErrorString(e), __FILE__, __LINE__); \ + exit(1); \ + } \ + } while(0) + +// ---------------------------------------------------------------- types +template +using S = ck_tile::sequence; + +// Block: gemm0 32x256x32 (warp 32x64x32 x 4), gemm1 32x64x256 (warp 32x16x256 x 4) +// config from the fork's fused_moe setup (flatmm_ex engine): O=fp32 + OAtomic=2, +// Gemm0NInterleave=4, swizzled=false, PadHiddenSize=true +using FmoeShape = ck_tile::FusedMoeGemmShape, + S<1, 4, 1>, + S<32, 64, 32>, + S<32, 64, 256>, + S<1, 4, 1>, + S<32, 16, 256>>; + +using FmoeTraits = ck_tile::FusedMoeGemmTraits, // gmem load vec + S<4>, // bridge smem store vec + S<8, 8, 8, 4>, // smem load vec + 1, // o gmem store vec + 4, // gemm0 interleave + 1, // gemm1 interleave + ck_tile::FusedMoeGemmWeightPermuteEnum::no_permute, + false, // swizzled + true, // pad hidden size + false>; // pad intermediate size + +using FmoeProblem = ck_tile::FusedMoeGemmPipelineProblem; + +using FmoePipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; +using FmoePartitioner = ck_tile::FusedMoeGemmTilePartitioner_Linear; +using FmoeKernel = ck_tile::FusedMoeGemmUKKernel; + +constexpr int BlockM0 = 32; + +// ---------------------------------------------------------------- host sort +struct MoeSortResult +{ + std::vector sorted_token_ids; // base-offset encoding + std::vector sorted_weight; + std::vector sorted_expert_ids; // per BlockM0 tile + int32_t num_sorted_tiles; // in tiles of BlockM0 + int padded_rows = 0; +}; + +static MoeSortResult host_sort(const std::vector>& topk_ids, + const std::vector>& topk_w, + int tokens, + int experts, + int topk, + int enc_mode) +{ + std::vector> rows_of_expert(experts); + std::vector> w_of_expert(experts); + std::vector> kidx_of_expert(experts); // topk slot, for enc=2 + for(int t = 0; t < tokens; t++) + for(int k = 0; k < topk; k++) + { + int e = topk_ids[t][k]; + rows_of_expert[e].push_back(t); + w_of_expert[e].push_back(topk_w[t][k]); + kidx_of_expert[e].push_back(k); + } + + MoeSortResult r; + for(int e = 0; e < experts; e++) + { + int n = static_cast(rows_of_expert[e].size()); + if(n == 0) + continue; + int padded = (n + BlockM0 - 1) / BlockM0 * BlockM0; + for(int i = 0; i < padded; i++) + { + if(i < n) + { + int tok = rows_of_expert[e][i]; + int kid = kidx_of_expert[e][i]; + // enc=0: base-offset ids; enc=2: repo mock id (token | k<<24), + // what moe_sorting emits when CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID=1 + // (the repo default); enc=1: raw token ids. + if(enc_mode == 0) + r.sorted_token_ids.push_back(tok - (r.padded_rows + i)); + else if(enc_mode == 2) + r.sorted_token_ids.push_back((tok & 0x00ffffff) | ((kid & 0xff) << 24)); + else + r.sorted_token_ids.push_back(tok); + r.sorted_weight.push_back(w_of_expert[e][i]); + } + else + { + // padding: invalid token sentinel + weight 0 -> zero contribution + if(enc_mode == 0) + r.sorted_token_ids.push_back(-(r.padded_rows + i)); + else if(enc_mode == 2) + r.sorted_token_ids.push_back((tokens & 0x00ffffff) | ((topk & 0xff) << 24)); + else + r.sorted_token_ids.push_back(tokens); + r.sorted_weight.push_back(0.0f); + } + } + for(int i = 0; i < padded / BlockM0; i++) r.sorted_expert_ids.push_back(e); + r.padded_rows += padded; + } + r.num_sorted_tiles = r.padded_rows / BlockM0; + return r; +} + +// ---------------------------------------------------------------- weight shuffle +// fork fused_moe contract (mfma_type=1): [b, n, k] -> [b, n/16, k/32, 4, 16, 8] +// (i.e. per (n16, k32) block: [k%32/8 : 4][n%16 : 16][k%8 : 8] contiguous) +template +static std::vector shuffle_moe_weight(const std::vector& w, int b_, int n_, int k_) +{ + std::vector out(w.size()); + for(int b = 0; b < b_; b++) + for(int n = 0; n < n_; n++) + for(int k = 0; k < k_; k++) + { + size_t src = (static_cast(b) * n_ + n) * k_ + k; + size_t dst = + (((static_cast(b) * (n_ / 16) + n / 16) * (k_ / 32) + k / 32) * 4 + + (k % 32) / 8) * 128 + + (n % 16) * 8 + (k % 8); + out[dst] = w[src]; + } + return out; +} + +// ---------------------------------------------------------------- main +int main(int argc, char* argv[]) +{ + int tokens = argc > 1 ? atoi(argv[1]) : 256; + int experts = argc > 2 ? atoi(argv[2]) : 8; + int topk = argc > 3 ? atoi(argv[3]) : 2; + int hidden = argc > 4 ? atoi(argv[4]) : 512; + int interm = argc > 5 ? atoi(argv[5]) : 512; + int do_valid = argc > 6 ? atoi(argv[6]) : 1; + int warmup = argc > 7 ? atoi(argv[7]) : 5; + int repeat = argc > 8 ? atoi(argv[8]) : 20; + int init_mode = argc > 9 ? atoi(argv[9]) : 0; + uint32_t seed = argc > 10 ? static_cast(atoi(argv[10])) : 11939; + int enc_mode = argc > 11 ? atoi(argv[11]) : 1; // 1: raw token ids (default), 0: base-offset ids + + printf("fused_moe flatmm_uk(HCU): t:%d e:%d k:%d h:%d i:%d v:%d init:%d\n", + tokens, + experts, + topk, + hidden, + interm, + do_valid, + init_mode); + printf("kernel: %s\n", FmoeKernel::GetName().c_str()); + + // Load-vector alignment contract: the engine's A/G/U global loads + // vectorize along hidden (Policy::GetAlignment_{A,G}() elements; + // 4/8 for this fp16 config), and ck_tile masks out-of-range accesses + // per WHOLE vector - a hidden not divisible by the vector length lets + // the first vector's tail read past the row end into the next row's + // data (verified: hidden=1/2/3/6 give silent wrong results, multiples + // of 4 are bit-clean). Same contract as ck_tile GEMM's K alignment: + // the caller pads. Everything device-visible below uses hidden_p as + // the hidden stride with zeros in [hidden, hidden_p); the reference + // still models and validates only the logical columns. + constexpr int align_a = FmoePipeline::Engine::Policy::template GetAlignment_A(); + constexpr int align_g = FmoePipeline::Engine::Policy::template GetAlignment_G(); + const int align_h = align_a > align_g ? align_a : align_g; + const int hidden_p = (hidden + align_h - 1) / align_h * align_h; + if(hidden_p != hidden) + printf("note: hidden %d padded to %d (load-vector alignment: A=%d G=%d)\n", + hidden, + hidden_p, + align_a, + align_g); + + // ---- host data (natural layouts, no shuffle) + lcg_seed(seed); + std::vector a(static_cast(tokens) * hidden); + if(init_mode == 0) + for(auto& v : a) v = h_fp16(lcg_float() * 2.0f); + else if(init_mode == 2) + for(auto& v : a) v = h_fp16(1.0f); + else + for(size_t i = 0; i < a.size(); i++) + a[i] = h_fp16(0.01f * static_cast(i % 97)); + + std::vector w1(static_cast(experts) * 2 * interm * hidden); + std::vector w2(static_cast(experts) * hidden * interm); + if(init_mode == 0) + { + for(auto& v : w1) v = lcg_float(); + for(auto& v : w2) v = lcg_float(); + } + else if(init_mode == 2) + { + for(auto& v : w1) v = 1.0f; + for(auto& v : w2) v = 1.0f; + } + else if(init_mode == 3) + { + // coordinate-decodable D: value*100 = n*8 + k%8 + 1 + for(auto& v : w1) v = 0.01f; + for(int e = 0; e < experts; e++) + for(int n = 0; n < hidden; n++) + for(int k = 0; k < interm; k++) + w2[(static_cast(e) * hidden + n) * interm + k] = + 0.01f * (n * 8 + k % 8 + 1); + } + else + { + for(size_t i = 0; i < w1.size(); i++) + w1[i] = 0.001f * static_cast((i * 7919) % 113) - 0.05f; + for(size_t i = 0; i < w2.size(); i++) + w2[i] = 0.001f * static_cast((i * 104729) % 127) - 0.06f; + } + if(init_mode == 4) + { + // full-chain decode: a=1, gate row n = 0.001*(n+1), up = 1 -> + // y[n] = silu(0.001*(n+1)*hidden) * hidden + // D selects interm column n == j for output column j -> + // o[t][j] = silu(0.001*(j+1)*hidden) * hidden (strictly increasing in j) + for(auto& v : a) v = h_fp16(1.0f); + for(int e = 0; e < experts; e++) + for(int n = 0; n < 2 * interm; n++) + for(int kk = 0; kk < hidden; kk++) + { + float g = (n < interm) ? 0.001f * (n + 1) : 1.0f; + w1[(static_cast(e) * 2 * interm + n) * hidden + kk] = g; + } + for(int e = 0; e < experts; e++) + for(int j = 0; j < hidden; j++) + for(int n = 0; n < interm; n++) + w2[(static_cast(e) * hidden + j) * interm + n] = + (n == j) ? 1.0f : 0.0f; + } + if(init_mode == 5) + { + // position-encoded D: value = 1000*row + col, so any in-kernel read of + // the D tile decodes to the exact (hidden row, interm col) it touched + for(int e = 0; e < experts; e++) + for(int j = 0; j < hidden; j++) + for(int n = 0; n < interm; n++) + w2[(static_cast(e) * hidden + j) * interm + n] = + 1000.0f * j + n; + } + if(init_mode == 6) + { + // separable decode: a[t] = 0.01*(t+1), gate row n = 0.001*(n+1), up = 1, + // D selects interm column n == j => o[t][j] ~ (t+1)^2 * (j+1): + // a wrong A-row gather permutes the (t+1)^2 factor, a wrong column + // mapping permutes the (j+1) factor + for(int t = 0; t < tokens; t++) + for(int kk = 0; kk < hidden; kk++) + a[static_cast(t) * hidden + kk] = h_fp16(0.01f * (t + 1)); + for(int e = 0; e < experts; e++) + for(int n = 0; n < 2 * interm; n++) + for(int kk = 0; kk < hidden; kk++) + { + float g = (n < interm) ? 0.001f * (n + 1) : 1.0f; + w1[(static_cast(e) * 2 * interm + n) * hidden + kk] = g; + } + for(int e = 0; e < experts; e++) + for(int j = 0; j < hidden; j++) + for(int n = 0; n < interm; n++) + w2[(static_cast(e) * hidden + j) * interm + n] = + (n == j) ? 1.0f : 0.0f; + } + + // topk assignment + std::vector> topk_ids(tokens, std::vector(topk)); + std::vector> topk_w(tokens, std::vector(topk)); + for(int t = 0; t < tokens; t++) + { + std::vector used(experts, false); + for(int k = 0; k < topk; k++) + { + int e = static_cast(std::abs(static_cast(lcg_float() * 1000))) % experts; + while(used[e]) e = (e + 1) % experts; + used[e] = true; + topk_ids[t][k] = e; + topk_w[t][k] = 0.5f + std::fabs(lcg_float()); + } + } + + MoeSortResult sorted = host_sort(topk_ids, topk_w, tokens, experts, topk, enc_mode); + printf("sorted: %d tiles (%d padded rows), experts used: %zu\n", + sorted.num_sorted_tiles, + sorted.padded_rows, + sorted.sorted_expert_ids.size()); + if(do_valid >= 2) + { + printf("sorted_ids:"); + for(int i = 0; i < sorted.padded_rows && i < 64; i++) + printf(" %d", sorted.sorted_token_ids[i]); + printf("\nsorted_w:"); + for(int i = 0; i < sorted.padded_rows && i < 64; i++) + printf(" %.2f", sorted.sorted_weight[i]); + printf("\ntopk_map:"); + for(int t = 0; t < tokens && t < 32; t++) + printf(" %d", topk_ids[t][0]); + printf("\n"); + } + + // output uses the padded hidden stride too (kernel writes zeros into the + // pad columns: D rows [hidden, hidden_p) are zero-padded below) + std::vector o(static_cast(tokens) * hidden_p, 0.0f); + + // fp16 device copies of weights, pre-shuffled to the fork's fused_moe + // mfma_type=1 layout (engine contract, see main.cpp of fork's fused_moe) + int shuf_mode = argc > 12 ? atoi(argv[12]) : 0; // 0:none(natural, default) 1:W1 2:W2 3:both + std::vector w1_h(w1.size()), w2_h(w2.size()); + for(size_t i = 0; i < w1.size(); i++) w1_h[i] = h_fp16(w1[i]); + for(size_t i = 0; i < w2.size(); i++) w2_h[i] = h_fp16(w2[i]); + // reference snapshots BEFORE shuffling (logical layout) + std::vector g1f(w1_h.size()), g2f(w2_h.size()); + for(size_t i = 0; i < w1_h.size(); i++) g1f[i] = h_f32(w1_h[i]); + for(size_t i = 0; i < w2_h.size(); i++) g2f[i] = h_f32(w2_h[i]); + { + // toolchain probe: the data path above no longer uses native _Float16 + // casts (see h_fp16 note); this one surviving cast only serves as a + // canary — native_roundtrip must read 0.770000, anything else means + // the DTK -c host miscompile is back and the run is not trustworthy. + ck_tile::fp16_t ph = h_fp16(0.77f); + uint16_t pbits; + std::memcpy(&pbits, &ph, 2); + printf("fp16 sanity: bits(0.77f)=0x%04x roundtrip=%.6f native_roundtrip=%.6f\n", + pbits, + h_f32(ph), + static_cast(static_cast(0.77f))); + } + if(shuf_mode & 1) + w1_h = shuffle_moe_weight(w1_h, experts, 2 * interm, hidden); + if(shuf_mode & 2) + w2_h = shuffle_moe_weight(w2_h, experts, hidden, interm); + + // padded upload staging: rows of A/W1 get hidden_p columns (zero tail), + // W2 gets hidden_p rows (zero tail). Pad values are zero so they add + // nothing to either gemm, and the zero W2 rows make the kernel write + // exact zeros into the pad columns of O. + std::vector a_p(static_cast(tokens) * hidden_p, + h_fp16(0.0f)); + std::vector w1_p(static_cast(experts) * 2 * interm * hidden_p, + h_fp16(0.0f)); + std::vector w2_p(static_cast(experts) * hidden_p * interm, + h_fp16(0.0f)); + for(int t = 0; t < tokens; t++) + for(int c = 0; c < hidden; c++) + a_p[static_cast(t) * hidden_p + c] = a[static_cast(t) * hidden + c]; + for(size_t i = 0; i < w1_h.size(); i++) + w1_p[(i / hidden) * static_cast(hidden_p) + i % hidden] = w1_h[i]; + for(int e = 0; e < experts; e++) + for(int r = 0; r < hidden; r++) + for(int c = 0; c < interm; c++) + w2_p[(static_cast(e) * hidden_p + r) * interm + c] = + w2_h[(static_cast(e) * hidden + r) * interm + c]; + + // ---- device buffers + void *d_a, *d_g, *d_d, *d_o, *d_sid, *d_sw, *d_seid, *d_nst; + HIP_CHECK(hipMalloc(&d_a, a_p.size() * 2)); + HIP_CHECK(hipMalloc(&d_g, w1_p.size() * 2)); + HIP_CHECK(hipMalloc(&d_d, w2_p.size() * 2)); + HIP_CHECK(hipMalloc(&d_o, o.size() * 4)); + HIP_CHECK(hipMalloc(&d_sid, sorted.sorted_token_ids.size() * 4)); + HIP_CHECK(hipMalloc(&d_sw, sorted.sorted_weight.size() * 4)); + HIP_CHECK(hipMalloc(&d_seid, sorted.sorted_expert_ids.size() * 4)); + HIP_CHECK(hipMalloc(&d_nst, 4)); + + HIP_CHECK(hipMemcpy(d_a, a_p.data(), a_p.size() * 2, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_g, w1_p.data(), w1_p.size() * 2, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_d, w2_p.data(), w2_p.size() * 2, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_sid, sorted.sorted_token_ids.data(), + sorted.sorted_token_ids.size() * 4, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_sw, sorted.sorted_weight.data(), + sorted.sorted_weight.size() * 4, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_seid, sorted.sorted_expert_ids.data(), + sorted.sorted_expert_ids.size() * 4, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_nst, &sorted.padded_rows, 4, hipMemcpyHostToDevice)); + + // TEMP-DEBUG probe (remove after diagnosis) + printf("DBG: hidden_p=%d a_p=%zu w1_p=%zu w2_p=%zu o=%zu\n", + hidden_p, a_p.size(), w1_p.size(), w2_p.size(), o.size()); + printf("DBG: a_p0=%.4f w1_p0=%.4f w2_p0=%.4f a_plast=%.4f\n", + h_f32(a_p[0]), h_f32(w1_p[0]), h_f32(w2_p[0]), h_f32(a_p.back())); + { + std::vector rt(8); + HIP_CHECK(hipMemcpy(rt.data(), d_g, 16, hipMemcpyDeviceToHost)); + printf("DBG: d_g[0..3]=%.4f %.4f %.4f %.4f\n", + h_f32(rt[0]), h_f32(rt[1]), h_f32(rt[2]), h_f32(rt[3])); + HIP_CHECK(hipMemcpy(rt.data(), d_a, 16, hipMemcpyDeviceToHost)); + printf("DBG: d_a[0..3]=%.4f %.4f %.4f %.4f\n", + h_f32(rt[0]), h_f32(rt[1]), h_f32(rt[2]), h_f32(rt[3])); + HIP_CHECK(hipMemcpy(rt.data(), d_d, 16, hipMemcpyDeviceToHost)); + printf("DBG: d_d[0..3]=%.4f %.4f %.4f %.4f\n", + h_f32(rt[0]), h_f32(rt[1]), h_f32(rt[2]), h_f32(rt[3])); + } + + // ---- kernel args (field order must match Kargs exactly) + FmoeKernel::Hargs hargs{}; + hargs.a_ptr = d_a; + hargs.g_ptr = d_g; + hargs.d_ptr = d_d; + hargs.o_ptr = d_o; + hargs.sorted_token_ids_ptr = d_sid; + hargs.sorted_weight_ptr = d_sw; + hargs.sorted_expert_ids_ptr = d_seid; + hargs.num_sorted_tiles_ptr = d_nst; + hargs.hidden_size = hidden_p; // padded: see load-vector note above + hargs.intermediate_size = interm; + hargs.num_tokens = tokens; + hargs.num_experts = experts; + hargs.topk = topk; + hargs.stride_token = hidden_p; + + auto kargs = FmoeKernel::MakeKargs(hargs); + auto grids = FmoeKernel::GridSize(hargs); + auto blk = FmoeKernel::BlockSize(); + + if(do_valid >= 2) + { + // debug: dump the y-dim structure of the C-block and O distributions + auto c_block = FmoePipeline::Engine::Policy::template MakeCBlockTile_Gemm1(); + auto o_dist = FmoePipeline::Engine::Policy::template MakeGlobalTileDistribution_O(); + auto c_lens = to_sequence( + c_block.get_tile_distribution().get_ys_to_d_descriptor().get_lengths()); + auto o_lens = to_sequence(o_dist.get_ys_to_d_descriptor().get_lengths()); + printf("C block y-lens:"); + ck_tile::static_for<0, c_lens.size(), 1>{}( + [&](auto j) { printf(" %d", static_cast(c_lens.at(j))); }); + printf("\nO dist y-lens:"); + ck_tile::static_for<0, o_lens.size(), 1>{}( + [&](auto j) { printf(" %d", static_cast(o_lens.at(j))); }); + printf("\n"); + } + + printf("grid: %ux%ux%u block: %u smem: %d bytes\n", + grids.x, + grids.y, + grids.z, + blk.x, + FmoePipeline::GetSmemSizeBytes()); + + ck_tile::stream_config sc; + auto launch = [&]() { + float t = ck_tile::launch_kernel( + sc, + ck_tile::make_kernel( + FmoeKernel{}, grids, blk, 0, kargs)); + (void)t; + }; + + for(int i = 0; i < warmup; i++) + { + HIP_CHECK(hipMemset(d_o, 0, o.size() * 4)); + launch(); + } + HIP_CHECK(hipDeviceSynchronize()); + + hipEvent_t e0, e1; + HIP_CHECK(hipEventCreate(&e0)); + HIP_CHECK(hipEventCreate(&e1)); + HIP_CHECK(hipEventRecord(e0)); + for(int i = 0; i < repeat; i++) + { + HIP_CHECK(hipMemset(d_o, 0, o.size() * 4)); + launch(); + } + HIP_CHECK(hipEventRecord(e1)); + HIP_CHECK(hipEventSynchronize(e1)); + float ms = 0; + HIP_CHECK(hipEventElapsedTime(&ms, e0, e1)); + ms /= repeat; + double flops = static_cast(tokens) * topk * (6.0 * hidden * interm); + printf("time: %.4f ms perf: %.2f TFlops\n", ms, flops / (ms * 1e-3) / 1e12); + + // ---- validation + int ret = 0; + if(do_valid) + { + std::vector o_h(o.size()); + HIP_CHECK(hipMemcpy(o_h.data(), d_o, o.size() * 4, hipMemcpyDeviceToHost)); + printf("DBG: o_h[0..3]=%.6f %.6f %.6f %.6f\n", o_h[0], o_h[1], o_h[2], o_h[3]); + + // compact the padded output back to logical columns; the pad columns + // are exact zeros on both sides, and everything below (reference, + // counters, dumps) works on the compact view + std::vector o_c(static_cast(tokens) * hidden); + for(int t = 0; t < tokens; t++) + for(int c = 0; c < hidden; c++) + o_c[static_cast(t) * hidden + c] = + o_h[static_cast(t) * hidden_p + c]; + + // fp32 reference + std::vector ref(static_cast(tokens) * hidden, 0.0f); + std::vector y_row(interm); + std::vector af(static_cast(tokens) * hidden); + for(size_t i = 0; i < af.size(); i++) af[i] = h_f32(a[i]); + // g1f/g2f already snapshotted from unshuffled fp16 weights + for(int t = 0; t < tokens; t++) + for(int k = 0; k < topk; k++) + { + int e = topk_ids[t][k]; + float w = topk_w[t][k]; + for(int n = 0; n < interm; n++) + { + float gacc = 0, uacc = 0; + size_t gr = (static_cast(e) * 2 * interm + n) * hidden; + size_t ur = (static_cast(e) * 2 * interm + interm + n) * hidden; + for(int kk = 0; kk < hidden; kk++) + { + gacc += af[static_cast(t) * hidden + kk] * g1f[gr + kk]; + uacc += af[static_cast(t) * hidden + kk] * g1f[ur + kk]; + } + // mirror the engine's fp16 rounding chain: acc->fp16 before + // activation, activation out ->fp16, product ->fp16 + float g0 = h_f32(h_fp16(gacc)); + float u0 = h_f32(h_fp16(uacc)); + float sg = g0 / (1.0f + std::exp(-g0)); + float sgh = h_f32(h_fp16(sg)); + y_row[n] = h_f32(h_fp16(sgh * u0)); + } + for(int nn = 0; nn < hidden; nn++) + { + float acc = 0; + size_t dr = (static_cast(e) * hidden + nn) * interm; + for(int kk = 0; kk < interm; kk++) acc += y_row[kk] * g2f[dr + kk]; + ref[static_cast(t) * hidden + nn] += w * acc; + } + } + + double max_abs = 0, max_rel = 0, sum_sq = 0, sum_ref = 0; + long bad = 0, nonfinite_o = 0, nonzero_ref = 0, nonzero_o = 0; + size_t n = ref.size(); + for(size_t i = 0; i < n; i++) + { + double d = std::fabs(static_cast(o_c[i]) - ref[i]); + double denom = std::fabs(ref[i]); + if(denom < 1e-2) + denom = 1e-2; + double r = d / denom; + if(d > max_abs) + max_abs = d; + if(r > max_rel) + max_rel = r; + sum_sq += d * d; + sum_ref += ref[i] * ref[i]; + if(d > 3e-2 + 3e-2 * std::fabs(ref[i])) + bad++; + // degeneracy counters: NaN comparisons above are all false, so a + // NaN output or an all-zero reference would otherwise pass vacuously + if(std::isnan(o_c[i])) + nonfinite_o++; + if(ref[i] != 0.0f) + nonzero_ref++; + if(o_c[i] != 0.0f) + nonzero_o++; + } + printf("verify: max_abs=%.5f max_rel=%.5f rmse=%.6f ref_rms=%.6f bad=%ld/%zu\n", + max_abs, + max_rel, + std::sqrt(sum_sq / n), + std::sqrt(sum_ref / n), + bad, + n); + if(do_valid >= 2) + { + // dump first rows to inspect error structure (index-guarded: + // degenerate shapes can be smaller than the fixed window) + for(int t = 0; t < 2; t++) + { + printf("row %d o :", t); + for(int c = 0; c < 24; c++) + if(static_cast(t) * hidden + c < o_c.size()) + printf(" %12.4f", o_c[static_cast(t) * hidden + c]); + printf("\nrow %d ref:", t); + for(int c = 0; c < 24; c++) + if(static_cast(t) * hidden + c < o_c.size()) + printf(" %12.4f", ref[static_cast(t) * hidden + c]); + printf("\n"); + } + // ratio stats over non-tiny refs + double rsum = 0; + long rcnt = 0; + for(size_t i = 0; i < n; i++) + if(std::fabs(ref[i]) > 1e-1 && std::fabs(o_c[i]) > 1e-3) + { + rsum += o_c[i] / ref[i]; + rcnt++; + } + if(rcnt > 0) + printf("mean o/ref ratio over %ld elems: %.5f\n", rcnt, rsum / rcnt); + else + printf("no ratio samples (output near zero?)\n"); + // per-row bad counts + for(int t = 0; t < tokens && t < 40; t++) + { + long rb = 0; + for(int c = 0; c < hidden; c++) + { + size_t i = static_cast(t) * hidden + c; + if(std::fabs(o_c[i] - ref[i]) > 3e-2 + 3e-2 * std::fabs(ref[i])) + rb++; + } + if(rb > 0) + printf("row %d: %d/%d bad, o[0]=%.4f ref[0]=%.4f\n", + t, + static_cast(rb), + hidden, + o_c[static_cast(t) * hidden], + ref[static_cast(t) * hidden]); + } + // column pattern of row 0: compact bad ranges + printf("row0 bad ranges:"); + int c0 = -1; + for(int c = 0; c <= hidden; c++) + { + size_t i = c; + bool badc = c < hidden && + std::fabs(o_c[i] - ref[i]) > 3e-2 + 3e-2 * std::fabs(ref[i]); + if(badc && c0 < 0) + c0 = c; + else if(!badc && c0 >= 0) + { + printf(" [%d,%d)", c0, c); + c0 = -1; + } + } + printf(" o[48]=%.3f o[80]=%.3f o[112]=%.3f\n", + o_c.size() > 48 ? o_c[48] : 0.0f, + o_c.size() > 80 ? o_c[80] : 0.0f, + o_c.size() > 112 ? o_c[112] : 0.0f); + } + printf("sanity: nan_o=%ld nonzero_ref=%ld nonzero_o=%ld\n", nonfinite_o, nonzero_ref, nonzero_o); + if(nonfinite_o > 0 || nonzero_ref == 0 || (nonzero_ref > 0 && nonzero_o == 0)) + { + printf("FAIL (degenerate output/reference; the tolerance check above is vacuous)\n"); + ret = 1; + return ret; + } + if(bad == 0) + { + printf("PASS\n"); + } + else + { + printf("FAIL (%ld elems exceed tol)\n", bad); + ret = 1; + } + } + + static_cast(hipFree(d_a)); + static_cast(hipFree(d_g)); + static_cast(hipFree(d_d)); + static_cast(hipFree(d_o)); + static_cast(hipFree(d_sid)); + static_cast(hipFree(d_sw)); + static_cast(hipFree(d_seid)); + static_cast(hipFree(d_nst)); + return ret; +} diff --git a/example_hcu/ck_tile/17_fused_moe/instances/fused_moegemm_api_internal.hpp b/example_hcu/ck_tile/17_fused_moe/instances/fused_moegemm_api_internal.hpp index 4f9b408..f04b475 100644 --- a/example_hcu/ck_tile/17_fused_moe/instances/fused_moegemm_api_internal.hpp +++ b/example_hcu/ck_tile/17_fused_moe/instances/fused_moegemm_api_internal.hpp @@ -126,7 +126,9 @@ float fused_moegemm_swizzled_lds_(const ck_tile::stream_config& s, fused_moegemm Ts_::Gemm0NInterleave, Ts_::Gemm1NInterleave, ck_tile::FusedMoeGemmWeightPermuteEnum::no_permute, - Ts_::swizzled>; + Ts_::swizzled, + true /* pad hidden size */, + false /* pad intermediate size */>; using f_shape = ck_tile::FusedMoeGemmShape +#include + +namespace ck_tile { + +template +struct FusedMoeGemmUKKernel +{ + using Partitioner = remove_cvref_t; + using Pipeline = remove_cvref_t; + using Epilogue = remove_cvref_t; // TODO: not used + + using BlockShape = typename Pipeline::BlockShape; + static constexpr index_t BlockSize_ = BlockShape::BlockSize; + + using ADataType = typename Pipeline::Problem::ADataType; + using GDataType = typename Pipeline::Problem::GDataType; + using DDataType = typename Pipeline::Problem::DDataType; + using AccDataType = typename Pipeline::Problem::AccDataType; + using ODataType = typename Pipeline::Problem::ODataType; + using AScaleDataType = typename Pipeline::Problem::AScaleDataType; + using GScaleDataType = typename Pipeline::Problem::GScaleDataType; + using DScaleDataType = typename Pipeline::Problem::DScaleDataType; + using YSmoothScaleDataType = typename Pipeline::Problem::YSmoothScaleDataType; + using TopkWeightDataType = typename Pipeline::Problem::TopkWeightDataType; + using IndexDataType = typename Pipeline::Problem::IndexDataType; + using YDataType = typename Pipeline::Problem::YDataType; + + using Traits = typename Pipeline::Problem::Traits; + static constexpr bool UseUK = true; + + static constexpr bool IsGateOnly = Traits::IsGateOnly; + static constexpr bool UseSmoothQuant = Traits::UseSmoothQuant; + static constexpr bool PadHiddenSize = Traits::PadHiddenSize; + static constexpr bool PadIntermediateSize = Traits::PadIntermediateSize; + + // clang-format off + template struct t2s; + template <> struct t2s { static constexpr const char * name = "fp32"; }; + template <> struct t2s { static constexpr const char * name = "fp16"; }; + template <> struct t2s { static constexpr const char * name = "bf16"; }; + // clang-format on + + CK_TILE_HOST static std::string GetName() + { +#define _SS_ std::string +#define _TS_ std::to_string + using S_ = BlockShape; + + return _SS_("fused_moe_") + _SS_(t2s::name) + "_uk_" + + _TS_(S_::Block_M0) + "x" + _TS_(S_::Block_N0) + "x" + _TS_(S_::Block_K0) + "x" + + _TS_(S_::Block_N1) + "_" + _TS_(S_::WarpPerBlock_M0) + "x" + + _TS_(S_::WarpPerBlock_N0) + "x" + _TS_(S_::WarpPerBlock_K0) + "_" + _TS_(S_::Warp_M0) + + "x" + _TS_(S_::Warp_N0) + "x" + _TS_(S_::Warp_K0) + "_" + _SS_(Pipeline::name); +#undef _SS_ +#undef _TS_ + } + + struct FusedMoeGemmKargs + { + const void* a_ptr; // [m, k], input token + const void* a_scale_ptr; // [m, 1], token scale + const void* g_ptr; // [e, n, k]/[e, 2*n, k], pre-shuffled + const void* d_ptr; // [e, n, k], pre-shuffled + const void* g_scale_ptr; // [e, 1, n], gate(up) scale + const void* d_scale_ptr; // [e, 1, k], down scale + const void* g_zp_ptr; // [e, 2*n, k/group], gate(up) zero-point + const void* d_zp_ptr; // [e, k, n/group], down zero-point + const void* y_smooth_scale_ptr; // [e, 1, n], smooth-quant-scale + void* o_ptr; // [m, k], output token + void* deterministic_tmp_ptr; // reserved + void* deterministic_counter_ptr; + + const void* sorted_token_ids_ptr; + const void* sorted_weight_ptr; + const void* sorted_expert_ids_ptr; + const void* num_sorted_tiles_ptr; + + index_t hidden_size; // k + index_t intermediate_size; // n / TP, for Gate/Up/Down + index_t num_tokens; // input number of tokens + index_t num_experts; // number of groups + index_t topk; + + index_t stride_token; // for input/output, stride per row + index_t block_shape_n; + index_t block_shape_k; + }; + + using Kargs = FusedMoeGemmKargs; + using Hargs = FusedMoeGemmHostArgs; + + CK_TILE_HOST static constexpr Kargs MakeKargs(const Hargs& hargs) + { + return bit_cast(hargs); + } + + CK_TILE_HOST static constexpr auto GridSize(const Hargs& hargs) + { + constexpr index_t block_m = BlockShape::Block_M0; + int max_num_tokens_padded = + hargs.topk * hargs.num_tokens + hargs.num_experts * block_m - hargs.topk; + return Partitioner::GridSize(max_num_tokens_padded, hargs.intermediate_size); + } + + CK_TILE_HOST static constexpr auto BlockSize() { return dim3(BlockSize_); } + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return Pipeline::GetSmemSize(); } + + CK_TILE_DEVICE void operator()(Kargs kargs) const + { + // GetSmemSize() is an element count, not bytes — allocating a char + // array of that size truncates the LDS buffer to half and pushes the + // warps-2..3 N-slices of the gemm1 D tile out of bounds. + __shared__ CK_TILE_LDS_ADDR char smem[Pipeline::GetSmemSizeBytes()]; + + IndexDataType num_sorted_tiles = __builtin_amdgcn_readfirstlane( + *reinterpret_cast(kargs.num_sorted_tiles_ptr)); + + num_sorted_tiles = num_sorted_tiles / BlockShape::Block_M0; + + const auto [sorted_tile_id, intermediate_tile_id] = + Partitioner{}(num_sorted_tiles, kargs.intermediate_size); + if(sorted_tile_id >= num_sorted_tiles) + return; + + Pipeline{}(kargs, smem, sorted_tile_id, intermediate_tile_id); + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_flatmm_uk_engine.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_flatmm_uk_engine.hpp new file mode 100644 index 0000000..cfd4fb2 --- /dev/null +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_flatmm_uk_engine.hpp @@ -0,0 +1,966 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2026 Hygon Info Technologies Ltd. +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp" + +namespace ck_tile { + +/* +This pipeline deal with a gemm(actually 2 gemm) with one very small(token), one very big(weight) +we need to design the pipeline such that all waves along gemm-N dim (gemm-m only 1 wave) + + <----- gemm-N ------> + +----+----+----+----+ + | w0 | w1 | w2 | w3 | gemm-m + +----+----+----+----+ +*/ +template +struct FusedMoeGemmPipeline_FlatmmUkEngine +{ + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + + using BlockShape = typename Problem::BlockShape; // this is FusedMoeGemmShape + + using ADataType = typename Problem::ADataType; + using GDataType = typename Problem::GDataType; + using DDataType = typename Problem::DDataType; + using AccDataType = typename Problem::AccDataType; + using ODataType = typename Problem::ODataType; + using AScaleDataType = typename Problem::AScaleDataType; + using GScaleDataType = typename Problem::GScaleDataType; + using DScaleDataType = typename Problem::DScaleDataType; + using YSmoothScaleDataType = typename Problem::YSmoothScaleDataType; + using TopkWeightDataType = typename Problem::TopkWeightDataType; + using IndexDataType = typename Problem::IndexDataType; + using YDataType = typename Problem::YDataType; + + using Traits = typename Problem::Traits; + + static constexpr bool IsGateOnly = Traits::IsGateOnly; + static constexpr bool UseSmoothQuant = Traits::UseSmoothQuant; + static constexpr bool PadHiddenSize = Traits::PadHiddenSize; + static constexpr bool PadIntermediateSize = Traits::PadIntermediateSize; + + static constexpr index_t kAlignmentA = Policy::template GetAlignment_A(); + static constexpr index_t kAlignmentG = Policy::template GetAlignment_G(); + static constexpr index_t kAlignmentD = Policy::template GetAlignment_D(); + static constexpr index_t kAlignmentO = Policy::template GetAlignment_O(); + static constexpr index_t NumPrefetch = Problem::NumPrefetch; + + static constexpr index_t SLD_A = static_cast(FusedMoeGemmPipelineSequencerEnum::SLD_A); + static constexpr index_t GLD_A = static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_A); + static constexpr index_t GLD_B = static_cast(FusedMoeGemmPipelineSequencerEnum::GLD_B); + static constexpr index_t GST_O = static_cast(FusedMoeGemmPipelineSequencerEnum::GST_O); + + static constexpr index_t kBlockPerCu = []() { + if constexpr(Problem::kBlockPerCu != -1) + return Problem::kBlockPerCu; + else + { + // minimize occupancy + return 2; + } + }(); + + static constexpr const char* name = "fused_moe_flatmm"; + static constexpr bool kSupportsDeterministicFp64Output = true; + + // TODO: there are multiple buffers + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize_A() + { + return Policy::template GetSmemSize_A(); + } + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + { + return Policy::template GetSmemSize(); + } + + // this is the thread-offset along row/col + CK_TILE_HOST_DEVICE static auto GetACoord() + { + constexpr auto a_dist = Policy::template MakeGlobalTileDistribution_A(); + const auto a_coord = a_dist.calculate_index(); // tile distribution index + return a_coord; + } + + // this is the thread-offset along row/col + CK_TILE_HOST_DEVICE static auto GetOCoord() + { + constexpr auto o_dist = Policy::template MakeGlobalTileDistribution_O(); + const auto o_coord = o_dist.calculate_index(); + return o_coord; + } + + template + CK_TILE_DEVICE auto + operator()(const AWindow& a_window_, + const GWindow& g_window_, + const DWindow& d_window_, + const Epilogue& epilogue_, + CK_TILE_LDS_ADDR void* smem, + index_t hidden_size, + index_t intermediate_size, + index_t intermediate_col_base) + { + _Pragma("clang diagnostic push") _Pragma("clang diagnostic ignored \"-Wc++20-extensions\""); + constexpr auto I0 = number<0>{}; + constexpr auto I1 = number<1>{}; + constexpr auto I2 = number<2>{}; + constexpr auto I3 = number<3>{}; + + CK_TILE_LDS_ADDR ADataType* smem_0 = reinterpret_cast(smem); + + // real weight rows/columns of this intermediate tile (gemm0 N / gemm1 K + // extent minus the tile's column offset). Partial tiles (intermediate + // size not a multiple of Block_N0) must zero-fill beyond this count: + // gate rows would otherwise read the up half / the next expert, and the + // manual D load would read the next row's leading columns (or past the + // W2 buffer for the last expert) - silently wrong results. + const index_t interm_valid = intermediate_size - intermediate_col_base; + + auto g_view = g_window_.get_bottom_tensor_view(); + + // + auto u_view = [&]() { + if constexpr(IsGateOnly) + { + // FIXME: + static_assert(false, "GateOnly is not supported yet!"); + return g_view; + } + else + { + const GDataType* g_ptr = + g_window_.get_bottom_tensor_view().get_buffer_view().p_data_; + + const GDataType* u_ptr = g_ptr + intermediate_size * hidden_size; + + // declare only the real up-half rows; the pad transform + // zero-fills the tail of a partial tile (see interm_valid note) + const auto u_view_ = make_naive_tensor_view( + u_ptr, + make_tuple(interm_valid, hidden_size), + make_tuple(hidden_size, 1), + number{}, + number<1>{}); + + const auto u_view_1_ = pad_tensor_view( + u_view_, + make_tuple(number{}, + number<2 * BlockShape::Block_K0>{}), + sequence{}); + return u_view_1_; + } + }(); + + // input global load to lds components + auto a_copy_dram_window = make_tile_window( + a_window_.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_window_.get_window_origin(), + Policy::template MakeGlobalTileDistribution_A()); + constexpr auto a_copy_lds_block_desc = Policy::template MakeLdsStoreDesc_A(); + constexpr auto a_lds_space_size = ck_tile::integer_least_multiple( + Policy::template GetSmemSize_A(), BlockShape::Block_K0); + // constexpr auto a_lds_bufs_elem_offset = 0; + auto a_copy_lds_window0 = make_tile_window( + make_tensor_view(smem_0, a_copy_lds_block_desc), + make_tuple(number{}, number{}), + {0, 0}); + auto a_copy_lds_window1 = make_tile_window( + make_tensor_view(smem_0 + a_lds_space_size, + a_copy_lds_block_desc), + make_tuple(number{}, number{}), + {0, 0}); + auto a_copy_lds_windows = make_tuple(a_copy_lds_window0, a_copy_lds_window1); + + // gate load from global to vgpr then to lds components + auto g_copy_dram_window = make_tile_window( + g_window_.get_bottom_tensor_view(), // [hidden_states, moe_intermediate_size] + make_tuple(number{}, + number{}), //[BLOCK_N0, BLOKC_K0] + g_window_.get_window_origin(), + Policy::template MakeGlobalTileDistribution_G()); + + /* + //if need write to lds, use this part + constexpr auto g_copy_lds_block_desc = Policy::template + MakeLdsStoreDesc_G(); constexpr auto g_lds_space_size = + Policy::template GetSmemSize_G(); constexpr auto g_lds_bufs_elem_offset = + a_lds_space_size * NumPrefetch; auto g_copy_lds_block = + make_tensor_view(smem_0 + a_lds_bufs_elem_offset, + g_copy_lds_block_desc); + auto g_copy_lds_window = make_tile_window( + g_copy_lds_block, make_tuple(number{}, + number{}), {0, 0}); + */ + + // up proj reuse gate LDS and same as gate proj, all vgpr and lds layout + auto u_copy_dram_window = make_tile_window( + u_view, + make_tuple(number{}, number{}), + {0, 0}, + Policy::template MakeGlobalTileDistribution_G()); + /* + // For LDS flow + constexpr auto u_copy_lds_block_desc = Policy::template + MakeLdsStoreDesc_G(); constexpr auto u_lds_space_size = + Policy::template GetSmemSize_G(); constexpr auto g_lds_bufs_elem_offset = + a_lds_space_size * NumPrefetch; auto g_copy_lds_block = + make_tensor_view(smem_0 + a_lds_bufs_elem_offset, + g_copy_lds_block_desc); + auto g_copy_lds_window = make_tile_window( + g_copy_lds_block, make_tuple(number{}, + number{}), {0, 0}); + */ + + // gemm0 from LDS to VGPR + // TODO: support get_sub_warp_id for WASP + using WarpGemm0 = decltype(Policy::template GetWarpGemm0()); + auto warp_gemm_0 = WarpGemm0{}; + // input + // tile_windows here are for warps + constexpr auto a_gemm_sld_block_desc = Policy::template MakeLdsLoadDesc_A(); + auto a_sld_win0 = [&]() { + using WG = WarpGemm0; + auto a_gemm_sld_block = + make_tensor_view(smem_0, a_gemm_sld_block_desc); + return make_tile_window( + a_gemm_sld_block, + a_gemm_sld_block_desc.get_lengths(), + {0, 0}, + make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + }(); + auto a_sld_win1 = [&]() { // Prefetch window + using WG = WarpGemm0; + auto a_gemm_sld_block = make_tensor_view( + smem_0 + a_lds_space_size, a_gemm_sld_block_desc); + return make_tile_window( + a_gemm_sld_block, + a_gemm_sld_block_desc.get_lengths(), + {0, 0}, + make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); + }(); + + // weight + /* + //TODO: Enable for LDS solution + constexpr auto gu_gemm_sld_block_desc = Policy::template MakeLdsLoadDesc_G(); + auto gu_gemm_sld_block = make_tensor_view(smem_0 + + g_lds_bufs_elem_offset, g_gemm_sld_block_desc); auto gu_gemm_sld_block_window = + make_tile_window(g_gemm_sld_block, g_gemm_sld_block_desc.get_lengths(), {0, 0}); warp window + auto gu_sld_win = [&](){ + using WG = WarpGemm0; + return make_tile_window( + gu_gemm_sld_block_window.get_bottom_tensor_view(), + make_tuple(number{}, number{}), + gu_gemm_sld_block_window.get_window_origin() + multi_index<2>{iNWarp * WG::kN, 0}, + make_static_tile_distribution(typename WG::BWarpDstrEncoding{})); + }(); + */ + + // gemm1 weight global to vgpr + // TODO: replace buffer vgpr to buffer lds + auto d_copy_dram_window = make_tile_window( + d_window_.get_bottom_tensor_view(), // [hidden_states, moe_intermediate_size] + make_tuple(number{}, + number{}), //[BLOCK_N1, BLOCK_K1] + d_window_.get_window_origin(), + Policy::template MakeGlobalTileDistribution_D()); + + // D tile moves with explicit pointer math (calculate_index() supplies + // the per-thread base): x0 = 8*y0 + (2*warp + lane/32), x1 = 8*(lane%32) + // + y1 for the global distribution; the LDS tile is packed [N1, K1] + // row-major; the warp-B read back is row = 16*warp + lane%16, col = + // 16*kiter + 4*(lane/16) + kper. (This bypasses tile-window + // load/store on this path; both layouts were verified on hw with + // position-encoded weights.) + const DDataType* d_gptr = d_window_.get_bottom_tensor_view().get_buffer_view().p_data_; + const auto d_idx_base = d_copy_dram_window.get_tile_distribution().calculate_index(); + + // gemm 1 vgpr to lds: the LDS tile is a packed [Block_N1, Block_K1] + // row-major block written/read with explicit pointer math below + CK_TILE_LDS_ADDR DDataType* smem_1 = reinterpret_cast(smem); + + // gemm1 lds to vgpr + using WarpGemm1 = decltype(Policy::template GetWarpGemm1()); + auto warp_gemm_1 = WarpGemm1{}; + + auto d_gemm_sld_block_desc = Policy::template MakeLdsLoadDesc_D(); + auto d_sld_win = [&]() { + using WG = WarpGemm1; + constexpr auto d_outer_dstr_enc = tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + constexpr auto d_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + d_outer_dstr_enc, typename WG::BWarpDstrEncoding{}); + auto d_gemm_sld_block = + make_tensor_view(smem_1, d_gemm_sld_block_desc); + auto d_gemm_sld_win = + make_tile_window(d_gemm_sld_block, + d_gemm_sld_block_desc.get_lengths(), + {0, 0}, + make_static_tile_distribution(d_block_dstr_encode)); + return d_gemm_sld_win; + }(); + + // output column offset of the current gemm1 N-iteration, advanced by + // atomic_add_o() after every [Block_M1, Block_N1] C tile + index_t n_col = 0; + + auto bridge_sst_win = [&]() { + return make_tile_window( + make_tensor_view( + reinterpret_cast(smem), + Policy::template MakeBridgeLdsStoreDesc()), + Policy::template MakeBridgeLdsStoreDesc().get_lengths(), + {0, 0}); + }(); + + auto bridge_sld_win = [&]() { + return make_tile_window(make_tensor_view( + reinterpret_cast(smem), + Policy::template MakeBridgeLdsLoadDesc()), + Policy::template MakeBridgeLdsLoadDesc().get_lengths(), + {0, 0}, + Policy::template MakeYTileDistribution()); + }(); + + // also OK with C array, 2 register buffer, prefetch u or prefetch g + // statically_indexed_array gs; + + // constexpr auto issues_a = number{}; + constexpr auto issues_g = number{}; + // constexpr auto issues_d = number{}; + // constexpr auto issues_o = number{}; + + // Both manually double-buffered tails below consume two tiles. Pad-enabled + // shapes must therefore round the loop count to an even number of tiles; + // merely rounding to one tile drops the last K0 tile for odd counts (for + // example 1232 / 32 -> 39, while the old loop executed only 38). + const index_t num_blocks_k0 = PadHiddenSize + ? ((hidden_size + 2 * BlockShape::Block_K0 - 1) / + (2 * BlockShape::Block_K0)) * + 2 + : (hidden_size + BlockShape::Block_K0 - 1) / + BlockShape::Block_K0; + const index_t num_blocks_n1 = PadHiddenSize + ? ((hidden_size + 2 * BlockShape::Block_N1 - 1) / + (2 * BlockShape::Block_N1)) * + 2 + : (hidden_size + BlockShape::Block_N1 - 1) / + BlockShape::Block_N1; + + using a_thread_type = decltype(load_tile(a_sld_win0)); + using g_thread_type = decltype(load_tile(g_copy_dram_window)); + using d_thread_type = decltype(load_tile(d_copy_dram_window)); + // using sld_d_thread_type = decltype(load_tile(d_sld_win)); + statically_indexed_array a; + statically_indexed_array gs; + statically_indexed_array gd; + // statically_indexed_array sd; + + auto gld_a_init = [&]() { a_copy_dram_window.init_raw(); }; + + auto gld_a = [&](auto& a_store_) { + async_load_tile_by_inline_asm(a_store_, a_copy_dram_window); + }; + auto move_a = [&]() { + move_tile_window(a_copy_dram_window, {number<0>{}, number{}}); + }; + auto sld_a = [&](auto& a_, auto& win_) { load_tile(a_, win_); }; + + auto gld_g_init = [&]() { g_copy_dram_window.init_raw(); }; + + auto gld_g = [&](auto& g_) { + if constexpr(IsGateOnly) + { + // TODO + } + load_tile(g_, g_copy_dram_window); + }; + auto move_g = [&]() { + move_tile_window(g_copy_dram_window, {number<0>{}, number{}}); + }; + + auto gld_u_init = [&]() { u_copy_dram_window.init_raw(); }; + + auto gld_u = [&](auto& u_) { + if constexpr(IsGateOnly) + { + load_tile(u_, g_copy_dram_window); + } + else + { + load_tile(u_, u_copy_dram_window); + } + }; + auto move_u = [&]() { + if constexpr(IsGateOnly) + { + move_tile_window(g_copy_dram_window, {number<0>{}, number{}}); + } + else + { + move_tile_window(u_copy_dram_window, {number<0>{}, number{}}); + } + }; + // statically_indexed_array ds; + auto gld_d_init = [&]() { d_copy_dram_window.init_raw(); }; + auto gld_d = [&](auto& d_) { +#ifdef FMOE_UK_DEBUG_G + if(get_warp_id() == 0 && (get_lane_id() == 0 || get_lane_id() == 32)) + printf("VEXP w=%d lane=%d org=(%d,%d) base=(%d,%d) is=%d hs=%d d0=%f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (int)d_copy_dram_window.get_window_origin()[number<0>{}], + (int)d_copy_dram_window.get_window_origin()[number<1>{}], + (int)d_idx_base[number<0>{}], + (int)d_idx_base[number<1>{}], + (int)intermediate_size, + (int)hidden_size, + (float)d_gptr[0]); +#endif + const auto org = d_copy_dram_window.get_window_origin(); + const index_t row0 = org[number<0>{}] + d_idx_base[number<0>{}]; + const index_t col0 = org[number<1>{}] + d_idx_base[number<1>{}]; + auto& buf = d_.get_thread_buffer(); + static_for<0, 64, 1>{}([&](auto j) { + constexpr index_t y0 = j / 8; + constexpr index_t y1 = j % 8; + const index_t r = row0 + 8 * y0; + const index_t c = col0 + y1; + // c is a tile-local column; guard it against the real interm + // extent (window base already carries intermediate_col_base) + buf[j] = (r < hidden_size && c < interm_valid) + ? d_gptr[static_cast(r) * intermediate_size + c] + : DDataType{0}; + }); +#ifdef FMOE_UK_DEBUG_G + if(get_warp_id() == 0 && (get_lane_id() == 0 || get_lane_id() == 32)) + printf("VEXP w=%d lane=%d org=(%d,%d) base=(%d,%d) is=%d hs=%d " + "d1=%f d256=%f d257=%f b0=%f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (int)d_copy_dram_window.get_window_origin()[number<0>{}], + (int)d_copy_dram_window.get_window_origin()[number<1>{}], + (int)d_idx_base[number<0>{}], + (int)d_idx_base[number<1>{}], + (int)intermediate_size, + (int)hidden_size, + (float)d_gptr[1], + (float)d_gptr[256], + (float)d_gptr[257], + (float)buf[number<0>{}]); +#endif + }; + // VGPR -> LDS, same (row, col) each element was gathered from + auto d_vgpr_to_lds = [&](auto& d_) { + auto& buf = d_.get_thread_buffer(); + static_for<0, 64, 1>{}([&](auto j) { + constexpr index_t y0 = j / 8; + constexpr index_t y1 = j % 8; + const index_t r = d_idx_base[number<0>{}] + 8 * y0; + const index_t c = d_idx_base[number<1>{}] + y1; + *(smem_1 + r * BlockShape::Block_K1 + c) = buf[j]; + }); + }; + // LDS -> VGPR in the embedded warp-B layout expected by gemm_1() + auto d_lds_to_vgpr = [&]() { + using DSl = decltype(load_tile(d_sld_win)); // type only, never executed + DSl d_tmp; + const index_t rr = 16 * get_warp_id() + get_lane_id() % 16; + const index_t c0 = 4 * (get_lane_id() / 16); + auto& buf = d_tmp.get_thread_buffer(); + static_for<0, 64, 1>{}([&](auto j) { + constexpr index_t kiter = j / 4; + constexpr index_t kper = j % 4; + buf[j] = *(smem_1 + rr * BlockShape::Block_K1 + 16 * kiter + c0 + kper); + }); + return d_tmp; + }; + auto move_d = [&]() { + // d move along gemm-n + move_tile_window(d_copy_dram_window, {number{}, number<0>{}}); + }; + + // Per-row topk-weighted fp32 atomic scatter of one [Block_M1, Block_N1] C tile. + // + // The fork's original epilogue drove this through [4, Block_N1] one-shot + // indexing-transform scatter windows plus a per-issue scalar weight, whose + // slot arithmetic (o_coord[0] + i*4) was derived for an MR1 C layout and + // the policy's output distribution hardcodes MRepeat=1. Under the HCU + // transposed-C warp layout (MR2: row = 16*MRepeat + 4*(lane/16) + CM1) + // that arithmetic misassigns rows to tokens. This port instead scatters + // explicitly: each thread owns 8 C elements with thread-buffer index + // d = MRepeat*4 + CM1PerLane-index, and the row's own sorted slot + // supplies both the destination token and its topk weight. + auto atomic_add_o = [&](auto& o_) { + using WarpGemm = typename remove_cvref_t::WarpGemmAttribute; + static_assert(WarpGemm::MRepeat == 2 && WarpGemm::NRepeat == 1 && + WarpGemm::MInterleave == 1 && WarpGemm::NInterleave == 1, + "epilogue row/col math locked to MR2/NR1/MI1/NI1"); + static_assert(BlockShape::Repeat_M1 == 1 && BlockShape::WarpPerBlock_M1 == 1 && + BlockShape::Repeat_N1 == 1 && BlockShape::WarpPerBlock_N1 == 4, + "epilogue row/col math locked to WarpPerBlock_1 = <1,4,1>"); + + const index_t cml = get_lane_id() / 16; + const index_t col = 16 * get_warp_id() + get_lane_id() % 16 + n_col; + const index_t row_b = epilogue_.row_base + cml; + auto& buf = o_.get_thread_buffer(); + static_for<0, 8, 1>{}([&](auto jj) { + constexpr index_t mrep = jj / 4; + constexpr index_t cm1 = jj % 4; + // HCU transposed-C layout (probed on hw): buffer index jj covers + // row 16*mrep + 4*cm1 + (lane/16); CMLane has stride 1 and the + // per-lane M1 repeat steps by 4 rows. + const index_t row = row_b + 16 * mrep + 4 * cm1; + index_t token = epilogue_.sorted_ids[row]; +#if CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID + // repo sorting packs token | (topk_idx << 24) when mock ids are + // enabled (CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID, default on in + // core/config.hpp); strip the tag exactly like + // FusedMoeGemmKernel does, otherwise every id looks out of + // range, the range guard drops ALL writes and o stays zero. + token &= 0xffffff; +#endif + const float weight = epilogue_.sorted_weights[row]; + if(token >= 0 && token < epilogue_.num_tokens && col < epilogue_.hidden_size) + { + atomicAdd(epilogue_.o_ptr + + static_cast(token) * epilogue_.stride_token + col, + static_cast(buf[jj]) * weight); + } + }); + n_col += BlockShape::Block_N1; + }; + + // tuple(gate, up) + auto acc_0 = make_tuple(Policy::template MakeCBlockTile_Gemm0(), + Policy::template MakeCBlockTile_Gemm0()); + // auto acc_1 = Policy::template MakeCBlockTile_Gemm1(); + auto acc_1s = generate_tuple( + [&](auto) { return Policy::template MakeCBlockTile_Gemm1(); }, number<2>{}); + + // clang-format off + //t_c refers to all gate/up_proj results in 1 thread.(VGPR), t_a, t_b have similar meaning. + auto gemm_0 = [&](auto& t_c, auto& t_a, auto& t_b) { + using WarpGemm = remove_cvref_t; + + using AWarpTensor = typename WarpGemm::AWarpTensor; + using BWarpTensor = typename WarpGemm::BWarpTensor; + using CWarpTensor = typename WarpGemm::CWarpTensor; + using CWarpDstr = typename WarpGemm::CWarpDstr; + + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + //w_a is same as t_a + AWarpTensor w_a; + w_a.get_thread_buffer() = t_a.get_thread_buffer(); + + //w_b is same as t_b due to t_b has N x K = 4 x 8 + BWarpTensor w_b; + w_b.get_thread_buffer() = t_b.get_thread_buffer(); + CWarpTensor w_c; + + //t_c is blockwise tensor, has repeat and warps. 0,0 here refers to m,n repeat for now only 1 is supported + // <0,0,0,2> -> offset + w_c.get_thread_buffer() = t_c.get_y_sliced_thread_data( + merge_sequences(sequence<0, 0>{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + warp_gemm_0(w_c, w_a, w_b); + + //direct copy, we may use set_thread_buffer? + //Note: Layout here is MRepeat=2, NRpeat=4(4interleave) + t_c.set_y_sliced_thread_data( + merge_sequences(sequence<0, 0>{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + w_c.get_thread_buffer()); + }; + // clang-format on + + // clang-format off + auto gemm_1 = [&](auto& t_c, auto& t_a, auto& t_b) { + using WarpGemm = remove_cvref_t; + + using AWarpTensor = typename WarpGemm::AWarpTensor; + using BWarpTensor = typename WarpGemm::BWarpTensor; + using CWarpTensor = typename WarpGemm::CWarpTensor; + using CWarpDstr = typename WarpGemm::CWarpDstr; + + constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t{}; + constexpr auto c_warp_y_lengths = to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths()); + + //w_a is same as t_a + AWarpTensor w_a; + w_a.get_thread_buffer() = t_a.get_thread_buffer(); + + //w_b is same as t_b due to t_b has N x K = 4 x 8 + BWarpTensor w_b; + w_b.get_thread_buffer() = t_b.get_thread_buffer(); + CWarpTensor w_c; + //t_c is blockwise tensor, has repeat and warps. 0,0 here refers to m,n repeat for now only 1 is supported + w_c.get_thread_buffer() = t_c.get_y_sliced_thread_data( + merge_sequences(sequence<0, 0>{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths)); + warp_gemm_1(w_c, w_a, w_b); + + //direct copy + t_c.set_y_sliced_thread_data( + merge_sequences(sequence<0, 0>{}, c_warp_y_index_zeros), + merge_sequences(sequence<1, 1>{}, c_warp_y_lengths), + w_c.get_thread_buffer()); + }; + // clang-format on + _Pragma("clang diagnostic pop"); + + auto pipeline_gemm0 = [&]() { + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_setprio(0); + block_sync_load_raw(2 * issues_g); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + sld_a(a[I0], a_sld_win0); + gld_a(a_copy_lds_windows(I1)); + __builtin_amdgcn_sched_barrier(0); + + move_a(); + gld_g(gs[I2]); + move_g(); + gld_u(gs[I3]); + move_u(); + + // pre-load a and next u/g + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_setprio(1); + + // first mmac group + gemm_0(acc_0[I0], a[I0], gs[I0]); + gemm_0(acc_0[I1], a[I0], gs[I1]); + + // second mmac group + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_setprio(0); + block_sync_load_raw(2 * issues_g); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + sld_a(a[I0], a_sld_win1); + gld_a(a_copy_lds_windows(I0)); + __builtin_amdgcn_sched_barrier(0); + + move_a(); + gld_g(gs[I0]); + move_g(); + gld_u(gs[I1]); + move_u(); + __builtin_amdgcn_s_setprio(1); + + gemm_0(acc_0[I0], a[I0], gs[I2]); + gemm_0(acc_0[I1], a[I0], gs[I3]); + }; + + auto pipeline_gemm0_tail = [&]() { + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_setprio(0); + block_sync_load_raw(2 * issues_g); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + sld_a(a[I0], a_sld_win0); + gld_a(a_copy_lds_windows(I1)); + __builtin_amdgcn_sched_barrier(0); + gld_g(gs[I2]); + gld_u(gs[I3]); + + // pre-load a and next u/g + __builtin_amdgcn_sched_barrier(0); + __builtin_amdgcn_s_setprio(1); + + // first mmac group + gemm_0(acc_0[I0], a[I0], gs[I0]); + gemm_0(acc_0[I1], a[I0], gs[I1]); + + // second mmac group + __builtin_amdgcn_s_setprio(0); + block_sync_load_raw(2 * issues_g); + __builtin_amdgcn_s_barrier(); + __builtin_amdgcn_sched_barrier(0); + sld_a(a[I0], a_sld_win1); + + // block_sync_load_raw(issues_a + 2 * issues_g); + __builtin_amdgcn_s_setprio(1); + gemm_0(acc_0[I0], a[I0], gs[I2]); + gemm_0(acc_0[I1], a[I0], gs[I3]); + __builtin_amdgcn_sched_barrier(0); + }; + + auto y = Policy::template MakeYBlockTile(); + + auto pipeline_bridge = [&]() { + using WarpGemm = remove_cvref_t; + using CDataType = typename WarpGemm::CDataType; + if constexpr(IsGateOnly) + { + static_assert(false, "not supported yet"); + } + else + { + block_sync_lds(); // ensure last input load have been read from lds + sweep_tile( + acc_0.at(number<0>{}), + [&](auto idx0, auto idx1) { + const YDataType x0 = + type_convert(acc_0.at(number<0>{})(idx0)); + const YDataType x1 = + type_convert(acc_0.at(number<0>{})(idx1)); + AccDataType y0_float; + AccDataType y1_float; + typename Problem::GateActivation{}( + y0_float, type_convert(x0)); + typename Problem::GateActivation{}( + y1_float, type_convert(x1)); + const YDataType y0 = type_convert(y0_float); + const YDataType y1 = type_convert(y1_float); + acc_0.at(number<0>{})(idx0) = type_convert(y0); + acc_0.at(number<0>{})(idx1) = type_convert(y1); + }, + sequence<1, 2>{}); // m=1, n=4 continuos element in thread_buffer(VGPR) + // Match the unfused FP16 contract: GEMM0 output, activation output, and the + // gate/up product each cross an FP16 rounding boundary before GEMM1. + auto reduce_acc_0 = tile_elementwise_in( + [&](const auto& a_, const auto& b_) { + const YDataType activated = type_convert(a_); + const YDataType up = type_convert(b_); + const YDataType product = type_convert(activated * up); + return type_convert(product); + }, + acc_0.at(number<0>{}), + acc_0.at(number<1>{})); + + constexpr auto bridge_tile_enc = + Policy::template MakeBridgeTileDistribution(); + + auto bridge_tile_tensor = + make_static_distributed_tensor(bridge_tile_enc); + static_assert(bridge_tile_tensor.get_thread_buffer_size() == + reduce_acc_0.get_thread_buffer_size(), + "bridge tensor is not same as reduce tensor!"); + + constexpr index_t thread_buf_size = bridge_tile_tensor.get_thread_buffer_size(); + // Note: here we assume that every mmac output 4 element, and NInterleave mmac will + // merge together. So here 1 group refers to 4 x NInterleave. + static_for<0, thread_buf_size, 1>{}([&](auto i) { + index_t reduce_index = + (i / (Problem::Gemm0NInterleave * 4)) * (Problem::Gemm0NInterleave * 4) + + (i % (Problem::Gemm0NInterleave * 4)) / Problem::Gemm0NInterleave + + i % Problem::Gemm0NInterleave * 4; + bridge_tile_tensor.get_thread_buffer()[i] = + reduce_acc_0.get_thread_buffer()[reduce_index]; + }); + + // TODO: support i8 quant here + // auto y_pre = cast_tile(bridge_tile_tensor); + store_tile(bridge_sst_win, cast_tile(bridge_tile_tensor)); + clear_tile(acc_1s(I0)); + block_sync_lds(); + // wave_barrier(); + load_tile(y, bridge_sld_win); + clear_tile(acc_1s(I1)); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + } + }; + + // note, gemm-1 start from idx-1 to N-2 (0, 1, 2....N-1) + auto pipeline_gemm1 = [&]() { + d_vgpr_to_lds(gd[I0]); + gld_d(gd[I1]); + move_d(); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + auto d_tmp = d_lds_to_vgpr(); +#ifdef FMOE_UK_DEBUG_D + if(n_col == 0 && get_lane_id() < 16) + printf("DPRB w=%d bn=%d buf: %f %f %f %f %f %f %f %f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (float)d_tmp.get_thread_buffer()[number<0>{}], + (float)d_tmp.get_thread_buffer()[number<1>{}], + (float)d_tmp.get_thread_buffer()[number<2>{}], + (float)d_tmp.get_thread_buffer()[number<3>{}], + (float)d_tmp.get_thread_buffer()[number<4>{}], + (float)d_tmp.get_thread_buffer()[number<5>{}], + (float)d_tmp.get_thread_buffer()[number<6>{}], + (float)d_tmp.get_thread_buffer()[number<7>{}]); +#endif + gemm_1(acc_1s[I0], y, d_tmp); + atomic_add_o(acc_1s[I0]); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + d_vgpr_to_lds(gd[I1]); + + // Note: if support pk_atomic_add, here will use 2interleave and f16 atomic + // move_o(); + gld_d(gd[I0]); + move_d(); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + + d_tmp = d_lds_to_vgpr(); + gemm_1(acc_1s[I1], y, d_tmp); + atomic_add_o(acc_1s[I1]); + clear_tile(acc_1s(I0)); + clear_tile(acc_1s(I1)); + block_sync_lds(); + asm volatile("s_waitcnt vmcnt(0)"); + __builtin_amdgcn_sched_barrier(0); + }; + + auto pipeline_gemm1_head = [&]() { + gld_d_init(); + gld_d(gd[I0]); + move_d(); + }; + auto pipeline_gemm1_tail = [&]() { +#ifdef FMOE_UK_DEBUG_G + if(get_lane_id() % 32 == 0) + { + auto ginx = d_copy_dram_window.get_tile_distribution().calculate_index(); + printf("GINX w=%d lane=%d x=(%d,%d) buf8to15: %f %f %f %f %f %f %f %f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (int)ginx[number<0>{}], + (int)ginx[number<1>{}], + (float)gd[I0].get_thread_buffer()[number<8>{}], + (float)gd[I0].get_thread_buffer()[number<9>{}], + (float)gd[I0].get_thread_buffer()[number<10>{}], + (float)gd[I0].get_thread_buffer()[number<11>{}], + (float)gd[I0].get_thread_buffer()[number<12>{}], + (float)gd[I0].get_thread_buffer()[number<13>{}], + (float)gd[I0].get_thread_buffer()[number<14>{}], + (float)gd[I0].get_thread_buffer()[number<15>{}]); + } +#endif + d_vgpr_to_lds(gd[I0]); +#ifdef FMOE_UK_DEBUG_G + if(get_lane_id() % 32 == 0) + printf("GPRB w=%d lane=%d buf: %f %f %f %f %f %f %f %f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (float)gd[I0].get_thread_buffer()[number<0>{}], + (float)gd[I0].get_thread_buffer()[number<1>{}], + (float)gd[I0].get_thread_buffer()[number<2>{}], + (float)gd[I0].get_thread_buffer()[number<3>{}], + (float)gd[I0].get_thread_buffer()[number<4>{}], + (float)gd[I0].get_thread_buffer()[number<5>{}], + (float)gd[I0].get_thread_buffer()[number<6>{}], + (float)gd[I0].get_thread_buffer()[number<7>{}]); +#endif + gld_d(gd[I1]); + move_d(); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + auto d_tmp = d_lds_to_vgpr(); +#ifdef FMOE_UK_DEBUG_D + if(n_col == 0 && get_lane_id() < 16) + printf("DPRB w=%d bn=%d buf: %f %f %f %f %f %f %f %f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (float)d_tmp.get_thread_buffer()[number<0>{}], + (float)d_tmp.get_thread_buffer()[number<1>{}], + (float)d_tmp.get_thread_buffer()[number<2>{}], + (float)d_tmp.get_thread_buffer()[number<3>{}], + (float)d_tmp.get_thread_buffer()[number<4>{}], + (float)d_tmp.get_thread_buffer()[number<5>{}], + (float)d_tmp.get_thread_buffer()[number<6>{}], + (float)d_tmp.get_thread_buffer()[number<7>{}]); +#endif + gemm_1(acc_1s[I0], y, d_tmp); +#ifdef FMOE_UK_DEBUG_C + if(n_col == 0 && (get_lane_id() % 16) < 2) + printf("CPROBE w=%d lane=%d c: %f %f %f %f %f %f %f %f\n", + (int)get_warp_id(), + (int)get_lane_id(), + (float)acc_1s[I0].get_thread_buffer()[number<0>{}], + (float)acc_1s[I0].get_thread_buffer()[number<1>{}], + (float)acc_1s[I0].get_thread_buffer()[number<2>{}], + (float)acc_1s[I0].get_thread_buffer()[number<3>{}], + (float)acc_1s[I0].get_thread_buffer()[number<4>{}], + (float)acc_1s[I0].get_thread_buffer()[number<5>{}], + (float)acc_1s[I0].get_thread_buffer()[number<6>{}], + (float)acc_1s[I0].get_thread_buffer()[number<7>{}]); +#endif + atomic_add_o(acc_1s[I0]); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + d_vgpr_to_lds(gd[I1]); + block_sync_lds(); + __builtin_amdgcn_sched_barrier(0); + + d_tmp = d_lds_to_vgpr(); + gemm_1(acc_1s[I1], y, d_tmp); + atomic_add_o(acc_1s[I1]); + __builtin_amdgcn_sched_barrier(0); + }; + + // start of pipeline + // clang-format off + gld_a_init(); + gld_a(a_copy_lds_windows(I0)); + move_a(); + //quant + gld_g_init(); + gld_u_init(); + gld_g(gs[I0]); + move_g(); + gld_u(gs[I1]); + move_u(); + //b scale load + clear_tile(acc_0.at(I0)); + clear_tile(acc_0.at(I1)); + + + // make sure a,g loaded + // lds_load_fence(); + // block_sync_load_raw(2 * issues_g); + + // we manually unroll double buffer inside hot loop + const index_t iters_0 = (num_blocks_k0 - 2) / 2; + index_t i_0 = 0; // (void)i_0; (void)iters_0; (void)pipeline_gemm0; + while(i_0++ < iters_0) + { + pipeline_gemm0(); + } + pipeline_gemm0_tail(); + + pipeline_bridge(); + + const index_t iters_1 = (num_blocks_n1 - 2) / 2; + index_t i_1 = 0; // (void) i_1; (void)iters_1; (void)pipeline_gemm1; + pipeline_gemm1_head(); + while(i_1++ < iters_1) + { + pipeline_gemm1(); + } + pipeline_gemm1_tail(); + // clang-format on + } +}; + +} // namespace ck_tile diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp new file mode 100644 index 0000000..e1fdd76 --- /dev/null +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_uk.hpp @@ -0,0 +1,230 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2026 Hygon Info Technologies Ltd. +// Modified by Hygon Information Technology Co., Ltd. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/fused_moe/pipeline/fused_moegemm_flatmm_uk_engine.hpp" + +namespace ck_tile { + +template +struct FusedMoeGemmPipeline_FlatmmUk +{ + using Problem = remove_cvref_t; + using Policy = remove_cvref_t; + + using BlockShape = typename Problem::BlockShape; + + using ADataType = typename Problem::ADataType; + using GDataType = typename Problem::GDataType; + using DDataType = typename Problem::DDataType; + using AccDataType = typename Problem::AccDataType; + using ODataType = typename Problem::ODataType; + using AScaleDataType = typename Problem::AScaleDataType; + using GScaleDataType = typename Problem::GScaleDataType; + using DScaleDataType = typename Problem::DScaleDataType; + using YSmoothScaleDataType = typename Problem::YSmoothScaleDataType; + using TopkWeightDataType = typename Problem::TopkWeightDataType; + using IndexDataType = typename Problem::IndexDataType; + using YDataType = typename Problem::YDataType; + + using Traits = typename Problem::Traits; + + static constexpr bool IsGateOnly = Traits::IsGateOnly; + static constexpr bool UseSmoothQuant = Traits::UseSmoothQuant; + static constexpr bool PadHiddenSize = Traits::PadHiddenSize; + static constexpr bool PadIntermediateSize = Traits::PadIntermediateSize; + + static_assert(!IsGateOnly, "flatmm_uk(HCU): gate+up only, gate-only not supported by engine"); + + static constexpr index_t kAlignmentA = Policy::template GetAlignment_A(); + static constexpr index_t kAlignmentG = Policy::template GetAlignment_G(); + static constexpr index_t kAlignmentD = Policy::template GetAlignment_D(); + static constexpr index_t kAlignmentO = Policy::template GetAlignment_O(); + + static constexpr index_t kBlockPerCu = 2; + + static constexpr const char* name = "flatmm_uk"; + + // compute engine: HCU MMAC implementation of the 2-gemm dataflow + // (faithful copy of the fork's flatmm_ex pipeline) + using Engine = FusedMoeGemmPipeline_FlatmmUkEngine; + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() + { + return Engine::GetSmemSize(); + } + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSizeBytes() + { + constexpr index_t kMaxElemSize = + (sizeof(ADataType) > sizeof(DDataType)) + ? (sizeof(ADataType) > sizeof(YDataType) ? sizeof(ADataType) : sizeof(YDataType)) + : (sizeof(DDataType) > sizeof(YDataType) ? sizeof(DDataType) : sizeof(YDataType)); + return Engine::GetSmemSize() * kMaxElemSize; + } + + CK_TILE_HOST_DEVICE static auto GetACoord() + { + constexpr auto a_dist = Policy::template MakeGlobalTileDistribution_A(); + const auto a_coord = a_dist.calculate_index(); + return a_coord; + } + + CK_TILE_HOST_DEVICE static auto GetOCoord() + { + constexpr auto o_dist = Policy::template MakeGlobalTileDistribution_O(); + const auto o_coord = o_dist.calculate_index(); + return o_coord; + } + + template + CK_TILE_DEVICE auto operator()(const Kargs& kargs, + CK_TILE_LDS_ADDR void* smem, + index_t sorted_tile_id, + index_t intermediate_tile_id) + { + static_assert(BlockShape::Block_K1 == BlockShape::Block_N0, + "flatmm_uk(HCU): gemm1 K tile must match gemm0 N tile"); + + constexpr index_t hidden_radio_0 = IsGateOnly ? 1 : 2; // gate_up = 2 + + index_t expert_stride_0 = kargs.intermediate_size * hidden_radio_0 * kargs.hidden_size; + index_t expert_stride_1 = kargs.hidden_size * kargs.intermediate_size; + + const IndexDataType expert_id = __builtin_amdgcn_readfirstlane( + reinterpret_cast(kargs.sorted_expert_ids_ptr)[sorted_tile_id]); + + index_t hidden_idx = __builtin_amdgcn_readfirstlane( + intermediate_tile_id * BlockShape::Block_N0); + + const auto a_coord = GetACoord(); + const auto sorted_token_id = + a_coord[number<0>{}] + sorted_tile_id * BlockShape::Block_M0; + + index_t token_id = + reinterpret_cast(kargs.sorted_token_ids_ptr)[sorted_token_id]; +#if CK_TILE_REFERENCE_MOE_SORTING_MOCK_ID || CK_TILE_REFERENCE_MOE_SORTING_MOCK_EXPERT_ID + // same decode contract as FusedMoeGemmKernel (A-gather side): strip the + // packed topk/expert tag before using the id as a token row + token_id &= 0xffffff; +#endif + + const auto a_window = [&]() { + const ADataType* a_ptr = reinterpret_cast(kargs.a_ptr); + const auto a_view_ = make_naive_tensor_view( + a_ptr, + make_tuple(kargs.num_tokens, kargs.hidden_size), + make_tuple(kargs.stride_token, 1), + number{}, + number<1>{}); + + const auto a_gather_view_ = transform_tensor_view( + a_view_, + make_tuple(make_indexing_transform(kargs.num_tokens, token_id), + make_pass_through_transform(kargs.hidden_size)), + make_tuple(sequence<0>{}, sequence<1>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + const auto a_gather_view_1_ = pad_tensor_view( + a_gather_view_, + make_tuple(number{}, number<2 * BlockShape::Block_K0>{}), + sequence{}); + + const auto a_window_ = make_tile_window( + a_gather_view_1_, + make_tuple(number{}, number{}), + {0, 0}); + return a_window_; + }(); + + const auto g_window = [&]() { + const GDataType* g_ptr = reinterpret_cast(kargs.g_ptr) + + static_cast(expert_id) * expert_stride_0 + + hidden_idx * kargs.hidden_size; + + // real gate rows left in this tile; rows beyond intermediate_size + // belong to the up half / next expert (or past W1 for the last + // expert), so declare the true extent and let the pad transform + // zero-fill the tail of a partial tile (gemm1 sees zero columns + // through the matching guard in the engine's manual D load). + const index_t g_rows_valid = kargs.intermediate_size - hidden_idx; + + const auto g_view_ = make_naive_tensor_view( + g_ptr, + make_tuple(g_rows_valid, kargs.hidden_size), + make_tuple(kargs.hidden_size, 1), + number{}, + number<1>{}); + + const auto g_view_1_ = pad_tensor_view( + g_view_, + make_tuple(number{}, number<2 * BlockShape::Block_K0>{}), + sequence{}); + + const auto g_window_ = make_tile_window( + g_view_1_, + make_tuple(number{}, number{}), + {0, 0}); + return g_window_; + }(); + + const auto d_window = [&]() { + const DDataType* d_ptr = reinterpret_cast(kargs.d_ptr) + + static_cast(expert_id) * expert_stride_1 + + hidden_idx; + + const auto d_view_ = make_naive_tensor_view( + d_ptr, + make_tuple(kargs.hidden_size, BlockShape::Block_K1), + make_tuple(kargs.intermediate_size, 1), + number{}, + number<1>{}); + + const auto d_view_1_ = pad_tensor_view( + d_view_, + make_tuple(number<2 * BlockShape::Block_N1>{}, number{}), + sequence{}); + + const auto d_window_ = make_tile_window( + d_view_1_, + make_tuple(number{}, number{}), + {0, 0}); + return d_window_; + }(); + + struct MoeEpilogue + { + ODataType* o_ptr; // [num_tokens, hidden] fp32 atomic accumulation buffer + const IndexDataType* sorted_ids; + const TopkWeightDataType* sorted_weights; + index_t row_base; // sorted_tile_id * Block_M0 + index_t stride_token; + index_t num_tokens; + index_t hidden_size; + }; + + Engine{}(a_window, + g_window, + d_window, + MoeEpilogue{reinterpret_cast(kargs.o_ptr), + reinterpret_cast(kargs.sorted_token_ids_ptr), + reinterpret_cast(kargs.sorted_weight_ptr), + sorted_tile_id * BlockShape::Block_M0, + kargs.stride_token, + kargs.num_tokens, + kargs.hidden_size}, + smem, + kargs.hidden_size, + kargs.intermediate_size, + hidden_idx); + } +}; + +} // namespace ck_tile From a616e396ede63df78020a4d9feefdc9e98e16322 Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Mon, 12 Oct 2026 20:09:35 +0800 Subject: [PATCH 04/11] fix(fused_moe): adapt for FusedMoeGemmPipelineProblem --- .../17_fused_moe/fused_moe_flatmm_uk_fp16.cpp | 2 - .../instances/fused_moegemm_api_internal.hpp | 46 +++++++++---------- .../fused_moegemm_pipeline_problem.hpp | 8 ++-- 3 files changed, 27 insertions(+), 29 deletions(-) diff --git a/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp b/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp index 013dc39..24b1d7f 100644 --- a/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp +++ b/example_hcu/ck_tile/17_fused_moe/fused_moe_flatmm_uk_fp16.cpp @@ -203,8 +203,6 @@ using FmoeProblem = ck_tile::FusedMoeGemmPipelineProblem; + f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType>; using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmEx; // using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; @@ -154,14 +154,14 @@ float fused_moegemm_swizzled_lds_(const ck_tile::stream_config& s, fused_moegemm typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, - f_traits>; + f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType>; using f_pipeline = ck_tile::FusedMoeGemmPipeline_Flatmm_GU_Regs_D_LDS; // using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; @@ -222,14 +222,14 @@ float fused_moegemm_wreg_(const ck_tile::stream_config& s, fused_moegemm_args a) typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, - f_traits>; + f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType>; using f_pipeline = ck_tile::FusedMoeGemmPipeline_Flatmm_Wreg; // using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; @@ -294,14 +294,14 @@ float fused_moegemm_wreg_shuffle_(const ck_tile::stream_config& s, fused_moegemm typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, - f_traits>; + f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType>; using f_pipeline = ck_tile::FusedMoeGemmPipeline_Flatmm_Wreg_shuffle; // using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; @@ -366,14 +366,14 @@ float fused_moegemm_wt_bd_shuffle_(const ck_tile::stream_config& s, fused_moegem typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, - f_traits>; + f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType>; using f_pipeline = ck_tile::FusedMoeGemmPipeline_Flatmm_Wt_shuffle; // using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmUk; @@ -438,14 +438,14 @@ float fused_moegemm_wt_shuffle_quant(const ck_tile::stream_config& s, fused_moeg typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType, Ts_::FusedQuant, Ts_::QuantBlockSizeN, Ts_::QuantBlockSizeK>; @@ -512,14 +512,14 @@ float fused_moegemm_wt_w8a8_quant(const ck_tile::stream_config& s, fused_moegemm typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType, Ts_::FusedQuant, Ts_::QuantBlockSizeN, Ts_::QuantBlockSizeK>; @@ -586,14 +586,14 @@ float fused_moegemm_wt_w4a16_quant(const ck_tile::stream_config& s, fused_moegem typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType, Ts_::FusedQuant, Ts_::QuantBlockSizeN, Ts_::QuantBlockSizeK>; @@ -659,14 +659,14 @@ float fused_moegemm_wt_w4a8_quant(const ck_tile::stream_config& s, fused_moegemm typename Ts_::AScaleDataType, typename Ts_::GScaleDataType, typename Ts_::DScaleDataType, - typename Ts_::GZeroPointDataType, - typename Ts_::DZeroPointDataType, typename Ts_::YSmoothScaleDataType, typename Ts_::TopkWeightDataType, typename Ts_::IndexDataType, f_act_, // TODO: hardcoded f_shape, f_traits, + typename Ts_::GZeroPointDataType, + typename Ts_::DZeroPointDataType, Ts_::FusedQuant, Ts_::QuantBlockSizeN, Ts_::QuantBlockSizeK>; diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_problem.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_problem.hpp index 21c669a..09717ff 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_problem.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_problem.hpp @@ -17,17 +17,17 @@ template struct FusedMoeGemmPipelineProblem { From a0714fd7234fc4fdb8f4759abd1bfb740453bcd8 Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Tue, 13 Oct 2026 16:59:39 +0800 Subject: [PATCH 05/11] fix(ds_read): fix pointer type mismatch --- include/ck/utility/hcu_dsreadm.hpp | 28 +++++++++---------- .../gemm/warp/warp_dsreadm_attribute_impl.hpp | 8 +++--- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/include/ck/utility/hcu_dsreadm.hpp b/include/ck/utility/hcu_dsreadm.hpp index 75432e4..7c29db1 100644 --- a/include/ck/utility/hcu_dsreadm.hpp +++ b/include/ck/utility/hcu_dsreadm.hpp @@ -13,7 +13,7 @@ __device__ float4_t intrin_ds_read_m32x8_f32(const float* ptr) { float4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x8_f32(ptr); + reg = __builtin_hcu_ds_read_m32x8_f32(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -24,7 +24,7 @@ __device__ int32x4_t intrin_ds_read_m32x8_tf32(const int* ptr) { int32x4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x8_tf32(ptr); + reg = __builtin_hcu_ds_read_m32x8_tf32(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -36,7 +36,7 @@ __device__ half8_t intrin_ds_read_m32x16_f16(const half_t* ptr) half8_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) - reg = __builtin_hcu_ds_read_m32x16_f16(ptr); + reg = __builtin_hcu_ds_read_m32x16_f16(reinterpret_cast<__fp16*>(const_cast(ptr))); #else swallow(ptr, reg); #endif @@ -48,7 +48,7 @@ __device__ half8_t intrin_ds_read_m32x16_f16_alt(const half_t* ptr) half8_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) - reg = __builtin_hcu_ds_read_m32x16_f16_alt(ptr); + reg = __builtin_hcu_ds_read_m32x16_f16_alt(reinterpret_cast<__fp16*>(const_cast(ptr))); #else swallow(ptr, reg); #endif @@ -59,7 +59,7 @@ __device__ short8_t intrin_ds_read_m32x16_bf16(const short* ptr) { short8_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x16_bf16(ptr); + reg = __builtin_hcu_ds_read_m32x16_bf16(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -70,7 +70,7 @@ __device__ short8_t intrin_ds_read_m32x16_bf16_alt(const short* ptr) { short8_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x16_bf16_alt(ptr); + reg = __builtin_hcu_ds_read_m32x16_bf16_alt(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -81,7 +81,7 @@ __device__ int32x4_t intrin_ds_read_m32x32_i8(const int* ptr) { int32x4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x32_i8(ptr); + reg = __builtin_hcu_ds_read_m32x32_i8(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -92,7 +92,7 @@ __device__ int32x4_t intrin_ds_read_m32x32_i8_alt2(const int* ptr) { int32x4_t reg; #if defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x32_i8_alt2(ptr); + reg = __builtin_hcu_ds_read_m32x32_i8_alt2(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -103,7 +103,7 @@ __device__ int32x4_t intrin_ds_read_m64x16_i8_alt4(const int* ptr) { int32x4_t reg; #if defined(__gfx938__) - reg = __builtin_hcu_ds_read_m64x16_i8_alt4(ptr); + reg = __builtin_hcu_ds_read_m64x16_i8_alt4(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -114,7 +114,7 @@ __device__ int32x4_t intrin_ds_read_m32x32_u8(const int* ptr) { int32x4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x32_u8(ptr); + reg = __builtin_hcu_ds_read_m32x32_u8(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -125,7 +125,7 @@ __device__ int32x4_t intrin_ds_read_m32x32_u8_alt2(const int* ptr) { int32x4_t reg; #if defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x32_u8_alt2(ptr); + reg = __builtin_hcu_ds_read_m32x32_u8_alt2(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -136,7 +136,7 @@ __device__ int32x4_t intrin_ds_read_m64x16_u8_alt4(const int* ptr) { int32x4_t reg; #if defined(__gfx938__) - reg = __builtin_hcu_ds_read_m64x16_u8_alt4(ptr); + reg = __builtin_hcu_ds_read_m64x16_u8_alt4(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -147,7 +147,7 @@ __device__ int32x4_t intrin_ds_read_m32x64_i4(const int* ptr) { int32x4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x64_i4(ptr); + reg = __builtin_hcu_ds_read_m32x64_i4(const_cast(ptr)); #else swallow(ptr, reg); #endif @@ -158,7 +158,7 @@ __device__ int32x4_t intrin_ds_read_m32x64_u4(const int* ptr) { int32x4_t reg; #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) - reg = __builtin_hcu_ds_read_m32x64_u4(ptr); + reg = __builtin_hcu_ds_read_m32x64_u4(const_cast(ptr)); #else swallow(ptr, reg); #endif diff --git a/include/ck_tile/ops/gemm/warp/warp_dsreadm_attribute_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_dsreadm_attribute_impl.hpp index 706033f..0a74bf1 100644 --- a/include/ck_tile/ops/gemm/warp/warp_dsreadm_attribute_impl.hpp +++ b/include/ck_tile/ops/gemm/warp/warp_dsreadm_attribute_impl.hpp @@ -48,7 +48,7 @@ struct WarpDsreadmAttributeImplF16MN32K16 { #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) - ret = __builtin_hcu_ds_read_m32x16_f16(lds_ptr); + ret = __builtin_hcu_ds_read_m32x16_f16(reinterpret_cast<__fp16*>(const_cast(lds_ptr))); #else detail::swallow{lds_ptr, ret}; #endif @@ -88,7 +88,7 @@ struct WarpDsreadmAttributeImplF16MN32K16Alt { #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) - ret = __builtin_hcu_ds_read_m32x16_f16_alt(lds_ptr); + ret = __builtin_hcu_ds_read_m32x16_f16_alt(reinterpret_cast<__fp16*>(const_cast(lds_ptr))); #else detail::swallow{lds_ptr, ret}; #endif @@ -108,7 +108,7 @@ struct WarpDsreadmAttributeImplBF16MN32K16 : WarpDsreadmAttributeImplF16MN32K16 #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) const auto raw = __builtin_hcu_ds_read_m32x16_f16( - reinterpret_cast(lds_ptr)); + reinterpret_cast<__fp16*>(const_cast(lds_ptr))); ret = bit_cast(raw); #else detail::swallow{lds_ptr, ret}; @@ -126,7 +126,7 @@ struct WarpDsreadmAttributeImplBF16MN32K16Alt : WarpDsreadmAttributeImplF16MN32K #if defined(__gfx928__) || defined(__gfx92a__) || defined(__gfx936__) || defined(__gfx938__) || \ defined(__gfx946__) const auto raw = __builtin_hcu_ds_read_m32x16_f16_alt( - reinterpret_cast(lds_ptr)); + reinterpret_cast<__fp16*>(const_cast(lds_ptr))); ret = bit_cast(raw); #else detail::swallow{lds_ptr, ret}; From a1a410fd2975193fdc7f0c8d4a5336c494aad7df Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Wed, 14 Oct 2026 10:38:56 +0800 Subject: [PATCH 06/11] fix(ck_tile): fix bit_cast ambiguous --- include/ck_tile/core/container/meta_data_buffer.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/include/ck_tile/core/container/meta_data_buffer.hpp b/include/ck_tile/core/container/meta_data_buffer.hpp index 7493b93..eba60fa 100644 --- a/include/ck_tile/core/container/meta_data_buffer.hpp +++ b/include/ck_tile/core/container/meta_data_buffer.hpp @@ -30,7 +30,7 @@ struct meta_data_buffer { constexpr index_t size = sizeof(T); - auto tmp = bit_cast>(data); + auto tmp = ck_tile::bit_cast>(data); for(int i = 0; i < size; i++) { @@ -66,7 +66,7 @@ struct meta_data_buffer pos++; } - data = bit_cast(tmp); + data = ck_tile::bit_cast(tmp); } return data; @@ -86,7 +86,7 @@ struct meta_data_buffer pos++; } - auto data = bit_cast(tmp); + auto data = ck_tile::bit_cast(tmp); return data; } From 378b2b90bffbec2c0e589f241f036f8898fa79e1 Mon Sep 17 00:00:00 2001 From: wugangsheng Date: Mon, 12 Oct 2026 17:19:47 +0800 Subject: [PATCH 07/11] feat(grouped_gemm): integrate v7/v8 uniform d2l bf16 NN/TN kernels as CK instances Port the standalone uniform Direct-to-LDS NN (dgrad) / TN (wgrad) kernels into the ck_tile framework as device-args instances: - gemm_pipeline_v7_uniform.hpp: uniform (non warp-specialized) pipeline, 256x256x16, 512 threads, 4-stage 16KB Direct-to-LDS ring, two-group wave phase interleave; V7UniformKind::NN walks a K-contiguous A, ::TN walks the [K, M] panel (both operands m32x16). - grouped_gemm_v7_uniform_epilogue.hpp: per-wave LDS slab staging, 16B vector stores, no cross-warp barrier. - instances 5000 (v7_bf16_nn) / 5001 (v8_bf16_tn), family v7, opt-in via CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN / _V8_TN (disable wins); selector gates: gfx936, bf16, K>=64 && K%16==0, N%256==0, G<=16, fixed K, no empty groups. - entries are NOT registry candidates: a dedicated archive carries the tuned codegen contract (-mllvm -support-768-vgprs et al.); the flat over-covered grid early-exits blocks past the last group's block_end (FindGroupId only converges in-range), and the LDS ring is DYNAMIC shared memory addressed from offset 0 per the reference kernels. - shared-header decoupling: kernel.hpp only includes the 10-line forward header (complete tile-parameter POD + enum); the pipeline body and the GemmConfigV7Uniform configs compile only inside the v7/v8 TUs, so tuning recompiles two translation units, not the archive. - grouped_gemm_v7_instance_test: correctness (CPU reference + default instance triangle) and perf through the versioned device-args ABI; selector/query tests extended to the 20-candidate registry. gfx936, groups=4 M=[8100,8100,8091,8091] N=4096 K=4096: NN 2.741 ms / 396.7 TF (default instance 6.508 ms / 167.1 TF) TN 2.700 ms / 402.7 TF --- .../ck_tile/19_grouped_gemm/CMakeLists.txt | 77 ++ .../ck_tile/19_grouped_gemm/grouped_gemm.hpp | 33 + .../grouped_gemm_device_args.cpp | 71 ++ .../grouped_gemm_device_args_query_test.cpp | 12 +- .../grouped_gemm_v7_instance_test.cpp | 787 ++++++++++++++++++ .../grouped_gemm_v7_uniform_epilogue.hpp | 206 +++++ .../grouped_gemm_device_args_provider.hpp | 8 +- .../grouped_gemm_v7_bf16_nn_device_args.cpp | 52 ++ .../grouped_gemm_v7_uniform_launch.hpp | 198 +++++ .../grouped_gemm_v8_bf16_tn_device_args.cpp | 51 ++ .../gemm/grouped_gemm_bw_family_selected.hpp | 49 ++ .../ops/gemm/grouped_gemm_device_args.h | 9 +- .../ops/gemm/grouped_gemm_selector.hpp | 72 +- .../ops/gemm/kernel/grouped_gemm_kernel.hpp | 105 ++- .../pipeline/gemm_pipeline_v7_uniform.hpp | 781 +++++++++++++++++ .../pipeline/gemm_pipeline_v7_uniform_fwd.hpp | 43 + 16 files changed, 2537 insertions(+), 17 deletions(-) create mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp create mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_uniform_epilogue.hpp create mode 100644 example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_bf16_nn_device_args.cpp create mode 100644 example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_uniform_launch.hpp create mode 100644 example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v8_bf16_tn_device_args.cpp create mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp create mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform_fwd.hpp diff --git a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt index 8d8ee66..c335ffa 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt +++ b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt @@ -339,6 +339,52 @@ set_source_files_properties( ${HCU_GROUPED_GEMM_GFX936_SELECTED_SOURCES} PROPERTIES COMPILE_OPTIONS "-Wno-unused-template;-Wno-unneeded-internal-declaration") +# --------------------------------------------------------------------------- +# v7 BF16 NN / TN kernels as precompiled CK instances. +# +# These two translation units are the instance entries (see their headers): the +# whole kernel body lives in ck_tile headers (gemm_pipeline_v7_uniform.hpp + +# grouped_gemm_v7_uniform_epilogue.hpp, instantiated through GroupedGemmKernel by +# grouped_gemm_v7_uniform_launch.hpp). They are NOT part of +# grouped_gemm_device_args_provider_sources.txt on purpose: the v7 pipelines +# were tuned under an explicit codegen contract -- the flags below (in +# particular the VGPR budget that a wrapping tile loop or a generic-pointer +# kargs read would break) -- and a consumer that compiles the manifest itself +# would silently drop that contract. The archive carries the flags with it, and +# any consumer of the manifest that wants these instances must extend its own +# compile options to match. +if(GPU_TARGETS MATCHES "gfx936") + add_library(ck_tile_hcu_grouped_gemm_v7_v8_bf16 STATIC EXCLUDE_FROM_ALL + instances/grouped_gemm_v7_bf16_nn_device_args.cpp + instances/grouped_gemm_v8_bf16_tn_device_args.cpp) + target_include_directories(ck_tile_hcu_grouped_gemm_v7_v8_bf16 + PRIVATE ${CMAKE_CURRENT_LIST_DIR}) + # No -D__gfx936__ here: the device pass defines the arch macro itself for + # --offload-arch=gfx936, and a fat binary building this target for gfx938 + # as well must see gfx938, not a forced gfx936. The pipeline's + # CACHE_SWIZZLE descriptor bit (word1 bit 30) keys off + # `__gfx936__ || __gfx938__`, so it stays live either way. + target_compile_options(ck_tile_hcu_grouped_gemm_v7_v8_bf16 + PRIVATE + "SHELL: -mllvm -support-768-vgprs=true" + "SHELL: -mllvm -disable-machine-sink" + "SHELL: -mllvm -disable-code-sink" + "SHELL: -mllvm -amdgpu-enable-rewrite-partial-reg-uses=false" + "SHELL: -mllvm -allow-gvn-convergent-call=true" + "SHELL: -mllvm -disallow-uniform-vmed3-combine=true" + "SHELL: -mllvm -hcu-pre-emit-load-store-opt=false" + "SHELL: -mllvm -amdgpu-early-inline-all=true" + "SHELL: -mllvm -amdgpu-function-calls=false" + "-fno-finite-math-only" + "-Wno-unused-result") + # -mcode-object-version=5 and --gpu-max-threads-per-block=256 from the + # reference harness are deliberately NOT copied: the first is a module ABI + # choice this build already makes globally, and the second states 256 + # threads while these kernels launch 512. + set_target_properties(ck_tile_hcu_grouped_gemm_v7_v8_bf16 + PROPERTIES POSITION_INDEPENDENT_CODE ON) +endif() + target_include_directories(ck_tile_hcu_grouped_gemm PRIVATE ${CMAKE_CURRENT_LIST_DIR}) target_compile_definitions(ck_tile_hcu_grouped_gemm PRIVATE @@ -381,6 +427,16 @@ set_target_properties(ck_tile_hcu_grouped_gemm PROPERTIES POSITION_INDEPENDENT_C ck_tile_hcu_grouped_gemm_apply_codegen(ck_tile_hcu_grouped_gemm) target_link_libraries(ck_tile_hcu_grouped_gemm PUBLIC ck_tile_hcu_grouped_gemm_device_args_providers) +if(TARGET ck_tile_hcu_grouped_gemm_v7_v8_bf16) + # Pair the dispatch (grouped_gemm_device_args.cpp) and the selector route + # with the entries this archive actually provides: an arch build without + # them must not name InstanceId 5000/5001, exactly like the promoted-NN + # instances above. + target_link_libraries(ck_tile_hcu_grouped_gemm + PUBLIC ck_tile_hcu_grouped_gemm_v7_v8_bf16) + target_compile_definitions(ck_tile_hcu_grouped_gemm + PRIVATE CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16) +endif() # Minimal consumer-independent library used by the public C ABI matrix. It # deliberately excludes all default-off Phase316/317 candidate TUs so generic @@ -430,6 +486,27 @@ if(GPU_TARGETS MATCHES "gfx936") PRIVATE ck_tile_hcu_grouped_gemm) endif() +# Correctness + performance driver for the v7 BF16 NN / TN instances +# (InstanceId 5000 / 5001). It goes through the versioned device-args ABI +# exactly as a consumer does (select + run), and it links the grouped-GEMM +# archive, so it also covers the selector route and the ABI dispatch that +# production uses. +if(TARGET ck_tile_hcu_grouped_gemm_v7_v8_bf16) + add_executable(tile_example_grouped_gemm_v7_instance_test EXCLUDE_FROM_ALL + grouped_gemm_v7_instance_test.cpp) + target_include_directories(tile_example_grouped_gemm_v7_instance_test + PRIVATE ${CMAKE_CURRENT_LIST_DIR}) + target_compile_definitions(tile_example_grouped_gemm_v7_instance_test + PRIVATE CK_TILE_GROUPED_GEMM_FAST_BUILD) + target_link_libraries(tile_example_grouped_gemm_v7_instance_test + PRIVATE ck_tile_hcu_grouped_gemm) + add_test(NAME tile_example_grouped_gemm_v7_instance_test + COMMAND $ + --layout=NN --m=512 --n=512 --k=256 --groups=4 --iters=5 --warmup=2) + set_tests_properties(tile_example_grouped_gemm_v7_instance_test + PROPERTIES LABELS "hcu;ck_tile;grouped_gemm;v7;correctness;perf") +endif() + # Host-only Phase-1 selector/contract test. This target deliberately does not # link a grouped-GEMM provider, so adding it cannot alter a device launch. add_executable(tile_example_grouped_gemm_selector_test EXCLUDE_FROM_ALL diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp index ab73c31..f70d12b 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp @@ -39,6 +39,13 @@ #define CK_TILE_PIPELINE_COMPUTE_V3_W8_OVERLAP 8 #define CK_TILE_PIPELINE_COMPUTE_V3_TRITON_OPERAND 9 #define CK_TILE_PIPELINE_COMPUTE_V3_DSREADM_STAGE 10 +// Uniform (non warp-specialized) Direct-to-LDS grouped GEMM for the NN (dgrad) +// and TN (wgrad) operand layouts. See +// include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp. +// Two ids: the layouts share the pipeline body but differ in the operand walk, +// and the config carries the layout that the host-side selector gated on. +#define CK_TILE_PIPELINE_COMPUTE_V7 12 +#define CK_TILE_PIPELINE_COMPUTE_V8 13 #ifndef CK_TILE_PIPELINE_DEFAULT #define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 @@ -2116,6 +2123,32 @@ struct PipelineTypeTraits using UniversalGemmPipeline = ck_tile::GemmPipelineAgBgCrMls; }; +// The v7/v8 uniform pipeline's forward declarations (enum + template names) +// arrive from grouped_gemm_kernel.hpp's inclusion of +// gemm_pipeline_v7_uniform_fwd.hpp; the pipeline body and the configs are +// completed only in the v7/v8 translation units (via the launch header). +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = + ck_tile::GemmPipelineV7Uniform; + template + using UniversalGemmPipeline = + ck_tile::GemmPipelineV7Uniform; +}; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = + ck_tile::GemmPipelineV7Uniform; + template + using UniversalGemmPipeline = + ck_tile::GemmPipelineV7Uniform; +}; + template struct GemmConfigComputeInt4 : public GemmConfigBase { diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args.cpp index 3dc1f84..bdaa4c5 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args.cpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args.cpp @@ -5,6 +5,9 @@ #include "ck_tile/ops/gemm/grouped_gemm_selector.hpp" #include "ck_tile/ops/gemm/grouped_gemm_bw_family_selected.hpp" #include "ck_tile/host/device_prop.hpp" +// Header-only: the tile bound the v7 instances are launched with. Nothing from +// it is instantiated unless the v7 route runs (see grouped_gemm_v7_uniform_launch). +#include "instances/grouped_gemm_v7_uniform_launch.hpp" #include @@ -45,6 +48,20 @@ ck_tile_hcu_grouped_gemm_bw_family_bf16_nn_rr_fast_vmem3_early_final_barrier_ck_ void*, int, std::uint32_t, hipStream_t); #endif +#if defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16) +// Declared in grouped_gemm_bw_family_selected.hpp (with the enable switches). +// The v7 instances take one extra argument the other families do not need: an +// upper bound on the (group, m tile, n tile) tile count, rounded up to a whole +// number of CTAs. Their config is non-persistent (one tile per CTA, by +// construction -- see gemm_pipeline_v7_uniform.hpp), so the launcher has to size the +// grid, and per-group M is device-resident. The bound is computed from the +// problem's host-visible lengths/capacities below. +extern "C" int ck_tile_hcu_grouped_gemm_v7_bf16_nn_device_args( + void*, int, std::uint32_t, hipStream_t, std::uint64_t); +extern "C" int ck_tile_hcu_grouped_gemm_v8_bf16_tn_device_args( + void*, int, std::uint32_t, hipStream_t, std::uint64_t); +#endif + namespace { constexpr std::uint32_t ArchGfx936 = @@ -110,6 +127,12 @@ static_assert(static_cast(ck_tile::GroupedGemmCandidateFamily::gfx CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_GFX938_MLS_V1); static_assert(static_cast(ck_tile::GroupedGemmCandidateFamily::v4) == CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V4_V1); +static_assert(static_cast(InstanceId::v7_bf16_nn) == + CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V7_BF16_NN_V1); +static_assert(static_cast(InstanceId::v8_bf16_tn) == + CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V8_BF16_TN_V1); +static_assert(static_cast(ck_tile::GroupedGemmCandidateFamily::v7) == + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V7_V1); constexpr std::uint32_t KnownProblemFlags = CK_TILE_HCU_GROUPED_GEMM_PROBLEM_MAY_HAVE_EMPTY_GROUPS_V1 | @@ -273,11 +296,25 @@ int convert_problem(const ck_tile_hcu_grouped_gemm_problem_v1* source, if(!ck_tile_hcu_grouped_gemm_bf16_tn_logical_k_tail_enabled()) policy.enable_bf16_tn_logical_k_tail = false; + // v7 NN / TN: opt-in, so the env switch is the ENABLE side here (the + // selector's own default is false). See + // grouped_gemm_bw_family_selected.hpp for the switch semantics. + policy.enable_bf16_v7_nn = ck_tile_hcu_grouped_gemm_bf16_v7_nn_enabled(); + policy.enable_bf16_v8_tn = ck_tile_hcu_grouped_gemm_bf16_v8_tn_enabled(); + #if !defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_GFX936_PROMOTED) // Do not name InstanceId 1001/1002 when this TU cannot dispatch them. policy.enable_gfx936_bf16_nn_promoted = false; #endif +#if !defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16) + // Same rule for InstanceId 5000/5001: this archive does not carry the v7 + // entries (gfx936-only, dedicated codegen contract), so the route must not + // be selectable. + policy.enable_bf16_v7_nn = false; + policy.enable_bf16_v8_tn = false; +#endif + #if defined(CK_TILE_GROUPED_GEMM_P1_FAMILY_A_V3_128_GFX938_DISABLE) policy.enable_p1_family_a_v3_128_gfx938 = false; #endif @@ -403,6 +440,30 @@ int dispatch_registry_candidate(const ck_tile::GroupedGemmProblem& problem, layout == ck_tile::GroupedGemmLayout::tn) return ck_tile_hcu_grouped_gemm_bw_family_bf16_tn_blas_transposed_logical_k_tail_device_args( device_args, groups, launch.num_cu, stream); +#endif + break; + case ck_tile::GroupedGemmInstanceId::v7_bf16_nn: +#if defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16) + if(dtype == ck_tile::GroupedGemmDataType::bf16 && + layout == ck_tile::GroupedGemmLayout::nn) + return ck_tile_hcu_grouped_gemm_v7_bf16_nn_device_args( + device_args, + groups, + launch.num_cu, + stream, + v7_uniform_launch::block_id_range(problem, launch.num_cu)); +#endif + break; + case ck_tile::GroupedGemmInstanceId::v8_bf16_tn: +#if defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16) + if(dtype == ck_tile::GroupedGemmDataType::bf16 && + layout == ck_tile::GroupedGemmLayout::tn) + return ck_tile_hcu_grouped_gemm_v8_bf16_tn_device_args( + device_args, + groups, + launch.num_cu, + stream, + v7_uniform_launch::block_id_range(problem, launch.num_cu)); #endif break; case ck_tile::GroupedGemmInstanceId::gfx936_v3_256_m_only_padding: @@ -514,6 +575,16 @@ void candidate_masks(ck_tile::GroupedGemmInstanceId id, data_type_mask = DTypeFp16; layout_mask = LayoutNn; break; + case ck_tile::GroupedGemmInstanceId::v7_bf16_nn: + architecture_mask = ArchGfx936; + data_type_mask = DTypeBf16; + layout_mask = LayoutNn; + break; + case ck_tile::GroupedGemmInstanceId::v8_bf16_tn: + architecture_mask = ArchGfx936; + data_type_mask = DTypeBf16; + layout_mask = LayoutTn; + break; } } diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args_query_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args_query_test.cpp index 41322f7..500c8a7 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args_query_test.cpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_device_args_query_test.cpp @@ -120,7 +120,7 @@ int main() constexpr std::uint32_t all_layout = CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_NT_V1 | CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_NN_V1 | CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_TN_V1; - constexpr std::array expected{{ + constexpr std::array expected{{ {CK_TILE_HCU_GROUPED_GEMM_INSTANCE_BW_FAMILY_SELECTED_V1, CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_BW_V1, CK_TILE_HCU_GROUPED_GEMM_ARCHITECTURE_MASK_GFX936_V1, @@ -216,6 +216,16 @@ int main() CK_TILE_HCU_GROUPED_GEMM_ARCHITECTURE_MASK_GFX936_V1, CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_MASK_FP16_V1, CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_NN_V1}, + {CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V7_BF16_NN_V1, + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V7_V1, + CK_TILE_HCU_GROUPED_GEMM_ARCHITECTURE_MASK_GFX936_V1, + CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_MASK_BF16_V1, + CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_NN_V1}, + {CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V8_BF16_TN_V1, + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V7_V1, + CK_TILE_HCU_GROUPED_GEMM_ARCHITECTURE_MASK_GFX936_V1, + CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_MASK_BF16_V1, + CK_TILE_HCU_GROUPED_GEMM_LAYOUT_MASK_TN_V1}, }}; std::set names; bool passed = true; diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp new file mode 100644 index 0000000..633e4bf --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp @@ -0,0 +1,787 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technologies Co., Ltd. +// +// Correctness + performance driver for the v7 BF16 NN / TN grouped-GEMM +// instances (InstanceId 5000 / 5001). +// +// WHAT THIS DRIVES +// The instances themselves, through the versioned device-args ABI Primus uses: +// ck_tile_hcu_grouped_gemm_select_device_args_v1 -> must return the v7 id +// ck_tile_hcu_grouped_gemm_run_device_args_v1 -> launches the kernel +// Nothing is called behind the ABI's back: the kargs array is built here exactly +// as the ABI's builders build it (one ck_tile::GemmTransKernelArg per group, +// uploaded once, consumed by the device-side tile walk). +// +// The v7 routes are OPT-IN, so the driver sets the same switch production +// would set: +// CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN / _TN +// and --ref=1 re-runs the same kargs through the DEFAULT instance for the same +// problem (switch disabled) to give a same-process, same-descriptor A/B. +// +// WHAT IT CHECKS +// * --check=1 (default): an independent CPU reference (fp32 accumulate over +// bf16 inputs, which are exact in fp32 -- only the summation order differs) +// and, with --ref=1, the default CK instance as a triangle. +// * always: selection identity, launch status, and the achieved time/TFLOPS. +// +// Layout contract (the one the v7 NN / TN entries consume): +// NN : A [M_g, K] row-major (stride_a = K) +// B [K, N] row-major (stride_b = N) +// TN : A [K, M_g] row-major panel, per-group a_ptr, stride_a = the panel pitch +// B [K, N] row-major (stride_b = N) +// D [M_g, N] row-major (stride_e = N) +// TN's A panel load covers a whole 256-wide column tile, so the A allocation +// carries slack past the last group (kAtailPad) -- the caller-side pad +// obligation the v7 TN entry's header documents. +// +// Usage: +// grouped_gemm_v7_instance_test [--layout=NN|TN] [--m=512] [--n=512] +// [--k=256] [--groups=2] [--iters=20] +// [--warmup=3] [--check=1] [--ref=1] + +#include + +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/grouped_gemm_device_args.h" +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Bf16 = ck_tile::bf16_t; + +constexpr std::uint32_t kV7NnInstance = 5000; +constexpr std::uint32_t kV8TnInstance = 5001; + +// Slack the TN A panel needs past the last group's columns (see the header). +constexpr std::size_t kAtailPad = 1024; + +struct Options +{ + std::string layout = "NN"; + int m = 512; + int n = 512; + int k = 256; + int groups = 2; + int iters = 20; + int warmup = 3; + bool check = true; + bool ref = true; + // The independent fp32 CPU reference the checked runs compare against (bf16 + // inputs are exact in fp32 -- only the summation order differs). --check=0 + // skips it for perf runs on production shapes. + bool cpu_ref = true; + // Ragged group lengths: the shapes the v7 harness benchmarks + // (--grouped-length=8100,8100,8091,8091). Empty => `groups` copies of `m`. + std::vector group_lengths; + bool dump = false; + // Grid override. 0 => the device's CU count (the production value). + int num_cu = 0; +}; + +std::size_t align_to(std::size_t value, std::size_t alignment) +{ + return (value + alignment - 1) / alignment * alignment; +} + +// Every hip* call below whose result carries [[nodiscard]] goes through this, so +// the driver both stays -Wunused-result clean and reports the first failure +// instead of running on with garbage. +void hip_check(hipError_t status, const char* what) +{ + if(status != hipSuccess) + { + std::fprintf(stderr, "HIP error %d (%s) at %s\n", static_cast(status), + hipGetErrorName(status), what); + std::exit(1); + } +} + +// Small multiples of 0.125: exact in bf16 (8 mantissa bits), so the only +// difference between the CPU reference and the kernel is the summation order. +// Both indices contribute, so no lane pattern of the kernels cancels a term. +float a_value(int row, int kk) +{ + return static_cast((row % 7) - 3) * 0.5f + + static_cast(kk % 3) * 0.25f; +} +float b_value(int kk, int col) +{ + return static_cast((col % 5) - 2) * 0.25f + + static_cast(kk % 4) * 0.125f; +} + +struct Case +{ + // Per-group row counts (ragged) plus their sum. `m` stays the first group's + // length for the single-group paths; every buffer is sized from group_m. + std::vector group_m; + int total_m = 0; + bool tn = false; + int m = 0; + int n = 0; + int k = 0; + int g = 0; + std::vector strides_a; + std::vector strides_b; + std::vector a_elements; + std::vector b_elements; + std::vector> host_a; + std::vector> host_b; + std::vector> host_d_ref; +}; + +int parse_int_arg(const char* text, const char* name) +{ + char* end = nullptr; + const long value = std::strtol(text, &end, 10); + if(end == text || value < 0) + { + std::fprintf(stderr, "bad --%s=%s\n", name, text); + std::exit(2); + } + return static_cast(value); +} + +void parse_args(int argc, char* argv[], Options& options) +{ + for(int i = 1; i < argc; ++i) + { + const std::string arg = argv[i]; + const auto value_of = [&](const char* name) -> const char* { + const std::size_t at = arg.find('='); + if(at == std::string::npos) + { + std::fprintf(stderr, "%s expects --%s=VALUE\n", arg.c_str(), name); + std::exit(2); + } + return argv[i] + at + 1; + }; + if(arg.rfind("--layout=", 0) == 0) + { + options.layout = value_of("layout"); + } + else if(arg.rfind("--m=", 0) == 0) + { + options.m = parse_int_arg(value_of("m"), "m"); + } + else if(arg.rfind("--n=", 0) == 0) + { + options.n = parse_int_arg(value_of("n"), "n"); + } + else if(arg.rfind("--k=", 0) == 0) + { + options.k = parse_int_arg(value_of("k"), "k"); + } + else if(arg.rfind("--groups=", 0) == 0) + { + options.groups = parse_int_arg(value_of("groups"), "groups"); + } + else if(arg.rfind("--iters=", 0) == 0) + { + options.iters = parse_int_arg(value_of("iters"), "iters"); + } + else if(arg.rfind("--warmup=", 0) == 0) + { + options.warmup = parse_int_arg(value_of("warmup"), "warmup"); + } + else if(arg.rfind("--check=", 0) == 0) + { + options.check = parse_int_arg(value_of("check"), "check") != 0; + } + else if(arg.rfind("--ref=", 0) == 0) + { + options.ref = parse_int_arg(value_of("ref"), "ref") != 0; + } + else if(arg.rfind("--group-lengths=", 0) == 0) + { + options.group_lengths.clear(); + const std::string list = value_of("group-lengths"); + std::size_t start = 0; + while(start <= list.size()) + { + const std::size_t comma = list.find(',', start); + const std::string item = list.substr(start, comma - start); + if(!item.empty()) + { + options.group_lengths.push_back(parse_int_arg(item.c_str(), "group-lengths")); + } + if(comma == std::string::npos) break; + start = comma + 1; + } + } + else if(arg.rfind("--cpu-ref=", 0) == 0) + { + options.cpu_ref = parse_int_arg(value_of("cpu-ref"), "cpu-ref") != 0; + } + else if(arg.rfind("--dump=", 0) == 0) + { + options.dump = parse_int_arg(value_of("dump"), "dump") != 0; + } + else if(arg.rfind("--num-cu=", 0) == 0) + { + options.num_cu = parse_int_arg(value_of("num-cu"), "num-cu"); + } + else + { + std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); + std::exit(2); + } + } + if(options.layout != "NN" && options.layout != "TN") + { + std::fprintf(stderr, "--layout must be NN or TN\n"); + std::exit(2); + } +} + +void build_case(const Options& options, Case& c) +{ + c.tn = options.layout == "TN"; + c.n = options.n; + c.k = options.k; + c.group_m = options.group_lengths.empty() + ? std::vector(static_cast(options.groups), options.m) + : options.group_lengths; + c.g = static_cast(c.group_m.size()); + c.m = c.group_m.front(); + c.total_m = 0; + for(int g = 0; g < c.g; ++g) c.total_m += c.group_m[g]; + // TN's A panel is [K, sum(M_g)]: it is addressed across the whole group set, + // so its pitch is the summed M (padded up to a whole tile, which is also the + // 256-element slack the panel fetch needs past the last group's columns). + const int a_pitch_tn = (c.total_m + 255) / 256 * 256; + + c.strides_a.resize(c.g); + c.strides_b.resize(c.g); + c.a_elements.resize(c.g); + c.b_elements.resize(c.g); + c.host_a.resize(c.g); + c.host_b.resize(c.g); + c.host_d_ref.resize(c.g); + + for(int g = 0; g < c.g; ++g) + { + // NN: A row-major [M_g, K]. TN: A is the [K, M_g] panel of the + // column-major [M_g, K] view, so its pitch is the M extent (the same + // choice the device-args smoke makes for its TN cases). + const int stride_a = c.tn ? a_pitch_tn : c.k; + const int stride_b = c.n; + c.strides_a[g] = stride_a; + c.strides_b[g] = stride_b; + c.a_elements[g] = + (c.tn ? c.k : c.group_m[g]) * stride_a + (c.tn ? static_cast(kAtailPad) : 0); + c.b_elements[g] = c.k * stride_b; + + c.host_a[g].assign(static_cast(c.a_elements[g]), + ck_tile::type_convert(0.0f)); + c.host_b[g].assign(static_cast(c.b_elements[g]), + ck_tile::type_convert(0.0f)); + c.host_d_ref[g].assign(static_cast(c.group_m[g]) * c.n, 0.0f); + + for(int row = 0; row < c.group_m[g]; ++row) + { + for(int kk = 0; kk < c.k; ++kk) + { + const std::size_t at = + c.tn ? static_cast(kk) * stride_a + row + : static_cast(row) * stride_a + kk; + c.host_a[g][at] = ck_tile::type_convert(a_value(row, kk)); + } + } + for(int kk = 0; kk < c.k; ++kk) + { + for(int col = 0; col < c.n; ++col) + { + c.host_b[g][static_cast(kk) * stride_b + col] = + ck_tile::type_convert(b_value(kk, col)); + } + } + } +} + +void cpu_reference(Case& c) +{ + for(int g = 0; g < c.g; ++g) + { + const int stride_a = c.strides_a[g]; + const int stride_b = c.strides_b[g]; + for(int col = 0; col < c.n; ++col) + { + for(int row = 0; row < c.m; ++row) + { + float acc = 0.0f; + for(int kk = 0; kk < c.k; ++kk) + { + const std::size_t at = + c.tn ? static_cast(kk) * stride_a + row + : static_cast(row) * stride_a + kk; + const float a = ck_tile::type_convert(c.host_a[g][at]); + const float b = ck_tile::type_convert( + c.host_b[g][static_cast(kk) * stride_b + col]); + acc += a * b; + } + c.host_d_ref[g][static_cast(row) * c.n + col] = acc; + } + } + } +} + +struct DeviceCase +{ + std::vector a_ptrs; + std::vector b_ptrs; + std::vector c_ptrs; + std::vector a_blocks; + std::vector b_blocks; + std::vector c_blocks; + void* kargs = nullptr; + std::size_t kargs_sz = 0; + bool allocated = false; + + ~DeviceCase() { release(); } + + void release() + { + if(!allocated) return; + for(auto* p : a_blocks) hip_check(hipFree(p), "hipFree"); + for(auto* p : b_blocks) hip_check(hipFree(p), "hipFree"); + for(auto* p : c_blocks) hip_check(hipFree(p), "hipFree"); + if(kargs != nullptr) hip_check(hipFree(kargs), "hipFree"); + a_blocks.clear(); b_blocks.clear(); c_blocks.clear(); + kargs = nullptr; + allocated = false; + } + + void allocate(const Case& c, const Options& options) + { + a_ptrs.resize(c.g); + b_ptrs.resize(c.g); + c_ptrs.resize(c.g); + a_blocks.resize(c.g); + b_blocks.resize(c.g); + c_blocks.resize(c.g); + + for(int g = 0; g < c.g; ++g) + { + const std::size_t a_bytes = + static_cast(c.a_elements[g]) * sizeof(Bf16); + const std::size_t b_bytes = + static_cast(c.b_elements[g]) * sizeof(Bf16); + const std::size_t c_bytes = + static_cast(c.m) * c.n * sizeof(Bf16); + hip_check(hipMalloc(&a_blocks[g], align_to(a_bytes, 256)), "hipMalloc(A)"); + hip_check(hipMalloc(&b_blocks[g], align_to(b_bytes, 256)), "hipMalloc(B)"); + hip_check(hipMalloc(&c_blocks[g], align_to(c_bytes, 256)), "hipMalloc(C)"); + hip_check(hipMemcpy(a_blocks[g], + c.host_a[g].data(), + a_bytes, + hipMemcpyHostToDevice), + "hipMemcpy(A)"); + hip_check(hipMemcpy(b_blocks[g], + c.host_b[g].data(), + b_bytes, + hipMemcpyHostToDevice), + "hipMemcpy(B)"); + hip_check(hipMemset(c_blocks[g], 0, c_bytes), "hipMemset(C)"); + a_ptrs[g] = a_blocks[g]; + b_ptrs[g] = b_blocks[g]; + c_ptrs[g] = c_blocks[g]; + } + + // block_start/block_end must be filled: the framework's non-persistent + // kernel path maps a 1D block id to (group, m tile, n tile) through them, + // which is how the v7 uniform instances resolve their tile. + const ck_tile::index_t n_tiles = (c.n + 255) / 256; + ck_tile::index_t block_cursor = 0; + std::vector host_kargs; + host_kargs.reserve(c.g); + for(int g = 0; g < c.g; ++g) + { + host_kargs.emplace_back(ck_tile::UniversalGemmKernelArgs<>{{a_ptrs[g]}, + {b_ptrs[g]}, + {}, + c_ptrs[g], + c.group_m[g], + c.n, + c.k, + {c.strides_a[g]}, + {c.strides_b[g]}, + {}, + c.n, + 1}, + block_cursor, + block_cursor + + ((c.group_m[g] + 255) / 256) * n_tiles); + block_cursor += ((c.group_m[g] + 255) / 256) * n_tiles; + } + (void)options; + kargs_sz = host_kargs.size() * sizeof(ck_tile::GemmTransKernelArg); + hip_check(hipMalloc(&kargs, kargs_sz), "hipMalloc(kargs)"); + hip_check(hipMemcpy(kargs, host_kargs.data(), kargs_sz, hipMemcpyHostToDevice), + "hipMemcpy(kargs)"); + allocated = true; + } +}; + +// --dump=1: the first few host operands, the CPU reference and the device +// output. Values, not aggregates: an aggregate cannot tell a broken mapping +// from a broken reference. +void dump_case(const Case& c, const std::vector& c_blocks) +{ + std::printf("dump host_a[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(ck_tile::type_convert(c.host_a[0][i]))); + std::printf("\ndump host_b[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(ck_tile::type_convert(c.host_b[0][i]))); + std::printf("\ndump host_d[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(c.host_d_ref[0][i])); + std::vector raw(16); + hip_check(hipMemcpy(raw.data(), + c_blocks[0], + raw.size() * sizeof(unsigned short), + hipMemcpyDeviceToHost), + "hipMemcpy(dump)"); + std::printf("\ndump device[0][0..7]:"); + for(int i = 0; i < 8; ++i) + { + const Bf16 value = *reinterpret_cast(&raw[i]); + std::printf(" %g(0x%04x)", + static_cast(ck_tile::type_convert(value)), + raw[i]); + } + std::printf("\n"); +} + +void build_problem(const Case& c, ck_tile_hcu_grouped_gemm_problem_v1& problem) +{ + problem = {}; + problem.struct_size = sizeof(problem); + problem.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + problem.data_type = CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_BF16_V1; + problem.layout = c.tn ? CK_TILE_HCU_GROUPED_GEMM_LAYOUT_TN_V1 + : CK_TILE_HCU_GROUPED_GEMM_LAYOUT_NN_V1; + problem.group_count = c.g; + problem.problem_flags = 0; + problem.policy_disable_flags = 0; + // Ragged per-group M: declared the way the ABI has it -- a device-lengths + // dimension with a host-visible capacity (the summed rows, plus one tile of + // rounding per group). A common M stays expressible too. + const bool common_m = (c.total_m == c.m * c.g); + problem.m = {common_m ? CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1 + : CK_TILE_HCU_GROUPED_GEMM_DIMENSION_DEVICE_LENGTHS_WITH_CAPACITY_V1, + 0, + common_m ? static_cast(c.m) + : static_cast(c.total_m) + 256 * c.g}; + problem.n = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.n}; + problem.k = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.k}; +} + +int select_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + std::uint32_t& instance_id, + std::string& name) +{ + ck_tile_hcu_grouped_gemm_selection_v1 selection{}; + selection.struct_size = sizeof(selection); + selection.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + const int status = ck_tile_hcu_grouped_gemm_select_device_args_v1(&problem, &selection); + if(status != 0) return status; + instance_id = selection.instance_id; + const char* instance_name = + ck_tile_hcu_grouped_gemm_instance_name_v1(selection.instance_id); + name = instance_name == nullptr ? std::string("?") : std::string(instance_name); + return 0; +} + +int run_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + void* kargs, + std::uint32_t instance_id, + std::uint32_t num_cu, + hipStream_t stream) +{ + ck_tile_hcu_grouped_gemm_launch_v1 launch{}; + launch.struct_size = sizeof(launch); + launch.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + launch.device_args = kargs; + launch.num_cu = num_cu; + launch.stream = reinterpret_cast(stream); + launch.expected_instance_id = instance_id; + ck_tile_hcu_grouped_gemm_selection_v1 launched{}; + launched.struct_size = sizeof(launched); + launched.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + return ck_tile_hcu_grouped_gemm_run_device_args_v1(&problem, &launch, &launched); +} + +double time_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + void* kargs, + std::uint32_t instance_id, + std::uint32_t num_cu, + hipStream_t stream, + int warmup, + int iters) +{ + for(int i = 0; i < warmup; ++i) + { + run_instance(problem, kargs, instance_id, num_cu, stream); + } + hipEvent_t start{}; + hipEvent_t stop{}; + hip_check(hipEventCreate(&start), "hipEventCreate"); + hip_check(hipEventCreate(&stop), "hipEventCreate"); + hip_check(hipEventRecord(start, stream), "hipEventRecord"); + for(int i = 0; i < iters; ++i) + { + run_instance(problem, kargs, instance_id, num_cu, stream); + } + hip_check(hipEventRecord(stop, stream), "hipEventRecord"); + hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); + float ms = 0.0f; + hip_check(hipEventElapsedTime(&ms, start, stop), "hipEventElapsedTime"); + hip_check(hipEventDestroy(start), "hipEventDestroy"); + hip_check(hipEventDestroy(stop), "hipEventDestroy"); + return static_cast(ms) / iters; +} + +struct CheckResult +{ + double max_err = 0.0; + std::size_t bad = 0; + std::size_t total = 0; + // First few mismatches, as (group, row, col, got, expected): the pattern is + // what identifies a broken tile/row/column mapping, and aggregate counts do + // not carry it. + struct Sample + { + int g; + int row; + int col; + float got; + float ref; + }; + std::vector samples; +}; + +CheckResult check_against_reference(const Case& c, + const std::vector& blocks, + int stride_c) +{ + CheckResult result; + for(int g = 0; g < c.g; ++g) + { + std::vector host_d(static_cast(c.m) * stride_c); + hip_check(hipMemcpy(host_d.data(), + blocks[g], + host_d.size() * sizeof(Bf16), + hipMemcpyDeviceToHost), + "hipMemcpy(D)"); + for(int row = 0; row < c.m; ++row) + { + for(int col = 0; col < c.n; ++col) + { + const std::size_t at = static_cast(row) * stride_c + col; + const double got = static_cast(ck_tile::type_convert(host_d[at])); + const double ref = + static_cast(c.host_d_ref[g][static_cast(row) * c.n + col]); + const double err = std::fabs(got - ref); + if(err > result.max_err) result.max_err = err; + ++result.total; + if(err > 0.5) + { + ++result.bad; + if(result.samples.size() < 8) + { + result.samples.push_back({g, + row, + col, + static_cast(got), + static_cast(ref)}); + } + } + } + } + } + return result; +} + +void print_check(const char* tag, const CheckResult& checked) +{ + std::printf("check %s max_err=%.6g bad=%zu/%zu\n", + tag, + checked.max_err, + checked.bad, + checked.total); + for(const auto& s : checked.samples) + { + std::printf(" sample g=%d row=%d col=%d got=%.6g ref=%.6g\n", + s.g, + s.row, + s.col, + static_cast(s.got), + static_cast(s.ref)); + } +} + +} // namespace + +int main(int argc, char* argv[]) +{ + Options options; + parse_args(argc, argv, options); + + Case c; + build_case(options, c); + if(options.check && options.cpu_ref) + { + cpu_reference(c); + } + + DeviceCase device; + device.allocate(c, options); + + ck_tile_hcu_grouped_gemm_problem_v1 problem{}; + build_problem(c, problem); + + const std::uint32_t num_cu = options.num_cu > 0 + ? static_cast(options.num_cu) + : static_cast(ck_tile::get_num_cu()); + hipStream_t stream = nullptr; + std::printf("grid num_cu=%u (device=%d)\n", + num_cu, + static_cast(ck_tile::get_num_cu())); + + bool ok = true; + + const char* enable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V8_TN" + : "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN"; + const char* disable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V8_TN" + : "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V7_NN"; + + // --- v7 instance ----------------------------------------------------- + setenv(enable_switch, "1", 1); + unsetenv(disable_switch); + + std::uint32_t v7_id = 0; + std::string v7_name; + if(select_instance(problem, v7_id, v7_name) != 0) + { + std::printf("layout=%s result=FAIL reason=select_status\n", options.layout.c_str()); + return 1; + } + const std::uint32_t expected = c.tn ? kV8TnInstance : kV7NnInstance; + std::printf("select layout=%s instance_id=%u expected=%u name=%s\n", + options.layout.c_str(), + v7_id, + expected, + v7_name.c_str()); + ok = ok && v7_id == expected; + + const int run_status = + run_instance(problem, device.kargs, v7_id, num_cu, stream); + if(run_status != 0) + { + std::printf("layout=%s result=FAIL reason=run_status=%d\n", options.layout.c_str(), run_status); + return 1; + } + const double v7_ms = + time_instance(problem, device.kargs, v7_id, num_cu, stream, options.warmup, options.iters); + + if(options.dump) + { + dump_case(c, device.c_ptrs); + } + + if(options.check) + { + if(options.cpu_ref) + { + const CheckResult checked = check_against_reference(c, device.c_ptrs, c.n); + print_check("port_vs_cpu", checked); + ok = ok && checked.bad == 0; + } + } + + const double flops = 2.0 * static_cast(c.m) * c.n * c.k * c.g; + std::printf("perf layout=%s m=%d n=%d k=%d groups=%d v7_ms=%.4f v7_tf=%.1f\n", + options.layout.c_str(), + c.m, + c.n, + c.k, + c.g, + v7_ms, + flops / (v7_ms * 1.0e-3) / 1.0e12); + + // --- default CK instance for the same kargs (same-process A/B) ---------- + if(options.ref) + { + setenv(disable_switch, "1", 1); + std::uint32_t ref_id = 0; + std::string ref_name; + if(select_instance(problem, ref_id, ref_name) != 0) + { + std::printf("ref layout=%s result=SKIP reason=select_status\n", options.layout.c_str()); + } + else if(ref_id == v7_id) + { + std::printf("ref layout=%s result=SKIP reason=same_instance\n", options.layout.c_str()); + } + else + { + // Contract check: the TN candidates in the registry consume the + // TRANSPOSED view of the problem (C^T = B^T A), i.e. a different + // kargs meaning than the direct view this driver builds for the v7 + // entry. Running them on these kargs would compare two different + // problems, so skip instead of printing a meaningless number. + ck_tile_hcu_grouped_gemm_candidate_info_v1 ref_info{}; + ref_info.struct_size = sizeof(ref_info); + ref_info.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + const bool ref_transposed_view = + ck_tile_hcu_grouped_gemm_query_instance_v1(ref_id, &ref_info) == 0 && + (ref_info.contract_flags & + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_CONTRACT_C_TRANSPOSED_VIEW_V1) != 0; + + const int ref_status = ref_transposed_view + ? 0 + : run_instance(problem, device.kargs, ref_id, num_cu, stream); + if(ref_transposed_view) + { + std::printf("ref layout=%s instance_id=%u name=%s result=SKIP " + "reason=transposed_view_contract\n", + options.layout.c_str(), + ref_id, + ref_name.c_str()); + } + else if(ref_status != 0) + { + std::printf("ref layout=%s instance_id=%u result=SKIP reason=run_status=%d\n", + options.layout.c_str(), + ref_id, + ref_status); + } + else + { + const double ref_ms = time_instance( + problem, device.kargs, ref_id, num_cu, stream, options.warmup, options.iters); + std::printf("ref layout=%s instance_id=%u name=%s ref_ms=%.4f ref_tf=%.1f\n", + options.layout.c_str(), + ref_id, + ref_name.c_str(), + ref_ms, + flops / (ref_ms * 1.0e-3) / 1.0e12); + } + } + unsetenv(disable_switch); + } + + std::printf("layout=%s result=%s\n", options.layout.c_str(), ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_uniform_epilogue.hpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_uniform_epilogue.hpp new file mode 100644 index 0000000..f7385f2 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_uniform_epilogue.hpp @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// +// Store epilogue for the uniform Direct-to-LDS NN/TN pipelines +// (gemm_pipeline_v7_uniform.hpp). Same writeback the reference kernels use, +// expressed in ck_tile terms. +// +// WHY NOT CShuffleEpilogue +// ------------------------ +// The MMAC accumulator puts one lane's four results at columns {q, q+4, q+8, +// q+12} (q = lane >> 4) across 16 rows, so direct scalar stores never coalesce: +// measured 33.2M VL1 write requests on the reference shape against hipBLASLt's +// 8.4M, and because two CTAs overlap per CU that write flood congests TA/TC for +// the WHOLE kernel. Staging each 16-row slice through the (by then idle) ring +// LDS and reading it back so every lane holds one contiguous 16 B run turns 128 +// scalar stores into 16 vector stores per wave: 33.2M -> ~1M requests, measured +// 1.100 ms / 247 TF -> 0.852 ms / 319 TF on the reference shape. +// +// Shuffle-free (per-wave LDS slabs, no block barrier) and 16 B stores are the two +// properties that matter; the layout below is the reference kernel's. +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp" + +// The address arithmetic below is a port of the reference kernels, which are +// built with -Wno-old-style-cast; every cast is a fixed-width unsigned/signed +// conversion in an address expression. CK's own arch/buffer-addressing headers +// handle it the same way (see cast_pointer_to_constant_address_space). The +// 16 B store path reinterprets the output element pointer as one dwordx4 +// vector -- an alignment-checked bit reinterpretation, not a type pun. +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wold-style-cast" +#pragma clang diagnostic ignored "-Wundefined-reinterpret-cast" + +namespace ck_tile { + +namespace v7u_epi_detail { + +CK_TILE_DEVICE void ds_write_ele(uint32_t addr, uint16_t v) +{ + asm volatile("ds_write_b16 %0, %1" ::"v"(addr), "v"((uint32_t)v) : "memory"); +} +CK_TILE_DEVICE void ds_write_ele(uint32_t addr, uint32_t v) +{ + asm volatile("ds_write_b32 %0, %1" ::"v"(addr), "v"(v) : "memory"); +} + +using u32x4 = int32_t __attribute__((ext_vector_type(4))); +CK_TILE_DEVICE u32x4 ds_read_b128(uint32_t addr) +{ + u32x4 v; + asm volatile("ds_read_b128 %0, %1" : "=v"(v) : "v"(addr) : "memory"); + return v; +} + +// fp32 -> output. The bf16 arm is the reference kernels' round-to-nearest-even +// shift (bit-identical to what the validated trees produce); fp16 goes through +// the hardware convert. +template +CK_TILE_DEVICE Element_O fp32_to(float f); +template <> +CK_TILE_DEVICE float fp32_to(float f) +{ + return f; +} +template <> +CK_TILE_DEVICE bf16_t fp32_to(float f) +{ + union + { + float fp32; + uint32_t u32; + } u = {f}; + u.u32 += 0x7fff + ((u.u32 >> 16) & 1); + bf16_t r; + *reinterpret_cast(&r) = (uint16_t)(u.u32 >> 16); + return r; +} +template <> +CK_TILE_DEVICE half_t fp32_to(float f) +{ + return __float2half(f); +} + +} // namespace v7u_epi_detail + +// --------------------------------------------------------------------------- +// The epilogue. `store_tile` is called with the accumulator the pipeline filled +// and the same LDS the ring used (which the k loop has released by then). +// --------------------------------------------------------------------------- +template +struct GroupedGemmV7UniformStoreEpilogue +{ + using Element_O = remove_cvref_t; + using AccTile = remove_cvref_t; + using G = Geometry; + + // ---- framework-facing surface (what GroupedGemmKernel reads) -------- + using ODataType = Element_O; + using DsDataType = ck_tile::tuple<>; // no D tensors: single-stage store + using DsLayout = ck_tile::tuple<>; + static constexpr index_t NumDTensor = 0; + static constexpr memory_operation_enum MemoryOperation = memory_operation_enum::set; + + static constexpr int kE = (int)sizeof(Element_O); // 2 | 4 + static constexpr int kRowB = (int)G::kWaveTileN * kE; // 256 | 512 + static constexpr int kUnits = kRowB / 16; // 16 | 32 + static constexpr int kSliceB = 16 * kRowB; // 4 KB | 8 KB + + // The epilogue reuses the pipeline's ring buffer, so it asks for none of its + // own: the framework's max(pipeline, epilogue) keeps the pipeline's 64 KB. + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return 0; } + + static_assert(kE == 2 || kE == 4, "2- or 4-byte output only"); + static_assert(kSliceB * (int)G::kWaves <= (int)G::kLdsBytes, + "the per-wave staging slabs must fit the ring LDS"); + + // The store ALSO drains whatever loads the k tail left outstanding: their + // LDS write-backs would otherwise land in the ring while the slabs are being + // staged (observed as a few hundred bad elements, varying per run). + CK_TILE_DEVICE static void store_tile(const V7UniformTileParams& p, + const AccTile& acc, + void* smem) + { + using namespace v7u_epi_detail; + (void)smem; + + Element_O* d_out = reinterpret_cast(p.d_ptr); + const index_t d_stride = p.d_stride; + + const int tid = (int)threadIdx.x; + const int wid = tid >> 6; + const int lane = tid & 63; + const int wave_m_start = (wid >> 1) * (int)G::kWaveTileM; + const int wave_n_start = (wid & 1) * (int)G::kWaveTileN; + + v7u::wait_all_buffer_data_arrived(0); + + const uint32_t epi_base = (uint32_t)(wid * kSliceB); + const uint32_t row_q = (uint32_t)(lane & 15); + const uint32_t row_xor = (uint32_t)(lane & 7); + const uint32_t col_q = (uint32_t)(lane >> 4); + +#pragma unroll + for(index_t mF = 0; mF < G::kMTilesPerWave; ++mF) + { +#pragma unroll + for(index_t nF = 0; nF < G::kNTilesPerWave; ++nF) + { + const index_t idx = mF * G::kNTilesPerWave + nF; +#pragma unroll + for(int c = 0; c < 4; ++c) + { + const uint32_t byte = (col_q + (uint32_t)(nF * 16 + 4 * c)) * (uint32_t)kE; + const uint32_t addr = epi_base + row_q * (uint32_t)kRowB + + (((byte >> 4) ^ row_xor) << 4) + (byte & 15u); + const Element_O val = fp32_to(acc.v[idx][c]); + if constexpr(kE == 2) + { + ds_write_ele(addr, *reinterpret_cast(&val)); + } + else + { + ds_write_ele(addr, *reinterpret_cast(&val)); + } + } + } + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + // Read the whole slice back FIRST, then store: the ds_read results + // are inline-asm VGPR defs, so the compiler will not insert an LGKM + // wait before the consuming global stores. + constexpr int kJ = kSliceB / 64 / 16; // 4 | 8 reads per lane + u32x4 vj[kJ]; +#pragma unroll + for(int j = 0; j < kJ; ++j) + { + const uint32_t t = (uint32_t)lane + 64u * (uint32_t)j; + const uint32_t r_in = t / (uint32_t)kUnits; + const uint32_t u = t % (uint32_t)kUnits; + const uint32_t addr = epi_base + r_in * (uint32_t)kRowB + + ((u ^ (r_in & 7u)) << 4); + vj[j] = ds_read_b128(addr); + } + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); +#pragma unroll + for(int j = 0; j < kJ; ++j) + { + const uint32_t t = (uint32_t)lane + 64u * (uint32_t)j; + const uint32_t r_in = t / (uint32_t)kUnits; + const uint32_t u = t % (uint32_t)kUnits; + const index_t row = p.m_start + wave_m_start + mF * 16 + (index_t)r_in; + if(row < p.m) + { + const index_t col0 = + p.n_start + wave_n_start + (index_t)u * (16 / kE); + *reinterpret_cast(&d_out[row * d_stride + col0]) = vj[j]; + } + } + } + } +}; + +} // namespace ck_tile + +#pragma clang diagnostic pop diff --git a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_device_args_provider.hpp b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_device_args_provider.hpp index b6cb387..17eafb2 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_device_args_provider.hpp +++ b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_device_args_provider.hpp @@ -154,7 +154,13 @@ int launch_registry_candidate(ck_tile::GroupedGemmInstanceId instance_id, case ck_tile::GroupedGemmInstanceId::bw_family_bf16_nn_fixed_srd_lds8: case ck_tile::GroupedGemmInstanceId::bw_family_bf16_nn_fast_vmem3: case ck_tile::GroupedGemmInstanceId::bw_family_blas_transposed: - case ck_tile::GroupedGemmInstanceId::bw_family_logical_k_tail: break; + case ck_tile::GroupedGemmInstanceId::bw_family_logical_k_tail: + // The v7 instances (5000/5001) are not registry candidates: they are + // dispatched by grouped_gemm_device_args.cpp straight to their device-args + // entries (see dispatch_registry_candidate), so reaching this switch with + // one of them is a caller error and falls through to unsupported. + case ck_tile::GroupedGemmInstanceId::v7_bf16_nn: + case ck_tile::GroupedGemmInstanceId::v8_bf16_tn: break; #if defined(CK_TILE_GROUPED_GEMM_DEVICE_ARGS_REGISTRY_GFX936) case ck_tile::GroupedGemmInstanceId::gfx936_v3_256_m_only_padding: if(layout == ck_tile::GroupedGemmLayout::nt) diff --git a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_bf16_nn_device_args.cpp b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_bf16_nn_device_args.cpp new file mode 100644 index 0000000..49303a4 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_bf16_nn_device_args.cpp @@ -0,0 +1,52 @@ +// grouped_gemm_v7_bf16_nn_device_args.cpp — CK grouped-GEMM instance entry +// (InstanceId 5000) for the uniform Direct-to-LDS NN (dgrad) kernel. +// +// The whole kernel lives in ck_tile: pipeline (gemm_pipeline_v7_uniform.hpp), epilogue +// (grouped_gemm_v7_uniform_epilogue.hpp) and config, instantiated through the +// framework's own GroupedGemmKernel by grouped_gemm_v7_uniform_launch.hpp. What this +// entry adds is the ABI shape: it receives the device kargs array the CK +// device-args ABI builds (one ck_tile::GemmTransKernelArg per group, built on +// the device, never read back to the host) plus the host-side tile bound, and +// launches one flat grid over it. +// +// CONTRACT (enforced by the CK selector, not here): bf16 in, bf16 out, gfx936, +// NN layout (A [M_g, K] row-major, B [K, N] row-major, D [M_g, N] row-major), +// K % 16 == 0, K >= 64 (the 4-stage ring), N % 256 == 0, +// 1 <= G <= kGroupedGemmV7MaxGroups, positive group lengths. A group's M may +// be ragged -- the epilogue masks rows past M_g. +// +// This TU is compiled under the v7 codegen contract (see CMakeLists.txt: the +// -mllvm switches and -fno-finite-math-only the kernel's register budget was +// tuned under). Do not move it into the provider manifest, whose consumers +// compile with their own options. +// +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// SPDX-License-Identifier: MIT + +#include "grouped_gemm_v7_uniform_launch.hpp" + +extern "C" int ck_tile_hcu_grouped_gemm_v7_bf16_nn_device_args( + void* device_args, + int group_count, + std::uint32_t num_cu, + hipStream_t stream, + std::uint64_t tile_slots) { + (void)num_cu; // the grid is sized by tile_slots, not by the CU count + if (device_args == nullptr || group_count <= 0 || + group_count > ck_tile::kGroupedGemmV7MaxGroups || tile_slots == 0) { + return -1; + } + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + return v7_uniform_launch::launch(device_args, + group_count, + tile_slots, + stream); +} diff --git a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_uniform_launch.hpp b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_uniform_launch.hpp new file mode 100644 index 0000000..244e3e5 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v7_uniform_launch.hpp @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// +// Launch helper for the v7 uniform (NN / TN) grouped-GEMM instances. +// +// This is the CK-kernel path: the whole tile chain +// (TileGemmShape -> TilePartitioner -> GroupedGemmKernel -> pipeline -> epilogue) +// is instantiated from the config, exactly as the other CK instances do it. What +// differs is the GRID: +// +// * the config is NON-persistent, so the kernel maps block id -> (group, +// m tile, n tile) from the kargs' block ranges and runs exactly ONE tile. +// The kernel cannot be made persistent: this pipeline's register budget +// (215 VGPR, 0 spill, the reference kernel's numbers) only survives when no +// loop wraps its k loop -- see the pipeline header. +// * the grid therefore has to cover the tile count, which the device-args +// provider cannot read (per-group M lives in device memory). The dispatch +// passes a host-side upper bound, rounded up to a whole number of CTAs: +// BlockIdRange() below is that bound, and the kernels' `t >= total` tiles +// simply exit early when the bound over-covers. +#pragma once + +#include + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/grouped_gemm_selector.hpp" // GroupedGemmProblem +#include "../grouped_gemm.hpp" +#include "../grouped_gemm_v7_uniform_epilogue.hpp" + +// ----------------------------------------------------------------------------- +// v7/v8 uniform NN / TN configs: 256x256x16 bf16, one tile per CTA, +// Direct-to-LDS. They live HERE (not in grouped_gemm.hpp) on purpose: the +// configs are consumed only by the v7/v8 translation units, so tuning them -- +// the common edit while debugging -- recompiles these two TUs, not the whole +// grouped-GEMM archive. +// +// GEOMETRY vs the CONFIG. The pipeline is fixed at 256x256x16 with a 4m x 2n +// wave grid (64M x 128N per wave, 32 MMACs per k-stage) and a 4-stage 16 KB LDS +// ring; the numbers below are what the FRAMEWORK needs to size its tile grid and +// to derive (M,N) tile origins from the block id. K_Tile is 64 because the +// framework only uses it for tile/loop bookkeeping -- the pipeline's own k +// loop is the 4-stage ring, advanced by the 16-wide k-stage. +// +// Persistent = false is load-bearing, not taste: the non-persistent kernel path +// maps block id -> (group, m tile, n tile) from the kargs' block ranges and then +// runs ONE tile, which is the only structure this pipeline compiles to the +// reference kernel's register budget (215 VGPR / 0 spill; a tile loop costs 40 +// VGPRs and ~20% of the throughput, see the pipeline header). +// +// kPadM/N/K stay false: ragged M and the N tail are handled by the epilogue's +// store mask and the selector's shape gates, not by padding the operands. +// ----------------------------------------------------------------------------- +template +struct GemmConfigV7Uniform : public GemmConfigBase +{ + static_assert(std::is_same_v, + "the v7/v8 uniform pipelines are bf16 kernels"); + + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 4; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool kPadM = false; + static constexpr bool kPadN = false; + static constexpr bool kPadK = false; + + static constexpr bool Persistent = false; + static constexpr bool DoubleSmemBuffer = false; + + static constexpr bool SupportsFastRowMajorB = true; + static constexpr bool SupportsFastColumnMajorA = true; + + static constexpr ck_tile::index_t TileParitionerGroupNum = 1; + static constexpr ck_tile::index_t TileParitionerM01 = 1; + + static constexpr int kBlockPerCu = 1; // 64 KB ring => one CTA per CU + + static constexpr ck_tile::index_t Pipeline = PipelineId; +}; + +using GemmConfigV7NnBf16 = + GemmConfigV7Uniform; +using GemmConfigV8TnBf16 = + GemmConfigV7Uniform; + +namespace v7_uniform_launch { + +// Upper bound on the (group, m tile, n tile) tile count for the v7 layouts, +// rounded up to a multiple of `num_cu` (one CTA per tile, so the grid must be a +// whole number of CTA slots). `problem.m`/`n` carry either a logical common +// length or a host-visible capacity -- never device-read values -- so this is +// computable here and stays graph-capture safe. +// +// Device-resident M: the capacity is the backing storage, i.e. an upper bound on +// the summed group rows, and each group can round up to its own tile, so +// ceil(capacity / 256) + group_count is the bound on the summed row tiles. +inline std::uint64_t block_id_range(const ck_tile::GroupedGemmProblem& problem, + std::uint32_t num_cu) +{ + const std::int64_t groups = problem.group_count > 0 ? problem.group_count : 1; + const std::int64_t n_tiles = (problem.n.extent + 255) / 256; + const bool device_m = + problem.m.kind == ck_tile::GroupedGemmDimensionKind::device_lengths_with_capacity; + const std::int64_t m_tiles = device_m ? (problem.m.extent + 255) / 256 + groups + : (problem.m.extent + 255) / 256 * groups; + const std::uint64_t tiles = static_cast(m_tiles * n_tiles); + const std::uint64_t slots = num_cu > 0 ? num_cu : 1; + return (tiles + slots - 1) / slots * slots; +} + +template +inline int launch(const void* kargs_dev, + int group_count, + std::uint64_t tile_slots, + hipStream_t stream) +{ + using Shape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + + using Traits = ck_tile::TileGemmUniversalTraits; + + using Problem = ck_tile::UniversalGemmPipelineProblem; + + using Pipeline = typename PipelineTypeTraits::template GemmPipeline; + using Epilogue = + ck_tile::GroupedGemmV7UniformStoreEpilogue; + using Partitioner = ck_tile::GemmSpatiallyLocalTilePartitioner; + using Kernel = ck_tile::GroupedGemmKernel; + + if(kargs_dev == nullptr || group_count <= 0 || tile_slots == 0) + { + return -1; + } + + const dim3 grid(static_cast(tile_slots)); + const dim3 block(Kernel::BlockSize()); + + const auto kernel = ck_tile::make_kernel<512, Config::kBlockPerCu>( + Kernel{}, + grid, + block, + Kernel::GetSmemSize(), + ck_tile::cast_pointer_to_constant_address_space(kargs_dev), + group_count); + ck_tile::launch_kernel(ck_tile::stream_config{stream}, kernel); + return 0; +} + +} // namespace v7_uniform_launch diff --git a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v8_bf16_tn_device_args.cpp b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v8_bf16_tn_device_args.cpp new file mode 100644 index 0000000..15b942a --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_v8_bf16_tn_device_args.cpp @@ -0,0 +1,51 @@ +// grouped_gemm_v8_bf16_tn_device_args.cpp — CK grouped-GEMM instance entry +// (InstanceId 5001) for the uniform Direct-to-LDS TN (wgrad) kernel. +// +// Same construction as grouped_gemm_v7_bf16_nn_device_args.cpp (read that +// file's header first: the ABI shape, the flat grid and the codegen contract +// are all there and are not repeated here). The TN-specific points: +// +// * Layout mapping. This entry consumes the CK TN description DIRECTLY: +// A (true A^T) : [k, sum(M_g)] row-major, per-group a_ptr into the panel, +// stride_a = the panel's row pitch in elements +// B (true B) : [k, n] row-major per group, stride_b = its row pitch +// D : [sum(M_g), n] row-major, stride_e = its row pitch +// i.e. the group's own column origin travels in a_ptr. The pipeline's A +// panel load covers a whole 256-wide column tile, so the caller must keep +// 256 elements of addressable slack past the last group's columns -- the +// pad obligation the selector documents (ragged M is allowed; the caller +// owns the pad). +// +// * The accumulation order is the plain K-serial one over the whole K extent +// of the problem; there is no k-split anywhere in the v7 path. +// +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// SPDX-License-Identifier: MIT + +#include "grouped_gemm_v7_uniform_launch.hpp" + +extern "C" int ck_tile_hcu_grouped_gemm_v8_bf16_tn_device_args( + void* device_args, + int group_count, + std::uint32_t num_cu, + hipStream_t stream, + std::uint64_t tile_slots) { + (void)num_cu; // the grid is sized by tile_slots, not by the CU count + if (device_args == nullptr || group_count <= 0 || + group_count > ck_tile::kGroupedGemmV7MaxGroups || tile_slots == 0) { + return -1; + } + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + return v7_uniform_launch::launch(device_args, + group_count, + tile_slots, + stream); +} diff --git a/include/ck_tile/ops/gemm/grouped_gemm_bw_family_selected.hpp b/include/ck_tile/ops/gemm/grouped_gemm_bw_family_selected.hpp index a69c067..cfa8df5 100644 --- a/include/ck_tile/ops/gemm/grouped_gemm_bw_family_selected.hpp +++ b/include/ck_tile/ops/gemm/grouped_gemm_bw_family_selected.hpp @@ -26,6 +26,55 @@ inline bool ck_tile_hcu_grouped_gemm_bf16_tn_logical_k_tail_enabled() return !(disable != nullptr && disable[0] == '1' && disable[1] == '\0'); } +// v7 BF16 NN / TN instances. These routes are OPT-IN (the selector policy +// fields default to false, see grouped_gemm_selector.hpp), so unlike the two +// switches above the enable side is the one that has to be set explicitly. The +// disable side still wins, which is what makes a same-process A/B able to pin +// either state: +// CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN=1 -> route BF16 NN to instance 5000 +// CK_TILE_GROUPED_GEMM_ENABLE_BF16_V8_TN=1 -> route BF16 TN to instance 5001 +// CK_TILE_GROUPED_GEMM_DISABLE_BF16_V7_V8*=1 -> force the route off +inline bool ck_tile_hcu_grouped_gemm_bf16_v7_nn_enabled() +{ + const char* disable = std::getenv("CK_TILE_GROUPED_GEMM_DISABLE_BF16_V7_NN"); + if(disable != nullptr && disable[0] == '1' && disable[1] == '\0') + return false; + const char* enable = std::getenv("CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN"); + return enable != nullptr && enable[0] == '1' && enable[1] == '\0'; +} + +inline bool ck_tile_hcu_grouped_gemm_bf16_v8_tn_enabled() +{ + const char* disable = std::getenv("CK_TILE_GROUPED_GEMM_DISABLE_BF16_V8_TN"); + if(disable != nullptr && disable[0] == '1' && disable[1] == '\0') + return false; + const char* enable = std::getenv("CK_TILE_GROUPED_GEMM_ENABLE_BF16_V8_TN"); + return enable != nullptr && enable[0] == '1' && enable[1] == '\0'; +} + +// The v7 device-args entries. Declared unconditionally (with C linkage) so the +// dispatch can name them; whether they exist at link time is what +// CK_TILE_GROUPED_GEMM_DEVICE_ARGS_HAS_V7_V8_BF16 guards. +// +// `tile_slots` is an upper bound on the (group, m tile, n tile) tile count, +// rounded up to a whole number of CTAs. The v7 configs are NON-persistent -- +// one tile per CTA, which is exactly what keeps their register budget at the +// reference kernels' level (a wrapping tile loop costs 40 VGPRs and ~20%) -- so +// the launcher has to size a flat grid, and per-group M is device-resident. The +// bound comes from the host-visible problem dimensions and is never read back. +extern "C" int ck_tile_hcu_grouped_gemm_v7_bf16_nn_device_args( + void* device_args, + int group_count, + std::uint32_t num_cu, + hipStream_t stream, + std::uint64_t tile_slots); +extern "C" int ck_tile_hcu_grouped_gemm_v8_bf16_tn_device_args( + void* device_args, + int group_count, + std::uint32_t num_cu, + hipStream_t stream, + std::uint64_t tile_slots); + // Launch the selected BW-family FP16/BF16 grouped-GEMM kernels from an // existing device array of ck_tile::GemmTransKernelArg. The validated family // currently covers gfx936 and gfx938 under the same aligned-shape contract. diff --git a/include/ck_tile/ops/gemm/grouped_gemm_device_args.h b/include/ck_tile/ops/gemm/grouped_gemm_device_args.h index d85a486..545f081 100644 --- a/include/ck_tile/ops/gemm/grouped_gemm_device_args.h +++ b/include/ck_tile/ops/gemm/grouped_gemm_device_args.h @@ -63,7 +63,11 @@ typedef enum ck_tile_hcu_grouped_gemm_instance_id_v1 CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V3_128_M4_FULL_PADDING_V1 = 4009, CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V3_128_M4_FULL_PADDING_GFX936_V1 = 4010, CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V4_64_M01_4_GFX936_V1 = 4011, - CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V6_M32_NONPADDING_MLE8_FAMILY_B_GFX936_V1 = 4104 + CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V6_M32_NONPADDING_MLE8_FAMILY_B_GFX936_V1 = 4104, + // v7: the uniform Direct-to-LDS NN / TN kernels driven through the + // device-args ABI. + CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V7_BF16_NN_V1 = 5000, + CK_TILE_HCU_GROUPED_GEMM_INSTANCE_V8_BF16_TN_V1 = 5001 } ck_tile_hcu_grouped_gemm_instance_id_v1; typedef enum ck_tile_hcu_grouped_gemm_candidate_family_v1 @@ -71,7 +75,8 @@ typedef enum ck_tile_hcu_grouped_gemm_candidate_family_v1 CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_BW_V1 = 0, CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_GFX936_V3_V1 = 1, CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_GFX938_MLS_V1 = 2, - CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V4_V1 = 3 + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V4_V1 = 3, + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_FAMILY_V7_V1 = 4 } ck_tile_hcu_grouped_gemm_candidate_family_v1; typedef enum ck_tile_hcu_grouped_gemm_architecture_mask_v1 diff --git a/include/ck_tile/ops/gemm/grouped_gemm_selector.hpp b/include/ck_tile/ops/gemm/grouped_gemm_selector.hpp index cde8c06..24a5807 100644 --- a/include/ck_tile/ops/gemm/grouped_gemm_selector.hpp +++ b/include/ck_tile/ops/gemm/grouped_gemm_selector.hpp @@ -101,8 +101,22 @@ struct GroupedGemmSelectorPolicy // Production default after independent gfx936/gfx938 correctness and focused fallback // comparisons. Callers can still set this field false for an explicit V4 rollback. bool enable_p3_fp16_tn_m1024_k4096_route = true; + // v7 BF16 NN / TN (the uniform Direct-to-LDS kernels as device-args + // instances; see instances/grouped_gemm_v7_v8_bf16_{nn,tn}_device_args.cpp). + // + // OPT-IN, unlike every route above: these are the first landings of the v7 + // uniform-pipeline kernels as CK instances, and the production NN/TN routes + // must not move until they are measured on the real shapes. The adapter also + // honours CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_V8 so the A/B does not + // need a rebuild; the disable side of those switches wins. + bool enable_bf16_v7_nn = false; + bool enable_bf16_v8_tn = false; }; +// The v7 kernels carry their group table in the kernel argument, and that +// table is fixed at 16. +constexpr std::int32_t kGroupedGemmV7MaxGroups = 16; + enum class GroupedGemmInstanceId { // Values are stable identity codes. Add new instances to a family range; @@ -126,6 +140,10 @@ enum class GroupedGemmInstanceId v3_128_m4_full_padding_gfx936 = 4010, v4_64_m01_4_gfx936 = 4011, v6_m32_nonpadding_mle8_family_b_gfx936 = 4104, + // v7 family (the uniform Direct-to-LDS NN/TN kernels as device-args + // instances). + v7_bf16_nn = 5000, + v8_bf16_tn = 5001, }; enum class GroupedGemmCandidateFamily @@ -134,6 +152,7 @@ enum class GroupedGemmCandidateFamily gfx936_v3, gfx938_mls, v4, + v7, }; struct GroupedGemmCandidateIdentity @@ -143,7 +162,7 @@ struct GroupedGemmCandidateIdentity GroupedGemmCandidateFamily family; }; -inline constexpr std::array +inline constexpr std::array grouped_gemm_candidate_registry{{ {GroupedGemmInstanceId::bw_family_selected, "bw_family.selected", @@ -199,6 +218,17 @@ inline constexpr std::array {GroupedGemmInstanceId::v6_m32_nonpadding_mle8_family_b_gfx936, "v6.32x128x64.nonpadding.selector_m_le_8.family_b.gfx936", GroupedGemmCandidateFamily::v4}, + // v7: 256x256x16 tile, 8 waves (4m x 2n), 4-stage 16 KB LDS ring fed + // by Direct-to-LDS buffer loads, m32x16 / b64 LDS reads, two-group phase + // interleave and the LDS-staged dwordx4 epilogue. The suffix names the + // operand form, because that is what separates the NN entry (only the + // second operand is [k, feature]) from the TN one (both are). + {GroupedGemmInstanceId::v7_bf16_nn, + "v7.nn_b.uniform_d2l.256x256x16", + GroupedGemmCandidateFamily::v7}, + {GroupedGemmInstanceId::v8_bf16_tn, + "v8.tn_ab.uniform_d2l.256x256x16", + GroupedGemmCandidateFamily::v7}, }}; constexpr const GroupedGemmCandidateIdentity* @@ -229,6 +259,8 @@ enum class GroupedGemmSelectionReason p2_bf16_g8_gfx936_fc2_nn_fast_vmem3, p2_bf16_g8_gfx938_v5_fc1_route, p3_fp16_tn_m1024_k4096_route, + v7_bf16_nn, + v8_bf16_tn, v4_n_aligned, v4_fallback, }; @@ -321,6 +353,44 @@ constexpr GroupedGemmSelection select_grouped_gemm_candidate( } } + // v7 BF16 NN / TN instances (opt-in; see policy.enable_bf16_v7_v8*). + // + // The gates are the kernels' own contract, not a shape whitelist: + // K % 16 == 0 and K >= 64 -- the 16-deep k-stage and the 4-stage ring + // N % 256 == 0 -- the LDS-staged epilogue stores whole 16 B + // row segments and carries no column + // predicate, so a partial N tile would write + // past the row + // common K -- the k-advance of both streams is per-launch + // non-empty groups -- not validated for the walk yet; keep the + // fallback until they are + // <= kGroupedGemmV7MaxGroups -- the device-side group table is what + // bounds the kargs scan at kernel entry + // gfx936 -- the kernels' ISA contract (v7 loads, the + // CACHE_SWIZZLE descriptor bit) + // TN additionally requires an M that the launcher can address: its A panel + // is [k, sum(M_g)] with a host-known pitch, and the panel load covers a whole + // 256-wide column tile, so a ragged last tile reads up to 256 columns past + // the group's last column. Reading into the NEXT group's columns is harmless + // (those rows are masked on the store), but the read past the LAST group's + // columns must land in storage the caller owns -- the reference kernel's + // documented tail-pad obligation, which the selector cannot verify. Ragged M + // is therefore allowed; the caller owns the pad. + const bool v7_contract = + problem.arch == GroupedGemmArch::gfx936 && !has_variable_k && + problem.group_count <= kGroupedGemmV7MaxGroups && + !problem.may_have_empty_groups && k >= 64 && k % 16 == 0 && n % 256 == 0; + if(v7_contract && is_bf16 && is_nn && policy.enable_bf16_v7_nn) + { + return {GroupedGemmInstanceId::v7_bf16_nn, + GroupedGemmSelectionReason::v7_bf16_nn}; + } + if(v7_contract && is_bf16 && is_tn && policy.enable_bf16_v8_tn) + { + return {GroupedGemmInstanceId::v8_bf16_tn, + GroupedGemmSelectionReason::v8_bf16_tn}; + } + const bool bw_arch_layout = problem.arch == GroupedGemmArch::gfx936 || (problem.arch == GroupedGemmArch::gfx938 && is_bf16_tn); 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 1ccd71d..95dfea7 100644 --- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp @@ -11,6 +11,11 @@ #include "ck_tile/host/stream_utils.hpp" // #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" +// Forward declarations only (the enum + the pipeline and tile-parameter +// names): the v7 uniform pipeline BODY is completed inside the v7 translation +// units via the launch header, so no other TU here pays for it. Keep this at +// file scope -- it declares namespace ck_tile names. +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform_fwd.hpp" // #include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp" #include "ck_tile/host.hpp" @@ -29,6 +34,25 @@ struct is_grouped_gemm_mls_pipeline> { }; +// Detects the uniform Direct-to-LDS NN/TN pipelines +// (include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp), which carry a +// static `IsV7UniformPipeline = true`. Routes them through RunGemmV7Uniform: the +// raw-pointer path that builds no tensor view and no tile window. That is not a +// simplification for its own sake -- wrapping this pipeline's k loop in ANY +// surrounding control flow (a persistent tile loop, a window) costs 40 VGPRs +// and a scratch spill, measured ~20% of the throughput. The entry is therefore +// a flat one-CTA-per-tile grid, sized by the launcher. +template +struct is_grouped_gemm_v7_uniform_pipeline : std::false_type +{ +}; + +template +struct is_grouped_gemm_v7_uniform_pipeline> + : std::bool_constant +{ +}; + /// @brief The Grouped GEMM kernel host arguments. /// /// @par Overview @@ -615,22 +639,35 @@ struct GroupedGemmKernel splitk_batch_offset.bs_k_split_offset[0]; CDataType* c_ptr = static_cast(kargs.e_ptr); - // allocate LDS - __shared__ char smem_ptr_0[GetSmemSize()]; - - if constexpr(is_grouped_gemm_mls_pipeline::value) + if constexpr(is_grouped_gemm_v7_uniform_pipeline::value) { - RunGemmMls(a_ptr, - b_ptr, - c_ptr, - smem_ptr_0, - kargs, - splitk_batch_offset, - i_m, - i_n); + // The v7 ring comes from DYNAMIC shared memory: the launcher + // passes GetSmemSize() bytes, and the pipeline's inline asm + // addresses LDS from offset 0 (the reference kernels allocate the + // ring exactly this way). The 64 KB static array below must stay + // OUT of this instantiation: it would sit between LDS 0 and the + // dynamic allocation, and static + dynamic together would not fit + // the CU's workgroup LDS budget. + RunGemmV7Uniform(nullptr, kargs, i_m, i_n); } else { + // allocate LDS + __shared__ char smem_ptr_0[GetSmemSize()]; + + if constexpr(is_grouped_gemm_mls_pipeline::value) + { + RunGemmMls(a_ptr, + b_ptr, + c_ptr, + smem_ptr_0, + kargs, + splitk_batch_offset, + i_m, + i_n); + } + else + { static_assert( GemmPipeline::DoubleSmemBuffer || !GemmPipeline::Preshuffle, "SingleSmemBuffer and Preshuffle cannot both be enabled simultaneously!"); @@ -741,6 +778,7 @@ struct GroupedGemmKernel i_n); } } + } } #if defined(CK_TILE_GROUPED_GEMM_FUSED_PACK_GRADX) @@ -764,6 +802,39 @@ struct GroupedGemmKernel #endif } + CK_TILE_DEVICE static void + RunGemmV7Uniform(void* p_smem, + const KernelArgs& kargs, + const index_t block_idx_m, + const index_t block_idx_n) + { + static_assert(is_grouped_gemm_v7_uniform_pipeline::value); + + // NN and TN differ only in which tensor plays which operand role; both + // take the group's raw base pointers and its row pitches, exactly as the + // reference launchers handed them to their kernels. + V7UniformTileParams p; + p.a_ptr = kargs.as_ptr[0]; + p.b_ptr = kargs.bs_ptr[0]; + p.d_ptr = kargs.e_ptr; + p.m = kargs.M; + p.n = kargs.N; + p.k = kargs.K; + p.a_stride = kargs.stride_As[0]; + p.b_stride = kargs.stride_Bs[0]; + p.d_stride = kargs.stride_E; + // Group-relative tile origins: the descriptor bound, the A row mapping + // and the D row mask are all group-relative, exactly as in the reference + // kernels (whose grid.y was the group's own M-tile index). + p.m_start = block_idx_m; + p.n_start = block_idx_n; + + using AccTile = typename GemmPipeline::AccTile; + AccTile acc; + GemmPipeline{}.run_tile(p, acc, p_smem); + EpiloguePipeline{}.store_tile(p, acc, p_smem); + } + CK_TILE_DEVICE static void RunGemmMls(const ADataType* a_ptr, const BDataType* b_ptr, @@ -1067,6 +1138,16 @@ struct GroupedGemmKernel const auto gemm_desc_ptr = reinterpret_cast( cast_pointer_to_generic_address_space(gemm_descs_const)); + // A flat over-covered grid (the v7 instances pad the host-side tile + // bound up to a whole number of CTAs because per-group M is + // device-resident) hands some CTAs block ids past the last group's + // block_end. FindGroupId only converges for in-range ids, so exit + // before it. Exact-sized standard launches never take this branch. + if(block_id >= gemm_desc_ptr[group_count - 1].block_end) + { + return; + } + const index_t group_id = FindGroupId(gemm_desc_ptr, block_id, group_count); const auto& kargs = gemm_desc_ptr[group_id]; diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp new file mode 100644 index 0000000..b1ff19a --- /dev/null +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform.hpp @@ -0,0 +1,781 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// +// Uniform (non warp-specialized) Direct-to-LDS grouped-GEMM pipeline for the +// gfx936 / gfx938 "v7" NN (dgrad) and TN (wgrad) kernels, ported into the +// ck_tile pipeline vocabulary. +// +// WHY A SEPARATE PIPELINE FAMILY +// ------------------------------ +// The two operand layouts this serves share one structure that no existing CK +// pipeline has: at least one operand is [K, feature]-major, so its per-lane MMAC +// fragment is NOT a contiguous 16 B run in global memory. The treatment (from +// the reference kernels, where it was measured as the only arrangement that fits +// this LDS budget) is +// +// * fetch 16 B (8 features) of one reduction row per lane into LDS with +// Direct-to-LDS (`buffer_load_dwordx4 ... offen lds`), +// * let the hardware perform the K/fragment redistribution inside the LDS READ +// (`ds_read_m32x16_b16`), and +// * interleave the per-lane store order so the read is bank-conflict-free. +// +// NN : A = [M, K] K-contiguous (cheap: loaded as a fragment, one +// `ds_read_b64` per m-tile) and B = [K, N] N-contiguous (m32x16 read). +// TN : BOTH operands are [K, feature]-major, so both go through m32x16; the one +// read instruction serves A (a 64-column window, 2 reads) and B (a +// 128-column window, 4 reads) because the A and B fragments of the MMAC +// are the same lane pattern. +// +// GEOMETRY (fixed by the reference kernels; the CK config must match) +// ------------------------------------------------------------------ +// block tile 256M x 256N x 16K, 512 threads = 8 waves of 64 +// wave grid 4m x 2n, wave tile 64M x 128N, 32 MMACs per k-stage +// LDS 4-stage ring of 16 KB (8 KB A + 8 KB B), the whole 64 KB budget +// phase 12-phase unroll, operand ring 3, prefetch distance 5 (ring + 1), +// two-group wave interleave (waves 0-3 mmac-first, waves 4-7 +// load-first) -- see the phase-discipline note below. +// +// WHY THE KERNEL MUST RUN ONE TILE PER CTA (measured, not stylistic) +// ----------------------------------------------------------------- +// Wrapping this pipeline's k-loop in a tile loop -- even a loop the compiler can +// see through -- costs 40 VGPRs (215 -> 255/256) and a scratch spill, which +// measured ~20% of the throughput. The scan over the group table and the per-tile +// workgroup barrier are both free; it is the loop itself, because the epilogue's +// live range then overlaps the next tile's prologue in the register allocator. +// The device-args entry therefore launches a FLAT grid (one CTA per tile, the +// grid sized from a host-side tile bound) and re-derives the (group, m, n) tile +// coordinates from the device kargs inside the kernel. Do not "simplify" this +// into a persistent loop. +// +// PHASE DISCIPLINE (carried over verbatim, with the measurements that set it) +// --------------------------------------------------------------------------- +// * Two-group interleave, not a uniform order: the LDS bursts of waves 0-3 and +// 4-7 stay one MMAC phase apart, so a load's LDS write-back always finds a +// quiet window. The uniform (blasLt-like) order measured -12%, the mirror of +// the interleave (group-B first) is throughput-neutral. +// * The load tail waits are vmcnt(4) / lgkmcnt(8): with 2 loads per step and a +// 4-slot ring, vmcnt(4) is 2 steps of lookback, which is the loosest wait +// that still guards the slot being re-filled. +// * s_setprio(3) around the MMAC burst, s_setprio(1) elsewhere (blasLt's +// scheme; measured inside the run-to-run spread, kept as the reference +// kernel has it). +// * The per-step workgroup barrier is NOT optional: dropping it measured +1.35% +// and corrupts the result (the ring's cross-wave visibility depends on it). +// +// TILING CONTRACT +// --------------- +// NN : A = [M, K] row major (a_stride), B = [K, N] row major (b_stride, the +// group's own panel), D = [M, N] row major (d_stride) +// TN : A = [K, M] row major (a_stride, M contiguous), B = [K, N] row major +// (b_stride), D = [M, N] row major (d_stride) +// Rows past M and columns past N are masked on the STORE only; the K extent must +// be a whole number of 16-wide k-stages and N a whole number of 256-wide tiles +// (both are contract-level requirements, enforced by the host-side selector). +// The A panel fetch reads a whole 256-column tile, so TN additionally requires +// addressable slack after the last group's A columns (the caller's obligation, +// documented by the standalone kernel's header). +#pragma once + +#include "ck_tile/core.hpp" + +// V7UniformKind, V7UniformTileParams (complete) and the pipeline template +// declaration: the shared copy every grouped-GEMM TU already carries. +#include "gemm_pipeline_v7_uniform_fwd.hpp" + +// The address arithmetic below is a port of the reference kernels, which are +// built with -Wno-old-style-cast; every cast is a fixed-width unsigned/signed +// conversion in an address expression. CK's own arch/buffer-addressing headers +// handle it the same way (see cast_pointer_to_constant_address_space). +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wold-style-cast" + +namespace ck_tile { + +// --------------------------------------------------------------------------- +// Geometry. +// --------------------------------------------------------------------------- +struct V7UniformGeometry +{ + static constexpr index_t kBlockM = 256; + static constexpr index_t kBlockN = 256; + static constexpr index_t kStageK = 16; + static constexpr index_t kThreads = 512; + static constexpr index_t kWaves = 8; + static constexpr index_t kWaveSize = 64; + static constexpr index_t kWaveGridM = 4; + static constexpr index_t kWaveGridN = 2; + static constexpr index_t kWaveTileM = kBlockM / kWaveGridM; // 64 + static constexpr index_t kWaveTileN = kBlockN / kWaveGridN; // 128 + static constexpr index_t kMfmaM = 16; + static constexpr index_t kMfmaN = 16; + static constexpr index_t kMfmaK = 16; + static constexpr index_t kMTilesPerWave = kWaveTileM / kMfmaM; // 4 + static constexpr index_t kNTilesPerWave = kWaveTileN / kMfmaN; // 8 + static constexpr index_t kMmacPerStage = kMTilesPerWave * kNTilesPerWave; // 32 + + static constexpr index_t kNumStages = 4; + static constexpr index_t kOpRing = 3; + static constexpr index_t kPhaseCycle = 12; + static constexpr index_t kWaveGroup = kWaves / 2; + static constexpr index_t kPrefetchDist = kNumStages + 1; // 5 + + static constexpr index_t kLdsBytesA = kBlockM * kStageK * 2; // 8 KB + static constexpr index_t kLdsBytesB = kStageK * kBlockN * 2; // 8 KB + static constexpr index_t kLdsBytesPerStage = kLdsBytesA + kLdsBytesB; // 16 KB + static constexpr index_t kLdsBytes = kLdsBytesPerStage * kNumStages; // 64 KB + + static constexpr index_t kAccVgprPerLane = kMmacPerStage * 4; // 128 + + static constexpr index_t kAccMTiles = kMTilesPerWave; + static constexpr index_t kAccNTiles = kNTilesPerWave; + + static_assert(kThreads == kWaves * kWaveSize, "512 = 8 x 64"); + static_assert(kMfmaK == kStageK, "one MMAC covers exactly one k-stage"); + static_assert(kLdsBytes == 64 * 1024, "the ring is the whole CU LDS budget"); + static_assert(kPhaseCycle % kOpRing == 0 && kPhaseCycle % kNumStages == 0, + "the unroll must span whole operand sets and whole LDS rings"); + static_assert(kPrefetchDist > kNumStages, "distance ring+1 is the whole point"); +}; + +// --------------------------------------------------------------------------- +// Accumulator tile: 4 m-iterations x 8 n-iterations x 4 fp32 = 128 VGPR/lane. +// A VECTOR quad per slot (not four scalars): that is the granularity the +// register allocator works with around every MMAC, and the arrangement the +// reference kernels use. +// --------------------------------------------------------------------------- +struct V7UniformAccTile +{ + using vec_t = float __attribute__((ext_vector_type(4))); + vec_t v[V7UniformGeometry::kAccMTiles * V7UniformGeometry::kAccNTiles]; +}; + +// V7UniformKind and the pipeline template declaration come from the forward +// header: it is the copy every grouped-GEMM TU already carries. +namespace v7u_detail { + +using vec4_bf16 = unsigned short __attribute__((ext_vector_type(4))); +using vec8_bf16 = unsigned short __attribute__((ext_vector_type(8))); +using vec4_fp32 = float __attribute__((ext_vector_type(4))); +using u32x4 = int32_t __attribute__((ext_vector_type(4))); + +CK_TILE_DEVICE void wait_all_warp_arrived() +{ + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_barrier\n"); + __builtin_amdgcn_sched_barrier(0); +} + +template +CK_TILE_DEVICE void wait_all_buffer_data_arrived(int wait_count = 0) +{ + __builtin_amdgcn_sched_barrier(0); + if constexpr(Sync) + { + asm volatile("s_waitcnt vmcnt(%0)\n\ts_barrier\n" ::"n"(wait_count)); + } + else + { + asm volatile("s_waitcnt vmcnt(%0)\n" ::"n"(wait_count)); + } + __builtin_amdgcn_sched_barrier(0); +} + +template +CK_TILE_DEVICE void wait_lds_data_arrived(int wait_count = 0) +{ + __builtin_amdgcn_sched_barrier(0); + if constexpr(Sync) + { + asm volatile("s_waitcnt lgkmcnt(%0)\n\ts_barrier\n" ::"n"(wait_count)); + } + else + { + asm volatile("s_waitcnt lgkmcnt(%0)\n" ::"n"(wait_count)); + } + __builtin_amdgcn_sched_barrier(0); +} + +template +CK_TILE_DEVICE void set_prio() +{ + static_assert(Prio >= 0 && Prio <= 3, "s_setprio operand must be 0..3"); + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_setprio %0" ::"n"(Prio)); + __builtin_amdgcn_sched_barrier(0); +} + +// MMAC. The 4-element bf16 operands and the 4-element fp32 accumulator are the +// instruction's own fragment layout; the redistribution from LDS is the +// ds_read_m32x16's job (see the header note). +CK_TILE_DEVICE vec4_fp32 mmac_bf16(const vec4_bf16& a, const vec4_bf16& b, const vec4_fp32& c) +{ + return __builtin_hcu_mmac_f32_16x16x16_bf16(a, b, c); +} + +using buffer_rsrc_t = int32_t __attribute__((ext_vector_type(4))); + +// Buffer resource descriptor. STRIDE is the CACHE_SWIZZLE rotation period (bit +// 30) on the A side of NN -- the one arch-dependent field, gated on the build's +// arch macros. Everything else is layout-neutral. +CK_TILE_DEVICE buffer_rsrc_t make_buffer_rsrc(const void* ptr, + uint32_t stride_bytes, + uint32_t size_bytes, + bool allow_swizzle) +{ + buffer_rsrc_t r; + const uint64_t addr = reinterpret_cast(ptr); + r[0] = (int32_t)(addr & 0xffffffffu); + r[1] = (int32_t)(addr >> 32); + r[2] = (int32_t)size_bytes; + r[3] = 0x00020000; + uint32_t w1 = (stride_bytes & 0x3fffu) << 16; +#if defined(__gfx936__) || defined(__gfx938__) + if(allow_swizzle && stride_bytes % 64 == 0) + { + const uint32_t s64 = stride_bytes / 64; + if((s64 & (s64 - 1)) == 0) w1 |= 1u << 30; + } +#else + (void)allow_swizzle; // CACHE_SWIZZLE does not exist on this part +#endif + r[1] |= (int32_t)w1; + return r; +} + +// Global -> LDS 16 B per lane. `idxen` is the pair {index lo, offset hi} form; +// `offen` keeps the whole address in the vaddr and the k-advance in the soffset. +CK_TILE_DEVICE void buffer_load_dwordx4_offen_lds(buffer_rsrc_t rsrc, + uint32_t vaddr, + uint32_t soffset, + uint32_t m0) +{ + asm volatile("s_mov_b32 m0, %[m0]\n\t" + "buffer_load_dwordx4 %[v], %[rsrc], %[so] offen lds\n" + : + : [v] "v"(vaddr), [rsrc] "s"(rsrc), [so] "s"(soffset), [m0] "s"(m0) + : "memory"); +} + +// SLC (evict-first) is set on the NN A load on purpose: it is the 32-row +// scattered load that overruns the VL1 address/alloc path, and marking its lines +// evict-first measured +0.8%. The B side must NOT carry it (-6.2%). +CK_TILE_DEVICE void buffer_load_dwordx4_idxen_lds_t(buffer_rsrc_t rsrc, + uint64_t pair, + uint32_t soffset, + uint32_t m0) +{ + asm volatile("s_mov_b32 m0, %[m0]\n\t" + "buffer_load_dwordx4 %[pair], %[rsrc], %[so] idxen offen slc lds\n" + : + : [pair] "v"(pair), [rsrc] "s"(rsrc), [so] "s"(soffset), [m0] "s"(m0) + : "memory"); +} + +CK_TILE_DEVICE void buffer_load_dwordx4_idxen_lds(buffer_rsrc_t rsrc, + uint64_t pair, + uint32_t soffset, + uint32_t m0) +{ + asm volatile("s_mov_b32 m0, %[m0]\n\t" + "buffer_load_dwordx4 %[pair], %[rsrc], %[so] idxen offen lds\n" + : + : [pair] "v"(pair), [rsrc] "s"(rsrc), [so] "s"(soffset), [m0] "s"(m0) + : "memory"); +} + +// LDS reads. b64 x4 = one b64 per m-tile from a single base + immediates (the +// NN A side, whose per-lane fragment IS contiguous); m32x16 x8 / x4 = the +// hardware redistribution used for every [K, feature]-major operand. +template +CK_TILE_DEVICE void ds_read_b64_x4( + uint32_t base, vec4_bf16& f0, vec4_bf16& f1, vec4_bf16& f2, vec4_bf16& f3) +{ + asm volatile("ds_read_b64 %0, %4 offset:%5\n\t" + "ds_read_b64 %1, %4 offset:%6\n\t" + "ds_read_b64 %2, %4 offset:%7\n\t" + "ds_read_b64 %3, %4 offset:%8\n" + : "=&v"(f0), "=&v"(f1), "=&v"(f2), "=&v"(f3) + : "v"(base), "n"(kOff), "n"(kOff + 512u), "n"(kOff + 1024u), + "n"(kOff + 1536u) + : "memory"); +} + +template +CK_TILE_DEVICE void ds_read_m32x16_x8(const uint32_t (&rb)[4], vec4_bf16 (&out)[8]) +{ + vec8_bf16 d0, d1, d2, d3; + asm volatile("ds_read_m32x16_b16 %0, %4 offset:%8\n\t" + "ds_read_m32x16_b16 %1, %5 offset:%8\n\t" + "ds_read_m32x16_b16 %2, %6 offset:%8\n\t" + "ds_read_m32x16_b16 %3, %7 offset:%8\n" + : "=&v"(d0), "=&v"(d1), "=&v"(d2), "=&v"(d3) + : "v"(rb[0]), "v"(rb[1]), "v"(rb[2]), "v"(rb[3]), "n"(kOff) + : "memory"); + out[0] = __builtin_shufflevector(d0, d0, 0, 1, 2, 3); + out[1] = __builtin_shufflevector(d0, d0, 4, 5, 6, 7); + out[2] = __builtin_shufflevector(d1, d1, 0, 1, 2, 3); + out[3] = __builtin_shufflevector(d1, d1, 4, 5, 6, 7); + out[4] = __builtin_shufflevector(d2, d2, 0, 1, 2, 3); + out[5] = __builtin_shufflevector(d2, d2, 4, 5, 6, 7); + out[6] = __builtin_shufflevector(d3, d3, 0, 1, 2, 3); + out[7] = __builtin_shufflevector(d3, d3, 4, 5, 6, 7); +} + +template +CK_TILE_DEVICE void ds_read_m32x16_x4(const uint32_t (&rb)[2], vec4_bf16 (&out)[4]) +{ + vec8_bf16 d0, d1; + asm volatile("ds_read_m32x16_b16 %0, %2 offset:%4\n\t" + "ds_read_m32x16_b16 %1, %3 offset:%4\n" + : "=&v"(d0), "=&v"(d1) + : "v"(rb[0]), "v"(rb[1]), "n"(kOff) + : "memory"); + out[0] = __builtin_shufflevector(d0, d0, 0, 1, 2, 3); + out[1] = __builtin_shufflevector(d0, d0, 4, 5, 6, 7); + out[2] = __builtin_shufflevector(d1, d1, 0, 1, 2, 3); + out[3] = __builtin_shufflevector(d1, d1, 4, 5, 6, 7); +} + +} // namespace v7u_detail + +// Namespace alias for the asm-level helpers, declared where C++ allows one. +namespace v7u = v7u_detail; + +// --------------------------------------------------------------------------- +// The pipeline. One instance serves one tile: `run_tile` runs the whole K loop +// and leaves the accumulator in `acc` for the epilogue. +// --------------------------------------------------------------------------- +template +struct GemmPipelineV7Uniform +{ + using G = V7UniformGeometry; + using AccTile = V7UniformAccTile; + + // The framework's trait for the raw (no tensor view, no tile window) kernel + // path: see GroupedGemmKernel's is_grouped_gemm_v7_uniform_pipeline. + static constexpr bool IsV7UniformPipeline = true; + + static constexpr V7UniformKind kind = kKind; + + // ---- framework-facing surface (what GroupedGemmKernel reads) -------- + // The kernel takes BlockSize as a VALUE, not a call: keep the member form. + static constexpr index_t BlockSize = G::kThreads; + + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + static constexpr index_t KPerBlock = remove_cvref_t::kK; + + // The one-CTA-per-tile structure the register budget depends on: the config + // must be non-persistent, and nothing double-buffers. + static constexpr bool UsePersistentKernel = false; + static constexpr bool DoubleSmemBuffer = false; + static constexpr bool Preshuffle = false; + + // The LDS the ring occupies (the whole CU budget). + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return G::kLdsBytes; } + + using AOpBank = v7u::vec4_bf16[G::kOpRing][G::kMTilesPerWave]; + using BOpBank = v7u::vec4_bf16[G::kOpRing][G::kNTilesPerWave]; + + // Per-thread scalar state: the descriptors / per-lane bases the k loop needs. + struct Addr + { + v7u::buffer_rsrc_t a_rsrc, b_rsrc; + uint32_t a_vaddr, b_vaddr; // per-lane global byte offsets (TN, NN-B) + uint64_t a_pair; // NN A: {index, offset} idxen pair + uint32_t a_rd; // NN A b64 read base + uint32_t a_rb[2]; // TN A m32x16 read bases + uint32_t b_rb[4]; // B m32x16 read bases + uint32_t m0_a_base, m0_b_base; + int b_stride_bytes; + int a_stride_bytes; // TN: the A panel's row pitch (the k advance); + // NN keeps the K-contiguous 2-byte advance and + // never reads this. + int n_stages; + int wid; + }; + + CK_TILE_DEVICE static void mmac32(AccTile& acc, + const AOpBank& a_bank, + const BOpBank& b_bank, + int bank) + { + // The pragma-unrolled loop body is written out for bank 0 so the MMAC + // burst starts with a `s_setprio 3` as in the reference kernel. + acc.v[0] = v7u::mmac_bf16(a_bank[bank][0], b_bank[bank][0], acc.v[0]); + v7u::set_prio<3>(); +#pragma unroll + for(index_t i = 1; i < G::kMmacPerStage; ++i) + { + const index_t m = i / G::kNTilesPerWave; + const index_t n = i % G::kNTilesPerWave; + acc.v[i] = v7u::mmac_bf16(a_bank[bank][m], b_bank[bank][n], acc.v[i]); + } + v7u::set_prio<1>(); + } + + template + CK_TILE_DEVICE static void issue_load_and_read(Addr& s, + AOpBank& a_bank, + BOpBank& b_bank, + int kt) + { + if((kt + 2) < s.n_stages) + { + if constexpr(kKind == V7UniformKind::NN) + { + v7u::ds_read_m32x16_x8(s.b_rb, b_bank[kBank]); + v7u::ds_read_b64_x4( + s.a_rd, a_bank[kBank][0], a_bank[kBank][1], a_bank[kBank][2], a_bank[kBank][3]); + } + else + { + v7u::ds_read_m32x16_x8(s.b_rb, b_bank[kBank]); + v7u::ds_read_m32x16_x4(s.a_rb, a_bank[kBank]); + } + } + // The k advance rides the soffsets in every form: the descriptor base is + // never rewritten inside the loop. + const int kStage = (kt + G::kPrefetchDist < s.n_stages) ? kt + G::kPrefetchDist : 0; + // A's k advance: NN walks a K-contiguous [M, K] operand (32 B per + // 16-wide stage); TN walks the [K, M] panel row-major, so one stage is + // a_stride_bytes (16 panel rows). The reference kernels advance + // exactly this way (bf16_gemm_tn_hipc issue_load_and_read). + const uint32_t so_a = + (uint32_t)(kStage * G::kStageK * + (kKind == V7UniformKind::NN ? 2u : (uint32_t)s.a_stride_bytes)); + const uint32_t so_b = (uint32_t)(kStage * G::kStageK * s.b_stride_bytes); + if constexpr(kKind == V7UniformKind::NN) + { + v7u::buffer_load_dwordx4_idxen_lds_t( + s.a_rsrc, s.a_pair, so_a, s.m0_a_base + kLoadSlot); + v7u::buffer_load_dwordx4_offen_lds( + s.b_rsrc, s.b_vaddr, so_b, s.m0_b_base + kLoadSlot); + } + else + { + v7u::buffer_load_dwordx4_offen_lds( + s.a_rsrc, s.a_vaddr, so_a, s.m0_a_base + kLoadSlot); + v7u::buffer_load_dwordx4_offen_lds( + s.b_rsrc, s.b_vaddr, so_b, s.m0_b_base + kLoadSlot); + } + } + + // One k-stage. kGroupA selects the phase order (compile-time: the divergent + // runtime branch per step compiled to a full exec-mask prologue every step); + // kBarrier=false on the last step of an unrolled cycle, where the caller + // emits the barrier after its loop control so the trip test overlaps the + // drain instead of being exposed behind it. + template + CK_TILE_DEVICE static void stage_step(Addr& s, AccTile& acc, AOpBank& a_bank, BOpBank& b_bank, int kt) + { + if constexpr(kGroupA) + { + mmac32(acc, a_bank, b_bank, kBankCur); + issue_load_and_read(s, a_bank, b_bank, kt); + v7u::wait_lds_data_arrived(8); + } + else + { + issue_load_and_read(s, a_bank, b_bank, kt); + v7u::wait_lds_data_arrived(8); + mmac32(acc, a_bank, b_bank, kBankCur); + } + v7u::wait_all_buffer_data_arrived(4); + if constexpr(kBarrier) + { + v7u::wait_all_warp_arrived(); + } + } + + CK_TILE_DEVICE static void run_tile(const V7UniformTileParams& p, AccTile& acc, void* smem); +}; + +// --------------------------------------------------------------------------- +// The tile body. Written once, with `if constexpr` on the operand layout; the +// instruction order, the waits and the ring arithmetic are the reference +// kernels' (see the header note for what each one is worth). +// --------------------------------------------------------------------------- +template +CK_TILE_DEVICE void GemmPipelineV7Uniform::run_tile( + const V7UniformTileParams& p, AccTile& acc, void* smem) +{ + using namespace v7u_detail; + (void)smem; // the ring is the launcher's dynamic LDS allocation + + const index_t tid = threadIdx.x; + const int wid = (int)(tid >> 6); + const int lane = (int)(tid & (G::kWaveSize - 1)); + const int wave_m = wid >> 1; + const int wave_n = wid & 1; + const int wave_m_start = wave_m * (int)G::kWaveTileM; + + const int m_block_start = p.m_start; + const int n_block_start = p.n_start; + +#pragma unroll + for(index_t i = 0; i < G::kMmacPerStage; ++i) acc.v[i] = vec4_fp32{0, 0, 0, 0}; + + Addr S; + S.wid = wid; + S.n_stages = (int)(p.k / G::kStageK); + S.b_stride_bytes = (int)p.b_stride * 2; + + const uint32_t m0_a_base = __builtin_amdgcn_readfirstlane((uint32_t)(wid * 1024)); + const uint32_t m0_b_base = + __builtin_amdgcn_readfirstlane((uint32_t)(G::kLdsBytesA + wid * 1024)); + S.m0_a_base = m0_a_base; + S.m0_b_base = m0_b_base; + + if constexpr(kKind == V7UniformKind::NN) + { + // ---- A = [M, K], K-contiguous ------------------------------------ + const int a_stride_bytes = (int)p.a_stride * 2; + // Lane pairs fetch one row's two 16 B k-halves; rows r = 4..7, 12..15 + // (mod 16) take the halves swapped, which spreads the b64 read phases + // over all 8 bank groups (a plain row-major 32 B stride reaches 4 and + // measured 2-way conflicts on every read phase). + const int a_row_global = m_block_start + wid * 32 + (lane >> 1); + const uint32_t a_kh = + ((((lane >> 1) & 15) >> 2) & 1) ^ (uint32_t)(lane & 1); + // Descriptor STRIDE = rotation period: the largest power-of-two multiple + // of 64 that fits the 14-bit field. Non-power-of-two pitches keep the + // rotation through the idxen split (index*STRIDE + row*rem + col is + // byte-identical to row*pitch + col). + S.a_stride_bytes = a_stride_bytes; // unused on the NN advance path + uint32_t a_stride_desc = 64u; + while(a_stride_desc * 2u <= (uint32_t)a_stride_bytes && a_stride_desc * 2u <= 0x3fffu) + { + a_stride_desc *= 2u; + } + const uint32_t a_off_rem = (uint32_t)a_stride_bytes - a_stride_desc; + S.a_pair = (uint64_t)(uint32_t)a_row_global | + ((uint64_t)((uint32_t)(a_row_global * (int)a_off_rem + a_kh * 16)) << 32); + const uint32_t a_size_bytes = (uint32_t)((int64_t)p.m * a_stride_bytes); + S.a_rsrc = make_buffer_rsrc(p.a_ptr, a_stride_desc, a_size_bytes, /*swizzle=*/true); + // Read base: row (lane&15) at +32*row, 8 B k-unit (lane>>4) at +du, with + // du compensating the swapped rows' [kh1|kh0] placement. + const uint32_t a_rho = (uint32_t)(lane & 15); + const uint32_t a_uu = (uint32_t)(lane >> 4); + const uint32_t a_t = (a_rho >> 2) & 1; + const uint32_t a_du = a_t ? ((a_uu < 2u) ? 16u + 8u * a_uu : 8u * a_uu - 16u) : 8u * a_uu; + S.a_rd = (uint32_t)(wave_m_start * 32 + a_rho * 32 + a_du); + } + else + { + // ---- A = [K, M], M-contiguous (the mirrored NN-B case) ------------ + const int a_stride_bytes = (int)p.a_stride * 2; + const int a_k_in_stage = wid * 2 + (lane >> 5); + const int a_m_global = + m_block_start + 8 * (((lane >> 3) & 3) * 8 + ((lane & 7) ^ (4 * ((lane >> 5) & 1)))); + S.a_vaddr = (uint32_t)(a_k_in_stage * a_stride_bytes + a_m_global * 2); + const uint32_t a_size_bytes = (uint32_t)((int64_t)p.k * a_stride_bytes); + S.a_rsrc = make_buffer_rsrc(p.a_ptr, (uint32_t)a_stride_bytes, a_size_bytes, false); + S.a_stride_bytes = a_stride_bytes; + } + + { + // ---- B = [K, N], N-contiguous (both layouts) ---------------------- + const int b_stride_bytes = S.b_stride_bytes; + // Monotonic walk: lanes 0-31 walk row 2w's full 512 B in address order, + // lanes 32-63 walk row 2w+1's; the odd row's lane->octet map is XORed by + // 4 so the two rows land on complementary bank halves. That is the L1 + // request-merge shape (the merged walk measured 310.8 -> 369.1 TF). + const int b_k_in_stage = wid * 2 + (lane >> 5); + const int b_n_global = + n_block_start + 8 * (((lane >> 3) & 3) * 8 + ((lane & 7) ^ (4 * ((lane >> 5) & 1)))); + S.b_vaddr = (uint32_t)(b_k_in_stage * b_stride_bytes + b_n_global * 2); + const uint32_t b_size_bytes = (uint32_t)((int64_t)p.k * b_stride_bytes); + S.b_rsrc = make_buffer_rsrc(p.b_ptr, (uint32_t)b_stride_bytes, b_size_bytes, false); + } + + // Read anchors, one per-lane vaddr per read: lane l supplies read r's anchor + // for octet o = base + 4r + (l&3), row parity rw = (l>>2)&1, window (l>>3): + // addr = 1024*(l>>3) + 512*rw + 16*(o ^ (4*rw)) + // (the ^4*rw is the same phase term the walk applies to the odd row, so the + // two reads hit complementary bank halves). One register per read is what + // makes "+4 octets" a free immediate; the per-lane anchors are why the read + // forms take arrays. + const uint32_t rw = (uint32_t)((lane >> 2) & 1); + const uint32_t ph = 4u * rw; + const uint32_t b_rd0 = + (uint32_t)G::kLdsBytesA + (uint32_t)(1024 * (lane >> 3)) + 512u * rw; +#pragma unroll + for(int r = 0; r < 4; ++r) + { + const uint32_t o = (uint32_t)(16 * wave_n) + 4u * (uint32_t)r + (uint32_t)(lane & 3); + S.b_rb[r] = b_rd0 + 16u * (o ^ ph); + } + const uint32_t a_rd0 = (uint32_t)(1024 * (lane >> 3)) + 512u * rw; +#pragma unroll + for(int r = 0; r < 2; ++r) + { + const uint32_t o = (uint32_t)(8 * wave_m) + 4u * (uint32_t)r + (uint32_t)(lane & 3); + S.a_rb[r] = a_rd0 + 16u * (o ^ ph); + } + + AOpBank a_bank; + BOpBank b_bank; + + // Prologue: fill the ring, consume banks 0/1, prime the pre-loop slot. + // so_a: NN walks K-contiguous A (32 B per stage); TN walks the [K, M] + // panel row-major (a_stride_bytes per stage) -- see issue_load_and_read. +#pragma unroll + for(index_t st = 0; st < G::kNumStages; ++st) + { + const uint32_t so_a = + (uint32_t)(st * G::kStageK * + (kKind == V7UniformKind::NN ? 2u : (uint32_t)S.a_stride_bytes)); + const uint32_t so_b = (uint32_t)(st * G::kStageK * S.b_stride_bytes); + if constexpr(kKind == V7UniformKind::NN) + { + buffer_load_dwordx4_idxen_lds_t( + S.a_rsrc, S.a_pair, so_a, S.m0_a_base + (uint32_t)(st * G::kLdsBytesPerStage)); + buffer_load_dwordx4_offen_lds( + S.b_rsrc, S.b_vaddr, so_b, S.m0_b_base + (uint32_t)(st * G::kLdsBytesPerStage)); + } + else + { + buffer_load_dwordx4_offen_lds( + S.a_rsrc, S.a_vaddr, so_a, S.m0_a_base + (uint32_t)(st * G::kLdsBytesPerStage)); + buffer_load_dwordx4_offen_lds( + S.b_rsrc, S.b_vaddr, so_b, S.m0_b_base + (uint32_t)(st * G::kLdsBytesPerStage)); + } + } + + wait_all_buffer_data_arrived(2 * (G::kNumStages - 1)); + if constexpr(kKind == V7UniformKind::NN) + { + ds_read_b64_x4<0>(S.a_rd, a_bank[0][0], a_bank[0][1], a_bank[0][2], a_bank[0][3]); + } + else + { + ds_read_m32x16_x4<0>(S.a_rb, a_bank[0]); + } + ds_read_m32x16_x8<0>(S.b_rb, b_bank[0]); + wait_all_buffer_data_arrived(2 * (G::kNumStages - 2)); + if constexpr(kKind == V7UniformKind::NN) + { + ds_read_b64_x4( + S.a_rd, a_bank[1][0], a_bank[1][1], a_bank[1][2], a_bank[1][3]); + } + else + { + ds_read_m32x16_x4(S.a_rb, a_bank[1]); + } + ds_read_m32x16_x8(S.b_rb, b_bank[1]); + wait_all_buffer_data_arrived(2); + // lgkmcnt(0), not 8: the pre-loop load below overwrites slot 0, so both + // prologue bursts must be fully drained first. + wait_lds_data_arrived(0); + wait_all_warp_arrived(); + { + constexpr uint32_t kPreSlot = + (uint32_t)(((G::kPrefetchDist - 1) % G::kNumStages) * G::kLdsBytesPerStage); + // Runtime, not constexpr: it wraps to 0 when the K extent is exactly the + // ring's own depth (K == kNumStages * kStageK), same as the reference. + const int kPre4 = (S.n_stages > (int)G::kPrefetchDist - 1) ? (int)G::kPrefetchDist - 1 : 0; + const uint32_t so_a = + (uint32_t)(kPre4 * (int)G::kStageK * + (kKind == V7UniformKind::NN ? 2u : (uint32_t)S.a_stride_bytes)); + const uint32_t so_b = (uint32_t)(kPre4 * (int)G::kStageK * S.b_stride_bytes); + if constexpr(kKind == V7UniformKind::NN) + { + buffer_load_dwordx4_idxen_lds_t(S.a_rsrc, S.a_pair, so_a, S.m0_a_base + kPreSlot); + buffer_load_dwordx4_offen_lds(S.b_rsrc, S.b_vaddr, so_b, S.m0_b_base + kPreSlot); + } + else + { + buffer_load_dwordx4_offen_lds(S.a_rsrc, S.a_vaddr, so_a, S.m0_a_base + kPreSlot); + buffer_load_dwordx4_offen_lds(S.b_rsrc, S.b_vaddr, so_b, S.m0_b_base + kPreSlot); + } + } + + // The k loop. One divergent branch for the whole kernel (not per step): the + // exec-mask save/restore is paid once, and each arm is straight-line because + // kGroupA is a compile-time parameter. The branch must stay DIVERGENT: the + // uniform-branch variant corrupts on this toolchain (the scheduler reorders + // the inline-asm waits). + int kt = 0; + if(S.wid < (int)G::kWaveGroup) + { +#pragma unroll 1 + for(bool go = (kt + (int)G::kPhaseCycle <= S.n_stages); go;) + { + stage_step<0, true>(S, acc, a_bank, b_bank, kt + 0); + stage_step<1, true>(S, acc, a_bank, b_bank, kt + 1); + stage_step<2, true>(S, acc, a_bank, b_bank, kt + 2); + stage_step<3, true>(S, acc, a_bank, b_bank, kt + 3); + stage_step<4, true>(S, acc, a_bank, b_bank, kt + 4); + stage_step<5, true>(S, acc, a_bank, b_bank, kt + 5); + stage_step<6, true>(S, acc, a_bank, b_bank, kt + 6); + stage_step<7, true>(S, acc, a_bank, b_bank, kt + 7); + stage_step<8, true>(S, acc, a_bank, b_bank, kt + 8); + stage_step<9, true>(S, acc, a_bank, b_bank, kt + 9); + stage_step<10, true>(S, acc, a_bank, b_bank, kt + 10); + stage_step<11, true, false>(S, acc, a_bank, b_bank, kt + 11); + kt += (int)G::kPhaseCycle; + go = (kt + (int)G::kPhaseCycle <= S.n_stages); + wait_all_warp_arrived(); + } + if(kt < S.n_stages) { stage_step<0, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<1, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<2, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<3, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<4, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<5, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<6, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<7, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<8, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<9, true>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<10, true>(S, acc, a_bank, b_bank, kt); ++kt; } + } + else + { +#pragma unroll 1 + for(bool go = (kt + (int)G::kPhaseCycle <= S.n_stages); go;) + { + stage_step<0, false>(S, acc, a_bank, b_bank, kt + 0); + stage_step<1, false>(S, acc, a_bank, b_bank, kt + 1); + stage_step<2, false>(S, acc, a_bank, b_bank, kt + 2); + stage_step<3, false>(S, acc, a_bank, b_bank, kt + 3); + stage_step<4, false>(S, acc, a_bank, b_bank, kt + 4); + stage_step<5, false>(S, acc, a_bank, b_bank, kt + 5); + stage_step<6, false>(S, acc, a_bank, b_bank, kt + 6); + stage_step<7, false>(S, acc, a_bank, b_bank, kt + 7); + stage_step<8, false>(S, acc, a_bank, b_bank, kt + 8); + stage_step<9, false>(S, acc, a_bank, b_bank, kt + 9); + stage_step<10, false>(S, acc, a_bank, b_bank, kt + 10); + stage_step<11, false, false>(S, acc, a_bank, b_bank, kt + 11); + kt += (int)G::kPhaseCycle; + go = (kt + (int)G::kPhaseCycle <= S.n_stages); + wait_all_warp_arrived(); + } + if(kt < S.n_stages) { stage_step<0, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<1, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<2, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<3, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<4, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<5, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<6, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<7, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<8, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<9, false>(S, acc, a_bank, b_bank, kt); ++kt; } + if(kt < S.n_stages) { stage_step<10, false>(S, acc, a_bank, b_bank, kt); ++kt; } + } +} + +} // namespace ck_tile + +#pragma clang diagnostic pop diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform_fwd.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform_fwd.hpp new file mode 100644 index 0000000..04fd4db --- /dev/null +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_v7_uniform_fwd.hpp @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. +// +// Shared declarations for the v7/v8 uniform pipeline +// (gemm_pipeline_v7_uniform.hpp). The grouped-GEMM kernel header includes +// THIS header instead of the pipeline body: the tile parameters must be a +// COMPLETE type for every TU that may instantiate RunGemmV7Uniform (the HIP +// device pass instantiates the body eagerly), while the pipeline itself is +// still completed only inside the v7/v8 translation units. The struct is a +// trivially-copyable parameter block; parsing it costs the other TUs nothing. +#pragma once + +namespace ck_tile { + +enum class V7UniformKind +{ + NN, // A = [M, K], B = [K, N] + TN, // A = [K, M], B = [K, N] +}; + +// Per-tile parameters, as the kernel path derives them from the device kargs. +// Raw pointers and explicit origins: the pipeline builds its own per-lane +// addresses from them (no tensor views, no windows), which is what keeps the +// register budget at the reference kernel's level. +struct V7UniformTileParams +{ + const void* a_ptr = nullptr; + const void* b_ptr = nullptr; + void* d_ptr = nullptr; + index_t m = 0; // this group's row count (M) + index_t n = 0; // common N + index_t k = 0; // common K + index_t a_stride = 0; // row pitch of the TRUE A under the active layout + index_t b_stride = 0; // row pitch of the TRUE B (n) + index_t d_stride = 0; // row pitch of D (n) + index_t m_start = 0; // this tile's global M origin (from the group base) + index_t n_start = 0; // this tile's global N origin +}; + +template +struct GemmPipelineV7Uniform; + +} // namespace ck_tile From 559098452defc58872093d47fb5cf4df852287fd Mon Sep 17 00:00:00 2001 From: wugangsheng Date: Tue, 13 Oct 2026 09:52:24 +0800 Subject: [PATCH 08/11] feat(grouped_gemm): integrate Wasp NT pipeline Add the CK Wasp producer-consumer pipeline, grouped single-group launch path, direct-store epilogue, and focused correctness/performance driver. Preserve the C/Cp stage schedule while hoisting producer wave-ID reads out of the K loop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ck_tile/19_grouped_gemm/CMakeLists.txt | 59 + .../ck_tile/19_grouped_gemm/grouped_gemm.hpp | 46 +- .../grouped_gemm_v7_instance_test.cpp | 649 +++++- .../grouped_gemm_wasp_epilogue.hpp | 565 +++++ .../grouped_gemm_wasp_test.cpp | 2068 +++++++++++++++++ .../instances/grouped_gemm_impl.hpp | 69 +- .../ops/gemm/kernel/grouped_gemm_kernel.hpp | 195 +- .../pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp | 1653 +++++++++++++ 8 files changed, 5292 insertions(+), 12 deletions(-) create mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp create mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp create mode 100644 include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp diff --git a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt index c335ffa..3b994cb 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt +++ b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt @@ -500,6 +500,24 @@ if(TARGET ck_tile_hcu_grouped_gemm_v7_v8_bf16) PRIVATE CK_TILE_GROUPED_GEMM_FAST_BUILD) target_link_libraries(tile_example_grouped_gemm_v7_instance_test PRIVATE ck_tile_hcu_grouped_gemm) + # Optional device-side functional reference: hipBLASLt, one plain GEMM per + # group (--hipblaslt-ref=1). Enabled only when the install the standalone + # harness builds against is present, so the test still builds everywhere. + set(HIPBLASLT_INSTALL "/root/workspace/hipblaslt-install-0820/hipblaslt-install" + CACHE PATH "hipBLASLt install root for the v7 instance test") + if(EXISTS "${HIPBLASLT_INSTALL}/lib/libhipblaslt.so") + # SYSTEM: the hipBLASLt public header is not -Weverything clean; as a + # system include clang does not warn inside it. + target_include_directories(tile_example_grouped_gemm_v7_instance_test + SYSTEM PRIVATE ${HIPBLASLT_INSTALL}/include) + target_link_libraries(tile_example_grouped_gemm_v7_instance_test + PRIVATE ${HIPBLASLT_INSTALL}/lib/libhipblaslt.so) + set_target_properties(tile_example_grouped_gemm_v7_instance_test + PROPERTIES BUILD_RPATH "${HIPBLASLT_INSTALL}/lib" + INSTALL_RPATH "${HIPBLASLT_INSTALL}/lib") + target_compile_definitions(tile_example_grouped_gemm_v7_instance_test + PRIVATE CK_TILE_V7_TEST_HAS_HIPBLASLT) + endif() add_test(NAME tile_example_grouped_gemm_v7_instance_test COMMAND $ --layout=NN --m=512 --n=512 --k=256 --groups=4 --iters=5 --warmup=2) @@ -507,6 +525,47 @@ if(TARGET ck_tile_hcu_grouped_gemm_v7_v8_bf16) PROPERTIES LABELS "hcu;ck_tile;grouped_gemm;v7;correctness;perf") endif() +# Explicit opt-in driver for the Wasp NT grouped-GEMM pipeline. It is not part +# of the default test suite: use the documented 8152x4096x4096 workload when +# validating correctness, resources, and performance. +add_executable(tile_example_grouped_gemm_wasp_test EXCLUDE_FROM_ALL + grouped_gemm_wasp_test.cpp) +target_include_directories(tile_example_grouped_gemm_wasp_test + PRIVATE ${CMAKE_CURRENT_LIST_DIR}) +target_compile_definitions(tile_example_grouped_gemm_wasp_test + PRIVATE CK_TILE_GROUPED_GEMM_FAST_BUILD) +if(NOT DEFINED HIPBLASLT_INSTALL) + set(HIPBLASLT_INSTALL "" CACHE PATH "hipBLASLt install root for the Wasp test") +endif() +if(EXISTS "${HIPBLASLT_INSTALL}/include/hipblaslt/hipblaslt.h" AND + EXISTS "${HIPBLASLT_INSTALL}/lib/libhipblaslt.so") + target_include_directories(tile_example_grouped_gemm_wasp_test + SYSTEM PRIVATE ${HIPBLASLT_INSTALL}/include) + target_link_libraries(tile_example_grouped_gemm_wasp_test + PRIVATE ${HIPBLASLT_INSTALL}/lib/libhipblaslt.so) + set_target_properties(tile_example_grouped_gemm_wasp_test + PROPERTIES BUILD_RPATH "${HIPBLASLT_INSTALL}/lib") + target_compile_definitions(tile_example_grouped_gemm_wasp_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) +endif() +option(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT + "Instrument per-consumer Wasp epilogue cycles in the standalone test" OFF) +if(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + target_compile_definitions(tile_example_grouped_gemm_wasp_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) +endif() +option(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT + "Instrument one steady-state Wasp producer/consumer VMEM wait per wave" OFF) +if(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + target_compile_definitions(tile_example_grouped_gemm_wasp_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) +endif() +option(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER + "Diagnostic: reverse Wasp CTA-to-tile mapping without changing tile work" OFF) +if(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + target_compile_definitions(tile_example_grouped_gemm_wasp_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) +endif() # Host-only Phase-1 selector/contract test. This target deliberately does not # link a grouped-GEMM provider, so adding it cannot alter a device launch. add_executable(tile_example_grouped_gemm_selector_test EXCLUDE_FROM_ALL diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp index f70d12b..5c57bde 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm.hpp @@ -23,6 +23,7 @@ #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3_bw_family_operand.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mls.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp" // #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" @@ -46,6 +47,7 @@ // and the config carries the layout that the host-side selector gated on. #define CK_TILE_PIPELINE_COMPUTE_V7 12 #define CK_TILE_PIPELINE_COMPUTE_V8 13 +#define CK_TILE_PIPELINE_COMPUTE_WASP 14 #ifndef CK_TILE_PIPELINE_DEFAULT #define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 @@ -434,6 +436,38 @@ struct GemmConfigComputeV3RegPrefetch256 : public GemmConfigBase static constexpr int kBlockPerCu = 1; }; +// Explicitly selected gfx936 BF16 NT warp-specialized path. The tile warp grid +// describes the eight consumer waves; the Wasp pipeline contributes four +// producer waves to the 12-wave workgroup. +template +struct GemmConfigWasp256 : public GemmConfigBase +{ + static_assert(std::is_same_v, + "The Wasp NT pipeline currently supports BF16 only"); + + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 8; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool kPadM = true; + static constexpr bool kPadN = false; + static constexpr bool kPadK = false; + static constexpr bool Persistent = true; + static constexpr bool DoubleSmemBuffer = false; + static constexpr bool SupportsFastRowMajorB = false; + static constexpr bool SupportsFastColumnMajorA = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_WASP; + static constexpr int kBlockPerCu = 1; +}; + template struct GemmConfigComputeV3RegPrefetch256W16Group8Aligned : public GemmConfigComputeV3RegPrefetch256 @@ -2149,6 +2183,15 @@ struct PipelineTypeTraits ck_tile::GemmPipelineV7Uniform; }; +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrWasp; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; +}; + template struct GemmConfigComputeInt4 : public GemmConfigBase { @@ -2659,4 +2702,5 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, void* kargs_ptr, const std::vector>& output_resets, bool splitk = false, - std::uint32_t requested_num_cu = 0); + std::uint32_t requested_num_cu = 0, + ck_tile::index_t flat_grid_size = 0); diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp index 633e4bf..ccdc343 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp @@ -38,6 +38,11 @@ // grouped_gemm_v7_instance_test [--layout=NN|TN] [--m=512] [--n=512] // [--k=256] [--groups=2] [--iters=20] // [--warmup=3] [--check=1] [--ref=1] +// [--hipblaslt-ref=0] [--hipblaslt-perf=0] +// --hipblaslt-ref=1: verify against one hipBLASLt GEMM per group, compared +// on the device (no CPU pass over the output) +// --hipblaslt-perf=1: time hipBLASLt (plans built once, bare Matmul in the +// loop) and print blas_ms / blas_tf next to the v7 numbers #include @@ -46,6 +51,11 @@ #include "ck_tile/host/device_prop.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" +#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) +#include +#include +#endif + #include #include #include @@ -85,6 +95,15 @@ struct Options bool dump = false; // Grid override. 0 => the device's CU count (the production value). int num_cu = 0; + // Device-side functional reference: run the same kargs through hipBLASLt + // (one plain GEMM per group) and compare on the GPU -- no CPU pass over the + // output, so production shapes check in seconds. Built only when the + // hipBLASLt install is present (see CMakeLists CK_TILE_V7_TEST_HAS_HIPBLASLT). + bool hipblaslt_ref = false; + // Performance reference: time hipBLASLt (plans built once, bare Matmul in + // the loop) with the same warmup/iters as the v7 timing, and print + // blas_ms / blas_tf next to v7_ms / v7_tf. + bool hipblaslt_perf = false; }; std::size_t align_to(std::size_t value, std::size_t alignment) @@ -230,6 +249,14 @@ void parse_args(int argc, char* argv[], Options& options) { options.num_cu = parse_int_arg(value_of("num-cu"), "num-cu"); } + else if(arg.rfind("--hipblaslt-ref=", 0) == 0) + { + options.hipblaslt_ref = parse_int_arg(value_of("hipblaslt-ref"), "hipblaslt-ref") != 0; + } + else if(arg.rfind("--hipblaslt-perf=", 0) == 0) + { + options.hipblaslt_perf = parse_int_arg(value_of("hipblaslt-perf"), "hipblaslt-perf") != 0; + } else { std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); @@ -335,6 +362,452 @@ void cpu_reference(Case& c) } } +#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) +// --------------------------------------------------------------------------- +// Device-side functional reference: one plain hipBLASLt GEMM per group, then a +// device-side element compare -- the output never touches the CPU, so +// production shapes verify in seconds. +// +// The col-major formulation (same construction bf16_gemm_hipblaslt.cpp uses): +// D_row[M_g, N] = A . B <=> C_col[N, M_g] = B_view[N, K] . A_view[K, M_g] +// with B_view the [N, K] col-major view of the [K, N] row-major B (ld = N), +// and A_view the operand buffer seen as +// NN: [K, M_g] col-major, ld = K (the [M_g, K] row-major A) -> opB = N +// TN: [M_g, K] col-major, ld = pitch (the [K, pitch] A panel) -> opB = T +// The panel leading dimension is the one thing the bf16_gemm_hipblaslt wrapper +// cannot express for the grouped TN operand (it hardcodes ldb = k); here it is +// passed explicitly. +// --------------------------------------------------------------------------- +void hipblaslt_check_status(hipblasStatus_t status, const char* what) +{ + if(status != HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "hipBLASLt error %d at %s\n", + static_cast(status), + what); + std::exit(1); + } +} + +// --------------------------------------------------------------------------- +// Tuned-solution table -- same file and semantics as bf16_gemm_hipblaslt.cpp +// (BF16GEMM_HIPBLASLT_SOLUTION_FILE selects it, BF16GEMM_NO_TUNING=1 disables; +// lines "m,n,k,transA,transB,solution_index", '*' = wildcard, '#' comments). +// The keys live in THIS test's hipblasLt parameter space (m = caller N, +// n = M_g, k = K; transA is always N here, transB = T for the TN panel form), +// which is the same space the wrapper tunes in. Resolution: best-scoring +// entry -> hipblaslt_ext getAlgosFromIndex + matmulIsAlgoSupported -> else the +// stock heuristic, warning once per rejected key. +// --------------------------------------------------------------------------- +struct TunedSolution +{ + long long m = -1, n = -1, k = -1; // -1 = wildcard + char ta = '?', tb = '?'; + int index = -1; +}; + +static long long tuned_num(const std::string& s) +{ + return (s == "*") ? -1 : std::strtoll(s.c_str(), nullptr, 10); +} + +static bool tuned_hit(long long key, long long want) +{ + return key < 0 || key == want; +} + +static bool tuned_hit(char key, char want) +{ + return key == '*' || key == '?' || key == want; +} + +static const std::vector& tuned_table() +{ + static const std::vector table = [] { + std::vector t; + const char* disable = std::getenv("BF16GEMM_NO_TUNING"); + if(disable != nullptr && std::atoi(disable) != 0) + { + return t; + } + const char* path = std::getenv("BF16GEMM_HIPBLASLT_SOLUTION_FILE"); + if(path == nullptr || *path == '\0') + { + return t; + } + FILE* f = std::fopen(path, "r"); + if(f == nullptr) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: cannot open '%s', using the stock heuristic\n", + path); + return t; + } + char line[512]; + while(std::fgets(line, sizeof(line), f) != nullptr) + { + char* p = line; + while(*p == ' ' || *p == '\t') + { + ++p; + } + if(*p == '#' || *p == '\n' || *p == '\r' || *p == '\0') + { + continue; + } + std::string fields[6]; + int nf = 0; + for(char* tok = std::strtok(p, ", \t\r\n"); tok != nullptr && nf < 6; + tok = std::strtok(nullptr, ", \t\r\n")) + { + fields[nf++] = tok; + } + if(nf < 6) + { + std::fprintf(stderr, "[hipblasLt] WARNING: malformed line ignored: %s", line); + continue; + } + TunedSolution e; + e.m = tuned_num(fields[0]); + e.n = tuned_num(fields[1]); + e.k = tuned_num(fields[2]); + e.ta = static_cast(std::toupper(static_cast(fields[3][0]))); + e.tb = static_cast(std::toupper(static_cast(fields[4][0]))); + e.index = static_cast(std::strtol(fields[5].c_str(), nullptr, 10)); + t.push_back(e); + } + std::fclose(f); + std::fprintf(stderr, + "[hipblasLt] tuned-solution table: %zu entries from %s\n", + t.size(), + path); + return t; + }(); + return table; +} + +static bool tuned_pick(long long m, long long n, long long k, char ta, char tb, int& index) +{ + int best_score = 6; + bool found = false; + for(const TunedSolution& e : tuned_table()) + { + if(!(tuned_hit(e.m, m) && tuned_hit(e.n, n) && tuned_hit(e.k, k) && + tuned_hit(e.ta, ta) && tuned_hit(e.tb, tb))) + { + continue; + } + const int score = (e.m < 0 ? 1 : 0) + (e.n < 0 ? 1 : 0) + (e.k < 0 ? 1 : 0) + + ((e.ta == '*' || e.ta == '?') ? 1 : 0) + + ((e.tb == '*' || e.tb == '?') ? 1 : 0); + if(score < best_score) + { + best_score = score; + index = e.index; + found = true; + if(score == 0) + { + break; // exact key: cannot do better + } + } + } + return found; +} + +// D = 1.0 * op(B) . op(A) + 0.0, all bf16 in / bf16 out, fp32 compute. +// Descriptor + algorithm selection is split from the execution: the perf +// reference builds each group's plan once and times bare hipblasLtMatmul +// calls, exactly what a production consumer would do. +struct HipblasltPlan +{ + hipblasLtMatmulDesc_t desc = nullptr; + hipblasLtMatrixLayout_t a_layout = nullptr; // the B operand, col-major [N, K], ld = N + hipblasLtMatrixLayout_t b_layout = nullptr; // the A operand + hipblasLtMatrixLayout_t c_layout = nullptr; // D, col-major [N, M_g], ld = N + hipblasLtMatmulAlgo_t algo{}; +}; + +bool hipblaslt_make_plan(HipblasltPlan& plan, + hipblasLtHandle_t handle, + void* workspace, + std::size_t workspace_size, + int m_g, + int n, + int k, + int ld_a, + bool tn) +{ + (void)workspace; // the workspace pointer only matters at Matmul time + hipblasLtMatmulPreference_t pref = nullptr; + + hipblasOperation_t op_a = HIPBLAS_OP_N; + hipblasOperation_t op_b = tn ? HIPBLAS_OP_T : HIPBLAS_OP_N; + + hipblaslt_check_status( + hipblasLtMatmulDescCreate(&plan.desc, HIPBLAS_COMPUTE_32F, HIP_R_32F), + "hipblasLtMatmulDescCreate"); + hipblaslt_check_status( + hipblasLtMatmulDescSetAttribute(plan.desc, + HIPBLASLT_MATMUL_DESC_TRANSA, + &op_a, + sizeof(op_a)), + "SetAttribute(TRANSA)"); + hipblaslt_check_status( + hipblasLtMatmulDescSetAttribute(plan.desc, + HIPBLASLT_MATMUL_DESC_TRANSB, + &op_b, + sizeof(op_b)), + "SetAttribute(TRANSB)"); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.a_layout, + HIP_R_16BF, + static_cast(n), + static_cast(k), + static_cast(n)), + "LayoutCreate(B view)"); + // opT stores the [K, M_g] logical matrix through the [M_g, K] panel: ld is + // the panel pitch, not K. + const std::size_t b_ld = tn ? static_cast(ld_a) + : static_cast(k); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.b_layout, + HIP_R_16BF, + static_cast(k), + static_cast(m_g), + b_ld), + "LayoutCreate(A view)"); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.c_layout, + HIP_R_16BF, + static_cast(n), + static_cast(m_g), + static_cast(n)), + "LayoutCreate(D)"); + + hipblaslt_check_status( + hipblasLtMatmulPreferenceCreate(&pref), + "hipblasLtMatmulPreferenceCreate"); + hipblaslt_check_status( + hipblasLtMatmulPreferenceSetAttribute( + pref, + HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, + sizeof(workspace_size)), + "PreferenceSetAttribute(WORKSPACE)"); + + // 1) tuned table (best-scoring entry, resolved through the ext API); + // 2) stock heuristic. A table entry the library rejects only warns. + const char ta_ch = (op_a == HIPBLAS_OP_N) ? 'N' : 'T'; + const char tb_ch = (op_b == HIPBLAS_OP_N) ? 'N' : 'T'; + int tuned_index = -1; + if(tuned_pick(n, m_g, k, ta_ch, tb_ch, tuned_index)) + { + std::vector resolved(1); + std::vector algo_index(1, tuned_index); + const float probe_alpha = 1.0f; + const float probe_beta = 0.0f; + std::size_t required = 0; + if(hipblaslt_ext::getAlgosFromIndex(handle, algo_index, resolved) != + HIPBLAS_STATUS_SUCCESS || + resolved[0].state != HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d does not exist; " + "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", + tuned_index, + n, + m_g, + k, + ta_ch, + tb_ch); + } + else if(hipblaslt_ext::matmulIsAlgoSupported(handle, + plan.desc, + &probe_alpha, + plan.a_layout, + plan.b_layout, + &probe_beta, + plan.c_layout, + plan.c_layout, + resolved[0].algo, + required) != + HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d unsupported here; " + "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", + tuned_index, + n, + m_g, + k, + ta_ch, + tb_ch); + } + else if(required > workspace_size) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d needs %zu bytes " + "workspace (have %zu); stock heuristic\n", + tuned_index, + required, + workspace_size); + } + else + { + plan.algo = resolved[0].algo; + std::fprintf(stderr, + "[hipblasLt] tuned: m=%d n=%d k=%d trans=%c%c -> index=%d\n", + n, + m_g, + k, + ta_ch, + tb_ch, + tuned_index); + hipblasLtMatmulPreferenceDestroy(pref); + return true; + } + } + + hipblasLtMatmulHeuristicResult_t algo{}; + int algo_count = 0; + hipblaslt_check_status( + hipblasLtMatmulAlgoGetHeuristic(handle, + plan.desc, + plan.a_layout, + plan.b_layout, + plan.c_layout, + plan.c_layout, + pref, + 1, + &algo, + &algo_count), + "hipblasLtMatmulAlgoGetHeuristic"); + const bool ok = algo_count > 0; + if(ok) + { + plan.algo = algo.algo; + } + hipblasLtMatmulPreferenceDestroy(pref); + return ok; +} + +void hipblaslt_run_plan(const HipblasltPlan& plan, + hipblasLtHandle_t handle, + void* workspace, + std::size_t workspace_size, + hipStream_t stream, + const void* b_dev, + const void* a_dev, + void* d_dev) +{ + const float alpha = 1.0f; + const float beta = 0.0f; + hipblaslt_check_status( + hipblasLtMatmul(handle, + plan.desc, + &alpha, + b_dev, + plan.a_layout, + a_dev, + plan.b_layout, + &beta, + d_dev, + plan.c_layout, + d_dev, + plan.c_layout, + &plan.algo, + workspace, + workspace_size, + stream), + "hipblasLtMatmul"); +} + +void hipblaslt_destroy_plan(HipblasltPlan& plan) +{ + hipblasLtMatrixLayoutDestroy(plan.c_layout); + hipblasLtMatrixLayoutDestroy(plan.b_layout); + hipblasLtMatrixLayoutDestroy(plan.a_layout); + hipblasLtMatmulDescDestroy(plan.desc); + plan = HipblasltPlan{}; +} + +// Element compare on the device: bad = |got - ref| > atol + rtol*|ref|, +// max_err tracked via the non-negative-float bit pattern (uint order == float +// order for non-negative values). +__global__ void blas_compare_kernel(const Bf16* __restrict__ got, + const Bf16* __restrict__ ref, + long long count, + float atol, + float rtol, + unsigned long long* bad, + unsigned int* max_bits) +{ + const long long i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if(i >= count) + { + return; + } + const float g = ck_tile::type_convert(got[i]); + const float r = ck_tile::type_convert(ref[i]); + const float d = fabsf(g - r); + if(d > atol + rtol * fabsf(r)) + { + atomicAdd(bad, 1ull); + } + atomicMax(max_bits, __float_as_uint(d)); +} + +struct BlasCheckResult +{ + long long bad = 0; + double max_err = 0.0; +}; + +BlasCheckResult blas_compare(hipStream_t stream, + const void* got, + const void* ref, + long long count) +{ + unsigned long long* bad_dev = nullptr; + unsigned int* max_bits_dev = nullptr; + hip_check(hipMalloc(&bad_dev, sizeof(*bad_dev)), "hipMalloc(bad)"); + hip_check(hipMalloc(&max_bits_dev, sizeof(*max_bits_dev)), "hipMalloc(max)"); + hip_check(hipMemset(bad_dev, 0, sizeof(*bad_dev)), "hipMemset(bad)"); + hip_check(hipMemset(max_bits_dev, 0, sizeof(*max_bits_dev)), "hipMemset(max)"); + + const int block = 256; + const long long grid = (count + block - 1) / block; + blas_compare_kernel<<(grid), block, 0, stream>>>( + static_cast(got), + static_cast(ref), + count, + 1.0f, + 0.01f, + bad_dev, + max_bits_dev); + hip_check(hipGetLastError(), "blas_compare_kernel launch"); + + unsigned long long bad = 0; + unsigned int max_bits = 0; + hip_check(hipMemcpyAsync(&bad, bad_dev, sizeof(bad), hipMemcpyDeviceToHost, stream), + "hipMemcpyAsync(bad)"); + hip_check(hipMemcpyAsync(&max_bits, max_bits_dev, sizeof(max_bits), hipMemcpyDeviceToHost, stream), + "hipMemcpyAsync(max)"); + hip_check(hipStreamSynchronize(stream), "blas_compare sync"); + hip_check(hipFree(bad_dev), "hipFree(bad)"); + hip_check(hipFree(max_bits_dev), "hipFree(max)"); + + BlasCheckResult result; + result.bad = static_cast(bad); + // The bit pattern came back over memcpy: rebuild the non-negative float. + static_assert(sizeof(max_bits) == sizeof(float), "bit transfer size"); + std::memcpy(&result.max_err, &max_bits, sizeof(result.max_err)); + return result; +} +#endif // CK_TILE_V7_TEST_HAS_HIPBLASLT + struct DeviceCase { std::vector a_ptrs; @@ -710,7 +1183,181 @@ int main(int argc, char* argv[]) } } - const double flops = 2.0 * static_cast(c.m) * c.n * c.k * c.g; + // --- hipBLASLt functional reference (device-side, no CPU pass) ---------- + if(options.hipblaslt_ref) + { +#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT + hipblasLtHandle_t handle = nullptr; + hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); + void* workspace = nullptr; + const std::size_t workspace_size = 256ull << 20; + hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); + + std::vector ref_blocks(static_cast(c.g), nullptr); + std::vector plans(static_cast(c.g)); + for(int g = 0; g < c.g; ++g) + { + const std::size_t bytes = + static_cast(c.group_m[g]) * c.n * sizeof(Bf16); + hip_check(hipMalloc(&ref_blocks[g], bytes), "hipMalloc(blas ref)"); + hip_check(hipMemset(ref_blocks[g], 0, bytes), "hipMemset(blas ref)"); + if(!hipblaslt_make_plan(plans[g], + handle, + workspace, + workspace_size, + c.group_m[g], + c.n, + c.k, + c.strides_a[g], + c.tn)) + { + std::fprintf(stderr, "hipBLASLt: no algorithm for group %d\n", g); + std::exit(1); + } + // One plain GEMM per group on the same stream the instance used. + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + ref_blocks[g]); + } + hip_check(hipStreamSynchronize(stream), "hipblaslt ref sync"); + + long long blas_bad = 0; + double blas_max_err = 0.0; + for(int g = 0; g < c.g; ++g) + { + const BlasCheckResult r = + blas_compare(stream, + device.c_ptrs[g], + ref_blocks[g], + static_cast(c.group_m[g]) * c.n); + blas_bad += r.bad; + blas_max_err = blas_max_err > r.max_err ? blas_max_err : r.max_err; + hip_check(hipFree(ref_blocks[g]), "hipFree(blas ref)"); + } + std::printf("blas_check layout=%s bad=%lld max_err=%.4f tolerance=1.0+0.01*|ref|\n", + options.layout.c_str(), + blas_bad, + blas_max_err); + ok = ok && blas_bad == 0; + for(auto& plan : plans) + { + hipblaslt_destroy_plan(plan); + } + hip_check(hipFree(workspace), "hipFree(blas workspace)"); + hipblasLtDestroy(handle); +#else + std::printf("blas_check result=FAIL reason=built_without_hipblaslt\n"); + ok = false; +#endif + } + + // --- hipBLASLt performance reference -------------------------------------- + // Plans are built once per group; the timed loop issues bare + // hipblasLtMatmul calls for all groups per iteration -- the launch pattern a + // production consumer would use against the same buffers. + if(options.hipblaslt_perf) + { +#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT + hipblasLtHandle_t handle = nullptr; + hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); + void* workspace = nullptr; + const std::size_t workspace_size = 256ull << 20; + hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); + + std::vector plans(static_cast(c.g)); + bool plans_ok = true; + for(int g = 0; g < c.g && plans_ok; ++g) + { + plans_ok = hipblaslt_make_plan(plans[g], + handle, + workspace, + workspace_size, + c.group_m[g], + c.n, + c.k, + c.strides_a[g], + c.tn); + } + if(!plans_ok) + { + std::printf("blas_perf layout=%s result=FAIL reason=no_algo\n", + options.layout.c_str()); + ok = false; + } + else + { + for(int i = 0; i < options.warmup; ++i) + { + for(int g = 0; g < c.g; ++g) + { + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + device.c_ptrs[g]); + } + } + hipEvent_t start{}; + hipEvent_t stop{}; + hip_check(hipEventCreate(&start), "hipEventCreate"); + hip_check(hipEventCreate(&stop), "hipEventCreate"); + hip_check(hipEventRecord(start, stream), "hipEventRecord"); + for(int i = 0; i < options.iters; ++i) + { + for(int g = 0; g < c.g; ++g) + { + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + device.c_ptrs[g]); + } + } + hip_check(hipEventRecord(stop, stream), "hipEventRecord"); + hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); + float elapsed_ms = 0.0f; + hip_check(hipEventElapsedTime(&elapsed_ms, start, stop), "hipEventElapsedTime"); + const double blas_ms = + static_cast(elapsed_ms) / static_cast(options.iters); + const double blas_tf = + 2.0 * static_cast(c.total_m) * c.n * c.k / + (blas_ms * 1.0e-3) / 1.0e12; + std::printf("blas_perf layout=%s m=%d n=%d k=%d groups=%d " + "blas_ms=%.4f blas_tf=%.1f\n", + options.layout.c_str(), + c.m, + c.n, + c.k, + c.g, + blas_ms, + blas_tf); + hip_check(hipEventDestroy(start), "hipEventDestroy"); + hip_check(hipEventDestroy(stop), "hipEventDestroy"); + } + for(auto& plan : plans) + { + hipblaslt_destroy_plan(plan); + } + hip_check(hipFree(workspace), "hipFree(blas workspace)"); + hipblasLtDestroy(handle); +#else + std::printf("blas_perf result=FAIL reason=built_without_hipblaslt\n"); + ok = false; +#endif + } + + const double flops = 2.0 * static_cast(c.total_m) * c.n * c.k; std::printf("perf layout=%s m=%d n=%d k=%d groups=%d v7_ms=%.4f v7_tf=%.1f\n", options.layout.c_str(), c.m, diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp new file mode 100644 index 0000000..cff2147 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp @@ -0,0 +1,565 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technologies Co., Ltd. +// +// Epilogue for the WASP (warp-specialized producer / consumer) NT pipeline. +// +// Ported line by line from the standalone warp-specialized reference +// implementation (`gemm_epilogue_writeback_via_lds_pc`). +// +// WHY THIS IS THE POINT OF THE PORT +// --------------------------------- +// The native MMAC output puts one lane's four results at columns +// {q, q+4, q+8, q+12} across 16 rows (q = lane >> 4), so direct scalar stores +// never coalesce: measured 33.2M VL1 write requests on m=8100 vs +// hipBLASLt's 8.4M, and because two CTAs overlap per CU that write flood +// congests TA/TC for the WHOLE kernel. Staging each 128-row slice through the +// (by then idle) 64 KB of LDS and reading it back with a wave-64-covers- +// 4-rows-x-128B layout turns 128 scalar stores into 16 vector stores per wave: +// write requests 33.2M -> ~1M, measured 1.100 ms / 247 TF -> 0.852 ms / 319 TF. +// +// WHAT THIS EPILOGUE MUST DO THAT CShuffleEpilogue CANNOT +// ------------------------------------------------------ +// 1. ROLE AWARE. Only the 8 consumer waves (warp id >= 4) own output. The 4 +// producer waves fall straight through. CShuffle's kBlockSize is +// MWave*NWave*warp_size = 512 and its thread-raked tile distribution +// indexes by warp id, so waves 8..11 of a 768 thread workgroup would +// compute out-of-range coordinates. +// 2. BARRIER FREE. CShuffle emits several unconditional block_sync_lds() +// (s_waitcnt lgkmcnt(0) + s_barrier) -- see cshuffle_epilogue.hpp:881, +// 883, 900, 917, 934, 944, 965. A workgroup barrier the producer waves +// never reach deadlocks, so this epilogue uses per-wave LDS slabs and only +// the LDS drain it actually needs. The cross-tile drain lives at the +// PIPELINE's entry instead (block_sync_lds()), where both roles arrive. +// +// LDS: 8 consumer waves x 8 KB (128 rows x 32 columns x 2 B) = the pipeline's +// whole 64 KB buffer, reused after the K loop drains. GetSmemSize() is 0 so the +// framework's max(GemmPipeline, EpiloguePipeline) keeps the pipeline's 65536. +// +// STAGE 3 of the port (see the pipeline header for stages 1-2): the writeback +// itself. Correctness and the coalescing win are in place; note the default +// LDS_XOR below is what removes the store-side bank conflicts. +#pragma once + +#include "ck_tile/core.hpp" + +#if defined(CK_TILE_GROUPED_GEMM_WASP_DEBUG_RECORD) +// Debug only: what each CTA's consumer waves actually saw when they reached the +// epilogue. Answers "did this CTA run at all" and "what origin/size did it +// think it had", which reasoning cannot settle. 5 slots per CTA. +// [0] epilogue entries [1] block_org[0] [2] M [3] N [4] ldd +// static: internal linkage, so -Wmissing-variable-declarations (part of +// -Weverything) stays quiet and each TU that opts in gets its own copy. +// per CTA: [0] reached [1] block_org[0] [2] M [3] N [4] ldd +// [5] calculated origin offset [6] p_data_ low [7] d_ptr low +// [8] col_base [9] consumer_id +__device__ static int g_wasp_dbg[128 * 10]; +// One slot per (CTA, wave): a plain store, so it cannot lose the increment that +// a shared counter or a shared OR would. +__device__ static int g_wasp_dbg_waves[128 * 16]; +#endif + +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) +__device__ static uint64_t g_wasp_epi_cycles[1024 * 16 * 9]; +#endif + +template +struct GroupedGemmWaspStoreEpilogue +{ + using AccDataType = ck_tile::remove_cvref_t; + using ODataType = ck_tile::remove_cvref_t; + using DsDataType = ck_tile::tuple<>; + using DsLayout = ck_tile::tuple<>; + + // The pipeline owns the WARP SPLIT; the epilogue only needs the count so + // that it can filter the consumer waves. + using Pipeline = ck_tile::remove_cvref_t; + + using AccTile = typename Pipeline::AccTile; + + static constexpr ck_tile::index_t NumDTensor = 0; + static constexpr ck_tile::memory_operation_enum MemoryOperation = + ck_tile::memory_operation_enum::set; + + // v1 ships 2-byte output only (bf16/fp16 share this path). fp32 output uses + // a different geometry entirely (EpiGeomF32/EpiPassF32) and is not + // ported yet. + static_assert(sizeof(ODataType) == 2, "WASP epilogue v1 is 2-byte output only"); + + // ------------------------------------------------------------------ + // Geometry, mirroring EpiCoalescePC. + // ------------------------------------------------------------------ + static constexpr ck_tile::index_t kWaveM = 256; // rows per consumer wave + static constexpr ck_tile::index_t kWaveN = 32; // columns per consumer wave + + static constexpr int kEpiHalvesPerWarp = 2; // 128-row slabs, two passes + static constexpr int kEpiElemsPerLane = 8; // 16 B per lane per store + + static constexpr int kColsPerRow = kWaveN / kEpiElemsPerLane; // 4 + static constexpr int kLog2ColsPerRow = 2; // log2(4) + static constexpr int kRowsPerPass = 64 / kColsPerRow; // 16 + + static constexpr ck_tile::index_t kRowsPerHalf = kWaveM / kEpiHalvesPerWarp; // 128 + static constexpr int kPasses = kRowsPerHalf / kRowsPerPass; // 8 + + static constexpr int kSubN = kWaveN / 16; // 2 n sub-tiles + static constexpr int kSubMPerHalf = kRowsPerHalf / 16; // 8 m sub-tiles + + static constexpr int kLdsRowBytes = kWaveN * sizeof(ODataType); // 64 + static constexpr int kLdsVectorBytes = 8; + static constexpr int kLdsVectorsPerRow = kLdsRowBytes / kLdsVectorBytes; + static constexpr ck_tile::index_t kSlabBytesPerWave = + kRowsPerHalf * kWaveN * sizeof(ODataType); // 8192 + + static_assert(kColsPerRow == 4 && kRowsPerPass == 16 && kPasses == 8, "epilogue geometry"); + static_assert(kLdsVectorsPerRow == 8, "the epilogue expects eight 8-byte chunks per row"); + // One half consumes kSubMPerHalf x kSubN = 16 accumulator slots; the two + // halves together must cover the whole 16 x 2 = 32 slot wave tile. + static_assert(kEpiHalvesPerWarp * kSubMPerHalf * kSubN == + AccTile::kMIters * AccTile::kNIters, + "the two halves must cover the whole accumulator tile"); + + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetVectorSizeC() + { + return kEpiElemsPerLane; + } + + template + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t + GetVectorSizeD(ck_tile::number) + { + return 1; + } + + // Reuse the pipeline's LDS buffer; the framework takes the max of the two. + CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; } + +#if defined(CK_TILE_GROUPED_GEMM_WASP_NO_BOUNDS_CHECK) + // Debug switch: store without any M/N masking. Used to tell "the second + // m-tile is not processed at all" apart from "the bounds check rejects every + // row of it". Safe for the shapes the test drives, because the rejected rows + // are still inside the allocated output buffer. + static constexpr bool kBoundsCheck = false; +#else + static constexpr bool kBoundsCheck = true; +#endif + + // One 16 B lane payload: two 8 B halves, or eight elements. + union EpiVec8 + { + int64_t q[2]; + int32_t d[4]; + ODataType e[kEpiElemsPerLane]; + }; + + CK_TILE_DEVICE static uint16_t FloatToBf16RneNonNan(float value) + { +#if defined(__HIP_DEVICE_COMPILE__) && defined(__gfx936__) + const uint32_t bits = ck_tile::bit_cast(value); + const uint32_t rounded = bits + 0x7fffU + ((bits >> 16) & 1U); + return static_cast(rounded >> 16); +#else + return ck_tile::bit_cast( + ck_tile::type_convert(value)); +#endif + } + + // ------------------------------------------------------------------ + // Keep the 8-byte XOR-swizzle used by the HIPC epilogue so the LDS reload + // produces output-order vectors directly, without a VGPR permutation. + // ------------------------------------------------------------------ + CK_TILE_DEVICE static int LdsOffset(int row_in_slab, int col) + { + const int vector = (col / 4) ^ (row_in_slab & (kLdsVectorsPerRow - 1)); + return row_in_slab * kLdsRowBytes + vector * kLdsVectorBytes + + (col & 3) * int(sizeof(ODataType)); + } + + // ------------------------------------------------------------------ + // Store one 128-row half of the pre-converted wave tile into the LDS slab. + // + // The MMAC C layout is MEASURED, not assumed: accumulator slot c of lane + // group g holds column 4*c + g (NOT 4*g + c -- the fp8-era convention was + // never valid on gfx936). A lane's four values are therefore four columns + // apart, so the four converted values are staged individually. + // ------------------------------------------------------------------ + template + CK_TILE_DEVICE static void StoreHalfToLds(ODataType* slab, + ck_tile::index_t lane, + const ODataType (*converted)[4]) + { + const int lane_and_15 = static_cast(lane) & 15; + const int lane_lsr_4 = static_cast(lane) >> 4; + char* const base = reinterpret_cast(slab); + + ck_tile::static_for<0, kSubMPerHalf, 1>{}([&](auto mm_local) { + constexpr int mm_global = HALF * kSubMPerHalf + mm_local.value; + const int row_in_slab = mm_local.value * 16 + lane_and_15; + ck_tile::static_for<0, kSubN, 1>{}([&](auto nn) { + constexpr int flat = mm_global * kSubN + nn.value; + ck_tile::static_for<0, 4, 1>{}([&](auto c_index) { + const int c = c_index.value; + const int col = nn.value * 16 + 4 * c + lane_lsr_4; + *reinterpret_cast(base + LdsOffset(row_in_slab, col)) = + converted[flat][c]; + }); + }); + }); + } + + // ------------------------------------------------------------------ + // Read the wave's slab back and write it to global with coalesced 16 B + // stores: a wave-64 now covers 4 rows x 128 B, so one store per lane is a + // whole cache line instead of 16 scattered line transactions. + // ------------------------------------------------------------------ + template + CK_TILE_DEVICE static void LoadLdsAndStoreGlobal(const ODataType* slab, + ODataType* d_ptr, + ck_tile::index_t ldd, + ck_tile::index_t row_base_global, + ck_tile::index_t col_base_global, + [[maybe_unused]] ck_tile::index_t M, + [[maybe_unused]] ck_tile::index_t N, + ck_tile::index_t lane) + { + const int row_in_pass = static_cast(lane) >> kLog2ColsPerRow; + const int chunk = static_cast(lane) & (kColsPerRow - 1); + const int col_in_wave = chunk * kEpiElemsPerLane; + const ck_tile::index_t col_global = col_base_global + col_in_wave; + + // Read two adjacent 8-byte chunks per pass. The LDS swizzle maps them + // directly to the final output order. + EpiVec8 buf[kPasses]; +#pragma unroll + for(int pass = 0; pass < kPasses; ++pass) + { + const int row_in_slab = pass * kRowsPerPass + row_in_pass; +#pragma unroll + for(int q = 0; q < 2; ++q) + { + const int col = col_in_wave + q * 4; + const int off = LdsOffset(row_in_slab, col); + const auto* addr = reinterpret_cast(slab) + off; + const auto lds_addr = static_cast( + static_cast(reinterpret_cast(addr))); + uint64_t loaded_pair; + ck_tile::smem_load<8>{}(loaded_pair, lds_addr, 0); + buf[pass].q[q] = static_cast(loaded_pair); + } + } + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + __builtin_amdgcn_sched_barrier(0); + + // Phase B: coalesced global stores. + bool all_rows_in_bounds = true; + if constexpr(kBoundsCheck) + { + const ck_tile::index_t last_row = row_base_global + HALF * kRowsPerHalf + + kRowsPerHalf - kRowsPerPass + row_in_pass; + all_rows_in_bounds = (last_row < M); + } + + // IsSupportedArgument requires N to be a multiple of the 256-column + // block tile. Each consumer owns 32 columns and stores aligned 8-BF16 + // vectors, so the column-tail/scalar-store path is unreachable. + if(all_rows_in_bounds) + { +#pragma unroll + for(int pass = 0; pass < kPasses; ++pass) + { + const ck_tile::index_t row_global = row_base_global + HALF * kRowsPerHalf + + pass * kRowsPerPass + row_in_pass; + *reinterpret_cast(d_ptr + row_global * ldd + col_global) = buf[pass]; + } + } + else + { +#pragma unroll + for(int pass = 0; pass < kPasses; ++pass) + { + const ck_tile::index_t row_global = row_base_global + HALF * kRowsPerHalf + + pass * kRowsPerPass + row_in_pass; + if constexpr(kBoundsCheck) + { + if(row_global < M) + { + *reinterpret_cast(d_ptr + row_global * ldd + col_global) = + buf[pass]; + } + } + else + { + *reinterpret_cast(d_ptr + row_global * ldd + col_global) = + buf[pass]; + } + } + } + } + + // ------------------------------------------------------------------ + // Entry point. Called by GroupedGemmKernel on EVERY wave of the workgroup + // (grouped_gemm_kernel.hpp:943-947, no role guard there), so the filter is + // ours. Contains no workgroup barrier on purpose. + // ------------------------------------------------------------------ +#if defined(CK_TILE_GROUPED_GEMM_WASP_DIRECT_STORE) + // ------------------------------------------------------------------ + // Debug path: store the accumulator straight to DRAM with no LDS staging. + // 128 scalar stores per lane, i.e. deliberately slow -- but it answers the + // one question that bisects this port: is the ACCUMULATOR itself right? + // Same failure pattern as the staged path -> the compute is wrong. + // Pattern disappears -> the staging is wrong. + // Not for production; no coalescing, no XOR swizzle. + // ------------------------------------------------------------------ + template + CK_TILE_DEVICE void operator()(OWindow& o_window, + const AccTile_& acc_tile, + const DWindow&, + void* = nullptr) const + { + const ck_tile::index_t wid = ck_tile::get_warp_id(); + if(wid < Pipeline::kNumProducerWaves) + { + return; + } + const ck_tile::index_t consumer_id = wid - Pipeline::kNumProducerWaves; + const ck_tile::index_t lane = ck_tile::get_lane_id(); + + auto o_view = o_window.get_bottom_tensor_view(); + const auto& o_desc = o_view.get_tensor_descriptor(); + const auto block_org = o_window.get_window_origin(); + + // Raw base, no window origin: see the staged path for why adding it + // again double counts the tile's row offset. + auto* d_ptr = o_view.get_buffer_view().p_data_; + const ck_tile::index_t ldd = o_desc.calculate_offset(ck_tile::make_multi_index(1, 0)) - + o_desc.calculate_offset(ck_tile::make_multi_index(0, 0)); + [[maybe_unused]] const ck_tile::index_t M = o_desc.get_lengths()[ck_tile::number<0>{}]; + [[maybe_unused]] const ck_tile::index_t N = o_desc.get_lengths()[ck_tile::number<1>{}]; + + const ck_tile::index_t row_in_tile = lane & 15; + const ck_tile::index_t col_quad = lane >> 4; + const ck_tile::index_t col_base = block_org[1] + consumer_id * kWaveN; + +#pragma unroll + for(int m = 0; m < static_cast(AccTile::kMIters); ++m) + { + const ck_tile::index_t row = block_org[0] + m * 16 + row_in_tile; + if constexpr(kBoundsCheck) + { + if(row >= M) + { + continue; + } + } +#pragma unroll + for(int n = 0; n < static_cast(AccTile::kNIters); ++n) + { +#pragma unroll + for(int c = 0; c < 4; ++c) + { + const ck_tile::index_t col = col_base + n * 16 + 4 * c + col_quad; + if(!kBoundsCheck || col < N) + { + d_ptr[row * ldd + col] = ck_tile::type_convert( + acc_tile.v[m * static_cast(AccTile::kNIters) + n][c]); + } + } + } + } + } +#else + template + CK_TILE_DEVICE void operator()(OWindow& o_window, + const AccTile_& acc_tile, + const DWindow&, + void* p_smem = nullptr) const + { + const ck_tile::index_t wid = ck_tile::get_warp_id(); + if(wid < Pipeline::kNumProducerWaves) + { + return; // producer waves own no output + } + if((ck_tile::wasp_detail::wasp_ablate() & 16) != 0) + { + // Ablation (see gemm_pipeline_ag_bg_cr_wasp.hpp): run the whole GEMM + // but throw the writeback away, so the epilogue's cost can be timed + // on its own. + return; + } + // Only used by the DEBUG_RECORD block below; store_raw recomputes its + // own. Keep for that diagnostic without tripping -Wunused otherwise. + [[maybe_unused]] const ck_tile::index_t consumer_id = + wid - Pipeline::kNumProducerWaves; + [[maybe_unused]] const ck_tile::index_t lane = ck_tile::get_lane_id(); + + auto o_view = o_window.get_bottom_tensor_view(); + const auto& o_desc = o_view.get_tensor_descriptor(); + const auto block_org = o_window.get_window_origin(); + + // d_ptr is deliberately the RAW tensor base, WITHOUT the window origin: + // row_global and col_global below are already GLOBAL (they start at + // block_org), so adding the origin again double counts it. That is + // invisible for the first tile (origin 0) and pushes every store of + // every later tile out of the buffer -- which then looks like "the + // second m-tile was never written". + auto* d_ptr = o_view.get_buffer_view().p_data_; + const ck_tile::index_t ldd = o_desc.calculate_offset(ck_tile::make_multi_index(1, 0)) - + o_desc.calculate_offset(ck_tile::make_multi_index(0, 0)); + const ck_tile::index_t M = o_desc.get_lengths()[ck_tile::number<0>{}]; + const ck_tile::index_t N = o_desc.get_lengths()[ck_tile::number<1>{}]; + +#if defined(CK_TILE_GROUPED_GEMM_WASP_DEBUG_RECORD) + { + const int cta = static_cast(ck_tile::get_block_id()); + if(lane == 0 && cta >= 0 && cta < 128) + { + const auto p_data_low = static_cast( + reinterpret_cast(o_view.get_buffer_view().p_data_) & 0xFFFFFFFFu); + const auto d_ptr_low = static_cast( + reinterpret_cast(d_ptr) & 0xFFFFFFFFu); + g_wasp_dbg_waves[cta * 16 + static_cast(wid)] = 1; + g_wasp_dbg[cta * 10 + 0] = 1; + g_wasp_dbg[cta * 10 + 1] = static_cast(block_org[0]); + g_wasp_dbg[cta * 10 + 2] = static_cast(M); + g_wasp_dbg[cta * 10 + 3] = static_cast(N); + g_wasp_dbg[cta * 10 + 4] = static_cast(ldd); + g_wasp_dbg[cta * 10 + 5] = static_cast( + o_desc.calculate_offset(ck_tile::make_multi_index(block_org[0], block_org[1]))); + g_wasp_dbg[cta * 10 + 6] = p_data_low; + g_wasp_dbg[cta * 10 + 7] = d_ptr_low; + g_wasp_dbg[cta * 10 + 8] = static_cast(block_org[1]); + g_wasp_dbg[cta * 10 + 9] = static_cast(consumer_id); + // Proof of life: can this CTA's epilogue reach its own output + // region at all? If the host does not see 777 here, the global + // store itself is being lost, and no amount of bound/offset + // reasoning matters. + d_ptr[0] = ck_tile::type_convert(777.0F); + } + } +#endif + + store_raw(d_ptr, ldd, M, N, block_org[0], block_org[1], acc_tile, p_smem); + } + + // Raw-pointer store core: the staged writeback fed by SCALARS instead of an + // output tile window. `d_ptr` is the RAW tensor base (NO tile origin folded + // in) -- row/col below are global and already start at row_base/col_base, so + // adding the origin again would double-count it. row_base/col_base are the + // tile's global row/col origin (block_org[0]/[1]). The grouped kernel's + // RunGemmWasp path calls this directly with kargs-derived scalars, so no + // output tile-window object is built or kept live. + // + // The role filter and the ablation gate live HERE (not only in the window + // operator() above): RunGemmWasp calls store_raw directly on EVERY wave of + // the workgroup, so producer waves (wid < kNumProducerWaves) must be dropped + // here, exactly as the window path drops them before extracting the window. +#endif // CK_TILE_GROUPED_GEMM_WASP_DIRECT_STORE + + CK_TILE_DEVICE static void store_raw(ODataType* d_ptr, + ck_tile::index_t ldd, + ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t row_base, + ck_tile::index_t col_base, + const AccTile& acc_tile, + void* p_smem) + { + const ck_tile::index_t wid = ck_tile::get_warp_id(); + if(wid < Pipeline::kNumProducerWaves) + { + return; // producer waves own no output + } + if((ck_tile::wasp_detail::wasp_ablate() & 16) != 0) + { + return; // ablation: run the GEMM but throw the writeback away + } + const ck_tile::index_t consumer_id = wid - Pipeline::kNumProducerWaves; + const ck_tile::index_t lane = ck_tile::get_lane_id(); + + auto* const slab = reinterpret_cast(static_cast(p_smem) + + consumer_id * kSlabBytesPerWave); + + const ck_tile::index_t row_base_global = row_base; + const ck_tile::index_t col_base_global = col_base + consumer_id * kWaveN; + +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + uint64_t epi_times[10]{}; + asm volatile("s_memtime %0" : "=s"(epi_times[0])); +#endif + // Convert the complete accumulator before LDS staging, as in the HIPC + // epilogue. This deliberately keeps the converted tile live across both + // halves instead of repeating conversion while issuing LDS stores. + ODataType converted[AccTile::kMIters * AccTile::kNIters][4]; + ck_tile::static_for<0, AccTile::kMIters * AccTile::kNIters, 1>{}([&](auto flat) { + ck_tile::static_for<0, 4, 1>{}([&](auto c) { + if constexpr(std::is_same_v) + { + converted[flat.value][c.value] = + ck_tile::bit_cast( + FloatToBf16RneNonNan(acc_tile.v[flat.value][c.value])); + } + else + { + converted[flat.value][c.value] = + ck_tile::type_convert(acc_tile.v[flat.value][c.value]); + } + }); + }); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[1])); +#endif + + StoreHalfToLds<0>(slab, lane, converted); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[2])); +#endif + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + __builtin_amdgcn_sched_barrier(0); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[3])); +#endif + LoadLdsAndStoreGlobal<0>( + slab, d_ptr, ldd, row_base_global, col_base_global, M, N, lane); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[4])); +#endif + + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + __builtin_amdgcn_sched_barrier(0); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[5])); +#endif + StoreHalfToLds<1>(slab, lane, converted); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[6])); +#endif + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + __builtin_amdgcn_sched_barrier(0); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[7])); +#endif + LoadLdsAndStoreGlobal<1>( + slab, d_ptr, ldd, row_base_global, col_base_global, M, N, lane); +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + asm volatile("s_memtime %0" : "=s"(epi_times[8])); + asm volatile("s_waitcnt vmcnt(0)\n" ::: "memory"); + asm volatile("s_memtime %0" : "=s"(epi_times[9])); + if(lane == 0) + { + const ck_tile::index_t cta = ck_tile::get_block_id(); + if(cta < 1024) + { + const ck_tile::index_t index = (cta * 16 + wid) * 9; + for(int stage = 0; stage < 9; ++stage) + { + g_wasp_epi_cycles[index + stage] = + epi_times[stage + 1] - epi_times[stage]; + } + } + } +#endif + } +}; diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp new file mode 100644 index 0000000..3138dc7 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp @@ -0,0 +1,2068 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technologies Co., Ltd. +// +// Correctness + performance driver for the WASP (warp-specialized +// producer / consumer) NT grouped GEMM pipeline. +// +// grouped_gemm_wasp_geometry_smoke.cpp guards the one structural assumption of +// the design (a 12 wave workgroup carrying a 1 x 8 consumer warp grid). THIS +// file answers the two questions that guard cannot: does it compute the right +// numbers, and how fast is it. +// +// Design notes +// ------------ +// * kargs come from the FRAMEWORK's own host-side builder, +// `GroupedGemmKernel<...>::MakeKargs(std::vector)`, so +// the descriptor the kernel sees is field-for-field what the device-args ABI +// would produce. Nothing is hand-filled, so a harness bug cannot hide behind +// a hand-copied layout. +// * A, B and C are laid out the way the ABI lays them out for NT: +// A [sum(M_g), K] row major, stride_A = K (group g at prefix_M[g]*K) +// B [G, N, K] row major, stride_B = K (group g at g*N*K) +// C [sum(M_g), N] row major, stride_E = N (group g at prefix_M[g]*N) +// Each group's a_ptr/e_ptr points INTO the shared buffer, which is exactly +// what lets the kernel recover the group table from the kargs alone. +// * The oracle is an independent CPU reference (fp32 accumulate), not another +// GPU kernel, so the check still means something when every GPU path is +// wrong. bf16 products are exact in fp32, so only the summation ORDER +// differs between the reference and the kernel: the tolerance can be tight. +// * `--ref=1` (default) additionally runs the same kargs through an existing CK +// configuration. That gives a triangle of evidence: if WASP and the CK +// reference both disagree with the CPU reference, the harness is wrong; if +// only WASP disagrees, it is the kernel. +// +// Current port limitations this driver enforces +// -------------------------------------------- +// The port does not yet mask A rows past M, B columns past N, or a K tail, so +// only aligned shapes are legal: per-group M % 256 == 0, N % 256 == 0, +// K % 64 == 0. Unaligned shapes are reported as SKIPPED (not FAIL) unless +// --allow-unaligned is given, which is what you want when debugging the masking +// itself. +// +// Usage: +// grouped_gemm_wasp_test [--m=...] [--n=...] [--k=...] [--groups=...] +// [--iters=N] [--warmup=N] [--ref=0|1] +// [--allow-unaligned] [--quiet] + +#include "instances/grouped_gemm_impl.hpp" + +#include + +#if defined(CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) +#include +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Bf16 = ck_tile::bf16_t; +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + +using WaspConfig = GemmConfigWasp256; + +// --------------------------------------------------------------------------- +// The type chain `grouped_gemm_tileloop` assembles internally (see +// instances/grouped_gemm_impl.hpp:920-1033), rebuilt here so that this driver +// can reach the framework's host-side kargs builder. Only the kargs builder is +// used from it; the launches below go through grouped_gemm_tileloop, which +// assembles its own kernel per configuration. +// --------------------------------------------------------------------------- +using WaspShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + +using WaspUniversalTraits = ck_tile::PersistentTileGemmUniversalTraits< + WaspConfig::kPadM, + WaspConfig::kPadN, + WaspConfig::kPadK, + WaspConfig::DoubleSmemBuffer, + Row, + Col, + Row, + WaspConfig::TransposeC, + false, + false, + true, + WaspConfig::SupportsFastRowMajorB, + WaspConfig::SupportsFastColumnMajorA>; + +using WaspProblem = ck_tile::UniversalGemmPipelineProblem; + +using WaspPipeline = + typename PipelineTypeTraits::template GemmPipeline; + +using WaspEpilogue = GroupedGemmWaspStoreEpilogue; + +using WaspPartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner< + WaspShape, + WaspConfig::TileParitionerGroupNum, + WaspConfig::TileParitionerM01>; + +using WaspKernel = ck_tile::GroupedGemmKernel; + +// Independent CK reference for the same 256x256x64 NT tile. +using RefConfig = GemmConfigComputeV3TritonOperand256ExactTraversal; + +// Only used to compare the two partitioners' grid derivation on the host. +using RefShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; + +// --------------------------------------------------------------------------- +// Host helpers +// --------------------------------------------------------------------------- + +struct Options +{ + // The option vocabulary mirrors the standalone reference harness + // (hipc_kernels/main.cpp) so that ONE shape script can drive both binaries + // and the two columns are directly comparable: + // --m / --n / --k / --grouped-length / --layout / --output-dtype / + // --bench / --persistent-kernel / --warmup-iters / --bench-iters + // The defaults match that harness too, so "run both with no arguments" means + // the same problem; the small correctness shapes are always passed + // explicitly by the script. + int m = 4096; // rows of the single group, when --grouped-length is absent + int n = 4096; + int k = 1024; + int groups = 1; // shorthand for N equal groups of --m rows (extension) + // One M per group, in group order. Empty => a single group of `m` rows. + std::vector group_lens; + int warmup = 5; // harness default + int iters = 20; // harness default + bool bench = false; // harness semantics: timing only, no verdict + bool ref = true; + bool allow_unaligned = false; + bool quiet = false; + bool hipblaslt_ref = false; + bool hipblaslt_perf = false; + bool persistent_kernel = true; // accepted for CLI parity; see main() + std::string layout = "NT"; + std::string out_dtype = "bf16"; + // bf16 operands: the products are exact in fp32, so the CPU oracle and the + // kernel differ only by summation order. In practice the observed noise + // floor of this comparison is ~one bf16 ULP (2^-8), so the default is two + // orders of magnitude above it -- tight enough to catch a wrong mapping, + // loose enough not to flag rounding. + float rel_tol = 1.0e-2F; + float abs_tol = 1.0e-2F; + // Input pattern. These exist to make the OUTPUT self-describing, which is + // what turns "26% of elements are wrong" into an attributable diagnosis: + // random : the real workload + // ones : every output must be exactly K. Any deviation is a value/K + // problem (a permutation of correctly computed values would still + // read K), so this separates "wrong value" from "wrong location". + // bindex : A[m,k] = (k == 0), B[n,k] = n. Then C[m,n] = n exactly, i.e. + // every output element PRINTS the B row it was computed from, so + // a wrong or shifted B stripe is visible directly in the values. + std::string pattern = "random"; + bool debug_ctas = false; + // When >= 0, print row `dump_row` of group 0 across n in steps of 8. With + // --pattern=bindex the printed numbers ARE the B row each output column was + // computed from, so a shifted/stale consumer stripe shows up as a step at + // the exact column where the mapping breaks. + int dump_row = -1; + // Ablation mask for the timing decomposition (see gemm_pipeline_ag_bg_cr_wasp.hpp): + // 1 = no B loads, 2 = no MMACs, 4 = every A read from row 0, 8 = no barriers. + // Needs CK_TILE_GROUPED_GEMM_WASP_ABLATION. Results become garbage by design. + int ablate = 0; +}; + +// Round-to-nearest-even fp32 -> bf16, returned as the exact float value the +// bf16 can hold. Doing it here (instead of relying on a conversion) means the +// CPU reference and the device buffer are guaranteed to see identical inputs. +float round_to_bf16(float x) +{ + std::uint32_t u = 0; + std::memcpy(&u, &x, sizeof(u)); + const std::uint32_t lsb = (u >> 16) & 1u; + if((u & 0x7FFFFFFFu) > 0x7F800000u) + { + return x; // NaN / inf: leave alone, the test data never contains one + } + u += 0x7FFFu + lsb; + u &= 0xFFFF0000u; + float r = 0.0F; + std::memcpy(&r, &u, sizeof(r)); + return r; +} + +// Deterministic, portable LCG. Values in [-0.5, 0.5). +float next_value(std::uint32_t& state) +{ + state = state * 1664525u + 1013904223u; + return (static_cast(state >> 8) / static_cast(1u << 24)) - 0.5F; +} + +template +class DevBuf +{ + public: + DevBuf() = default; + DevBuf(const DevBuf&) = delete; + DevBuf& operator=(const DevBuf&) = delete; + ~DevBuf() + { + if(ptr_ != nullptr) + { + (void)hipFree(ptr_); + } + } + + bool Alloc(std::size_t count) + { + return hipMalloc(reinterpret_cast(&ptr_), count * sizeof(T)) == hipSuccess; + } + T* Data() { return ptr_; } + const T* Data() const { return ptr_; } + bool Upload(const std::vector& src, std::size_t count) + { + return hipMemcpy(ptr_, src.data(), count * sizeof(T), hipMemcpyHostToDevice) == hipSuccess; + } + bool Download(std::vector& dst, std::size_t count) const + { + return hipMemcpy(dst.data(), ptr_, count * sizeof(T), hipMemcpyDeviceToHost) == hipSuccess; + } + + private: + T* ptr_ = nullptr; +}; + +template +bool UploadRaw(T* dst, const void* src, std::size_t bytes) +{ + return hipMemcpy(dst, src, bytes, hipMemcpyHostToDevice) == hipSuccess; +} + +int parse_int(const char* text, int fallback) +{ + char* end = nullptr; + const long v = std::strtol(text, &end, 10); + if(end == text || v < 0 || v > 100000000L) + { + return fallback; + } + return static_cast(v); +} + +float parse_float(const char* text, float fallback) +{ + char* end = nullptr; + const float v = std::strtof(text, &end); + if(end == text || !(v >= 0.0F)) + { + return fallback; + } + return v; +} + +// --------------------------------------------------------------------------- +// The problem shape, resolved from Options. Per-group M is a first-class thing +// here (not an equal split) because every real NT shape in the reference +// harness is ragged: 8152,8152,8257,8257 / 1923,2047,2179,2043 / 5120,1536,... . +// `prefix` carries the row offsets that A and C use, so group g's rows live at +// [prefix[g], prefix[g+1]) of the concatenated operand -- the same convention +// the reference harness and the device-args ABI use. +// --------------------------------------------------------------------------- +struct ShapeInfo +{ + std::vector lens; + std::vector prefix; // size == lens.size() + 1 + long sum_m = 0; + int n = 0; + int k = 0; + + int groups() const { return static_cast(lens.size()); } + + // Group owning a GLOBAL output row. Ragged groups mean this cannot be a + // division by a per-group M. + int group_of(long row) const + { + for(int g = static_cast(lens.size()) - 1; g >= 0; --g) + { + if(row >= prefix[static_cast(g)]) + { + return g; + } + } + return 0; + } + int row_in_group(long row) const + { + const int g = group_of(row); + return static_cast(row - prefix[static_cast(g)]); + } +}; + +ShapeInfo make_shape(const Options& opt) +{ + ShapeInfo sh; + if(!opt.group_lens.empty()) + { + sh.lens = opt.group_lens; + } + else + { + sh.lens.assign(static_cast(opt.groups > 0 ? opt.groups : 1), opt.m); + } + sh.prefix.resize(sh.lens.size() + 1); + long off = 0; + for(std::size_t g = 0; g < sh.lens.size(); ++g) + { + sh.prefix[g] = off; + off += sh.lens[g]; + } + sh.prefix[sh.lens.size()] = off; + sh.sum_m = off; + sh.n = opt.n; + sh.k = opt.k; + return sh; +} + +// "512,512" / "8152,8152,8257,8257". Returns false on a malformed list. +bool parse_group_lens(const char* text, std::vector& out) +{ + out.clear(); + const std::string s(text); + std::size_t begin = 0; + for(;;) + { + const std::size_t comma = s.find(',', begin); + const std::string tok = + s.substr(begin, comma == std::string::npos ? std::string::npos : comma - begin); + if(tok.empty()) + { + return false; + } + char* end = nullptr; + const long v = std::strtol(tok.c_str(), &end, 10); + if(end == tok.c_str() || *end != '\0' || v <= 0 || v > 100000000L) + { + return false; + } + out.push_back(static_cast(v)); + if(comma == std::string::npos) + { + break; + } + begin = comma + 1; + } + return !out.empty(); +} + +void print_usage() +{ + std::printf( + "usage: grouped_gemm_wasp_test [--m=N] [--n=N] [--k=N] [--grouped-length=a,b,c]\n" + " [--layout=NT] [--output-dtype=bf16]\n" + " [--bench=0|1] [--persistent-kernel=1]\n" + " [--warmup-iters=N] [--bench-iters=N]\n" + " [--ref=0|1] [--rel-tol=F] [--abs-tol=F]\n" + " [--hipblaslt-ref=0|1] [--hipblaslt-perf=0|1]\n" + " [--allow-unaligned] [--quiet]\n" + " extras not present in the reference harness:\n" + " [--groups=N] equal groups of --m rows\n" + " [--warmup=N] [--iters=N] aliases\n" + " [--pattern=random|ones|bindex|aindex]\n" + " [--dump-row=N] [--debug-ctas]\n"); +} + +// --------------------------------------------------------------------------- +// One problem instance +// --------------------------------------------------------------------------- + +struct CompareStats +{ + double max_abs = 0.0; + double max_rel = 0.0; + std::size_t bad = 0; +}; + +// One failing output element, with its grouped coordinates. The PATTERN of these +// is what localises a broken kernel: a wrong consumer stripe shows up as a whole +// n range, a wrong MMAC fragment map as a stride-16 row pattern, a missing k tail +// as a uniform small error. +struct Mismatch +{ + int g = 0; + int m = 0; + int n = 0; + float ref = 0.0F; + float got = 0.0F; +}; + +struct CaseResult +{ + bool skipped = false; + bool passed = false; + double ms = 0.0; + double tflops = 0.0; + bool cpu_ran = false; + double wasp_max_abs = 0.0; + double wasp_max_rel = 0.0; + std::size_t wasp_bad = 0; + bool ref_ran = false; + double ref_max_abs = 0.0; + double ref_max_rel = 0.0; + std::size_t ref_bad = 0; + // WASP vs the CK reference configuration, run on the same kargs. Once a shape + // outgrows the CPU oracle's budget this is the ONLY oracle available, so it + // is what gives a large shape a verdict at all. + bool wasp_vs_ref_ran = false; + double wasp_vs_ref_max_abs = 0.0; + std::size_t wasp_vs_ref_bad = 0; + bool wasp_vs_hipblaslt_ran = false; + double wasp_vs_hipblaslt_max_abs = 0.0; + std::size_t wasp_vs_hipblaslt_bad = 0; + double hipblaslt_ms = 0.0; + std::vector wasp_examples; + std::vector ref_examples; + // Per-32-column-stripe and per-(row%16) failure counts. These two + // histograms localise almost every class of bug in this kernel: a bad + // consumer stripe shows up in the first, a bad MMAC fragment mapping or a + // bad LDS row stride in the second. + std::size_t bad_per_stripe[8] = {}; + std::size_t bad_per_row16[16] = {}; + // Per 16-row block, i.e. per m-iteration of a consumer's 256-row stripe: + // isolates a bad MMAC fragment map from a bad store. + std::size_t bad_per_miter[16] = {}; + // Outputs still equal to the sentinel, i.e. NEVER WRITTEN by either kernel. + // Distinguishes "the epilogue stored the wrong value" from "the epilogue + // never reached this output". + std::size_t wasp_not_written = 0; + // First element of the second half of the output, used only with + // --debug-ctas: the epilogue writes a 777 probe there. + float probe_cell = 0.0F; + std::vector dump_ref; + std::vector dump_got; +}; + +CompareStats compare(const std::vector& ref, + const std::vector& got, + const ShapeInfo& sh, + float rel_tol, + float abs_tol, + std::vector* examples, + std::size_t max_examples) +{ + CompareStats s{}; + for(std::size_t i = 0; i < ref.size(); ++i) + { + const double r = ref[i]; + const double g = got[i]; + const double diff = std::fabs(r - g); + if(diff > s.max_abs) + { + s.max_abs = diff; + } + const double rel = diff / (std::fabs(r) + 1e-6); + if(rel > s.max_rel) + { + s.max_rel = rel; + } + if(diff > static_cast(rel_tol) * std::fabs(r) + static_cast(abs_tol)) + { + ++s.bad; + if(examples != nullptr && examples->size() < max_examples) + { + const long row = static_cast(i) / sh.n; + const int col = static_cast(i) % sh.n; + examples->push_back(Mismatch{sh.group_of(row), + sh.row_in_group(row), + col, + ref[i], + got[i]}); + } + } + } + return s; +} + +#if defined(CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) +struct TunedSolution +{ + long long m = -1; + long long n = -1; + long long k = -1; + char ta = '?'; + char tb = '?'; + int index = -1; +}; + +const std::vector& tuned_solutions() +{ + static const std::vector table = [] { + std::vector entries; + const char* disable = std::getenv("BF16GEMM_NO_TUNING"); + if(disable != nullptr && std::atoi(disable) != 0) + { + return entries; + } + + const char* path = std::getenv("BF16GEMM_HIPBLASLT_SOLUTION_FILE"); + if(path == nullptr || *path == '\0') + { + path = "/root/workspace/bf16_gemm/tune/hipblaslt_solutions.conf"; + } + FILE* file = std::fopen(path, "r"); + if(file == nullptr) + { + std::fprintf(stderr, + "[hipBLASLt] WARNING: cannot open tune file '%s'; using heuristic\n", + path); + return entries; + } + + char line[512]; + while(std::fgets(line, sizeof(line), file) != nullptr) + { + char* p = line; + while(*p == ' ' || *p == '\t') + { + ++p; + } + if(*p == '#' || *p == '\n' || *p == '\r' || *p == '\0') + { + continue; + } + + std::string fields[6]; + int field_count = 0; + for(char* token = std::strtok(p, ", \t\r\n"); + token != nullptr && field_count < 6; + token = std::strtok(nullptr, ", \t\r\n")) + { + fields[field_count++] = token; + } + if(field_count < 6 || fields[3].empty() || fields[4].empty()) + { + std::fprintf(stderr, "[hipBLASLt] WARNING: malformed tune entry ignored\n"); + continue; + } + + auto parse_dimension = [](const std::string& value, long long& output) { + if(value == "*") + { + output = -1; + return true; + } + char* end = nullptr; + output = std::strtoll(value.c_str(), &end, 10); + return end != value.c_str() && *end == '\0'; + }; + TunedSolution entry; + if(!parse_dimension(fields[0], entry.m) || + !parse_dimension(fields[1], entry.n) || + !parse_dimension(fields[2], entry.k)) + { + std::fprintf(stderr, "[hipBLASLt] WARNING: malformed tune dimensions ignored\n"); + continue; + } + entry.ta = static_cast( + std::toupper(static_cast(fields[3][0]))); + entry.tb = static_cast( + std::toupper(static_cast(fields[4][0]))); + char* end = nullptr; + entry.index = static_cast(std::strtol(fields[5].c_str(), &end, 10)); + if(end == fields[5].c_str() || *end != '\0' || entry.index < 0) + { + std::fprintf(stderr, "[hipBLASLt] WARNING: malformed tune index ignored\n"); + continue; + } + entries.push_back(entry); + } + std::fclose(file); + std::fprintf(stderr, + "[hipBLASLt] loaded %zu tuned solutions from %s\n", + entries.size(), + path); + return entries; + }(); + return table; +} + +bool tuned_solution_matches(const TunedSolution& entry, + long long m, + long long n, + long long k, + char ta, + char tb) +{ + return (entry.m < 0 || entry.m == m) && (entry.n < 0 || entry.n == n) && + (entry.k < 0 || entry.k == k) && + (entry.ta == '*' || entry.ta == '?' || entry.ta == ta) && + (entry.tb == '*' || entry.tb == '?' || entry.tb == tb); +} + +bool pick_tuned_solution(long long m, + long long n, + long long k, + char ta, + char tb, + int& index) +{ + int best_score = 6; + bool found = false; + for(const TunedSolution& entry : tuned_solutions()) + { + if(!tuned_solution_matches(entry, m, n, k, ta, tb)) + { + continue; + } + const int score = (entry.m < 0 ? 1 : 0) + (entry.n < 0 ? 1 : 0) + + (entry.k < 0 ? 1 : 0) + + ((entry.ta == '*' || entry.ta == '?') ? 1 : 0) + + ((entry.tb == '*' || entry.tb == '?') ? 1 : 0); + if(score < best_score) + { + best_score = score; + index = entry.index; + found = true; + if(score == 0) + { + break; + } + } + } + return found; +} +#endif + +CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStream_t stream) +{ + CaseResult res{}; + + const int groups = sh.groups(); + const int n = sh.n; + const int k = sh.k; + const long sum_m = sh.sum_m; + const double macs = static_cast(sum_m) * n * k; + + // STAGE 4 (M tail) is in place: a ragged per-group M is masked -- the + // producer zero-fills the A rows past M and the epilogue drops the output + // rows past M. N and K tails are NOT masked yet, so N must still be a whole + // number of N-tiles and K a whole number of K-tiles. The real NT reference + // shapes (N=4096/2048, K=4096/2048) all satisfy that; only their M is + // ragged (8152, 8257, 1923, ...), which is now legal. + bool aligned = + (n % WaspConfig::N_Tile == 0) && (k % WaspConfig::K_Tile == 0); + if(!aligned && !opt.allow_unaligned) + { + res.skipped = true; + return res; + } + + // ---- host data ---- + std::vector a_host(static_cast(sum_m) * k); + std::vector b_host(static_cast(groups) * n * k); + std::vector c_ref(static_cast(sum_m) * n, 0.0F); + // CK-reference output, hoisted to function scope so the diagnostics below + // can use it as the oracle once a shape outgrows the CPU reference. + std::vector ref_out; + + if(opt.pattern == "ones") + { + for(auto& v : a_host) + { + v = 1.0F; + } + for(auto& v : b_host) + { + v = 1.0F; + } + } + else if(opt.pattern == "bindex") + { + for(std::size_t i = 0; i < a_host.size(); ++i) + { + a_host[i] = (i % static_cast(k) == 0) ? 1.0F : 0.0F; + } + for(int g = 0; g < groups; ++g) + { + for(int j = 0; j < n; ++j) + { + for(int kk = 0; kk < k; ++kk) + { + b_host[(static_cast(g) * n + j) * k + kk] = + static_cast(j); + } + } + } + } + else if(opt.pattern == "aindex") + { + // Mirror of bindex for the A side: A[row][k] = row (bf16-exact up to + // 255) and B[col][k] = (k == 0), so C[row][col] = A[row][0] = row. + // Every output element therefore PRINTS the A row it was actually + // computed from -- a wrong row shows up as a wrong (but valid) number, + // an out-of-bounds read as garbage. + // The printed value is the row WITHIN the group, which bf16 represents + // exactly only up to 256 -- larger groups round, and the readout is then + // approximate. Fine for the debugger it is. + for(int g = 0; g < groups; ++g) + { + const long row0 = sh.prefix[static_cast(g)]; + const int m_g = sh.lens[static_cast(g)]; + for(int i = 0; i < m_g; ++i) + { + for(int kk = 0; kk < k; ++kk) + { + a_host[(static_cast(row0) + static_cast(i)) * k + + kk] = static_cast(i); + } + } + } + for(auto& v : b_host) + { + v = 0.0F; + } + for(int g = 0; g < groups; ++g) + { + for(int j = 0; j < n; ++j) + { + b_host[(static_cast(g) * n + j) * k] = 1.0F; + } + } + } + else + { + std::uint32_t state = 12345u; + for(auto& v : a_host) + { + v = round_to_bf16(next_value(state)); + } + for(auto& v : b_host) + { + v = round_to_bf16(next_value(state)); + } + } + + // ---- device buffers ---- + DevBuf a_dev; + DevBuf b_dev; + DevBuf c_dev; + if(!a_dev.Alloc(a_host.size()) || !b_dev.Alloc(b_host.size()) || !c_dev.Alloc(c_ref.size())) + { + std::printf(" ERROR: device allocation failed\n"); + return res; + } + + // The host arrays hold exact bf16 values, so this conversion is exact. + std::vector a_bf16(a_host.size()); + std::vector b_bf16(b_host.size()); + for(std::size_t i = 0; i < a_host.size(); ++i) + { + a_bf16[i] = ck_tile::type_convert(a_host[i]); + } + for(std::size_t i = 0; i < b_host.size(); ++i) + { + b_bf16[i] = ck_tile::type_convert(b_host[i]); + } + if(!a_dev.Upload(a_bf16, a_bf16.size()) || !b_dev.Upload(b_bf16, b_bf16.size())) + { + std::printf(" ERROR: device upload failed\n"); + return res; + } + + // Poison the output so that "never written" is distinguishable from "written + // with a wrong value". 1024 is exactly representable in bf16, and no GEMM in + // this test can produce it from unit-magnitude inputs. + constexpr float kSentinel = 1024.0F; + { + std::vector poison(c_ref.size(), ck_tile::type_convert(kSentinel)); + if(!c_dev.Upload(poison, poison.size())) + { + std::printf(" ERROR: output poison upload failed\n"); + return res; + } + } + + // ---- kargs via the framework builder ---- + std::vector hdescs; + hdescs.reserve(static_cast(groups)); + for(int g = 0; g < groups; ++g) + { + const long row0 = sh.prefix[static_cast(g)]; + const void* a_g = a_dev.Data() + static_cast(row0) * k; + const void* b_g = b_dev.Data() + static_cast(g) * n * k; + void* c_g = c_dev.Data() + static_cast(row0) * n; + hdescs.emplace_back(a_g, + b_g, + c_g, + 1, // k_batch + sh.lens[static_cast(g)], // M of this group + n, + k, + k, // stride_A: A is [sumM, K] row major + k, // stride_B: B is [N, K] with K contiguous (NT) + n); // stride_E + } + const auto kargs = WaspKernel::MakeKargs(hdescs); + if(kargs.size() != static_cast(groups)) + { + std::printf(" ERROR: MakeKargs produced %zu kargs for %d groups\n", + kargs.size(), + groups); + return res; + } + if(!WaspKernel::IsSupportedArgument(kargs)) + { + std::printf(" ERROR: Wasp requires gfx936 BF16 NT with aligned N/K\n"); + return res; + } + const ck_tile::index_t flat_grid_size = WaspKernel::GridSize(hdescs).x; + + DevBuf kargs_dev; + if(!kargs_dev.Alloc(kargs.size()) || + !UploadRaw(kargs_dev.Data(), kargs.data(), kargs.size() * sizeof(kargs[0]))) + { + std::printf(" ERROR: kargs upload failed\n"); + return res; + } + + // ---- warmup + timed launches ---- + const ck_tile::stream_config sc{stream}; + for(int i = 0; i < opt.warmup; ++i) + { + (void)grouped_gemm_tileloop(sc, + groups, + kargs_dev.Data(), + {}, + false, + 0, + flat_grid_size); + } + (void)hipDeviceSynchronize(); + + hipEvent_t ev0 = nullptr; + hipEvent_t ev1 = nullptr; + if(hipEventCreate(&ev0) != hipSuccess || hipEventCreate(&ev1) != hipSuccess) + { + std::printf(" ERROR: hipEventCreate failed\n"); + return res; + } + (void)hipEventRecord(ev0, stream); + for(int i = 0; i < opt.iters; ++i) + { + (void)grouped_gemm_tileloop(sc, + groups, + kargs_dev.Data(), + {}, + false, + 0, + flat_grid_size); + } + (void)hipEventRecord(ev1, stream); + const hipError_t sync_err = hipEventSynchronize(ev1); + float elapsed_ms = 0.0F; + (void)hipEventElapsedTime(&elapsed_ms, ev0, ev1); + { + // Every HIP call in this driver was previously unchecked, which hides a + // device-side fault behind a plausible-looking timing number. + const hipError_t last = hipGetLastError(); + if(sync_err != hipSuccess) + { + std::printf(" ERROR: hipEventSynchronize: %s\n", hipGetErrorString(sync_err)); + } + if(last != hipSuccess) + { + std::printf(" ERROR: hipGetLastError after launches: %s\n", + hipGetErrorString(last)); + } + } + (void)hipEventDestroy(ev0); + (void)hipEventDestroy(ev1); + + res.ms = static_cast(elapsed_ms) / opt.iters; + res.tflops = res.ms > 0.0 ? (2.0 * macs / (res.ms * 1.0e-3)) / 1.0e12 : 0.0; + + // ---- read back and compare ---- + std::vector c_host(c_ref.size()); + if(!c_dev.Download(c_host, c_host.size())) + { + std::printf(" ERROR: device download failed\n"); + return res; + } + std::vector wasp_out(c_ref.size()); + for(std::size_t i = 0; i < c_ref.size(); ++i) + { + wasp_out[i] = ck_tile::type_convert(c_host[i]); + } + + // ---- CPU reference ---- + const bool run_cpu = macs <= 2.0e9; + if(run_cpu) + { + for(int g = 0; g < groups; ++g) + { + const long row0 = sh.prefix[static_cast(g)]; + const int m_g = sh.lens[static_cast(g)]; + for(int i = 0; i < m_g; ++i) + { + const float* a_row = &a_host[static_cast(row0 + i) * k]; + const float* b_mat = &b_host[static_cast(g) * n * k]; + for(int j = 0; j < n; ++j) + { + const float* b_row = b_mat + static_cast(j) * k; + double acc = 0.0; + for(int kk = 0; kk < k; ++kk) + { + acc += static_cast(a_row[kk]) * static_cast(b_row[kk]); + } + c_ref[static_cast(row0 + i) * n + j] = static_cast(acc); + } + } + } + const CompareStats ws = compare(c_ref, + wasp_out, + sh, + opt.rel_tol, + opt.abs_tol, + &res.wasp_examples, + 8); + res.cpu_ran = true; + res.wasp_max_abs = ws.max_abs; + res.wasp_max_rel = ws.max_rel; + res.wasp_bad = ws.bad; + } + + // ---- optional CK reference on the same kargs ---- + if(opt.ref) + { + DevBuf c_ref_dev; + if(c_ref_dev.Alloc(c_ref.size())) + { + std::vector ref_descs; + ref_descs.reserve(static_cast(groups)); + for(int g = 0; g < groups; ++g) + { + const long row0 = sh.prefix[static_cast(g)]; + const void* a_g = a_dev.Data() + static_cast(row0) * k; + const void* b_g = b_dev.Data() + static_cast(g) * n * k; + void* c_g = c_ref_dev.Data() + static_cast(row0) * n; + ref_descs.emplace_back( + a_g, b_g, c_g, 1, sh.lens[static_cast(g)], n, k, k, k, n); + } + const auto ref_kargs = WaspKernel::MakeKargs(ref_descs); + DevBuf ref_kargs_dev; + if(ref_kargs_dev.Alloc(ref_kargs.size()) && + UploadRaw(ref_kargs_dev.Data(), + ref_kargs.data(), + ref_kargs.size() * sizeof(ref_kargs[0]))) + { + (void)grouped_gemm_tileloop(sc, + groups, + ref_kargs_dev.Data(), + {}, + false, + static_cast(num_cu)); + std::vector ref_host(c_ref.size()); + if(c_ref_dev.Download(ref_host, ref_host.size())) + { + ref_out.assign(c_ref.size(), 0.0F); + for(std::size_t i = 0; i < ref_out.size(); ++i) + { + ref_out[i] = ck_tile::type_convert(ref_host[i]); + } + res.ref_ran = true; + if(run_cpu) + { + const CompareStats rs = compare(c_ref, + ref_out, + sh, + opt.rel_tol, + opt.abs_tol, + &res.ref_examples, + 4); + res.ref_max_abs = rs.max_abs; + res.ref_max_rel = rs.max_rel; + res.ref_bad = rs.bad; + } + { + const CompareStats wr = compare(wasp_out, + ref_out, + sh, + opt.rel_tol, + opt.abs_tol, + nullptr, + 0); + res.wasp_vs_ref_ran = true; + res.wasp_vs_ref_max_abs = wr.max_abs; + res.wasp_vs_ref_bad = wr.bad; + } + } + } + } + } + + if(opt.hipblaslt_ref || opt.hipblaslt_perf) + { +#if defined(CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) + hipblasLtHandle_t handle = nullptr; + hipblasLtMatmulDesc_t desc = nullptr; + hipblasLtMatrixLayout_t b_layout = nullptr; + hipblasLtMatrixLayout_t a_layout = nullptr; + hipblasLtMatrixLayout_t d_layout = nullptr; + hipblasLtMatmulPreference_t preference = nullptr; + hipblasLtMatmulAlgo_t algo{}; + constexpr std::size_t workspace_size = 256ull << 20; + void* workspace = nullptr; + DevBuf blas_out; + + auto blas_check = [](hipblasStatus_t status, const char* operation) { + if(status != HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "hipBLASLt error %d at %s\n", + static_cast(status), + operation); + std::exit(1); + } + }; + auto hip_runtime_check = [](hipError_t status, const char* operation) { + if(status != hipSuccess) + { + std::fprintf(stderr, + "HIP error %d (%s) at %s\n", + static_cast(status), + hipGetErrorName(status), + operation); + std::exit(1); + } + }; + blas_check(hipblasLtCreate(&handle), "hipblasLtCreate"); + blas_check(hipblasLtMatmulDescCreate(&desc, HIPBLAS_COMPUTE_32F, HIP_R_32F), + "hipblasLtMatmulDescCreate"); + const hipblasOperation_t op_a = HIPBLAS_OP_T; + const hipblasOperation_t op_b = HIPBLAS_OP_N; + blas_check(hipblasLtMatmulDescSetAttribute( + desc, HIPBLASLT_MATMUL_DESC_TRANSA, &op_a, sizeof(op_a)), + "hipblasLtMatmulDescSetAttribute(TRANSA)"); + blas_check(hipblasLtMatmulDescSetAttribute( + desc, HIPBLASLT_MATMUL_DESC_TRANSB, &op_b, sizeof(op_b)), + "hipblasLtMatmulDescSetAttribute(TRANSB)"); + // NT input storage is B[N,K] and A[M,K], both row-major. Viewed as + // column-major matrices, the product is B^T * A^T = C^T. + blas_check(hipblasLtMatrixLayoutCreate( + &b_layout, HIP_R_16BF, k, n, k), + "hipblasLtMatrixLayoutCreate(B)"); + blas_check(hipblasLtMatrixLayoutCreate( + &a_layout, HIP_R_16BF, k, sh.lens[0], k), + "hipblasLtMatrixLayoutCreate(A)"); + blas_check(hipblasLtMatrixLayoutCreate( + &d_layout, HIP_R_16BF, n, sh.lens[0], n), + "hipblasLtMatrixLayoutCreate(D)"); + blas_check(hipblasLtMatmulPreferenceCreate(&preference), + "hipblasLtMatmulPreferenceCreate"); + blas_check(hipblasLtMatmulPreferenceSetAttribute( + preference, + HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, + sizeof(workspace_size)), + "hipblasLtMatmulPreferenceSetAttribute"); + bool selected_tuned = false; + int tuned_index = -1; + if(pick_tuned_solution(n, sh.lens[0], k, 'T', 'N', tuned_index)) + { + std::vector resolved(1); + std::vector indices{tuned_index}; + if(hipblaslt_ext::getAlgosFromIndex(handle, indices, resolved) == + HIPBLAS_STATUS_SUCCESS && + resolved[0].state == HIPBLAS_STATUS_SUCCESS) + { + const float probe_alpha = 1.0F; + const float probe_beta = 0.0F; + std::size_t required_workspace = 0; + const hipblasStatus_t supported = + hipblaslt_ext::matmulIsAlgoSupported(handle, + desc, + &probe_alpha, + b_layout, + a_layout, + &probe_beta, + d_layout, + d_layout, + resolved[0].algo, + required_workspace); + if(supported == HIPBLAS_STATUS_SUCCESS && + required_workspace <= workspace_size) + { + algo = resolved[0].algo; + selected_tuned = true; + std::fprintf(stderr, + "[hipBLASLt] tuned NT shape m=%d n=%d k=%d " + "transA=T transB=N -> solution_index=%d " + "(workspace=%zu B)\n", + n, + sh.lens[0], + k, + tuned_index, + required_workspace); + } + else + { + std::fprintf(stderr, + "[hipBLASLt] WARNING: tuned solution %d is unsupported " + "or needs %zu bytes workspace (available %zu); " + "using heuristic\n", + tuned_index, + required_workspace, + workspace_size); + } + } + else + { + std::fprintf(stderr, + "[hipBLASLt] WARNING: tuned solution %d is unavailable; " + "using heuristic\n", + tuned_index); + } + } + if(!selected_tuned) + { + hipblasLtMatmulHeuristicResult_t heuristic{}; + int algo_count = 0; + blas_check(hipblasLtMatmulAlgoGetHeuristic(handle, + desc, + b_layout, + a_layout, + d_layout, + d_layout, + preference, + 1, + &heuristic, + &algo_count), + "hipblasLtMatmulAlgoGetHeuristic"); + if(algo_count == 0) + { + std::fprintf(stderr, "hipBLASLt did not return an algorithm for NT shape\n"); + std::exit(1); + } + algo = heuristic.algo; + std::fprintf(stderr, "[hipBLASLt] using heuristic algorithm\n"); + } + blas_check(hipblasLtMatmulPreferenceDestroy(preference), + "hipblasLtMatmulPreferenceDestroy"); + if(hipMalloc(&workspace, workspace_size) != hipSuccess || + !blas_out.Alloc(c_ref.size())) + { + std::fprintf(stderr, "hipBLASLt workspace/output allocation failed\n"); + std::exit(1); + } + + const float alpha = 1.0F; + const float beta = 0.0F; + auto run_blas = [&](void* output) { + blas_check(hipblasLtMatmul(handle, + desc, + &alpha, + b_dev.Data(), + b_layout, + a_dev.Data(), + a_layout, + &beta, + output, + d_layout, + output, + d_layout, + &algo, + workspace, + workspace_size, + stream), + "hipblasLtMatmul"); + }; + + if(opt.hipblaslt_ref) + { + run_blas(blas_out.Data()); + if(hipStreamSynchronize(stream) != hipSuccess) + { + std::fprintf(stderr, "hipBLASLt reference synchronization failed\n"); + std::exit(1); + } + std::vector blas_host(c_ref.size()); + std::vector blas_values(c_ref.size()); + if(!blas_out.Download(blas_host, blas_host.size())) + { + std::fprintf(stderr, "hipBLASLt output download failed\n"); + std::exit(1); + } + for(std::size_t i = 0; i < blas_values.size(); ++i) + { + blas_values[i] = ck_tile::type_convert(blas_host[i]); + } + const CompareStats compared = compare(wasp_out, + blas_values, + sh, + opt.rel_tol, + opt.abs_tol, + nullptr, + 0); + res.wasp_vs_hipblaslt_ran = true; + res.wasp_vs_hipblaslt_max_abs = compared.max_abs; + res.wasp_vs_hipblaslt_bad = compared.bad; + } + if(opt.hipblaslt_perf) + { + for(int i = 0; i < opt.warmup; ++i) + { + run_blas(blas_out.Data()); + } + hipEvent_t start = nullptr; + hipEvent_t stop = nullptr; + if(hipEventCreate(&start) != hipSuccess || hipEventCreate(&stop) != hipSuccess) + { + std::fprintf(stderr, "hipBLASLt timing event creation failed\n"); + std::exit(1); + } + hip_runtime_check(hipEventRecord(start, stream), "hipEventRecord(start)"); + for(int i = 0; i < opt.iters; ++i) + { + run_blas(blas_out.Data()); + } + hip_runtime_check(hipEventRecord(stop, stream), "hipEventRecord(stop)"); + hip_runtime_check(hipEventSynchronize(stop), "hipEventSynchronize(stop)"); + float elapsed = 0.0F; + hip_runtime_check(hipEventElapsedTime(&elapsed, start, stop), + "hipEventElapsedTime"); + res.hipblaslt_ms = static_cast(elapsed) / opt.iters; + hip_runtime_check(hipEventDestroy(start), "hipEventDestroy(start)"); + hip_runtime_check(hipEventDestroy(stop), "hipEventDestroy(stop)"); + } + + hip_runtime_check(hipFree(workspace), "hipFree(hipBLASLt workspace)"); + hipblasLtMatrixLayoutDestroy(d_layout); + hipblasLtMatrixLayoutDestroy(a_layout); + hipblasLtMatrixLayoutDestroy(b_layout); + hipblasLtMatmulDescDestroy(desc); + hipblasLtDestroy(handle); +#else + std::fprintf(stderr, + "hipBLASLt requested but this target was built without hipBLASLt\n"); + std::exit(2); +#endif + } + + // ----------------------------------------------------------------------- + // Diagnostics against whichever oracle actually ran. The CPU reference is + // preferred (it is independent of every GPU path), but a real benchmark shape + // costs far more MACs than it can afford -- so the CK reference is the + // fallback, otherwise the histograms below would be silently absent on + // exactly the shapes that matter. + // ----------------------------------------------------------------------- + const std::vector* oracle = nullptr; + if(res.cpu_ran) + { + oracle = &c_ref; + } + else if(res.wasp_vs_ref_ran) + { + oracle = &ref_out; + } + if(oracle != nullptr) + { + const std::size_t len = oracle->size(); + for(std::size_t i = 0; i < len; ++i) + { + const double r = (*oracle)[i]; + const double g = wasp_out[i]; + if(std::fabs(g - static_cast(kSentinel)) < 1.0) + { + ++res.wasp_not_written; + } + if(std::fabs(r - g) > static_cast(opt.rel_tol) * std::fabs(r) + + static_cast(opt.abs_tol)) + { + const int col = static_cast(i % static_cast(n)); + const long grow = static_cast(i / static_cast(n)); + const int row = sh.row_in_group(grow); // ragged: not a modulo + const int stripe = col / 32; + if(stripe >= 0 && stripe < 8) + { + ++res.bad_per_stripe[stripe]; + } + ++res.bad_per_row16[row % 16]; + // Bucket per m-tile within the group, so the array stays in range + // for any M. + ++res.bad_per_miter[(row % 256) / 16]; + // On a real shape there can be millions of mismatches; keeping + // eight coordinates is enough to see the pattern. + if(res.wasp_examples.size() < 8) + { + res.wasp_examples.push_back( + Mismatch{sh.group_of(grow), row, col, (*oracle)[i], wasp_out[i]}); + } + } + } + // First element of the second half of the whole output. With a ragged + // group list this straddles groups, so it is a coarse "did the second + // half get written at all" probe, nothing finer. + const std::size_t half_idx = + static_cast(sum_m) * static_cast(n) / 2; + if(len > half_idx) + { + res.probe_cell = wasp_out[half_idx]; + } + if(opt.dump_row >= 0 && static_cast(opt.dump_row) < sum_m) + { + const std::size_t base = static_cast(opt.dump_row) * n; + for(int j = 0; j < n; j += 8) + { + res.dump_ref.push_back((*oracle)[base + j]); + res.dump_got.push_back(wasp_out[base + j]); + } + } + } + + // Verdict: prefer the CPU oracle, fall back to the CK reference. If neither + // ran there is nothing to assert, and claiming a pass would be a lie -- main + // reports that state explicitly instead. + if(res.cpu_ran) + { + res.passed = (res.wasp_bad == 0); + } + else if(res.wasp_vs_ref_ran) + { + res.passed = (res.wasp_vs_ref_bad == 0); + } + else + { + res.passed = true; + } + if(res.wasp_vs_hipblaslt_ran) + { + res.passed = res.passed && res.wasp_vs_hipblaslt_bad == 0; + } + return res; +} + +} // namespace + +int main(int argc, char* argv[]) +{ + Options opt{}; + for(int i = 1; i < argc; ++i) + { + const std::string arg = argv[i]; + if(arg == "--help") + { + print_usage(); + return 0; + } + else if(arg == "--allow-unaligned") + { + opt.allow_unaligned = true; + } + else if(arg == "--quiet") + { + opt.quiet = true; + } + else if(arg == "--debug-ctas") + { + opt.debug_ctas = true; + } + else if(arg.rfind("--m=", 0) == 0) + { + opt.m = parse_int(arg.c_str() + 4, opt.m); + } + else if(arg.rfind("--n=", 0) == 0) + { + opt.n = parse_int(arg.c_str() + 4, opt.n); + } + else if(arg.rfind("--k=", 0) == 0) + { + opt.k = parse_int(arg.c_str() + 4, opt.k); + } + else if(arg.rfind("--groups=", 0) == 0) + { + opt.groups = parse_int(arg.c_str() + 9, opt.groups); + } + else if(arg.rfind("--warmup=", 0) == 0) + { + opt.warmup = parse_int(arg.c_str() + 9, opt.warmup); + } + else if(arg.rfind("--iters=", 0) == 0) + { + opt.iters = parse_int(arg.c_str() + 8, opt.iters); + } + else if(arg.rfind("--ref=", 0) == 0) + { + opt.ref = parse_int(arg.c_str() + 6, 1) != 0; + } + else if(arg.rfind("--hipblaslt-ref=", 0) == 0) + { + opt.hipblaslt_ref = parse_int(arg.c_str() + 16, 0) != 0; + } + else if(arg.rfind("--hipblaslt-perf=", 0) == 0) + { + opt.hipblaslt_perf = parse_int(arg.c_str() + 17, 0) != 0; + } + else if(arg.rfind("--rel-tol=", 0) == 0) + { + opt.rel_tol = parse_float(arg.c_str() + 10, opt.rel_tol); + } + else if(arg.rfind("--abs-tol=", 0) == 0) + { + opt.abs_tol = parse_float(arg.c_str() + 10, opt.abs_tol); + } + else if(arg.rfind("--pattern=", 0) == 0) + { + opt.pattern = arg.substr(10); + } + else if(arg.rfind("--dump-row=", 0) == 0) + { + opt.dump_row = parse_int(arg.c_str() + 11, opt.dump_row); + } + else if(arg.rfind("--ablate=", 0) == 0) + { + opt.ablate = parse_int(arg.c_str() + 9, opt.ablate); + } + // ---- reference-harness vocabulary ---- + else if(arg.rfind("--grouped-length=", 0) == 0) + { + if(!parse_group_lens(arg.c_str() + 17, opt.group_lens)) + { + std::printf("bad --grouped-length: %s\n", arg.c_str()); + return 2; + } + } + else if(arg.rfind("--layout=", 0) == 0) + { + opt.layout = arg.substr(9); + } + else if(arg.rfind("--output-dtype=", 0) == 0) + { + opt.out_dtype = arg.substr(15); + } + else if(arg.rfind("--bench=", 0) == 0) + { + opt.bench = parse_int(arg.c_str() + 8, 1) != 0; + } + else if(arg == "--bench") + { + opt.bench = true; + } + else if(arg.rfind("--persistent-kernel=", 0) == 0) + { + opt.persistent_kernel = parse_int(arg.c_str() + 20, 1) != 0; + } + else if(arg.rfind("--warmup-iters=", 0) == 0) + { + opt.warmup = parse_int(arg.c_str() + 15, opt.warmup); + } + else if(arg.rfind("--bench-iters=", 0) == 0) + { + opt.iters = parse_int(arg.c_str() + 14, opt.iters); + } + else + { + std::printf("unknown argument: %s\n", arg.c_str()); + print_usage(); + return 2; + } + } + if(opt.iters < 1) + { + opt.iters = 1; + } + + // ---- Options this port cannot honour yet are refused LOUDLY. Accepting them + // and quietly doing something else is how a benchmark stops meaning + // anything. ---- + if(opt.layout != "NT") + { + std::printf("wasp test: layout=%s is not ported (NT only so far)\n", + opt.layout.c_str()); + return 2; + } + if(opt.out_dtype != "bf16") + { + std::printf("wasp test: output-dtype=%s is not ported (bf16 only so far)\n", + opt.out_dtype.c_str()); + return 2; + } + if((opt.hipblaslt_ref || opt.hipblaslt_perf) && + (opt.groups != 1 || !opt.group_lens.empty())) + { + std::printf("wasp test: hipBLASLt comparison/timing currently requires one group\n"); + return 2; + } + if(!opt.persistent_kernel) + { + // The CK grouped kernel walks a grid-stride tile list inside the device + // kernel; there is no per-group launch path to select. + std::printf("wasp test: --persistent-kernel=0 is not representable here (the CK " + "grouped kernel is always persistent); ignoring\n"); + } + const ShapeInfo sh = make_shape(opt); + + hipDeviceProp_t prop{}; + int dev = 0; + if(hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&prop, dev) != hipSuccess) + { + std::printf("wasp test: no HIP device -> SKIP\n"); + return 0; + } + const int num_cu = prop.multiProcessorCount; + hipStream_t stream = nullptr; + + std::printf("wasp test: device CU=%d, tile=%dx%dx%d, block=%d threads, LDS=%d B\n", + num_cu, + static_cast(WaspConfig::M_Tile), + static_cast(WaspConfig::N_Tile), + static_cast(WaspConfig::K_Tile), + static_cast(WaspPipeline::BlockSize), + static_cast(WaspPipeline::kWaspLdsBytes)); + // The whole design assumes 12 waves of 64 lanes. If the framework ever + // launches this kernel with wave32, BlockSize() halves and the role split + // silently collapses, so print what the framework actually resolved. + std::printf("wasp test: is_wave32=%d launched_block_x=%d (pipeline BlockSize=%d)\n", + static_cast(ck_tile::is_wave32()), + static_cast(WaspKernel::BlockSize().x), + static_cast(WaspPipeline::BlockSize)); + // The persistent tile loop derives its total work from the PARTITIONER's + // GridSize, not from the host. Print it, and the reference config's, so a + // "tile never processed" symptom can be attributed to the partitioner or + // ruled out immediately. + { + using RefPartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner< + RefShape, + RefConfig::TileParitionerGroupNum, + RefConfig::TileParitionerM01>; + // Per group: the grid is derived from each group's own M, so this is + // printed for group 0 only as a sanity check. + std::printf("wasp test: partitioner GridSize(m=%d n=%d) group0: wasp=%d " + "(MPerBlock=%d NPerBlock=%d M01=%d GroupNum=%d) ref=%d (M01=%d)\n", + sh.lens[0], + sh.n, + static_cast(WaspPartitioner::GridSize(sh.lens[0], sh.n)), + static_cast(WaspPartitioner::MPerBlock), + static_cast(WaspPartitioner::NPerBlock), + static_cast(WaspConfig::TileParitionerM01), + static_cast(WaspConfig::TileParitionerGroupNum), + static_cast(RefPartitioner::GridSize(sh.lens[0], sh.n)), + static_cast(RefConfig::TileParitionerM01)); + } + std::printf("wasp test: layout=%s groups=%d sum(M)=%ld M=[", + opt.layout.c_str(), + sh.groups(), + sh.sum_m); + for(int g = 0; g < sh.groups(); ++g) + { + std::printf("%s%d", g ? "," : "", sh.lens[static_cast(g)]); + } + std::printf("] N=%d K=%d persistent=1 (warmup=%d iters=%d ck_ref=%s)\n", + sh.n, + sh.k, + opt.warmup, + opt.iters, + opt.ref ? "on" : "off"); + +#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) + if(opt.ablate != 0) + { + const int mask = opt.ablate; + if(hipMemcpyToSymbol(ck_tile::wasp_detail::g_wasp_ablate, + &mask, + sizeof(mask), + 0, + hipMemcpyHostToDevice) != hipSuccess) + { + std::printf("wasp test: ERROR: could not set the ablation mask\n"); + return 2; + } + std::printf("wasp test: ABLATION mask=%d (1=no B loads, 2=no MMACs, " + "4=A read row0, 8=no barriers) -- results are garbage by design\n", + mask); + } +#endif + + const CaseResult r = run_case(opt, sh, num_cu, stream); + + if(r.skipped) + { + std::printf("wasp test: SKIPPED (shape outside the ported alignment contract)\n"); + return 0; + } + +#if defined(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + { + std::vector cycles(1024 * 16 * 9); + if(hipMemcpyFromSymbol(cycles.data(), + g_wasp_epi_cycles, + cycles.size() * sizeof(uint64_t), + 0, + hipMemcpyDeviceToHost) != hipSuccess) + { + std::fprintf(stderr, "failed to read Wasp epilogue cycle counters\n"); + return 2; + } + constexpr const char* stage_names[] = { + "bf16-convert", + "half0-lds-issue", + "half0-lds-drain", + "half0-lds-read+global-issue", + "inter-half-lds-drain", + "half1-lds-issue", + "half1-lds-drain", + "half1-lds-read+global-issue", + "global-store-drain"}; + for(std::size_t stage = 0; stage < 9; ++stage) + { + uint64_t min_cycles = UINT64_MAX; + uint64_t max_cycles = 0; + uint64_t sum_cycles = 0; + std::size_t sample_count = 0; + for(std::size_t i = stage; i < cycles.size(); i += 9) + { + const uint64_t value = cycles[i]; + if(value != 0) + { + min_cycles = value < min_cycles ? value : min_cycles; + max_cycles = value > max_cycles ? value : max_cycles; + sum_cycles += value; + ++sample_count; + } + } + if(sample_count != 0) + { + std::printf("wasp epilogue %s cycles: min=%llu avg=%.1f max=%llu samples=%zu\n", + stage_names[stage], + static_cast(min_cycles), + static_cast(sum_cycles) / static_cast(sample_count), + static_cast(max_cycles), + sample_count); + } + } + } +#endif + +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + { + std::vector samples(1024 * 16); + if(hipMemcpyFromSymbol(samples.data(), + ck_tile::wasp_detail::g_wasp_vmem_samples, + samples.size() * sizeof(samples[0]), + 0, + hipMemcpyDeviceToHost) != hipSuccess) + { + std::fprintf(stderr, "failed to read Wasp VMEM wait samples\n"); + return 2; + } + std::printf("wasp VMEM wait samples (stage %u; cycles, aggregated across CTAs):\n", + static_cast(ck_tile::wasp_detail::kVmemSampleStage)); + for(int wave = 0; wave < 12; ++wave) + { + uint64_t issue_min = UINT64_MAX; + uint64_t issue_max = 0; + uint64_t issue_sum = 0; + uint64_t wait_min = UINT64_MAX; + uint64_t wait_max = 0; + uint64_t wait_sum = 0; + uint64_t vmcnt4_min = UINT64_MAX; + uint64_t vmcnt4_max = 0; + uint64_t vmcnt4_sum = 0; + uint64_t vmcnt2_min = UINT64_MAX; + uint64_t vmcnt2_max = 0; + uint64_t vmcnt2_sum = 0; + std::size_t count = 0; + for(int cta = 0; cta < 1024; ++cta) + { + const auto& sample = samples[cta * 16 + wave]; + if(sample.valid == 0) + { + continue; + } + const uint64_t issue_cycles = sample.issue_to_wait_start; + const uint64_t wait_cycles = sample.wait_cycles; + issue_min = std::min(issue_min, issue_cycles); + issue_max = std::max(issue_max, issue_cycles); + issue_sum += issue_cycles; + wait_min = std::min(wait_min, wait_cycles); + wait_max = std::max(wait_max, wait_cycles); + wait_sum += wait_cycles; + vmcnt4_min = std::min(vmcnt4_min, sample.vmcnt4_cycles); + vmcnt4_max = std::max(vmcnt4_max, sample.vmcnt4_cycles); + vmcnt4_sum += sample.vmcnt4_cycles; + vmcnt2_min = std::min(vmcnt2_min, sample.vmcnt2_cycles); + vmcnt2_max = std::max(vmcnt2_max, sample.vmcnt2_cycles); + vmcnt2_sum += sample.vmcnt2_cycles; + ++count; + } + if(count != 0) + { + if(wave < 4) + { + std::printf( + " producer wave %d: issue-to-wait-start min/avg/max=%llu/%.1f/%llu, " + "vmcnt<=4 elapsed min/avg/max=%llu/%.1f/%llu, " + "vmcnt<=2 elapsed min/avg/max=%llu/%.1f/%llu, " + "vmcnt(0)-wait min/avg/max=%llu/%.1f/%llu, CTAs=%zu\n", + wave, + static_cast(issue_min), + static_cast(issue_sum) / static_cast(count), + static_cast(issue_max), + static_cast(vmcnt4_min), + static_cast(vmcnt4_sum) / static_cast(count), + static_cast(vmcnt4_max), + static_cast(vmcnt2_min), + static_cast(vmcnt2_sum) / static_cast(count), + static_cast(vmcnt2_max), + static_cast(wait_min), + static_cast(wait_sum) / static_cast(count), + static_cast(wait_max), + count); + } + else + { + std::printf(" consumer wave %d: issue-to-wait-start min/avg/max=%llu/%.1f/%llu, " + "vmcnt-wait min/avg/max=%llu/%.1f/%llu, CTAs=%zu\n", + wave - 4, + static_cast(issue_min), + static_cast(issue_sum) / static_cast(count), + static_cast(issue_max), + static_cast(wait_min), + static_cast(wait_sum) / static_cast(count), + static_cast(wait_max), + count); + } + } + } + + const int n_tiles = + (sh.n + static_cast(WaspPartitioner::NPerBlock) - 1) / + static_cast(WaspPartitioner::NPerBlock); + const int m_tiles = + (sh.lens[0] + static_cast(WaspPartitioner::MPerBlock) - 1) / + static_cast(WaspPartitioner::MPerBlock); +#if defined(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + const int tile_count = m_tiles * n_tiles; +#endif + std::printf("producer A wait by M tile (N-tile-fastest tile order%s):\n", +#if defined(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + ", reversed" +#else + "" +#endif + ); + for(int m_tile = 0; m_tile < m_tiles; ++m_tile) + { + std::vector waits; + uint64_t vmcnt4_sum = 0; + uint64_t vmcnt2_sum = 0; + uint64_t wait_sum = 0; + for(int n_tile = 0; n_tile < n_tiles; ++n_tile) + { + const int tile_id = m_tile * n_tiles + n_tile; +#if defined(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + const int cta = tile_count - 1 - tile_id; +#else + const int cta = tile_id; +#endif + for(int wave = 0; wave < 4; ++wave) + { + const auto& sample = samples[cta * 16 + wave]; + if(sample.valid != 0) + { + waits.push_back(sample.wait_cycles); + wait_sum += sample.wait_cycles; + vmcnt4_sum += sample.vmcnt4_cycles; + vmcnt2_sum += sample.vmcnt2_cycles; + } + } + } + if(waits.empty()) + { + continue; + } + std::sort(waits.begin(), waits.end()); + const std::size_t p95_index = (waits.size() - 1) * 95 / 100; + const int row_start = m_tile * static_cast(WaspPartitioner::MPerBlock); + const int row_end = std::min( + sh.lens[0], row_start + static_cast(WaspPartitioner::MPerBlock)); + const double sample_count = static_cast(waits.size()); + std::printf(" M tile %d rows [%d,%d): vmcnt(0) avg/p95/max=%0.1f/%llu/%llu, " + "vmcnt<=4/<=2 avg=%0.1f/%0.1f, samples=%zu\n", + m_tile, + row_start, + row_end, + static_cast(wait_sum) / sample_count, + static_cast(waits[p95_index]), + static_cast(waits.back()), + static_cast(vmcnt4_sum) / sample_count, + static_cast(vmcnt2_sum) / sample_count, + waits.size()); + } + } +#endif + + if(!opt.quiet) + { + std::printf("wasp test: %.3f ms %.2f TFLOPS\n", r.ms, r.tflops); + if(r.hipblaslt_ms > 0.0) + { + const double macs = static_cast(sh.sum_m) * sh.n * sh.k; + const double blas_tflops = 2.0 * macs / (r.hipblaslt_ms * 1.0e-3) / 1.0e12; + std::printf("hipblaslt: %.3f ms %.2f TFLOPS\n", r.hipblaslt_ms, blas_tflops); + } + } + if(opt.bench) + { + // Same semantics as the reference harness: --bench is timing only. + return 0; + } + if(r.cpu_ran) + { + std::printf("wasp test: wasp vs CPU max_abs=%.6g max_rel=%.3g bad=%zu\n", + r.wasp_max_abs, + r.wasp_max_rel, + r.wasp_bad); + } + else + { + std::printf("wasp test: wasp vs CPU (skipped: shape too large for the naive " + "oracle; the ck_ref comparison below is the oracle instead)\n"); + } + if(r.ref_ran) + { + if(r.cpu_ran) + { + std::printf("wasp test: ck_ref vs CPU max_abs=%.6g max_rel=%.3g bad=%zu\n", + r.ref_max_abs, + r.ref_max_rel, + r.ref_bad); + } + else + { + // Printing the zero-initialised fields here would read as a pass; the + // comparison simply never happened. + std::printf("wasp test: ck_ref vs CPU (not compared: the CPU oracle was " + "skipped)\n"); + } + } + if(r.wasp_vs_ref_ran) + { + std::printf("wasp test: wasp vs ck_ref max_abs=%.6g bad=%zu\n", + r.wasp_vs_ref_max_abs, + r.wasp_vs_ref_bad); + } + if(r.wasp_vs_hipblaslt_ran) + { + std::printf("wasp test: wasp vs hipBLASLt max_abs=%.6g bad=%zu\n", + r.wasp_vs_hipblaslt_max_abs, + r.wasp_vs_hipblaslt_bad); + } + // The histograms are oracle independent: print them whenever they were + // filled, not only when the CPU oracle happened to run -- otherwise they go + // missing on exactly the real shapes this is meant to debug. + if(!r.wasp_examples.empty()) + { + std::printf("wasp test: first failing outputs (g, m, n): ref -> got\n"); + for(const auto& e : r.wasp_examples) + { + std::printf("wasp test: (%d, %d, %d): %.6g -> %.6g\n", + e.g, + e.m, + e.n, + static_cast(e.ref), + static_cast(e.got)); + } + std::printf("wasp test: bad per 32-col stripe:"); + for(int s = 0; s < 8; ++s) + { + std::printf(" %zu", r.bad_per_stripe[s]); + } + std::printf("\nwasp test: bad per (row%%16):"); + for(int m = 0; m < 16; ++m) + { + std::printf(" %zu", r.bad_per_row16[m]); + } + std::printf("\nwasp test: bad per (row/16):"); + for(int m = 0; m < 16; ++m) + { + std::printf(" %zu", r.bad_per_miter[m]); + } + std::printf("\nwasp test: outputs never written by the kernel (sentinel left): %zu\n", + r.wasp_not_written); + } + if(!r.cpu_ran && !r.wasp_vs_ref_ran) + { + std::printf("wasp test: no oracle ran (--ref=0 and the shape is too large for the " + "CPU reference); nothing was verified\n"); + return 1; + } + +#if defined(CK_TILE_GROUPED_GEMM_WASP_DEBUG_RECORD) + if(opt.debug_ctas) + { + std::vector dbg(128 * 10); + std::vector waves(128 * 16); + const bool ok_a = + hipMemcpyFromSymbol(dbg.data(), g_wasp_dbg, dbg.size() * sizeof(int), 0, + hipMemcpyDeviceToHost) == hipSuccess; + const bool ok_b = + hipMemcpyFromSymbol(waves.data(), g_wasp_dbg_waves, waves.size() * sizeof(int), 0, + hipMemcpyDeviceToHost) == hipSuccess; + if(ok_a && ok_b) + { + int active = 0; + for(int c = 0; c < 128; ++c) + { + if(dbg[static_cast(c) * 10] > 0) + { + ++active; + } + } + std::printf("wasp test: CTAs that reached the epilogue: %d\n", active); + std::printf("wasp test: probe cell (first element of the second half) = %g " + "(777 = the store reached it, 1024 = still the sentinel)\n", + static_cast(r.probe_cell)); + for(int c = 0; c < 4; ++c) + { + const std::size_t b = static_cast(c) * 10; + if(dbg[b] == 0) + { + continue; + } + unsigned mask = 0; + for(int w = 0; w < 16; ++w) + { + if(waves[static_cast(c) * 16 + static_cast(w)] != 0) + { + mask |= (1U << static_cast(w)); + } + } + std::printf( + "wasp test: cta %d: waves=0x%03x block_org=(%d,%d) M=%d N=%d ldd=%d " + "calc_off=%d d_ptr-p_data_=%d elems\n", + c, + mask, + dbg[b + 1], + dbg[b + 8], + dbg[b + 2], + dbg[b + 3], + dbg[b + 4], + dbg[b + 5], + dbg[b + 7] - dbg[b + 6]); + } + } + } +#endif + if(opt.dump_row >= 0 && !r.dump_ref.empty()) + { + std::printf("wasp test: row %d dump, every 8th column\n", opt.dump_row); + for(std::size_t i = 0; i < r.dump_ref.size(); ++i) + { + std::printf("wasp test: n=%-4zu ref=%-14.6g got=%-14.6g\n", + i * 8, + static_cast(r.dump_ref[i]), + static_cast(r.dump_got[i])); + } + } + +#if defined(CK_TILE_GROUPED_GEMM_WASP_DUMP_A_STAGE) + // A staging verification. The stage is expected to hold the 256x64 A tile in + // the documented layout; only the VALUES are checked here (with --pattern=ones + // every entry must be bf16 1.0 == 0x3F80), which is exactly the property the + // consumer depends on. A wrong layout would leave far more than the ~5 % the + // runtime shows. + { + std::vector stage(2 * 256 * 64); + if(hipMemcpyFromSymbol(stage.data(), + ck_tile::wasp_detail::g_wasp_a_dump, + stage.size() * sizeof(unsigned short), + 0, + hipMemcpyDeviceToHost) == hipSuccess) + { + for(int s = 0; s < 2; ++s) + { + std::size_t good = 0; + std::size_t per_m[16] = {}; + for(int rr = 0; rr < 256; ++rr) + { + for(int cc = 0; cc < 64; ++cc) + { + if(stage[static_cast(s) * 256 * 64 + + static_cast(rr) * 64 + cc] == 0x3F80U) + { + ++good; + } + else + { + ++per_m[rr / 16]; + } + } + } + std::printf("wasp test: A stage snapshot slot %d: %zu / 16384 entries are bf16 1.0\n", + s, + good); + std::printf("wasp test: not-1.0 per 16-row m-iteration:"); + for(int m = 0; m < 16; ++m) + { + std::printf(" %zu", per_m[m]); + } + std::printf("\n"); + } + } + } +#endif + + if(!r.passed) + { + std::printf("wasp test: FAIL\n"); + return 1; + } + std::printf("wasp test: PASS\n"); + return 0; +} diff --git a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_impl.hpp b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_impl.hpp index 6efdfae..600f224 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_impl.hpp +++ b/example_hcu/ck_tile/19_grouped_gemm/instances/grouped_gemm_impl.hpp @@ -24,6 +24,7 @@ #include "ck_tile/ops/gemm.hpp" #include "ck_tile/host.hpp" #include "../grouped_gemm.hpp" +#include "../grouped_gemm_wasp_epilogue.hpp" namespace { @@ -514,6 +515,14 @@ int run_grouped_gemm_c_impl(const ck_tile_hcu_grouped_gemm_desc* descs, grouped_gemm_update_plain_desc_cache(descs, group_count); } + ck_tile::index_t flat_grid_size = 0; + for(const auto& arg : gemm_descs_vec) + { + flat_grid_size += + ck_tile::integer_divide_ceil(arg.M, GemmConfig::M_Tile) * + ck_tile::integer_divide_ceil(arg.N, GemmConfig::N_Tile) * arg.k_batch; + } + grouped_gemm_tileloop(gemm_descs_vec), - splitk); + splitk, + 0, + flat_grid_size); } else { @@ -762,12 +773,12 @@ float grouped_gemm(const std::vector::template GemmPipeline; #if defined(CK_TILE_GROUPED_GEMM_DSREADM_OUTPUT_TILE_DIRECT_STORE) - using GemmEpilogue = + using GroupedGemmRegularEpilogue = GroupedGemmOutputTileDirectStoreEpilogue; #elif defined(CK_TILE_GROUPED_GEMM_TRITON_OPERAND_NO_D) - using GemmEpilogue = + using GroupedGemmRegularEpilogue = GroupedGemmDirectStoreEpilogue>; #if defined(CK_TILE_GROUPED_GEMM_DIRECT_STORE_EPILOGUE) - using GemmEpilogue = std::conditional_t< + using GroupedGemmRegularEpilogue = std::conditional_t< std::tuple_size_v == 0 && memory_operation == ck_tile::memory_operation_enum::set, GroupedGemmDirectStoreEpilogue; #else - using GemmEpilogue = CShuffleGemmEpilogue; + using GroupedGemmRegularEpilogue = CShuffleGemmEpilogue; #endif #endif + using GemmEpilogue = std::conditional_t< + ck_tile::is_grouped_gemm_wasp_pipeline::value, + GroupedGemmWaspStoreEpilogue, + GroupedGemmRegularEpilogue>; using KernelBase = ck_tile::GroupedGemmKernel; #if defined(CK_TILE_GROUPED_GEMM_UNIQUE_KERNEL_TAG) @@ -915,7 +930,8 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, void* kargs_ptr, const std::vector>& output_resets, bool splitk, - std::uint32_t requested_num_cu) + std::uint32_t requested_num_cu, + ck_tile::index_t flat_grid_size) { using GemmShape = ck_tile::TileGemmShape< ck_tile::sequence, @@ -969,12 +985,12 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, using GemmPipeline = typename PipelineTypeTraits< GemmConfig::Pipeline>::template GemmPipeline; #if defined(CK_TILE_GROUPED_GEMM_DSREADM_OUTPUT_TILE_DIRECT_STORE) - using GemmEpilogue = + using GroupedGemmRegularEpilogue = GroupedGemmOutputTileDirectStoreEpilogue; #elif defined(CK_TILE_GROUPED_GEMM_TRITON_OPERAND_NO_D) - using GemmEpilogue = + using GroupedGemmRegularEpilogue = GroupedGemmDirectStoreEpilogue; #endif #else - using GemmEpilogue = ck_tile::CShuffleEpilogue< + using GroupedGemmRegularEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem, @@ -1021,6 +1037,10 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, ? 1 : GemmConfig::CShuffleWarpGemmNInterleave>>; #endif + using GemmEpilogue = std::conditional_t< + ck_tile::is_grouped_gemm_wasp_pipeline::value, + GroupedGemmWaspStoreEpilogue, + GroupedGemmRegularEpilogue>; using KernelBase = ck_tile::GroupedGemmKernel>; + + if constexpr(ck_tile::is_grouped_gemm_wasp_pipeline::value) + { + static_assert(memory_operation == ck_tile::memory_operation_enum::set, + "Wasp grouped GEMM only supports direct output stores"); + if(num_groups <= 0 || splitk || requested_num_cu != 0 || flat_grid_size <= 0) + { + throw std::runtime_error( + "Wasp grouped GEMM requires non-split groups and an exact flat grid"); + } + + constexpr dim3 blocks = Kernel::BlockSize(); + const dim3 grids{static_cast(flat_grid_size), 1, 1}; + if(s.log_level_ > 0) + { + std::cout << "Launching Wasp grouped GEMM with exact grid: {" << grids.x + << ", 1, 1}, blocks: {" << blocks.x << ", 1, 1}" << std::endl; + } + + using WaspKernel = ck_tile::GroupedGemmWaspKernel; + const auto kernel = ck_tile::make_kernel( + WaspKernel{}, + grids, + blocks, + 0, + ck_tile::cast_pointer_to_constant_address_space(kargs_ptr), + num_groups); + ave_time = ck_tile::launch_kernel(s, kernel); + return ave_time; + } + constexpr dim3 blocks = Kernel::BlockSize(); dim3 grids = Kernel::MaxOccupancyGridSize(s); if(requested_num_cu > 0) 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 95dfea7..4c09614 100644 --- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp @@ -53,6 +53,17 @@ struct is_grouped_gemm_v7_uniform_pipeline +struct is_grouped_gemm_wasp_pipeline : std::false_type +{ +}; + +template +struct is_grouped_gemm_wasp_pipeline> + : std::bool_constant +{ +}; + /// @brief The Grouped GEMM kernel host arguments. /// /// @par Overview @@ -391,6 +402,38 @@ struct GroupedGemmKernel CK_TILE_HOST static bool IsSupportedArgument(const std::vector& kargs) { + if constexpr(is_grouped_gemm_wasp_pipeline::value) + { + if(kargs.empty()) + { + return false; + } + if(get_hcu_target_enum() != hcu_target_enum::gfx936) + { + return false; + } + if constexpr(!std::is_same_v || + !std::is_same_v || + !std::is_same_v || + !std::is_same_v || + !std::is_same_v || + !std::is_same_v) + { + return false; + } + for(const auto& karg : kargs) + { + const auto& args = karg.group_karg; + if(args.k_batch != 1 || args.K % GemmPipeline::KPerBlock != 0 || + args.N % TilePartitioner::NPerBlock != 0 || + args.stride_As[0] < args.K || args.stride_Bs[0] < args.K || + args.stride_E < args.N) + { + return false; + } + } + } + if constexpr(is_grouped_gemm_mls_pipeline::value) { if(get_hcu_target_enum() != hcu_target_enum::gfx938) @@ -639,7 +682,19 @@ struct GroupedGemmKernel splitk_batch_offset.bs_k_split_offset[0]; CDataType* c_ptr = static_cast(kargs.e_ptr); - if constexpr(is_grouped_gemm_v7_uniform_pipeline::value) + if constexpr(is_grouped_gemm_wasp_pipeline::value) + { + __shared__ char smem_ptr_0[GetSmemSize()]; + RunGemmWasp(a_ptr, + b_ptr, + c_ptr, + smem_ptr_0, + kargs, + splitk_batch_offset, + i_m, + i_n); + } + else if constexpr(is_grouped_gemm_v7_uniform_pipeline::value) { // The v7 ring comes from DYNAMIC shared memory: the launcher // passes GetSmemSize() bytes, and the pipeline's inline asm @@ -835,6 +890,55 @@ struct GroupedGemmKernel EpiloguePipeline{}.store_tile(p, acc, p_smem); } + CK_TILE_DEVICE static void + RunGemmWasp(const ADataType* a_ptr, + const BDataType* b_ptr, + CDataType* d_ptr, + void* p_smem, + const KernelArgs& kargs, + const typename Base::SplitKBatchOffset& splitk_batch_offset, + const index_t block_idx_m, + const index_t block_idx_n) + { + static_assert(is_grouped_gemm_wasp_pipeline::value); + static_assert(std::is_same_v && + std::is_same_v && + std::is_same_v, + "The Wasp grouped path is specialized for NT GEMM"); + static_assert(GemmPipeline::kNumProducerWaves == 4 && + GemmPipeline::kNumConsumerWaves == 8, + "The Wasp grouped path requires 4 producer + 8 consumer waves"); + + const index_t num_loop = __builtin_amdgcn_readfirstlane( + TilePartitioner::GetLoopNum(splitk_batch_offset.splitted_k)); + const index_t row_base = block_idx_m; + const index_t col_base = block_idx_n; + + const auto* const a_tile = a_ptr + row_base * kargs.stride_As[0]; + const auto* const b_tile = b_ptr + col_base * kargs.stride_Bs[0]; + using AccTile = typename GemmPipeline::AccTile; + AccTile acc; + GemmPipeline{}.run_raw_into(a_tile, + kargs.stride_As[0], + kargs.M - row_base, + b_tile, + kargs.stride_Bs[0], + kargs.N, + col_base, + num_loop, + acc, + p_smem); + + EpiloguePipeline::store_raw(d_ptr, + kargs.stride_E, + kargs.M, + kargs.N, + row_base, + col_base, + acc, + p_smem); + } + CK_TILE_DEVICE static void RunGemmMls(const ADataType* a_ptr, const BDataType* b_ptr, @@ -1227,4 +1331,93 @@ struct GroupedGemmKernel } }; +template +struct GroupedGemmWaspKernel +{ + using TilePartitioner = typename GroupedKernel::TilePartitioner; + using GemmPipeline = typename GroupedKernel::GemmPipeline; + using EpiloguePipeline = typename GroupedKernel::EpiloguePipeline; + using KernelArgs = typename GroupedKernel::KernelArgs; + using TransKernelArg = typename GroupedKernel::TransKernelArg; + using ADataType = typename GroupedKernel::ADataType; + using BDataType = typename GroupedKernel::BDataType; + using CDataType = typename GroupedKernel::CDataType; + using OffsetTile1DPartitioner = typename GroupedKernel::OffsetTile1DPartitioner; + + static constexpr index_t kBlockSize = GroupedKernel::kBlockSize; + + CK_TILE_DEVICE void + operator()(const void CK_TILE_CONSTANT_ADDRESS_SPACE* gemm_desc_const, + index_t group_count) const + { + static_assert(is_grouped_gemm_wasp_pipeline::value); + + const auto* const descs = reinterpret_cast( + cast_pointer_to_generic_address_space(gemm_desc_const)); + const index_t block_id = get_block_1d_id(); + + // MakeKargs assigns each group a contiguous interval in the flat grid. + index_t first_group = 0; + index_t last_group = group_count; + while(first_group < last_group) + { + const index_t mid = first_group + (last_group - first_group) / 2; + if(block_id >= descs[mid].block_end) + { + first_group = mid + 1; + } + else + { + last_group = mid; + } + } + if(first_group >= group_count || block_id < descs[first_group].block_start) + { + return; + } + + const auto& desc = descs[first_group]; + const KernelArgs& kargs = desc.group_karg; + + index_t tile_id = block_id - desc.block_start; +#if defined(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + tile_id = TilePartitioner::GridSize(kargs.M, kargs.N) - 1 - tile_id; +#endif + const auto [i_m, i_n] = + OffsetTile1DPartitioner::GetOffsetedTileIndex(0, kargs.M, kargs.N, tile_id); + const index_t row_base = i_m * TilePartitioner::MPerBlock; + const index_t col_base = i_n * TilePartitioner::NPerBlock; + + const auto* const a_tile = + static_cast(kargs.as_ptr[0]) + + row_base * kargs.stride_As[0]; + const auto* const b_tile = + static_cast(kargs.bs_ptr[0]) + + col_base * kargs.stride_Bs[0]; + auto* const d_ptr = static_cast(kargs.e_ptr); + + __shared__ char smem[GroupedKernel::GetSmemSize()]; + using AccTile = typename GemmPipeline::AccTile; + AccTile acc; + GemmPipeline{}.run_raw_into(a_tile, + kargs.stride_As[0], + kargs.M - row_base, + b_tile, + kargs.stride_Bs[0], + kargs.N, + col_base, + TilePartitioner::GetLoopNum(kargs.K), + acc, + smem); + EpiloguePipeline::store_raw(d_ptr, + kargs.stride_E, + kargs.M, + kargs.N, + row_base, + col_base, + acc, + smem); + } +}; + } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp new file mode 100644 index 0000000..bb818b0 --- /dev/null +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp @@ -0,0 +1,1653 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Info Technologies Co., Ltd. +// +// Warp-specialized (producer / consumer) grouped-GEMM pipeline, ported from the +// standalone warp-specialized reference design into the ck_tile pipeline vocabulary. +// +// The structural fact that makes this expressible at all: +// +// UniversalGemmKernel::kBlockSize == GemmPipeline::BlockSize +// +// The number of waves in the workgroup is the PIPELINE's business, while the +// warp grid that partitions the output tile is the CONFIG's business. So a +// 12-wave workgroup (4 producers + 8 consumers) carries a 1 x 8 wave grid that +// describes the CONSUMERS ONLY: the producer waves are invisible to +// TileGemmShape / TilePartitioner / the epilogue geometry. +// +// block tile 256M x 256N x 64K +// wave grid M_Warp=1, N_Warp=8, warp tile 16x16x16 +// -> MIterPerWarp=16, NIterPerWarp=2 (32 MMAC, 128 acc VGPR) +// threads 12 waves x 64 = 768 +// producers waves 0..3 : stage A, 64 rows each, into a 2-deep ring +// consumers waves 4..11 : each owns a 256M x 32N output stripe +// LDS 64 KB, hand managed: 2 x 32 KB A ring, rounded out by the +// epilogue reusing the same buffer once the K loop drains. +// +// ----------------------------------------------------------------------------- +// STAGE 1 (this revision): structure only. +// +// The role split, the A ring ownership, the barrier cadence and the fragment / +// accumulator mapping are in place. Deliberately NOT in place yet, each a +// separate measured step: +// +// 1. XOR K swizzle on the A slab. The current layout is plain row-major +// (row stride 128 B), which is CORRECT but bank-conflicted: 128 B = 32 +// dwords = 0 mod 32 banks, so the 16 rows a wave reads in one ds_read +// phase all land on the same bank (measured max_bank = 16 vs 8 with +// the XOR). Expect this to be slow; it is the first thing to fix. +// 2. buffer_load_dwordx4 ... lds (SRD + m0) instead of a plain global load +// followed by ds_write. Same bytes, fewer VGPRs, no register round trip. +// 3. The C / C' 1-quarter-stage phase stagger between the two consumer half +// groups, and the peeled k-tail with its LDS zero-fill. +// 4. M / N tail masking: the A rows past M and the B columns past N are read +// unguarded right now. The reference kernel handles them with a bit-31 vaddr (hardware +// OOB zero-fill) and a register mask; the selector must therefore not +// route ragged shapes here until that lands. +// ----------------------------------------------------------------------------- +// +// ----------------------------------------------------------------------------- +// ISSUE-SLOT BUDGET (this revision): the k loops carried a lot of VALU. +// +// A VALU instruction and an MMAC issued by DIFFERENT waves on the same SIMD +// cannot share a cycle, so every address the producer and consumer loops +// recomputed was an issue slot taken from somebody's MMAC. Both loops are now +// VALU-free by construction: per-lane addresses are built once per TILE, and +// every stage-varying term rides either the uniform SGPR soffset (the producer's +// K column and ring slot, the B K column) or the ds_read's own immediate field +// (the A ring slot and the m-iteration stride). +// +// Measured off the ISA, per k-stage per wave: +// +// consumer : 31 VALU -> 0 (211 -> 175 instructions, 128 of them MMAC) +// producer : 42.5 VALU -> 0 (the loop body is 8 DMA loads + ~20 scalar ops) +// +// Wall clock on the ported shape: 245.6 -> 285.0 TFLOPS, VGPRs 254 -> 241, still +// 0 spills. What it took, beyond hoisting the arithmetic, was making the hoisted +// values NON-REMARKABLE to LLVM's rematerialisation heuristic -- it prefers +// rebuilding a cheap loop-invariant address inside the loop over holding it in a +// register. See AProdSrc (8 panel vaddrs) and AReadBases (2 K-half bases). +// ----------------------------------------------------------------------------- +// +// ----------------------------------------------------------------------------- +// THE 1/4-PHASE STAGGER (the standalone kernel's C / C' personas). +// +// WHAT IT IS. Both consumer bodies execute the same items per k-stage, but C' +// issues them rotated by ONE ITEM: C pairs each read with its MMAC inside one +// interval, C' lets one read cross each barrier and consumes it at the top of +// the next interval. In items (one item = one 8-unit m-chunk = 32 MMACs): +// +// C I1 = [R c0 M c0 R c1 M c1] C I2 = [R c2 M c2 R c3 M c3] +// C' I1 = [M c3(prev) R c0 M c0 R c1] C' I2 = [M c1 R c2 M c2 R c3] +// +// with a k-stage split into two intervals by BAR(k).mid, and c0/c1 the h0 chunks, +// c2/c3 the h1 chunks. C leads an interval with its ds_reads (LDS busy) while C' +// leads with its MMACs (matrix busy), and they swap at the tail. This is the +// same shape as the reference kernel's stage_body_C / stage_body_Cp; the +// granularity is one m-chunk instead of one K-half because the fragment has to +// fit (see below). +// +// MEASURED, this shape, two cadences: +// +// cadence body split TFLOPS spill/scratch +// (a) 1 work interval C - 286.4 0 / 0 +// (a) 1 work interval C' 0 299.4 5 / 24 B +// (a) 1 work interval C/C' 4+4 188.1 16 / 52 B +// (b) 2 intervals C - 247.7 16 / 52 B +// (b) 2 intervals C' 0 264.8 16 / 52 B +// (b) 2 intervals C/C' 4+4 172.0 16 / 52 B +// +// (b) is hipc's cadence and is what the rotation needs -- without the mid +// barrier there is no interval to rotate within. Read across the rows: the C' +// ORDER is worth +4..7% in both cadences, the mid-stage rendezvous costs 12%, +// and the rotation itself costs another 35% on top. The rotation is the only +// configuration that loses; every uniform body is better than any mix, and the +// loss is not codegen (the spill is 16 / 52 B in all three (b) rows). +// +// The mechanism is the rendezvous: a barrier releases when the SLOWEST wave +// arrives, and the two personas have different interval costs (C drains in +// place twice per interval, C' gets one cheap drain and one free one), so the +// workgroup runs at C's pace while C' idles -- and the LDS queue interleaving of +// two out-of-phase read streams makes both of them slower than either uniform +// stream. There is no phase offset to be had once a hard sync re-times every +// wave twice a stage. +// +// THE REGISTER SIDE of the same experiment: the reference kernel keeps a whole +// K-half fragment (64 VGPR) live across the barrier. Here that lands on the +// 256-VGPR cap with 260 registers spilled / 372 B/lane of scratch and measures +// 263.9 against 299.3 for the 8-unit chunk (spill 5 / 24 B). The reference +// kernel itself compiles at VGPR 256 / SGPR 104 / scratch 108 B/lane (read out +// of its own code object) -- it spills too, a quarter as much, because its body +// carries no framework state beyond the params struct. Its occupancy hint +// (amdgpu_waves_per_eu(1,1), mirrored here as CK_TILE_WAVES_PER_EU_HINT in +// host/kernel_launch.hpp) was tried and is NEUTRAL: the cap stays 256 and the +// fragment still spills. Chunking the fragment is what buys the difference. +// ----------------------------------------------------------------------------- +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp" + +namespace ck_tile { + +// --------------------------------------------------------------------------- +// Accumulator tile (16 +// m-iterations x 2 n-iterations x 4 fp32 = 128 VGPR per lane. +// --------------------------------------------------------------------------- +struct WaspAccTile +{ + static constexpr index_t kMIters = 16; // 256 M / 16 + static constexpr index_t kNIters = 2; // 32 N / 16 + + // One 4-dword quad per (m,n) slot, as a VECTOR: the backend keeps the quad + // as one value (the reference kernel's union_vec4_fp32 acc_o[]), instead of + // promoting it into four independent scalars. That granularity is what the + // register allocator works with around every MMAC. + using acc_vec_t = float __attribute__((ext_vector_type(4))); + acc_vec_t v[kMIters * kNIters]; +}; + +namespace wasp_detail { + +constexpr index_t kVmemSampleStage = 16; + +struct WaspVmemSample +{ + uint64_t issue_to_wait_start = 0; + uint64_t wait_cycles = 0; + uint64_t vmcnt4_cycles = 0; + uint64_t vmcnt2_cycles = 0; + uint32_t valid = 0; +}; + +struct WaspVmemDiag +{ + uint64_t issued_at = 0; + WaspVmemSample sample{}; + bool issued = false; + bool valid = false; +}; + +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) +__device__ static WaspVmemSample g_wasp_vmem_samples[1024 * 16]; +#endif + +CK_TILE_DEVICE uint64_t wasp_memtime() +{ + return static_cast(__builtin_amdgcn_s_memtime()); +} + +using frag_a_t = unsigned short __attribute__((ext_vector_type(4))); // 4 bf16 = 8 B +using frag_c_t = float __attribute__((ext_vector_type(4))); // 4 fp32 +using dword4_t = int32_t __attribute__((ext_vector_type(4))); // 16 B + +// 16 B staging unit: aliased as one dword4 or as two 8 B MMAC operand halves. +union alignas(16) stage16_t +{ + dword4_t d4; + frag_a_t frag[2]; +}; +static_assert(sizeof(stage16_t) == 16, "16 B staging unit"); + +// NOTE: a plain `ds_read_b128(byte_addr)` helper used to live here, taking the +// address as a "v" operand. It is deliberately gone: that form makes the +// compiler materialise a fresh per-lane address for every ds_read, i.e. 2 v_add +// per m-iteration, 32 VALU per k-stage per consumer wave -- and on this part a +// VALU and an MMAC from different waves on the same SIMD cannot issue in the +// same cycle. Use the pipeline's a_read_half (one base VGPR + immediates). + +#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) +// --------------------------------------------------------------------------- +// Ablation switch, set from the host with hipMemcpyToSymbol. Results are WRONG +// by construction -- the only thing being measured is the timing decomposition, +// because three separate hypotheses about where this kernel's time goes were +// each falsified by measurement and modelling is clearly not working. +// +// bit 0 : skip the B global loads +// bit 1 : skip the MMACs +// bit 2 : read every A fragment from row 0's address (same LDS traffic count, +// degenerate address pattern -- isolates "16 distinct addresses" from +// "the reads themselves") +// bit 3 : skip the workgroup barriers +// --------------------------------------------------------------------------- +__device__ static int g_wasp_ablate = 0; +#endif + +CK_TILE_DEVICE int wasp_ablate() +{ +#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) + return g_wasp_ablate; +#else + return 0; +#endif +} + +CK_TILE_DEVICE void ds_write_b128(uint32_t byte_addr, const dword4_t& v) +{ + asm volatile("ds_write_b128 %0, %1" ::"v"(byte_addr), "v"(v) : "memory"); +} + +CK_TILE_DEVICE frag_c_t mmac_bf16(const frag_a_t& a, + const frag_a_t& b, + const frag_c_t& c) +{ + return __builtin_hcu_mmac_f32_16x16x16_bf16(a, b, c); +} + +// Raw global 16 B read. STAGE 1 uses this instead of buffer_load ... lds; the +// bytes fetched are identical, only the path to LDS differs. +CK_TILE_DEVICE dword4_t global_load_b128(const void* p) +{ + return *reinterpret_cast(p); +} + +// --------------------------------------------------------------------------- +// STAGE 2: buffer_load_dwordx4 Direct-to-LDS. +// +// Ported from the reference producer (buffer_load_dwordx4_x8_to_lds_vaddr): +// each `buffer_load_dwordx4 vaddr, srsrc, soffset offen lds` streams one PANEL +// -- 64 lanes x 16 B = 1024 B -- straight from global into LDS at m0 + lane*16, +// with m0 advanced by 1024 between the 8 panels. This removes the VGPR round +// trip the STAGE-1 producer paid (8x global_load into 32 VGPR, then 8x +// ds_write): the memory system writes LDS directly, so the producer no longer +// holds the whole slab live in registers and the global->LDS hand-off is one +// instruction stream instead of two with a barrier-ordered clobber between. +// +// gfx936 / gfx938 only (the `lds` modifier and the CACHE_SWIZZLE descriptor +// bit). The STAGE-1 load+ds_write path is kept as the portable fallback. +// --------------------------------------------------------------------------- +using buffer_rsrc_t = int32_t __attribute__((ext_vector_type(4))); + +// V4 buffer resource descriptor. STRIDE (word1 bits [29:16]) is programmed for +// the CACHE_SWIZZLE rotation period only -- `offen` does NOT scale the index by +// it. CACHE_SWIZZLE (bit 30) exists on gfx936/938 and needs stride/64 to be a +// power of two, which K=4096 (pitch 8192 = 64 x 128) and K=2048 (4096 = 64 x +// 64) both satisfy; it is left off otherwise. +CK_TILE_DEVICE buffer_rsrc_t wasp_make_buffer_rsrc(const void* ptr, + uint32_t stride_bytes, + uint32_t size_bytes, + bool allow_swizzle = true) +{ + buffer_rsrc_t r; + const uint64_t addr = reinterpret_cast(ptr); + r[0] = static_cast(addr & 0xffffffffu); + r[1] = static_cast(addr >> 32); + r[2] = static_cast(size_bytes); + r[3] = 0x00020000; + + uint32_t val = (stride_bytes & 0x3fffu) << 16; +#if defined(__gfx936__) || defined(__gfx938__) + if(allow_swizzle && (stride_bytes % 64u) == 0u) + { + const uint32_t s64 = stride_bytes / 64u; + if((s64 & (s64 - 1u)) == 0u) + { + val |= 1u << 30; + } + } +#else + (void)allow_swizzle; +#endif + r[1] |= static_cast(val); + return r; +} + +// 8 panels, each a dwordx4 (16 B) per lane, Direct-to-LDS. `lds_m0` is the LDS +// BYTE address of the first panel (p_smem is at an LDS offset whose numeric +// pointer value IS that address); `soffset` is the stage's K column byte offset +// shared by all 8 loads; each `v_off[i]` is this lane's per-panel byte offset +// (row * pitch + swizzled K within the stage), with bit 31 set for an +// out-of-range row so the hardware zero-fills it. +CK_TILE_DEVICE void wasp_buffer_load_x8_to_lds(buffer_rsrc_t rsrc, + uint32_t lds_m0, + uint32_t soffset, + uint32_t v0, + uint32_t v1, + uint32_t v2, + uint32_t v3, + uint32_t v4, + uint32_t v5, + uint32_t v6, + uint32_t v7) +{ + asm volatile( + "s_mov_b32 m0, %1\n\t" + "buffer_load_dwordx4 %0, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %2, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %3, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %4, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %5, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %6, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %7, %9, %10 offen lds\n\t" + "s_add_u32 m0, m0, 0x400\n\t" + "buffer_load_dwordx4 %8, %9, %10 offen lds\n" + : + : "v"(v0), "s"(lds_m0), "v"(v1), "v"(v2), "v"(v3), "v"(v4), "v"(v5), "v"(v6), + "v"(v7), "s"(rsrc), "s"(soffset) + : "memory"); +} + +// One 16 B consumer-B read through the CACHE_SWIZZLE descriptor (gfx936/938). +// +// `pair` packs {index in the LOW dword, byte offset in the HIGH dword}; with +// `idxen offen` the hardware forms addr = base + index * STRIDE + offset + +// soffset. STRIDE (the descriptor's row pitch, programmed in +// wasp_make_buffer_rsrc) also sets the CACHE_SWIZZLE VL1 tag-bank rotation +// period, which is the whole point: the plain pointer dereference the STAGE-1/2 +// consumer used suffered a large VL1-Tag RAM stall on the B stream (every lane +// of a wave hitting the same tag bank). The index is the global B row (output +// column), the offset the in-row K byte (col_group*16 + k_half*64), and the +// stage's K column base rides the uniform SGPR `soffset`. +// FOUR buffer_load_dwordx4 in ONE asm block -- the reference kernel's B load +// form verbatim (buffer_load_dwordx4_x4_idxen_soff2 in bf16_gemm_nt_hipc.cpp). +// +// What matters about the shape: +// * two IDXEN index pairs only (one per 16-wide n-row), NOT one per load; +// * the K-half difference rides the instruction's IMMEDIATE field (+64 B = +// kPassK * 2 B), so it costs no register at all; +// * the stage's K column rides the uniform SGPR soffset; +// * the per-lane in-row byte offset rides the top half of each pair. +// The port used to issue four separate loads, each with its own 64-bit pair +// carrying the k-half inside the offset: same bytes, 4x the index registers and +// 4 asm blocks instead of 1. +CK_TILE_DEVICE void wasp_buffer_load_b128_x4_idxen_soff2(buffer_rsrc_t rsrc, + uint64_t p0, + uint64_t p1, + uint32_t s_col_base, + dword4_t& o0, + dword4_t& o1, + dword4_t& o2, + dword4_t& o3) +{ + asm volatile("buffer_load_dwordx4 %0, %4, %6, %7 idxen offen offset:0\n\t" + "buffer_load_dwordx4 %1, %4, %6, %7 idxen offen offset:64\n\t" + "buffer_load_dwordx4 %2, %5, %6, %7 idxen offen offset:0\n\t" + "buffer_load_dwordx4 %3, %5, %6, %7 idxen offen offset:64\n" + : "=&v"(o0), "=&v"(o1), "=&v"(o2), "=&v"(o3) + : "v"(p0), "v"(p1), "s"(rsrc), "s"(s_col_base) + : "memory"); +} + +CK_TILE_DEVICE dword4_t wasp_buffer_load_b128_idxen(buffer_rsrc_t rsrc, + uint64_t pair, + uint32_t soffset) +{ + dword4_t o; + asm volatile("buffer_load_dwordx4 %0, %1, %2, %3 idxen offen\n" + : "=v"(o) + : "v"(pair), "s"(rsrc), "s"(soffset) + : "memory"); + return o; +} + +// Workgroup barrier that actually publishes LDS. +// +// ck_tile::block_sync_lds() is NOT usable here: in its default configuration +// (CK_TILE_EXPERIMENTAL_BLOCK_SYNC_LDS_WITHOUT_SYNC_VMEM undefined) it expands +// to __syncthreads(), i.e. a bare s_barrier -- see core/arch/arch.hpp:121-134. +// s_barrier synchronizes EXECUTION, it does not drain posted LDS stores, so a +// producer's ds_write_b128 can still be in flight when a consumer's +// ds_read_b128 executes. The result is a race that is invisible to a single +// wave, intermittent across runs, and biased towards whichever lanes win the +// dispatch -- exactly the "same shape, different bad count" signature this port +// showed before the fix. LDS stores must be retired with s_waitcnt lgkmcnt(0) +// on the writer side BEFORE the barrier. +// Bare workgroup barrier -- the reference kernel's wait_all_warp_arrived() +// verbatim: sched_barrier(0), s_barrier, sched_barrier(0). +// +// NO LDS DRAIN. This is load bearing for the k-stage cadence: the C' persona +// issues its A(k,1) read at the very end of an interval and consumes it on the +// far side of the barrier. A draining barrier (s_waitcnt lgkmcnt(0) + s_barrier) +// forces that read to complete before ANY wave may pass, which is exactly the +// serialisation the 1/4-phase offset exists to remove -- and it also puts every +// wave's LDS latency on the critical path to the barrier. Each consumer waits +// for its own A data explicitly (the waits inside Q1/Q3), so nothing is lost by +// not draining here. +// +// The sched_barrier pair is the reference kernel's, not decoration: it stops the +// machine scheduler from hoisting a read across the barrier (which would break +// the deferral's ordering) or sinking an MMAC past it. +CK_TILE_DEVICE void wasp_barrier() +{ + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_barrier\n"); + __builtin_amdgcn_sched_barrier(0); +} + +CK_TILE_DEVICE void wasp_barrier_maybe() +{ + if((wasp_ablate() & 8) == 0) + { + wasp_barrier(); + } +} + +CK_TILE_DEVICE void wasp_sync_lds() +{ + // Only the LGKM counter is drained here, matching the reference kernel. + // + // An earlier revision also drained vmcnt(0) on the theory that the A stage + // (written with ds_write, an LDS-unit store) might be ordered by the global + // counter. The ablation disproves the cost model behind that: the barriers + // are ~20% of the runtime, and vmcnt(0) makes every wave -- including the 8 + // consumers, whose global loads are its own B operands and have nothing to + // do with publishing A -- stall on its own outstanding global traffic at + // each of the num_loop barriers. The producer's A writes are LDS stores and + // are covered by lgkmcnt. + asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier\n" ::: "memory"); +} + +// Barrier with the ablation bit honored (see wasp_ablate). Uniform branch, so +// it costs one scalar compare when the ablation is off. +CK_TILE_DEVICE void wasp_sync_lds_maybe() +{ + if((wasp_ablate() & 8) == 0) + { + wasp_sync_lds(); + } +} + +// Retire outstanding direct-to-LDS DMAs (vmcnt). The producer calls this before +// the barrier that publishes a stage, so the overlapping cadence can issue the +// next stage's DMA straight after a publish and let it run under the consumers' +// compute. All of a stage's 8 panels are in flight by the time this is reached, +// so it is one 8-wide MLP wait, not a per-panel serialization. +// s_setprio N (0..3): the CU's wave scheduler prefers a higher-priority wave +// when it picks which one to issue from next cycle; the hardware default is 1. +// The standalone kernel brackets every 64-MMAC chain with <3> ... <0>: a +// consumer held at priority 3 while its MMACs drain keeps the matrix pipe fed, +// and dropping back to 0 afterwards lets the producers' DMA and the consumers' +// own ds_reads through. The sched_barrier(0) pair pins the priority change in +// program order -- same form as the reference kernel's flash::set_prio. +template +CK_TILE_DEVICE void wasp_set_prio() +{ + static_assert(PRIO >= 0 && PRIO <= 3, "s_setprio operand must be 0..3"); + __builtin_amdgcn_sched_barrier(0); + asm volatile("s_setprio %0" ::"n"(PRIO)); + __builtin_amdgcn_sched_barrier(0); +} + +CK_TILE_DEVICE void wasp_wait_vmcnt0() +{ + asm volatile("s_waitcnt vmcnt(0)\n" ::: "memory"); +} + +CK_TILE_DEVICE void wasp_wait_vmcnt0(WaspVmemDiag& diag, index_t stage) +{ +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + if(stage == kVmemSampleStage && diag.issued) + { + const uint64_t wait_start = wasp_memtime(); + asm volatile("s_waitcnt vmcnt(0)\n" ::: "memory"); + const uint64_t wait_end = wasp_memtime(); + diag.sample.issue_to_wait_start = wait_start - diag.issued_at; + diag.sample.wait_cycles = wait_end - wait_start; + diag.sample.valid = 1; + diag.valid = true; + return; + } +#else + (void)diag; + (void)stage; +#endif + wasp_wait_vmcnt0(); +} + +CK_TILE_DEVICE void wasp_wait_vmcnt0_producer(WaspVmemDiag& diag, index_t stage) +{ +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + if(stage == kVmemSampleStage && diag.issued) + { + const uint64_t wait_start = wasp_memtime(); + asm volatile("s_waitcnt vmcnt(4)\n" ::: "memory"); + diag.sample.vmcnt4_cycles = wasp_memtime() - wait_start; + asm volatile("s_waitcnt vmcnt(2)\n" ::: "memory"); + diag.sample.vmcnt2_cycles = wasp_memtime() - wait_start; + asm volatile("s_waitcnt vmcnt(0)\n" ::: "memory"); + const uint64_t wait_end = wasp_memtime(); + diag.sample.issue_to_wait_start = wait_start - diag.issued_at; + diag.sample.wait_cycles = wait_end - wait_start; + diag.sample.valid = 1; + diag.valid = true; + return; + } +#else + (void)diag; + (void)stage; +#endif + wasp_wait_vmcnt0(); +} + +#if defined(CK_TILE_GROUPED_GEMM_WASP_DUMP_A_STAGE) +// --------------------------------------------------------------------------- +// DIAGNOSTIC: two snapshots of the raw A staging ring, taken by producer wave 0 +// right after the barrier that publishes each of the first two stages. +// +// Comparing them on the host against the expected [256 rows x 64 k] layout +// settles the one question runtime evidence cannot: is the A operand wrong +// because the STAGE was never written correctly, or because the CONSUMER reads +// a correctly written stage wrongly? +// +// Raw 16-bit patterns, decoded on the host, so nothing is lost in translation. +// The snapshot has to happen HERE and not later: the epilogue borrows the whole +// 64 KB for its output staging and would overwrite both ring slots. +// --------------------------------------------------------------------------- +__device__ static unsigned short g_wasp_a_dump[2 * 256 * 64]; + +// One wave cooperatively copies a whole 32 KB stage: 64 lanes x 256 bf16 each. +// +// Deliberately NOT gated by a "first caller wins" flag: the copy is cooperative, +// so a flag that stops 63 of the 64 lanes would leave their regions of the dump +// at their zero-initialised value and read as an empty stage. Every wave-0 visit +// simply overwrites the previous snapshot, which is harmless because a correctly +// staged tile looks the same every time. +CK_TILE_DEVICE void take_a_stage_snapshot(int slot, const char* lds) +{ + const index_t lane = ck_tile::get_lane_id(); + const unsigned short* src = + reinterpret_cast(lds + slot * 32768 + lane * 512); + for(index_t i = 0; i < 256; ++i) + { + g_wasp_a_dump[slot * 256 * 64 + lane * 256 + i] = src[i]; + } +} +#endif + +} // namespace wasp_detail + +template +struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 +{ + using Base = GemmPipelineAgBgCrCompV3; + + using BlockGemmShape = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + + // ------------------------------------------------------------------ + // Wave split. The consumer count must equal the configured warp grid, + // because those are the waves that own output. + // ------------------------------------------------------------------ + static constexpr index_t kNumProducerWaves = 4; + static constexpr index_t kNumConsumerWaves = 8; + // The 1/4-phase stagger: consumer waves 0..3 walk the four quarters in the C + // order, waves 4..7 in the C' order (the same four quarters rotated by one + // quarter), so the two halves lead an interval with different pipes -- the + // reference kernel's split verbatim (is_C_prime = consumer_id >= 4). + // + // MEASURED on this port, m=8152, with the k-stage structure exact: the C' + // ORDER run uniformly is worth +4..7%; the ROTATION itself is a net loss, + // because a barrier releases on the slowest wave and re-times the offset + // twice a stage. Both personas are kept because the rotation is the point. + static constexpr index_t kConsumerSplitWave = kNumConsumerWaves / 2; + + static constexpr index_t kNumWavesTotal = kNumProducerWaves + kNumConsumerWaves; + + // SHADOWS the base's `Problem::kBlockSize` (== NumWarps * warp size == 512 + // for this geometry). UniversalGemmKernel reads the pipeline member, so the + // launched workgroup is 768 threads while the tile grid stays 1 x 8. + static constexpr index_t BlockSize = kNumWavesTotal * get_warp_size(); + + static_assert(BlockGemmShape::NumWarps == kNumConsumerWaves, + "WASP: the tile warp grid must be exactly the consumer waves"); + static_assert(BlockSize == 768, "WASP workgroup is 12 waves of 64"); + + // ------------------------------------------------------------------ + // LDS budget, hand managed (does not follow the ComputeV3 policy). + // A stage: 4 producer slabs x 64 rows x 64 K x 2 B = 32 KB + // 2 stages = 64 KB + // ------------------------------------------------------------------ + static constexpr index_t kStageK = 64; + static constexpr index_t kPassK = 32; + static constexpr index_t kNumStages = 2; + static constexpr index_t kProdSlabRows = 64; + // Columns owned by one consumer wave. Every consumer reads/writes its OWN + // stripe, so this offset has to be applied on both the B operand and the + // store side -- forgetting it makes every consumer but the first compute the + // wrong columns. + static constexpr index_t kWaveN = BlockGemmShape::kN / kNumConsumerWaves; // 256/8 = 32 + static constexpr index_t kProdSlabBytes = kProdSlabRows * kStageK * sizeof(uint16_t); // 8192 + static constexpr index_t kAStageBytes = kNumProducerWaves * kProdSlabBytes; // 32768 + static constexpr index_t kWaspLdsBytes = kNumStages * kAStageBytes; // 65536 + + static_assert(kProdSlabBytes * kNumProducerWaves == kAStageBytes, "A stage budget"); + static_assert(BlockGemmShape::kK == kStageK, "block K must equal the A-stage depth"); + static_assert(BlockGemmShape::kM == kNumProducerWaves * kProdSlabRows, + "A rows are split evenly across the producer waves"); + static_assert(BlockGemmShape::kN == kNumConsumerWaves * 32, + "each consumer wave owns a 32 wide N stripe"); + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return kWaspLdsBytes; } + + using AccTile = WaspAccTile; + + // Marker for GroupedGemmKernel's compile-time dispatch (is_grouped_gemm_wasp + // _pipeline): routes this pipeline through RunGemmWasp, the raw-pointer path + // that feeds run_raw directly and keeps no tile-window state live across the + // k loop. Absent on every other pipeline, so they keep the generic path. + static constexpr bool IsWaspPipeline = true; + + // ------------------------------------------------------------------ + // Producer: waves 0..3, one 64-row slab each. + // + // STAGE 1 addressing (plain, see the header note): + // global row r of the tile, k element -> a_base + r*lda + k + // LDS slab p, local row r, 16 B chunk c -> slab + r*128 + c*16 + // lane l owns row_group = l>>3, chunk = l&7 (no XOR yet) + // ------------------------------------------------------------------ + // Producer A addressing: what is per-TILE and what is per-K-STAGE. + // + // The 8 panel vaddrs are `row * pitch + (chunk ^ row_group)*16` with + // bit 31 set for a row past m_valid. Every term in that expression is + // per-tile and per-lane only -- the k-stage does not appear anywhere in it + // (the stage's K column rides the uniform SGPR `soffset`, and the LDS ring + // slot rides `m0`). Rebuilding them inside the k loop therefore costs + // 8 v_add + 8 v_cmp + 8 v_cndmask + 11 v_or per stage per producer wave, + // measured off the ISA: 42.5 VALU per stage, against 8 useful DMA loads. + // + // That is the wrong trade on this part: an MMAC and a VALU instruction from + // different waves on the SAME SIMD cannot issue in the same cycle, so 4 + // producer waves x ~40 VALU ride on the issue slots the 8 consumer waves' + // MMACs want. (The producer's own register budget is tiny -- 9 VGPRs -- so + // hoisting the values costs nothing there either.) Same idea as the + // standalone kernel, which builds its vaddr array once before the loop and + // then only bumps m0 and soffset, both scalar. + // ------------------------------------------------------------------ + struct AProdSrc + { + wasp_detail::buffer_rsrc_t rsrc; // SGPR pair set, built once + uint32_t v_off[8]; // per-lane panel byte offsets + }; + + // a_base and lda are per-TILE, not per k-stage, so they are passed in rather + // than re-derived here. The first version called calculate_offset on every + // one of the num_loop calls -- 64 times per tile per wave -- and the ablation + // says this kernel's time is dominated by exactly that kind of scalar + // bookkeeping, not by the MMACs (removing every MMAC saves ~10%). + template + CK_TILE_DEVICE AProdSrc make_a_prod_src(const ADataTypePtr* a_base, + index_t lda, + index_t m_valid, + index_t wid) const + { + using namespace wasp_detail; + + const index_t lane = get_lane_id(); + + const index_t row_group = lane >> 3; // 0..7 + const index_t chunk = lane & 7; // 0..7, 16 B units along K +#if defined(CK_TILE_GROUPED_GEMM_WASP_REVERSE_PRODUCER_SLAB) + // Debug switch: hand the slabs out in reverse wave order. If the broken + // rows follow producer wave 3 the fault is in that wave's execution; if + // they stay at the same LDS addresses the fault is address dependent. + const index_t slab_id = kNumProducerWaves - 1 - wid; +#else + const index_t slab_id = wid; +#endif + const index_t row_base = slab_id * kProdSlabRows; + + // One buffer descriptor over this tile's A. The size field bounds valid + // rows to [0, m_valid): a row >= m_valid overruns num_records and reads + // as 0, which (with the bit-31 mask below) is the zero-fill the + // epilogue's row drop already assumes. + const uint32_t pitch_bytes = static_cast(lda) * sizeof(uint16_t); + const uint32_t size_bytes = + static_cast((m_valid > 0 ? m_valid : 0)) * pitch_bytes; + + AProdSrc s; + s.rsrc = wasp_make_buffer_rsrc(a_base, pitch_bytes, size_bytes); + + const uint32_t k_in_panel = + static_cast((chunk ^ row_group) * 16); // swizzled K byte offset +#pragma unroll + for(index_t i = 0; i < 8; ++i) + { + const index_t row = row_base + i * 8 + row_group; + const uint32_t v = static_cast(row) * pitch_bytes + k_in_panel; + s.v_off[i] = (row < m_valid) ? v : (v | 0x80000000u); + } + // ... and they have to SURVIVE the k loop rather than be rebuilt in it. + // Left alone, LLVM's rematerialisation heuristic decides that 5 ALU ops + // per panel per iteration are cheaper than 8 live VGPRs -- which is + // exactly backwards for this kernel: the producer's own budget is ~20 + // VGPRs, and every one of those ALU ops is a VALU instruction competing + // with a consumer wave's MMAC for the same SIMD issue slot. Measured on + // the ISA before this: the loop body carried 11 v_add + 11 v_or + 8 + // v_cmp + 8 v_cndmask per stage, against 8 useful DMA loads. An asm + // boundary makes the values non-rematerialisable, so they are computed + // once per tile and then only read. +#pragma unroll + for(index_t i = 0; i < 8; ++i) + { + asm volatile("" : "+v"(s.v_off[i])); + } + return s; + } + + // One k-stage's DMA issue: 8 panels, 16 B per lane each, straight to LDS. + // LDS m0 is a segment-relative byte offset, so its base is zero; passing + // the shared-memory pointer here would force a VGPR-to-SGPR readfirstlane. + // The only per-stage values are `soffset` (the stage's K column, uniform) + // and `m0` (the ring slot plus producer slab, uniform). + CK_TILE_DEVICE void producer_fill(const AProdSrc& src, index_t k_stage, index_t wid) const + { + using namespace wasp_detail; + + if((wasp_ablate() & 64) != 0) + { + return; // ablation: idle producer (see wasp_ablate) + } + + const index_t slot = (k_stage % kNumStages) * kAStageBytes; + + const uint32_t soffset = static_cast(k_stage * kStageK) * sizeof(uint16_t); + const uint32_t m0 = + static_cast(slot + wid * kProdSlabBytes); + + wasp_buffer_load_x8_to_lds(src.rsrc, + m0, + soffset, + src.v_off[0], + src.v_off[1], + src.v_off[2], + src.v_off[3], + src.v_off[4], + src.v_off[5], + src.v_off[6], + src.v_off[7]); + // NOTE: the DMA is only ISSUED here, not retired. Direct-to-LDS + // completion is on vmcnt, and the overlapping cadence in operator() + // drains it (wasp_wait_vmcnt0) right before the publishing barrier -- so + // that the NEXT stage's DMA, issued straight after a publish, overlaps + // the consumers' compute of the stage just published. + } + + // ------------------------------------------------------------------ + // Consumer: waves 4..11, one 256M x 32N stripe each. + // + // MMAC fragment contract (measured on gfx936, see the standalone reference kernel header): + // A lane l holds rows m = l&15 within its 16-row sub-tile and + // k = 4*(l>>4) + 0..3 of each K16 operand + // C accumulator slot c of lane group g holds column 4*c + g + // The 16 B LDS unit for lane group g, K-half kh holds global + // k = 8*(g + 4*kh) + 0..8, i.e. plain-layout offset (g + 4*kh)*16. + // ------------------------------------------------------------------ + // One B buffer: [k_half][n_tile]. + using BBuf = wasp_detail::stage16_t[2][WaspAccTile::kNIters]; + + // Prefetch a whole stage's B (2 K-halves x 2 n) into one buffer through the + // CACHE_SWIZZLE descriptor. In-row K byte of unit (kh,n) is 16*col_group + + // 64*kh; the stage K column base rides the SGPR soffset. + // The per-lane addressing values the k-stage consumes, computed ONCE per + // item -- the reference kernel's shape (a_base_h0 / a_base_h1 / pair_n0 / + // pair_n1, bf16_gemm_nt_hipc.cpp:1272-1373) -- and handed to the interval + // bodies. The k loop then carries them in registers and does NO addressing + // arithmetic: no VALU in the quarters, the base/pairs are read where the + // ds_read / buffer_load_dwordx4 need them. + // + // Per-lane because both are: the A read base is + // lds + row_in_tile*128 + ((col_group ^ (row_in_tile & 7)) * 16) + // and the B IDXEN pair is {index = global B row, offset = in-row bytes}. + struct AReadBases + { + uint32_t h0; + uint32_t h1; + }; + + struct BPairs + { + uint64_t p0; + uint64_t p1; + }; + + // The B prefetch: four buffer_load_dwordx4 in ONE asm block, the reference + // kernel's form -- two IDXEN pairs, the K-half riding the instruction's + // immediate (+0 / +64 B), the stage's K column riding the uniform SGPR + // soffset. The pairs are the per-item values built in operator(). + CK_TILE_DEVICE void load_B_stage(wasp_detail::buffer_rsrc_t b_rsrc, + const BPairs& bp, + index_t k_stage, + BBuf& buf, + wasp_detail::WaspVmemDiag& vmem_diag) const + { + using namespace wasp_detail; + const uint32_t s_col_base = static_cast(k_stage * kStageK) * sizeof(uint16_t); + if((wasp_ablate() & 1) != 0) + { + buf[0][0].d4 = 0; + buf[1][0].d4 = 0; + buf[0][1].d4 = 0; + buf[1][1].d4 = 0; + return; + } + // Destination order is the reference kernel's: (n0,kh0) (n0,kh1) + // (n1,kh0) (n1,kh1) -- buf[kh][n], the immediate 0/64 choosing the half. + wasp_buffer_load_b128_x4_idxen_soff2(b_rsrc, bp.p0, bp.p1, s_col_base, + buf[0][0].d4, // n0, kh0 (offset:0) + buf[1][0].d4, // n0, kh1 (offset:64) + buf[0][1].d4, // n1, kh0 (offset:0) + buf[1][1].d4); // n1, kh1 (offset:64) +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + if(k_stage == wasp_detail::kVmemSampleStage) + { + vmem_diag.issued_at = wasp_detail::wasp_memtime(); + vmem_diag.issued = true; + } +#else + (void)vmem_diag; +#endif + } + + // The 16 B bytes one lane reads per (n-tile, K-half) live at global K = + // 8*col + 32*khalf, i.e. col*16 B into the row; the K-half rides the + // instruction immediate (+0 / +64 B) and the stage's K column rides the + // uniform SGPR soffset -- the reference kernel's form. + static constexpr uint32_t kBHalfBytes = 64; + + // ------------------------------------------------------------------ + // The A fragment is ONE WHOLE K-HALF: 16 x 16 B per lane = 64 VGPR, read with + // ONE base register plus compile-time immediates, and consumed by ONE + // 64-MMAC burst. This is the standalone kernel's shape AND grouping + // (ds_read_a_half + mmac_half), and the k-stage body around it is unchanged: + // two intervals per stage, one 16-unit read burst and one 64-MMAC burst in + // each, in opposite order for the two consumer personas. See c_I1/c_I2 (C) + // and cp_I1/cp_I2 (C'). + // ------------------------------------------------------------------ + + // 16 consecutive ds_read_b128 off ONE base VGPR, with the m-iteration row + // offsets (16 rows x 128 B = 0x800) AND the 2-deep A ring slot as + // compile-time immediates. 32768 + 30720 = 63488 fits the 16-bit unsigned + // offset field. + // + // The immediates are load bearing. The natural form -- a generic + // ds_read_b128(byte_addr) helper taking the address as a "v" operand -- makes + // the compiler materialise a fresh per-lane address for every ds_read, i.e. + // 2 v_add per m-iteration, 32 VALU per k-stage per consumer wave, and those + // ride on the same SIMD issue slots as the other waves' MMACs. + // (row & 7) -- the XOR swizzle selector -- is invariant under +16 rows, so + // the same base serves all 16 reads of a half. + // + // NOT SPLIT INTO 16 SINGLE-OUTPUT BLOCKS, and that was measured: the + // per-statement "memory" clobber serialises the issue and it is 21% SLOWER + // (171.7 -> 135.7 TFLOPS) with the SAME spill, so it buys the allocator + // nothing either. + // `token` is an UNUSED input operand (the accumulator element the previous + // MMAC chain wrote). It exists only to create a data edge: the MMACs are + // builtins, so nothing else orders them against these ds_reads, and without + // the edge the reads can be issued first -- which makes the old fragment and + // the new one overlap, and the allocator answers with a 64-register copy of + // the old fragment. That copy is what pushes the window over the 256-VGPR + // cap and gets the per-item A bases / B pairs spilled into `buffer_load_dword` + // reloads inside the quarters. Zero instructions, one dependency. + // 16 consecutive ds_read_b128 off ONE base VGPR, with the m-iteration row + // offsets (16 rows x 128 B = 0x800) AND the 2-deep A ring slot as + // compile-time immediates. 32768 + 30720 = 63488 fits the 16-bit unsigned + // offset field. + // + // The outputs are dword4_t REFERENCES into the fragment's union objects (one + // per m-tile), not copies: the MMACs read the same objects through their + // bf16 view, so the fragment has ONE storage identity and the allocator + // reuses the same 64 registers instead of materialising a second fragment at + // the C' seam. That is the reference kernel's shape + // (union_vec16_bf16 a_regs[16] + reinterpret_cast(a_regs[i].i32x4), + // bf16_gemm_nt_hipc.cpp:1247 and :1281-1300). + template + CK_TILE_DEVICE static void + ds_read_b128_x16_imm(uint32_t base, + wasp_detail::dword4_t& o0, wasp_detail::dword4_t& o1, wasp_detail::dword4_t& o2, wasp_detail::dword4_t& o3, + wasp_detail::dword4_t& o4, wasp_detail::dword4_t& o5, wasp_detail::dword4_t& o6, wasp_detail::dword4_t& o7, + wasp_detail::dword4_t& o8, wasp_detail::dword4_t& o9, wasp_detail::dword4_t& o10, wasp_detail::dword4_t& o11, + wasp_detail::dword4_t& o12, wasp_detail::dword4_t& o13, wasp_detail::dword4_t& o14, wasp_detail::dword4_t& o15) + { + static_assert(SLOT + 30720 <= 65535, + "ds_read_b128 imm offset overflows the 16-bit unsigned field"); + asm volatile("ds_read_b128 %0, %16 offset:%c17\n\t" + "ds_read_b128 %1, %16 offset:%c18\n\t" + "ds_read_b128 %2, %16 offset:%c19\n\t" + "ds_read_b128 %3, %16 offset:%c20\n\t" + "ds_read_b128 %4, %16 offset:%c21\n\t" + "ds_read_b128 %5, %16 offset:%c22\n\t" + "ds_read_b128 %6, %16 offset:%c23\n\t" + "ds_read_b128 %7, %16 offset:%c24\n\t" + "ds_read_b128 %8, %16 offset:%c25\n\t" + "ds_read_b128 %9, %16 offset:%c26\n\t" + "ds_read_b128 %10, %16 offset:%c27\n\t" + "ds_read_b128 %11, %16 offset:%c28\n\t" + "ds_read_b128 %12, %16 offset:%c29\n\t" + "ds_read_b128 %13, %16 offset:%c30\n\t" + "ds_read_b128 %14, %16 offset:%c31\n\t" + "ds_read_b128 %15, %16 offset:%c32\n" + : "=&v"(o0), "=&v"(o1), "=&v"(o2), "=&v"(o3), + "=&v"(o4), "=&v"(o5), "=&v"(o6), "=&v"(o7), + "=&v"(o8), "=&v"(o9), "=&v"(o10), "=&v"(o11), + "=&v"(o12), "=&v"(o13), "=&v"(o14), "=&v"(o15) + : "v"(base), + "n"(SLOT + 0), "n"(SLOT + 2048), "n"(SLOT + 4096), "n"(SLOT + 6144), + "n"(SLOT + 8192), "n"(SLOT + 10240), "n"(SLOT + 12288), "n"(SLOT + 14336), + "n"(SLOT + 16384), "n"(SLOT + 18432), "n"(SLOT + 20480), "n"(SLOT + 22528), + "n"(SLOT + 24576), "n"(SLOT + 26624), "n"(SLOT + 28672), "n"(SLOT + 30720) + : "memory"); + } + + // One K-half of the 256M x K16 A operand for the whole wave tile. SLOT is + // the A-ring slot ((k_stage % 2) * kAStageBytes), which the k loop unrolled + // by parity makes compile-time. + template + CK_TILE_DEVICE void + a_read_half(const AReadBases& b, + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters]) const + { + ds_read_b128_x16_imm(H1 ? b.h1 : b.h0, + reinterpret_cast(a_regs[0]), + reinterpret_cast(a_regs[1]), + reinterpret_cast(a_regs[2]), + reinterpret_cast(a_regs[3]), + reinterpret_cast(a_regs[4]), + reinterpret_cast(a_regs[5]), + reinterpret_cast(a_regs[6]), + reinterpret_cast(a_regs[7]), + reinterpret_cast(a_regs[8]), + reinterpret_cast(a_regs[9]), + reinterpret_cast(a_regs[10]), + reinterpret_cast(a_regs[11]), + reinterpret_cast(a_regs[12]), + reinterpret_cast(a_regs[13]), + reinterpret_cast(a_regs[14]), + reinterpret_cast(a_regs[15])); + + } + + // 64 MMACs for one K-half: kk(2) x m(16) x n(2). The standalone kernel's + // mmac_half. + // + // kk stays OUTERMOST so consecutive MMACs write DIFFERENT accumulator + // registers: the two writes to one slot end up 16 MMACs apart, which is what + // hides the matrix pipe's accumulator pass-through behind 15 independent + // MMACs. A (m, n, kk) order has a 4-deep RAW chain per slot instead. + // + // The wave-priority bracketing the standalone kernel wraps around this chain + // (s_setprio 3 after one MMAC, 0 after the burst) was tried and it costs 4% + // here: 293.5 -> 281.2 TFLOPS on this shape. + // + // A NOTE ON PARTIAL DRAINS: splitting this into two 16-MMAC halves fed by + // halves of the A read (wait lgkmcnt(4) -> 16 MMACs -> wait lgkmcnt(0) -> + // the other 16) was measured and it is SLOWER -- 256 TFLOPS against 299 with + // the single drain. The LDS queue is shared with the other seven consumers, + // so "wait for my first four reads" does not return any sooner; it just + // takes two stalls instead of one. + // The Q1 drain -- the reference kernel's `wait lgkmcnt(0); wait vmcnt(0)`: + // lgkmcnt for A's LDS read, vmcnt for the B prefetch that Q1's MMACs consume. + CK_TILE_DEVICE void wasp_wait_q1(wasp_detail::WaspVmemDiag& vmem_diag, + index_t stage) const + { + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + wasp_detail::wasp_wait_vmcnt0(vmem_diag, stage); + } + + CK_TILE_DEVICE void mmac_half(const wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + const wasp_detail::stage16_t (&b_half)[WaspAccTile::kNIters], + AccTile& acc) const + { + using namespace wasp_detail; + if((wasp_ablate() & 2) != 0) + { + return; // ablation: skip MMACs + } + // The reference kernel peels the first MMAC out of the chain so that + // set_prio<3> can be raised AFTER one has issued and set_prio<0> dropped + // after the very last, with no per-iteration branch in the chain. + // + // (kk = 0, m = 0, n = 0) + mmac_one<0, 0, 0>(a_regs, b_half, acc); + wasp_set_prio<3>(); + + // Keep every accumulator index compile-time. A dynamically indexed + // vector array is lowered through private scratch by LLVM at this + // register pressure, even when the loops are marked unrolled. + static_for<1, WaspAccTile::kNIters, 1>{}([&](auto n) { + mmac_one<0, 0, n.value>(a_regs, b_half, acc); + }); + static_for<1, WaspAccTile::kMIters, 1>{}([&](auto m) { + static_for<0, WaspAccTile::kNIters, 1>{}([&](auto n) { + mmac_one<0, m.value, n.value>(a_regs, b_half, acc); + }); + }); + static_for<0, WaspAccTile::kMIters, 1>{}([&](auto m) { + static_for<0, WaspAccTile::kNIters, 1>{}([&](auto n) { + mmac_one<1, m.value, n.value>(a_regs, b_half, acc); + }); + }); + wasp_set_prio<0>(); + } + + // One 16x16x16 bf16 MMAC into acc.v[m*kNIters + n]. KK is literal so the + // fragment indices stay compile-time; the peeling above only changes WHERE + // the priority changes sit, never the order of the MMACs themselves. + template + CK_TILE_DEVICE void mmac_one(const wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + const wasp_detail::stage16_t (&b_half)[WaspAccTile::kNIters], + AccTile& acc) const + { + using namespace wasp_detail; + constexpr index_t slot = M * WaspAccTile::kNIters + N; + acc.v[slot] = mmac_bf16(a_regs[M].frag[KK], b_half[N].frag[KK], acc.v[slot]); + } + + // ------------------------------------------------------------------ + // PERSONA C's k-stage: (Q0,Q1) | BAR k.mid | (Q2,Q3) | BAR k.end. + // Every read is drained inside the interval that issued it, so every A slot + // is provably dead a full stage before the producer refills it. + // ------------------------------------------------------------------ + template + CK_TILE_DEVICE void c_I1(index_t consumer_id, + index_t k_stage, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const + { + using namespace wasp_detail; + (void)consumer_id; + (void)k_stage; + (void)num_loop; + constexpr index_t SLOT = CUR * kAStageBytes; + // This interval issues no B prefetch and no stage-dependent addressing, + // but keeps the same signature as its I2 sibling by design. + (void)k_stage; + (void)num_loop; + if((wasp_ablate() & 32) != 0) + { + return; // ablation: idle consumer + } + // Q0: A(k, h0). + a_read_half(a_rb, a_regs); + // Q1: A(h0) landed (lgkmcnt) + B(k) landed (vmcnt) -- B(k) was prefetched + // in the previous stage and B(k+1) has not been issued yet. + wasp_wait_q1(vmem_diag, k_stage); + mmac_half(a_regs, b_regs[CUR][0], acc); + } + + template + CK_TILE_DEVICE void c_I2(index_t consumer_id, + index_t k_stage, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::buffer_rsrc_t b_rsrc, + const BPairs& b_pairs, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const + { + using namespace wasp_detail; + constexpr index_t SLOT = CUR * kAStageBytes; + (void)consumer_id; + if((wasp_ablate() & 32) != 0) + { + return; // ablation: idle consumer + } + // Q2: issue the B prefetch -- 4 x buffer_load_dwordx4 in ONE asm block, + // the reference kernel's form (two IDXEN pairs, k-half on the immediate, + // the stage's K column on the uniform SGPR soffset) -- and read A(k, h1). + // + // The MMACs of the next interval consume these bytes, so the Q1 drain + // below is the reference kernel's `lgkmcnt(0); vmcnt(0)`. + if(k_stage + 1 < num_loop) + { + load_B_stage(b_rsrc, b_pairs, k_stage + 1, b_regs[CUR ^ 1], vmem_diag); + } + a_read_half(a_rb, a_regs); + // Q3: A(h1) landed. B(k)'s second half is resident, so NO vmcnt here -- + // draining it would stall on the B(k+1) prefetch just issued. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[CUR][1], acc); + } + + // ------------------------------------------------------------------ + // PERSONA C's k-stage, rotated by ONE QUARTER: + // (Q3 of k-1, Q0) | BAR k.mid | (Q1, Q2) | BAR k.end. + // + // What crosses a barrier is the MMAC and the A fragment registers it + // consumes. Each interval carries the same work as C's (one 16-unit read, one + // 64-MMAC burst, one real drain): a rendezvous releases when the SLOWEST wave + // arrives, so a heavier persona would just be paid for by the other half. + // ------------------------------------------------------------------ + template + CK_TILE_DEVICE void cp_I1(index_t consumer_id, + index_t k_stage, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + AccTile& acc) const + { + using namespace wasp_detail; + (void)consumer_id; + (void)k_stage; + (void)num_loop; + constexpr index_t SLOT = CUR * kAStageBytes; + (void)k_stage; + (void)num_loop; + if((wasp_ablate() & 32) != 0) + { + return; // ablation: idle consumer + } + if constexpr(!IS_FIRST) + { + // Q3 of the PREVIOUS stage: the quarter that crosses the barrier. + // Its data was read at the end of the previous stage and is drained + // HERE, so the wait is for a read issued a whole barrier ago -- that + // is what the deferral buys. B(k-1) is still resident (its buffer is + // CUR^1 here), which is why cp_I2's prefetch must not move earlier. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[CUR ^ 1][1], acc); + } + // Q0: A(k, h0). + a_read_half(a_rb, a_regs); + } + + template + CK_TILE_DEVICE void cp_I2(index_t consumer_id, + index_t k_stage, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::buffer_rsrc_t b_rsrc, + const BPairs& b_pairs, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const + { + using namespace wasp_detail; + constexpr index_t SLOT = CUR * kAStageBytes; + (void)consumer_id; + if((wasp_ablate() & 32) != 0) + { + return; // ablation: idle consumer + } + // Q1: A(k, h0) landed -- it was read before the mid barrier, so this wait + // is cheap by construction -- and B(k) landed (prefetched one stage ago). + wasp_wait_q1(vmem_diag, k_stage); + mmac_half(a_regs, b_regs[CUR][0], acc); + + // Q2: issue the B prefetch -- 4 x buffer_load_dwordx4 in ONE asm block, + // the reference kernel's form (two IDXEN pairs, k-half on the immediate, + // the stage's K column on the uniform SGPR soffset) -- and read A(k, h1). + // The bytes land in b_regs[CUR ^ 1] and are consumed by the next + // interval's Q1 MMACs; that is what wasp_wait_q1()'s vmcnt half waits for. + if(k_stage + 1 < num_loop) + { + load_B_stage(b_rsrc, b_pairs, k_stage + 1, b_regs[CUR ^ 1], vmem_diag); + } + a_read_half(a_rb, a_regs); + } + + // Persona C driver: two intervals per k-stage, mid barrier between them. + CK_TILE_DEVICE void consumer_run_c(index_t consumer_id, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::buffer_rsrc_t b_rsrc, + const BPairs& b_pairs, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const + { + for(index_t k = 0; k < num_loop; k += 2) + { + c_I1<0>(consumer_id, k, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid (bare, like the reference) + c_I2<0>(consumer_id, + k, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).end + if(k + 1 < num_loop) + { + c_I1<1>( + consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).mid + c_I2<1>(consumer_id, + k + 1, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).end + } + } + } + + // Persona C' driver: the same barrier sequence, the rotated item order. The + // first stage is peeled so IS_FIRST is compile-time (there is no stage -1 to + // consume), and the rotation leaves exactly ONE MMAC owed after the loop. It + // is executed here, without a barrier: the producer owes no extra rendezvous + // for it either, so the cadence stays balanced. + CK_TILE_DEVICE void consumer_run_cp(index_t consumer_id, + index_t num_loop, + const AReadBases& a_rb, + wasp_detail::buffer_rsrc_t b_rsrc, + const BPairs& b_pairs, + wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const + { + if(num_loop == 0) + { + return; + } + if((num_loop & 1) != 0) + { + cp_I1<0, /*IS_FIRST=*/true>(consumer_id, 0, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(0).mid + cp_I2<0>( + consumer_id, 0, num_loop, a_rb, b_rsrc, b_pairs, b_regs, a_regs, vmem_diag, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(0).end + for(index_t k = 1; k < num_loop; k += 2) + { + cp_I1<1, false>(consumer_id, k, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid + cp_I2<1>(consumer_id, + k, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).end + cp_I1<0, false>(consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).mid + cp_I2<0>(consumer_id, + k + 1, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).end + } + if((wasp_detail::wasp_ablate() & 32) == 0) + { + // Odd num_loop leaves the final stage at CUR=0. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[0][1], acc); + } + } + else + { + cp_I1<0, /*IS_FIRST=*/true>(consumer_id, 0, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(0).mid + cp_I2<0>( + consumer_id, 0, num_loop, a_rb, b_rsrc, b_pairs, b_regs, a_regs, vmem_diag, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(0).end + + for(index_t k = 1; k < num_loop - 1; k += 2) + { + cp_I1<1, false>(consumer_id, k, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid + cp_I2<1>(consumer_id, + k, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).end + cp_I1<0, false>(consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).mid + cp_I2<0>(consumer_id, + k + 1, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k+1).end + } + cp_I1<1, false>(consumer_id, num_loop - 1, num_loop, a_rb, b_regs, a_regs, acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid + cp_I2<1>(consumer_id, + num_loop - 1, + num_loop, + a_rb, + b_rsrc, + b_pairs, + b_regs, + a_regs, + vmem_diag, + acc); + wasp_detail::wasp_barrier_maybe(); // BAR(k).end + if((wasp_detail::wasp_ablate() & 32) == 0) + { + // Even num_loop leaves the final stage at CUR=1. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[1][1], acc); + } + } + } + // ------------------------------------------------------------------ + // Role-split entry point. + // + // Overlapping double-barrier cadence (two barriers per stage, A then B): + // + // producer : issue(0) + // for k: wait vmcnt | BAR_A(k) | issue(k+1) | BAR_B(k) + // consumer : for k: BAR_A(k) | consume(k) | BAR_B(k) + // + // BAR_A(k) publishes stage k: the producer has drained the stage-k DMA + // (vmcnt) before it, so the consumer reads a complete slab after it. + // issue(k+1) is started straight after BAR_A(k), into the OTHER ring slot + // ((k+1)%2 != k%2), so its global->LDS DMA runs UNDER the consumers' compute + // of stage k. BAR_B(k) then proves the consumers are done reading slot k%2 + // before the producer refills it at issue(k+2) next iteration. Both roles + // run exactly 2*num_loop barriers, in the same A,B order, so the split stays + // legal. + // + // The two slots are the whole point of the ring: with a single slot the + // producer could not stage k+1 while k is being read. The barrier carries + // the double-buffer handshake; the DMA's own vmcnt carries completion. + // + // The entry barrier drains the PREVIOUS tile's epilogue, which borrows the + // whole 64 KB buffer: without it this tile's issue(0) would overwrite A + // staging while the previous consumers are still reading it back for stores. + // ------------------------------------------------------------------ + // NOTE: GroupedGemmKernel invokes the pipeline with the SIX argument form + // (num_loop, has_hot_loop, tail_number, smem) -- see + // grouped_gemm_kernel.hpp:936-941 -- not the 4 argument one used by + // UniversalGemmKernel's plain RunGemm. The hot-loop flag and the tail + // number belong to the ComputeV3 TailHandler, which a role-split body does + // not use: the k-tail is folded into producer_fill's mask instead (TODO, + // see the header note). Both forms are provided so either caller works. + template + CK_TILE_DEVICE AccTile operator()(const ADramBlockWindow& a_dram_block_window, + const BDramBlockWindow& b_dram_block_window, + index_t num_loop, + bool /*has_hot_loop*/, + TailNumber /*tail_number*/, + void* p_smem) const + { + // ---- per-TILE tensor plumbing, hoisted out of the k loop ---- + // The windows are pure scalar carriers here: extract the base pointers, + // strides, M/N lengths and the tile origin, then hand them to run_raw. + // (The grouped kernel's RunGemmWasp path computes the SAME scalars + // straight from kargs + tile origin and calls run_raw directly, so no + // tile-window object lives across the k loop -- see run_raw's note.) + const auto a_view = a_dram_block_window.get_bottom_tensor_view(); + const auto& a_desc = a_view.get_tensor_descriptor(); + const auto a_org = a_dram_block_window.get_window_origin(); + const auto* a_base = a_view.get_buffer_view().p_data_ + + a_desc.calculate_offset(make_multi_index(a_org[0], a_org[1])); + const index_t lda = a_desc.calculate_offset(make_multi_index(1, 0)) - + a_desc.calculate_offset(make_multi_index(0, 0)); + // M tail: a_org[0] is this tile's first GLOBAL row within the group's A, + // get_lengths()[0] is the group's M, so valid tile-local rows = M - org. + const index_t a_m = a_desc.get_lengths()[number<0>{}]; + const index_t m_valid = a_m - a_org[0]; + + const auto b_view = b_dram_block_window.get_bottom_tensor_view(); + const auto& b_desc = b_view.get_tensor_descriptor(); + const auto b_org = b_dram_block_window.get_window_origin(); + const auto* b_base = b_view.get_buffer_view().p_data_ + + b_desc.calculate_offset(make_multi_index(b_org[0], b_org[1])); + // B is ColumnMajor for the NT layout, so the view is (N, K): index 1 is + // the n stride, i.e. the row pitch the reference kernel calls `b_stride`. + const index_t b_n_stride = b_desc.calculate_offset(make_multi_index(1, 0)) - + b_desc.calculate_offset(make_multi_index(0, 0)); + const index_t b_n_len = b_desc.get_lengths()[number<0>{}]; + + return run_raw(a_base, + lda, + m_valid, + b_base, + b_n_stride, + b_n_len, + static_cast(b_org[0]), + num_loop, + p_smem); + } + + // Raw-pointer core: the whole k-loop body, fed by SCALARS instead of tile + // windows. The framework's generic RunGemm builds tile-window objects for + // A/B/C/D and keeps the C/D ones live across the entire k loop (the epilogue + // needs them afterwards) -- ~30-40 VGPR of dead-weight live state stacked on + // the 224-VGPR accumulator/fragment floor, which is what pushed this kernel + // over the 256-VGPR cap and spilled the A fragment to scratch. The grouped + // kernel's RunGemmWasp path computes these scalars from kargs + tile origin + // and calls this directly, so no window state lives across the loop. + // + // a_base/lda/m_valid : A row-major base (origin folded in), row stride, + // valid tile-local rows + // b_base/b_n_stride : B col-major base (origin folded in), n row pitch + // b_n_len/b_n_origin : B's N length and this tile's n origin (for the SRD + // size bound (b_n_len - b_n_origin) * pitch) + CK_TILE_DEVICE AccTile run_raw(const ADataType* a_base, + index_t lda, + index_t m_valid, + const BDataType* b_base, + index_t b_n_stride, + index_t b_n_len, + index_t b_n_origin, + index_t num_loop, + void* p_smem) const + { + // Preserve the value-returning pipeline interface for generic CK + // callers. The specialized grouped Wasp path uses run_raw_into so + // producer waves never need a defined accumulator value. + AccTile acc{}; + run_raw_into(a_base, + lda, + m_valid, + b_base, + b_n_stride, + b_n_len, + b_n_origin, + num_loop, + acc, + p_smem); + return acc; + } + + CK_TILE_DEVICE void run_raw_into(const ADataType* a_base, + index_t lda, + index_t m_valid, + const BDataType* b_base, + index_t b_n_stride, + index_t b_n_len, + index_t b_n_origin, + index_t num_loop, + AccTile& acc, + void* p_smem) const + { + char* const lds = static_cast(p_smem); + const index_t wid = get_warp_id(); + + wasp_detail::WaspVmemDiag vmem_diag{}; + + // CACHE_SWIZZLE descriptor over this tile's B (built once per tile, not + // per k-stage). STRIDE = the row pitch in bytes; size bounds the SRD to + // the group's remaining N (b_n_len - b_n_origin). + const uint32_t b_pitch_bytes = static_cast(b_n_stride) * sizeof(uint16_t); + const uint32_t b_size_bytes = + static_cast(b_n_len - b_n_origin) * b_pitch_bytes; + const wasp_detail::buffer_rsrc_t b_rsrc = + wasp_detail::wasp_make_buffer_rsrc(b_base, b_pitch_bytes, b_size_bytes, + /*allow_swizzle=*/true); + // Built once per tile, not per k-stage: the prefetch's SRD. Its data is + // consumed by the k loop's MMACs one stage later (see the interval + // bodies' Q1/Q2 ordering). + + wasp_detail::wasp_sync_lds_maybe(); // drain the previous tile's epilogue (both roles) + + if(wid < kNumProducerWaves) + { + // The 8 panel vaddrs and the SRD are per-TILE: build them once, + // outside the k loop (see AProdSrc / make_a_prod_src). + const auto a_src = make_a_prod_src(a_base, lda, m_valid, wid); + + // hipc's producer shape: issue first, then the barrier pair -- + // BAR(k).mid is a rendezvous, BAR(k).end is what PUBLISHES A(k+1) + // (this role has waited for that DMA in front of it). The consumers' + // stage k+1 starts on the far side of that barrier. + producer_fill(a_src, 0, wid); // issue stage 0 + wasp_detail::wasp_barrier_maybe(); // BAR prologue.mid + wasp_detail::wasp_wait_vmcnt0(); // A(0) landed + wasp_detail::wasp_barrier_maybe(); // BAR prologue.end: publish A(0) + + for(index_t k = 0; k < num_loop; ++k) + { + // Issuing k+1 here instead of after the barrier buys its DMA a + // whole extra interval to land (it is confirmed below), and the + // slot it writes -- (k+1)%2 -- was last read by the consumers in + // stage k-1, i.e. one full interval before this issue. + if(k + 1 < num_loop) + { + producer_fill(a_src, k + 1, wid); // issue k+1 -> slot (k+1)%2 +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + if(k + 1 == wasp_detail::kVmemSampleStage) + { + vmem_diag.issued_at = wasp_detail::wasp_memtime(); + vmem_diag.issued = true; + } +#endif + } + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid + wasp_detail::wasp_wait_vmcnt0_producer(vmem_diag, k + 1); // A(k+1) landed + wasp_detail::wasp_barrier_maybe(); // BAR(k).end: publish A(k+1) + } + } + else + { + // Accumulators are consumed only by consumer waves. Keeping this + // initialization inside their role branch avoids issuing 128 + // accumulator-zeroing moves on producer waves. + acc = AccTile{}; + const index_t consumer_id = wid - kNumProducerWaves; + // B double buffer [parity][k_half][n]. Prime stage 0's B, then walk + // the k loop unrolled by parity so the live buffer index (CUR = k&1) + // is a compile-time template arg. + // The A fragment file is owned HERE, not by the per-stage body: the + // C' persona's deferred MMAC consumes it one stage after the read + // that filled it (see cp_I1). + wasp_detail::stage16_t a_regs[WaspAccTile::kMIters]; // one whole K-half + // Per-lane addressing values, computed once per item and shared by + // the interval bodies, so the K loop carries only the reference + // kernel's A bases and B index pairs. + const uint32_t lane = static_cast(get_lane_id()); + const uint32_t row = lane & 15u; + const uint32_t r7 = row & 7u; + const uint32_t colg = lane >> 4; + const uint32_t rowbase = static_cast(reinterpret_cast(lds)) + + row * (kStageK * sizeof(uint16_t)); + const AReadBases a_rb{rowbase + ((colg ^ r7) * 16u), + rowbase + (((colg + 4u) ^ r7) * 16u)}; + const uint32_t nrow = static_cast(consumer_id) * 32u + row; + const uint32_t inrow = colg * 16u; + const BPairs b_pr{(static_cast(inrow) << 32) | nrow, + (static_cast(inrow) << 32) | (nrow + 16u)}; + // Destination of the prefetch, double buffered [parity][k-half][n]. + wasp_detail::stage16_t b_regs[2][2][WaspAccTile::kNIters]; + // Prologue, mirroring the producer's: prime B(0), barrier, wait for + // it, barrier. The second barrier is the one that publishes A(0) -- + // the producer reaches it only after waiting for that DMA. + // Prime B(0); the first stage's MMACs consume it (see the wait below). + if(num_loop > 0) + { + load_B_stage(b_rsrc, b_pr, 0, b_regs[0], vmem_diag); + } + wasp_detail::wasp_barrier_maybe(); // BAR prologue.mid + wasp_detail::wasp_wait_vmcnt0(); // B(0) landed + wasp_detail::wasp_barrier_maybe(); // BAR prologue.end + + // waves 0..3 -> C, 4..7 -> C' + if(consumer_id < kConsumerSplitWave) + { + consumer_run_c( + consumer_id, num_loop, a_rb, b_rsrc, b_pr, b_regs, a_regs, vmem_diag, acc); + } + else + { + consumer_run_cp( + consumer_id, num_loop, a_rb, b_rsrc, b_pr, b_regs, a_regs, vmem_diag, acc); + } + } + + // Drain barrier before the epilogue. + // + // The epilogue borrows the whole 64 KB buffer to stage its output: its + // 8 per-consumer slabs tile the ENTIRE LDS, so consumer 0..3 write over + // the very region (the first A stage) that consumer 4..7 are still + // reading A from. Without this barrier a consumer that finishes its K + // loop early corrupts the A operands of the consumers that have not, + // which shows up as the second half of the N range being garbage while + // the first half is exact. + // + // It has to live HERE rather than in the epilogue: the epilogue runs + // only for the consumer waves, so a barrier inside it would never be + // reached by the producers. Every wave arrives here, the producers then + // fall through into the epilogue and return immediately. + wasp_detail::wasp_sync_lds_maybe(); + +#if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + if(get_lane_id() == 0) + { + const index_t cta = get_block_id(); + if(cta < 1024 && wid < 16) + { + wasp_detail::g_wasp_vmem_samples[cta * 16 + wid] = vmem_diag.sample; + } + } +#endif + + } + + // 4 argument form, forwarded. Kept so that a caller that only knows the + // plain RunGemm convention does not silently fall back to the inherited + // ComputeV3 body (same-arity overloads in a base class are hidden, not + // merged, by a derived declaration). + template + CK_TILE_DEVICE AccTile operator()(const ADramBlockWindow& a_dram_block_window, + const BDramBlockWindow& b_dram_block_window, + index_t num_loop, + void* p_smem) const + { + return (*this)(a_dram_block_window, + b_dram_block_window, + num_loop, + true, + TailNumber::Full, + p_smem); + } +}; + +} // namespace ck_tile From e13d6fecb7c475c8d4be59acbf69916ae517295c Mon Sep 17 00:00:00 2001 From: wugangsheng Date: Thu, 15 Oct 2026 12:50:18 +0800 Subject: [PATCH 09/11] Optimize WASP grouped GEMM pipeline Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ck_tile/19_grouped_gemm/CMakeLists.txt | 110 +- .../bench_grouped_gemm_fused_kernels.sh | 229 +++ .../grouped_gemm_fused_kernels_perf.md | 138 ++ ...p_test.cpp => grouped_gemm_fused_test.cpp} | 1806 +++++++++++++++-- .../grouped_gemm_v7_instance_test.cpp | 1434 ------------- .../grouped_gemm_wasp_epilogue.hpp | 13 +- .../pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp | 403 ++-- 7 files changed, 2175 insertions(+), 1958 deletions(-) create mode 100755 example_hcu/ck_tile/19_grouped_gemm/bench_grouped_gemm_fused_kernels.sh create mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_kernels_perf.md rename example_hcu/ck_tile/19_grouped_gemm/{grouped_gemm_wasp_test.cpp => grouped_gemm_fused_test.cpp} (52%) delete mode 100644 example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp diff --git a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt index 3b994cb..9a22a5a 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt +++ b/example_hcu/ck_tile/19_grouped_gemm/CMakeLists.txt @@ -486,85 +486,61 @@ if(GPU_TARGETS MATCHES "gfx936") PRIVATE ck_tile_hcu_grouped_gemm) endif() -# Correctness + performance driver for the v7 BF16 NN / TN instances -# (InstanceId 5000 / 5001). It goes through the versioned device-args ABI -# exactly as a consumer does (select + run), and it links the grouped-GEMM -# archive, so it also covers the selector route and the ABI dispatch that -# production uses. +# Unified correctness + performance driver for the v7/v8 BF16 NN/TN instances +# and the WASP BF16 NT pipeline. --layout selects the corresponding kernel +# family: NN -> v7, TN -> v8, NT -> WASP. if(TARGET ck_tile_hcu_grouped_gemm_v7_v8_bf16) - add_executable(tile_example_grouped_gemm_v7_instance_test EXCLUDE_FROM_ALL - grouped_gemm_v7_instance_test.cpp) - target_include_directories(tile_example_grouped_gemm_v7_instance_test + add_executable(tile_example_grouped_gemm_fused_test EXCLUDE_FROM_ALL + grouped_gemm_fused_test.cpp) + target_include_directories(tile_example_grouped_gemm_fused_test PRIVATE ${CMAKE_CURRENT_LIST_DIR}) - target_compile_definitions(tile_example_grouped_gemm_v7_instance_test + target_compile_definitions(tile_example_grouped_gemm_fused_test PRIVATE CK_TILE_GROUPED_GEMM_FAST_BUILD) - target_link_libraries(tile_example_grouped_gemm_v7_instance_test + target_link_libraries(tile_example_grouped_gemm_fused_test PRIVATE ck_tile_hcu_grouped_gemm) - # Optional device-side functional reference: hipBLASLt, one plain GEMM per - # group (--hipblaslt-ref=1). Enabled only when the install the standalone - # harness builds against is present, so the test still builds everywhere. + # Optional hipBLASLt correctness/performance references for both backends. set(HIPBLASLT_INSTALL "/root/workspace/hipblaslt-install-0820/hipblaslt-install" - CACHE PATH "hipBLASLt install root for the v7 instance test") - if(EXISTS "${HIPBLASLT_INSTALL}/lib/libhipblaslt.so") - # SYSTEM: the hipBLASLt public header is not -Weverything clean; as a - # system include clang does not warn inside it. - target_include_directories(tile_example_grouped_gemm_v7_instance_test + CACHE PATH "hipBLASLt install root for grouped-GEMM fused tests") + if(EXISTS "${HIPBLASLT_INSTALL}/include/hipblaslt/hipblaslt.h" AND + EXISTS "${HIPBLASLT_INSTALL}/lib/libhipblaslt.so") + target_include_directories(tile_example_grouped_gemm_fused_test SYSTEM PRIVATE ${HIPBLASLT_INSTALL}/include) - target_link_libraries(tile_example_grouped_gemm_v7_instance_test + target_link_libraries(tile_example_grouped_gemm_fused_test PRIVATE ${HIPBLASLT_INSTALL}/lib/libhipblaslt.so) - set_target_properties(tile_example_grouped_gemm_v7_instance_test + set_target_properties(tile_example_grouped_gemm_fused_test PROPERTIES BUILD_RPATH "${HIPBLASLT_INSTALL}/lib" INSTALL_RPATH "${HIPBLASLT_INSTALL}/lib") - target_compile_definitions(tile_example_grouped_gemm_v7_instance_test - PRIVATE CK_TILE_V7_TEST_HAS_HIPBLASLT) + target_compile_definitions(tile_example_grouped_gemm_fused_test + PRIVATE + CK_TILE_V7_TEST_HAS_HIPBLASLT + CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) endif() - add_test(NAME tile_example_grouped_gemm_v7_instance_test - COMMAND $ - --layout=NN --m=512 --n=512 --k=256 --groups=4 --iters=5 --warmup=2) - set_tests_properties(tile_example_grouped_gemm_v7_instance_test + add_test(NAME tile_example_grouped_gemm_fused_test + COMMAND $ + --layout=NN --m=512 --n=512 --k=256 + --groups=4 --iters=5 --warmup=2) + set_tests_properties(tile_example_grouped_gemm_fused_test PROPERTIES LABELS "hcu;ck_tile;grouped_gemm;v7;correctness;perf") -endif() -# Explicit opt-in driver for the Wasp NT grouped-GEMM pipeline. It is not part -# of the default test suite: use the documented 8152x4096x4096 workload when -# validating correctness, resources, and performance. -add_executable(tile_example_grouped_gemm_wasp_test EXCLUDE_FROM_ALL - grouped_gemm_wasp_test.cpp) -target_include_directories(tile_example_grouped_gemm_wasp_test - PRIVATE ${CMAKE_CURRENT_LIST_DIR}) -target_compile_definitions(tile_example_grouped_gemm_wasp_test - PRIVATE CK_TILE_GROUPED_GEMM_FAST_BUILD) -if(NOT DEFINED HIPBLASLT_INSTALL) - set(HIPBLASLT_INSTALL "" CACHE PATH "hipBLASLt install root for the Wasp test") -endif() -if(EXISTS "${HIPBLASLT_INSTALL}/include/hipblaslt/hipblaslt.h" AND - EXISTS "${HIPBLASLT_INSTALL}/lib/libhipblaslt.so") - target_include_directories(tile_example_grouped_gemm_wasp_test - SYSTEM PRIVATE ${HIPBLASLT_INSTALL}/include) - target_link_libraries(tile_example_grouped_gemm_wasp_test - PRIVATE ${HIPBLASLT_INSTALL}/lib/libhipblaslt.so) - set_target_properties(tile_example_grouped_gemm_wasp_test - PROPERTIES BUILD_RPATH "${HIPBLASLT_INSTALL}/lib") - target_compile_definitions(tile_example_grouped_gemm_wasp_test - PRIVATE CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) -endif() -option(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT - "Instrument per-consumer Wasp epilogue cycles in the standalone test" OFF) -if(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) - target_compile_definitions(tile_example_grouped_gemm_wasp_test - PRIVATE CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) -endif() -option(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT - "Instrument one steady-state Wasp producer/consumer VMEM wait per wave" OFF) -if(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) - target_compile_definitions(tile_example_grouped_gemm_wasp_test - PRIVATE CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) -endif() -option(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER - "Diagnostic: reverse Wasp CTA-to-tile mapping without changing tile work" OFF) -if(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) - target_compile_definitions(tile_example_grouped_gemm_wasp_test - PRIVATE CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + # Explicit diagnostic instrumentation options for the WASP backend. + option(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT + "Instrument per-consumer Wasp epilogue cycles in the fused test" OFF) + if(CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + target_compile_definitions(tile_example_grouped_gemm_fused_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_EPI_CYCLE_COUNT) + endif() + option(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT + "Instrument one steady-state Wasp producer/consumer VMEM wait per wave" OFF) + if(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + target_compile_definitions(tile_example_grouped_gemm_fused_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) + endif() + option(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER + "Diagnostic: reverse Wasp CTA-to-tile mapping without changing tile work" OFF) + if(CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + target_compile_definitions(tile_example_grouped_gemm_fused_test + PRIVATE CK_TILE_GROUPED_GEMM_WASP_REVERSE_CTA_ORDER) + endif() endif() # Host-only Phase-1 selector/contract test. This target deliberately does not # link a grouped-GEMM provider, so adding it cannot alter a device launch. diff --git a/example_hcu/ck_tile/19_grouped_gemm/bench_grouped_gemm_fused_kernels.sh b/example_hcu/ck_tile/19_grouped_gemm/bench_grouped_gemm_fused_kernels.sh new file mode 100755 index 0000000..ec656f1 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/bench_grouped_gemm_fused_kernels.sh @@ -0,0 +1,229 @@ +#!/usr/bin/env bash +# Reproduce the NT/NN/TN commands recorded in grouped_gemm_fused_kernels_perf.md. +# +# This script runs the four columns sequentially on one HCU: +# 1. tuned hipBLASLt + HIPC implementation +# 2. CK original device-args implementation +# 3. CK fused implementation (WASP for NT, v7/v8 for NN/TN) +# +# The workload source of truth is hipc_kernels/bench_grouped_shapes.sh. Do not +# duplicate the 27 shape strings here: this script extracts its SHAPES array. + +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "$SCRIPT_DIR/../../.." && pwd) + +DEVICE=${DEVICE:-5} +WARMUP=${WARMUP:-5} +ITERS=${ITERS:-20} +SUITE=${SUITE:-all} +LAYOUT=${LAYOUT:-all} +LOG_DIR=${LOG_DIR:-"$SCRIPT_DIR/perf-logs/$(date +%Y%m%d-%H%M%S)"} + +SHAPES_SCRIPT="$SCRIPT_DIR/hipc_kernels/bench_grouped_shapes.sh" +HIPC_BIN=${HIPC_BIN:-"$SCRIPT_DIR/hipc_kernels/bf16_gemm"} +CK_ORIGINAL_BIN=${CK_ORIGINAL_BIN:-"$REPO_ROOT/build/bin/tile_example_grouped_gemm_device_args_benchmark"} +CK_FUSED_BIN=${CK_FUSED_BIN:-"$REPO_ROOT/build/bin/tile_example_grouped_gemm_fused_test"} +HIPBLASLT_TUNE_FILE=${HIPBLASLT_TUNE_FILE:-"/root/workspace/bf16_gemm/tune/hipblaslt_solutions.conf"} + +usage() { + cat <&2; usage >&2; exit 2 ;; + esac +done + +case "$SUITE" in + all|hipc|ck-original|ck-fused) ;; + *) echo "--suite must be all, hipc, ck-original, or ck-fused" >&2; exit 2 ;; +esac +case "$LAYOUT" in + all|NT|NN|TN) ;; + *) echo "--layout must be all, NT, NN, or TN" >&2; exit 2 ;; +esac + +require_executable() { + local path=$1 + if [[ ! -x "$path" ]]; then + echo "Required benchmark binary is not executable: $path" >&2 + exit 2 + fi +} + +if [[ ! -f "$SHAPES_SCRIPT" ]]; then + echo "Shape source does not exist: $SHAPES_SCRIPT" >&2 + exit 2 +fi + +if [[ "$SUITE" == all || "$SUITE" == hipc ]]; then + require_executable "$HIPC_BIN" + if [[ ! -f "$HIPBLASLT_TUNE_FILE" ]]; then + echo "Tuned hipBLASLt solution file does not exist: $HIPBLASLT_TUNE_FILE" >&2 + exit 2 + fi +fi +if [[ "$SUITE" == all || "$SUITE" == ck-original ]]; then + require_executable "$CK_ORIGINAL_BIN" +fi +if [[ "$SUITE" == all || "$SUITE" == ck-fused ]]; then + require_executable "$CK_FUSED_BIN" +fi + +mapfile -t SHAPES < <( + awk ' + /^SHAPES=\(/ { in_shapes = 1; next } + in_shapes && /^\)/ { exit } + in_shapes && /^[[:space:]]*"/ { + line = $0 + sub(/^[[:space:]]*"/, "", line) + sub(/".*$/, "", line) + print line + } + ' "$SHAPES_SCRIPT" +) + +if [[ ${#SHAPES[@]} -ne 27 ]]; then + echo "Expected 27 shapes from $SHAPES_SCRIPT, found ${#SHAPES[@]}" >&2 + exit 2 +fi + +mkdir -p "$LOG_DIR" + +run_logged() { + local log=$1 + shift + + printf '\n$' + printf ' %q' "$@" + printf '\n' + printf '$' + printf ' %q' "$@" + printf '\n' >>"$log" + "$@" 2>&1 | tee -a "$log" +} + +run_hipc() { + local layout=$1 tag=$2 n=$3 k=$4 lengths=$5 + local log="$LOG_DIR/hipc_hipblaslt.log" + + printf '\n--- %s %s N=%s K=%s M=%s ---\n' \ + "$layout" "$tag" "$n" "$k" "$lengths" | tee -a "$log" + BF16GEMM_HIPBLASLT_SOLUTION_FILE="$HIPBLASLT_TUNE_FILE" \ + CUDA_VISIBLE_DEVICES="$DEVICE" \ + run_logged "$log" "$HIPC_BIN" \ + "--layout=$layout" "--grouped-length=$lengths" "--n=$n" "--k=$k" \ + --bench=1 "--warmup-iters=$WARMUP" "--bench-iters=$ITERS" +} + +run_ck_original() { + local layout=$1 tag=$2 n=$3 k=$4 lengths=$5 + local total_m=0 part + local log="$LOG_DIR/ck_original.log" + + IFS=',' read -r -a group_lengths <<<"$lengths" + for part in "${group_lengths[@]}"; do total_m=$((total_m + part)); done + + # The device-args interface models TN group lengths as per-group K. The + # report therefore used this equal-FLOP proxy (M=4096) for TN, not an + # exact-shape CK-original comparison. + local m=1 + if [[ "$layout" == TN ]]; then m=4096; fi + + printf '\n--- %s %s N=%s K=%s lengths=%s ---\n' \ + "$layout" "$tag" "$n" "$k" "$lengths" | tee -a "$log" + CUDA_VISIBLE_DEVICES="$DEVICE" \ + run_logged "$log" "$CK_ORIGINAL_BIN" \ + "--m=$m" "--n=$n" "--k=$k" "--group-lengths=$lengths" \ + "--variable-capacity=$total_m" --dtype=bf16 "--layout=$layout" \ + --backend=ck --time=1 "--warmup=$WARMUP" "--samples=$ITERS" \ + --validate=0 --include-args-build=0 +} + +run_ck_fused() { + local layout=$1 tag=$2 n=$3 k=$4 lengths=$5 + local first_m group_count + + IFS=',' read -r -a group_lengths <<<"$lengths" + first_m=${group_lengths[0]} + group_count=${#group_lengths[@]} + + if [[ "$layout" == NT ]]; then + local log="$LOG_DIR/ck_fused_nt.log" + printf '\n--- %s %s N=%s K=%s M=%s ---\n' \ + "$layout" "$tag" "$n" "$k" "$lengths" | tee -a "$log" + CUDA_VISIBLE_DEVICES="$DEVICE" \ + run_logged "$log" "$CK_FUSED_BIN" \ + --layout=NT "--grouped-length=$lengths" "--n=$n" "--k=$k" \ + --bench=1 "--warmup-iters=$WARMUP" "--bench-iters=$ITERS" + return + fi + + local expected_instance=5000 + if [[ "$layout" == TN ]]; then expected_instance=5001; fi + local log="$LOG_DIR/ck_fused_${layout,,}.log" + printf '\n--- %s %s N=%s K=%s M=%s ---\n' \ + "$layout" "$tag" "$n" "$k" "$lengths" | tee -a "$log" + local log_start + log_start=$(wc -l <"$log") + CUDA_VISIBLE_DEVICES="$DEVICE" \ + run_logged "$log" "$CK_FUSED_BIN" \ + "--layout=$layout" "--m=$first_m" "--n=$n" "--k=$k" \ + "--groups=$group_count" "--group-lengths=$lengths" \ + "--iters=$ITERS" "--warmup=$WARMUP" --check=0 --ref=0 --cpu-ref=0 + + if ! tail -n "+$((log_start + 1))" "$log" | grep -q \ + "select layout=$layout instance_id=$expected_instance expected=$expected_instance"; then + echo "Expected fused $layout instance $expected_instance was not selected for $tag" >&2 + exit 1 + fi +} + +printf 'device=%s warmup=%s iters=%s suite=%s layout=%s\n' \ + "$DEVICE" "$WARMUP" "$ITERS" "$SUITE" "$LAYOUT" | tee "$LOG_DIR/run_configuration.txt" +printf 'shape_source=%s\nhipblaslt_tune=%s\n' \ + "$SHAPES_SCRIPT" "$HIPBLASLT_TUNE_FILE" | tee -a "$LOG_DIR/run_configuration.txt" + +for entry in "${SHAPES[@]}"; do + IFS='|' read -r layout tag n k lengths description <<<"$entry" + if [[ "$LAYOUT" != all && "$LAYOUT" != "$layout" ]]; then continue; fi + + if [[ "$SUITE" == all || "$SUITE" == hipc ]]; then + run_hipc "$layout" "$tag" "$n" "$k" "$lengths" + fi + if [[ "$SUITE" == all || "$SUITE" == ck-original ]]; then + run_ck_original "$layout" "$tag" "$n" "$k" "$lengths" + fi + if [[ "$SUITE" == all || "$SUITE" == ck-fused ]]; then + run_ck_fused "$layout" "$tag" "$n" "$k" "$lengths" + fi +done + +echo "Completed selected sweep. Raw logs: $LOG_DIR" diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_kernels_perf.md b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_kernels_perf.md new file mode 100644 index 0000000..2a528c1 --- /dev/null +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_kernels_perf.md @@ -0,0 +1,138 @@ +# Grouped-GEMM performance comparison + +Performance comparison for every NT/NN/TN workload in +[`hipc_kernels/bench_grouped_shapes.sh`](./hipc_kernels/bench_grouped_shapes.sh). +All values are TFLOPS. A positive percentage means higher throughput than the +tuned hipBLASLt baseline. + +| 测事项 | shape | hipblaslt tuned | hipc | CK original | CK fused | hipc vs hipblaslt tuned | CK original vs hipblaslt tuned | CK fused vs hipblaslt tuned | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| NT tp1-fc1-fwd | M=[8152,8152,8257,8257], N=4096, K=4096 | 324.0 | 343.1 | 250.1 | 346.1 | +5.9% | -22.8% | +6.8% | +| NT tp1-fc2-fwd | M=[8152,8152,8257,8257], N=2048, K=4096 | 280.5 | 343.0 | 247.1 | 346.4 | +22.3% | -11.9% | +23.5% | +| NT tp1-fc1-dgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | 320.2 | 339.1 | 247.7 | 342.2 | +5.9% | -22.6% | +6.9% | +| NT tp1-fc2-dgrad | M=[8100,8100,8091,8091], N=4096, K=2048 | 261.4 | 303.8 | 222.7 | 308.3 | +16.2% | -14.8% | +17.9% | +| NT tp1-wgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | 320.1 | 339.0 | 247.6 | 342.1 | +5.9% | -22.6% | +6.9% | +| NT tp4-real | M=[1923,2047,2179,2043], N=4096, K=4096 | 265.2 | 318.9 | 231.5 | 320.6 | +20.3% | -12.7% | +20.9% | +| NT tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | 286.2 | 313.9 | 231.3 | 316.4 | +9.7% | -19.2% | +10.6% | +| NT tp4-mean-2048-fc2 | M=[2048,2048,2048,2048], N=2048, K=4096 | 290.0 | 291.1 | 200.6 | 280.7 | +0.4% | -30.8% | -3.2% | +| NT tp4-step-2112 | M=[2112,2112,2112,2112], N=4096, K=4096 | 269.3 | 290.4 | 208.0 | 292.7 | +7.8% | -22.8% | +8.7% | +| NT tp4-collapse-1875 | M=[1875,1875,1875,1875], N=4096, K=4096 | 243.4 | 294.2 | 58.2 | 295.9 | +20.9% | -76.1% | +21.6% | +| NT tp4-dgrad-a-16384 | M=[16384], N=4096, K=4096 | 352.7 | 342.6 | 249.6 | 344.8 | -2.8% | -29.2% | -2.2% | +| NT ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 352.4 | 344.0 | 250.5 | 345.9 | -2.4% | -28.9% | -1.9% | +| NT ctrl-8group-4096 | M=[4096,4096,4096,4096,4096,4096,4096,4096], N=4096, K=4096 | 288.6 | 341.1 | 250.3 | 341.8 | +18.2% | -13.3% | +18.4% | +| NT ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | 286.3 | 315.7 | 230.2 | 318.1 | +10.3% | -19.6% | +11.1% | +| NN nn-tp1-fc1 | M=[8100,8100,8091,8091], N=4096, K=4096 | 335.1 | 371.1 | 167.0 | 396.3 | +10.8% | -50.2% | +18.3% | +| NN nn-tp1-fc1-k2048 | M=[8100,8100,8091,8091], N=4096, K=2048 | 277.1 | 331.8 | 156.5 | 353.7 | +19.7% | -43.5% | +27.7% | +| NN nn-tp1-fc2 | M=[8100,8100,8091,8091], N=2048, K=4096 | 287.4 | 326.1 | 166.0 | 396.6 | +13.5% | -42.3% | +38.0% | +| NN nn-tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | 301.9 | 327.4 | 156.1 | 365.7 | +8.5% | -48.3% | +21.1% | +| NN nn-ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | 302.0 | 327.3 | 155.9 | 368.6 | +8.4% | -48.4% | +22.1% | +| NN nn-ctrl-16group | M=[2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048], N=4096, K=4096 | 301.9 | 327.6 | 168.7 | 395.2 | +8.5% | -44.1% | +30.9% | +| NN nn-ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 372.9 | 407.2 | 168.8 | 403.2 | +9.2% | -54.7% | +8.1% | +| TN tn-tp1-fc1-wgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | N/A | 345.4 | 348.7* | 402.7 | N/A | N/A* | N/A | +| TN tn-tp1-fc2-wgrad | M=[8100,8100,8091,8091], N=2048, K=4096 | N/A | 325.6 | 321.2* | 400.0 | N/A | N/A* | N/A | +| TN tn-tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | N/A | 329.7 | 298.6* | 372.3 | N/A | N/A* | N/A | +| TN tn-ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | N/A | 341.5 | 273.6* | 371.7 | N/A | N/A* | N/A | +| TN tn-ctrl-16group | M=[2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048], N=4096, K=4096 | N/A | 329.7 | 301.9* | 405.3 | N/A | N/A* | N/A | +| TN tn-ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 394.2 | 401.1 | 296.0* | 408.1 | +1.8% | N/A* | +3.5% | + +## Measurement details + +- Device: HCU 5 (80 CUs), selected with `CUDA_VISIBLE_DEVICES=5`; BF16. +- All measurements use 5 warmup iterations and 20 timed iterations. +- Reproduce the full sequential sweep with + `./bench_grouped_gemm_fused_kernels.sh`. It extracts the same 27-shape list + from `hipc_kernels/bench_grouped_shapes.sh`, writes raw logs beneath + `perf-logs/`, and runs tuned hipBLASLt/HIPC, CK original, and CK fused in the + same order. Use `--suite=hipc|ck-original|ck-fused` or + `--layout=NT|NN|TN` to run a subset. +- HIPC and hipBLASLt were measured by `hipc_kernels/bf16_gemm` with + `BF16GEMM_HIPBLASLT_SOLUTION_FILE=/root/workspace/bf16_gemm/tune/hipblaslt_solutions.conf`. + The tune table selected solution 4095 for the NT baseline and 4396 for NN; + grouped TN has no hipBLASLt baseline because that API cannot express its + per-group leading dimension. The single-group TN baseline used solution 2692. +- CK original used `tile_example_grouped_gemm_device_args_benchmark` with + `--backend=ck --validate=0 --include-args-build=0`, so the reported timing + excludes the device-args builder kernel. Its values are the reported median + across 20 samples. +- CK fused workloads use `tile_example_grouped_gemm_fused_test`; `--layout` + selects WASP for NT and the v7/v8 test for NN/TN. NN and TN verified selected + instances 5000 and 5001, respectively. The sweep is performance-only; use the functional + tests separately for correctness validation. +- Large outliers in the first HIPC pass were repeated three times on HCU 5. + The table uses the median of those repeats for NT `tp1-wgrad`, NT `tp4-real`, + TN `tn-ctrl-16group`, and TN `tn-ctrl-1group-sum32818`; other HIPC rows use + the initial sweep value. +- Percentage formula: `(kernel TFLOPS / hipBLASLt tuned TFLOPS - 1) * 100%`. + TN grouped workloads have no hipBLASLt denominator, so their relative values + are `N/A`. +- `*` CK original TN values are an equal-FLOP proxy, not exact-shape results: + the device-args CK interface represents TN `group-lengths` as per-group K, + whereas this HIPC shape list uses them as per-group M. For these proxy runs, + CK original used M=4096 and the listed values as per-group K. The absolute + TFLOPS are shown for reference only; no CK-original-vs-hipBLASLt ratio is + reported for TN because those operands are not dimension-equivalent. + +## Raw logs + +- HIPC and tuned hipBLASLt sweep: `/tmp/grouped_gemm_full_hipc_gpu5_20261014.log` +- CK original sweep: `/tmp/grouped_gemm_full_ck_original_gpu5_20261014.log` +- CK fused NT sweep: `/tmp/grouped_gemm_full_ck_fused_nt_gpu5_20261014.log` +- CK fused NN/TN sweep: `/tmp/grouped_gemm_full_ck_fused_nntn_gpu5_20261014.log` +- HIPC anomaly repeats: `/tmp/grouped_gemm_hipc_outlier_repeats_gpu5_20261014.log` + +## Refresh — 2026-10-15 + +This appended section records the 2026-10-15 rerun; the original table above +is preserved unchanged. Device, dtype, warmup, and timing parameters remain +HCU 5, BF16, 5 warmup iterations, and 20 timed iterations/samples. + +| 测事项 | shape | hipblaslt tuned | hipc | CK original | CK fused | hipc vs hipblaslt tuned | CK original vs hipblaslt tuned | CK fused vs hipblaslt tuned | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| NT tp1-fc1-fwd | M=[8152,8152,8257,8257], N=4096, K=4096 | 324.0 | 343.3 | 249.9 | 346.3 | +6.0% | -22.9% | +6.9% | +| NT tp1-fc2-fwd | M=[8152,8152,8257,8257], N=2048, K=4096 | 280.5 | 343.0 | 247.0 | 346.9 | +22.3% | -11.9% | +23.7% | +| NT tp1-fc1-dgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | 320.1 | 339.1 | 247.7 | 342.6 | +5.9% | -22.6% | +7.0% | +| NT tp1-fc2-dgrad | M=[8100,8100,8091,8091], N=4096, K=2048 | 261.7 | 303.7 | 222.4 | 311.2 | +16.0% | -15.0% | +18.9% | +| NT tp1-wgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | 320.2 | 339.1 | 247.2 | 342.6 | +5.9% | -22.8% | +7.0% | +| NT tp4-real | M=[1923,2047,2179,2043], N=4096, K=4096 | 265.2 | 318.9 | 231.6 | 323.6 | +20.3% | -12.7% | +22.0% | +| NT tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | 286.1 | 313.9 | 231.4 | 318.5 | +9.7% | -19.1% | +11.3% | +| NT tp4-mean-2048-fc2 | M=[2048,2048,2048,2048], N=2048, K=4096 | 290.0 | 291.1 | 200.5 | 287.3 | +0.4% | -30.9% | -1.0% | +| NT tp4-step-2112 | M=[2112,2112,2112,2112], N=4096, K=4096 | 269.5 | 290.3 | 207.5 | 292.1 | +7.7% | -23.0% | +8.4% | +| NT tp4-collapse-1875 | M=[1875,1875,1875,1875], N=4096, K=4096 | 243.3 | 294.3 | 57.9 | 295.7 | +21.0% | -76.2% | +21.6% | +| NT tp4-dgrad-a-16384 | M=[16384], N=4096, K=4096 | 350.8 | 342.6 | 249.1 | 349.2 | -2.3% | -29.0% | -0.5% | +| NT ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 352.2 | 344.0 | 250.1 | 349.6 | -2.3% | -29.0% | -0.7% | +| NT ctrl-8group-4096 | M=[4096,4096,4096,4096,4096,4096,4096,4096], N=4096, K=4096 | 288.5 | 341.0 | 250.2 | 343.8 | +18.2% | -13.3% | +19.2% | +| NT ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | 286.1 | 315.7 | 230.3 | 322.9 | +10.4% | -19.5% | +12.9% | +| NN nn-tp1-fc1 | M=[8100,8100,8091,8091], N=4096, K=4096 | 335.1 | 371.1 | 166.9 | 396.4 | +10.7% | -50.2% | +18.3% | +| NN nn-tp1-fc1-k2048 | M=[8100,8100,8091,8091], N=4096, K=2048 | 277.1 | 331.9 | 156.4 | 353.4 | +19.8% | -43.5% | +27.5% | +| NN nn-tp1-fc2 | M=[8100,8100,8091,8091], N=2048, K=4096 | 287.6 | 326.1 | 165.9 | 396.7 | +13.4% | -42.3% | +37.9% | +| NN nn-tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | 301.7 | 327.4 | 156.1 | 365.4 | +8.5% | -48.3% | +21.1% | +| NN nn-ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | 301.7 | 327.1 | 155.8 | 368.3 | +8.4% | -48.4% | +22.1% | +| NN nn-ctrl-16group | M=[2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048], N=4096, K=4096 | 301.9 | 327.6 | 168.7 | 395.5 | +8.5% | -44.1% | +31.0% | +| NN nn-ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 372.4 | 407.2 | 168.8 | 376.8† | +9.3% | -54.7% | N/A† | +| TN tn-tp1-fc1-wgrad | M=[8100,8100,8091,8091], N=4096, K=4096 | N/A | 345.6 | 348.6* | 402.4 | N/A | N/A* | N/A | +| TN tn-tp1-fc2-wgrad | M=[8100,8100,8091,8091], N=2048, K=4096 | N/A | 325.7 | 321.0* | 399.4 | N/A | N/A* | N/A | +| TN tn-tp4-mean-2048 | M=[2048,2048,2048,2048], N=4096, K=4096 | N/A | 329.7 | 298.6* | 371.1 | N/A | N/A* | N/A | +| TN tn-ctrl-skew-10x | M=[5120,1536,1024,512], N=4096, K=4096 | N/A | 341.4 | 273.5* | 371.4 | N/A | N/A* | N/A | +| TN tn-ctrl-16group | M=[2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048,2048], N=4096, K=4096 | N/A | 329.8 | 301.7* | 405.5 | N/A | N/A* | N/A | +| TN tn-ctrl-1group-sum32818 | M=[32818], N=4096, K=4096 | 394.2 | 400.7 | 296.0* | 407.9 | +1.7% | N/A* | +3.5% | + +### Refresh notes + +- The initial refresh measurements for NT `tp4-mean-2048` and TN + `tn-ctrl-skew-10x` were severe low-throughput outliers. Each implementation + for those shapes was run three additional times; the table uses the median + of those stable repeats. The raw first-pass records remain in the log. +- `†` The first-pass CK fused result for NN + `nn-ctrl-1group-sum32818` was 376.8 TFLOPS, but three immediate repeats + varied from 54.3 to 71.1 TFLOPS (median 59.8). The measurement is therefore + unstable and its CK-fused-vs-hipBLASLt ratio is intentionally not reported. + It should be remeasured on an idle device before drawing a comparison. +- `*` has the same equal-FLOP TN proxy meaning described above; it is not an + exact-shape CK-original result. +- Refresh logs: + `/tmp/grouped_gemm_fused_refresh_20261015/hipc_hipblaslt.log`, + `/tmp/grouped_gemm_fused_refresh_20261015/ck_original.log`, + `/tmp/grouped_gemm_fused_refresh_20261015/ck_fused_nt.log`, + `/tmp/grouped_gemm_fused_refresh_20261015/ck_fused_nn.log`, + `/tmp/grouped_gemm_fused_refresh_20261015/ck_fused_tn.log`, and + `/tmp/grouped_gemm_fused_refresh_20261015/outlier_repeats.log`. diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_test.cpp similarity index 52% rename from example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp rename to example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_test.cpp index 3138dc7..9a86ff6 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_test.cpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_fused_test.cpp @@ -1,53 +1,25 @@ // SPDX-License-Identifier: MIT // Copyright (c) 2026 Hygon Information Technologies Co., Ltd. // -// Correctness + performance driver for the WASP (warp-specialized -// producer / consumer) NT grouped GEMM pipeline. -// -// grouped_gemm_wasp_geometry_smoke.cpp guards the one structural assumption of -// the design (a 12 wave workgroup carrying a 1 x 8 consumer warp grid). THIS -// file answers the two questions that guard cannot: does it compute the right -// numbers, and how fast is it. -// -// Design notes -// ------------ -// * kargs come from the FRAMEWORK's own host-side builder, -// `GroupedGemmKernel<...>::MakeKargs(std::vector)`, so -// the descriptor the kernel sees is field-for-field what the device-args ABI -// would produce. Nothing is hand-filled, so a harness bug cannot hide behind -// a hand-copied layout. -// * A, B and C are laid out the way the ABI lays them out for NT: -// A [sum(M_g), K] row major, stride_A = K (group g at prefix_M[g]*K) -// B [G, N, K] row major, stride_B = K (group g at g*N*K) -// C [sum(M_g), N] row major, stride_E = N (group g at prefix_M[g]*N) -// Each group's a_ptr/e_ptr points INTO the shared buffer, which is exactly -// what lets the kernel recover the group table from the kargs alone. -// * The oracle is an independent CPU reference (fp32 accumulate), not another -// GPU kernel, so the check still means something when every GPU path is -// wrong. bf16 products are exact in fp32, so only the summation ORDER -// differs between the reference and the kernel: the tolerance can be tight. -// * `--ref=1` (default) additionally runs the same kargs through an existing CK -// configuration. That gives a triangle of evidence: if WASP and the CK -// reference both disagree with the CPU reference, the harness is wrong; if -// only WASP disagrees, it is the kernel. -// -// Current port limitations this driver enforces -// -------------------------------------------- -// The port does not yet mask A rows past M, B columns past N, or a K tail, so -// only aligned shapes are legal: per-group M % 256 == 0, N % 256 == 0, -// K % 64 == 0. Unaligned shapes are reported as SKIPPED (not FAIL) unless -// --allow-unaligned is given, which is what you want when debugging the masking -// itself. +// Unified correctness and performance driver for the v7/v8 NN/TN instances +// and the WASP NT grouped-GEMM pipeline. --layout selects the implementation; +// the remaining test options are forwarded to that implementation. // // Usage: -// grouped_gemm_wasp_test [--m=...] [--n=...] [--k=...] [--groups=...] -// [--iters=N] [--warmup=N] [--ref=0|1] -// [--allow-unaligned] [--quiet] +// grouped_gemm_fused_test --layout=NN|TN|NT [shape and test options] #include "instances/grouped_gemm_impl.hpp" - #include +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/grouped_gemm_device_args.h" +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" + +#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) +#include +#include +#endif #if defined(CK_TILE_GROUPED_GEMM_WASP_HAS_HIPBLASLT) #include #include @@ -64,7 +36,1388 @@ #include #include -namespace { +namespace v7_backend { + +using Bf16 = ck_tile::bf16_t; + +constexpr std::uint32_t kV7NnInstance = 5000; +constexpr std::uint32_t kV8TnInstance = 5001; + +// Slack the TN A panel needs past the last group's columns (see the header). +constexpr std::size_t kAtailPad = 1024; + +struct Options +{ + std::string layout = "NN"; + int m = 512; + int n = 512; + int k = 256; + int groups = 2; + int iters = 20; + int warmup = 3; + bool check = true; + bool ref = true; + // The independent fp32 CPU reference the checked runs compare against (bf16 + // inputs are exact in fp32 -- only the summation order differs). --check=0 + // skips it for perf runs on production shapes. + bool cpu_ref = true; + // Ragged group lengths: the shapes the v7 harness benchmarks + // (--grouped-length=8100,8100,8091,8091). Empty => `groups` copies of `m`. + std::vector group_lengths; + bool dump = false; + // Grid override. 0 => the device's CU count (the production value). + int num_cu = 0; + // Device-side functional reference: run the same kargs through hipBLASLt + // (one plain GEMM per group) and compare on the GPU -- no CPU pass over the + // output, so production shapes check in seconds. Built only when the + // hipBLASLt install is present (see CMakeLists CK_TILE_V7_TEST_HAS_HIPBLASLT). + bool hipblaslt_ref = false; + // Performance reference: time hipBLASLt (plans built once, bare Matmul in + // the loop) with the same warmup/iters as the v7 timing, and print + // blas_ms / blas_tf next to v7_ms / v7_tf. + bool hipblaslt_perf = false; +}; + +std::size_t align_to(std::size_t value, std::size_t alignment) +{ + return (value + alignment - 1) / alignment * alignment; +} + +// Every hip* call below whose result carries [[nodiscard]] goes through this, so +// the driver both stays -Wunused-result clean and reports the first failure +// instead of running on with garbage. +void hip_check(hipError_t status, const char* what) +{ + if(status != hipSuccess) + { + std::fprintf(stderr, "HIP error %d (%s) at %s\n", static_cast(status), + hipGetErrorName(status), what); + std::exit(1); + } +} + +// Small multiples of 0.125: exact in bf16 (8 mantissa bits), so the only +// difference between the CPU reference and the kernel is the summation order. +// Both indices contribute, so no lane pattern of the kernels cancels a term. +float a_value(int row, int kk) +{ + return static_cast((row % 7) - 3) * 0.5f + + static_cast(kk % 3) * 0.25f; +} +float b_value(int kk, int col) +{ + return static_cast((col % 5) - 2) * 0.25f + + static_cast(kk % 4) * 0.125f; +} + +struct Case +{ + // Per-group row counts (ragged) plus their sum. `m` stays the first group's + // length for the single-group paths; every buffer is sized from group_m. + std::vector group_m; + int total_m = 0; + bool tn = false; + int m = 0; + int n = 0; + int k = 0; + int g = 0; + std::vector strides_a; + std::vector strides_b; + std::vector a_elements; + std::vector b_elements; + std::vector> host_a; + std::vector> host_b; + std::vector> host_d_ref; +}; + +int parse_int_arg(const char* text, const char* name) +{ + char* end = nullptr; + const long value = std::strtol(text, &end, 10); + if(end == text || value < 0) + { + std::fprintf(stderr, "bad --%s=%s\n", name, text); + std::exit(2); + } + return static_cast(value); +} + +void parse_args(int argc, char* argv[], Options& options) +{ + for(int i = 1; i < argc; ++i) + { + const std::string arg = argv[i]; + const auto value_of = [&](const char* name) -> const char* { + const std::size_t at = arg.find('='); + if(at == std::string::npos) + { + std::fprintf(stderr, "%s expects --%s=VALUE\n", arg.c_str(), name); + std::exit(2); + } + return argv[i] + at + 1; + }; + if(arg.rfind("--layout=", 0) == 0) + { + options.layout = value_of("layout"); + } + else if(arg.rfind("--m=", 0) == 0) + { + options.m = parse_int_arg(value_of("m"), "m"); + } + else if(arg.rfind("--n=", 0) == 0) + { + options.n = parse_int_arg(value_of("n"), "n"); + } + else if(arg.rfind("--k=", 0) == 0) + { + options.k = parse_int_arg(value_of("k"), "k"); + } + else if(arg.rfind("--groups=", 0) == 0) + { + options.groups = parse_int_arg(value_of("groups"), "groups"); + } + else if(arg.rfind("--iters=", 0) == 0 || + arg.rfind("--bench-iters=", 0) == 0) + { + const char* name = + arg.rfind("--iters=", 0) == 0 ? "iters" : "bench-iters"; + options.iters = parse_int_arg(value_of(name), name); + } + else if(arg.rfind("--warmup=", 0) == 0 || + arg.rfind("--warmup-iters=", 0) == 0) + { + const char* name = + arg.rfind("--warmup=", 0) == 0 ? "warmup" : "warmup-iters"; + options.warmup = parse_int_arg(value_of(name), name); + } + else if(arg.rfind("--check=", 0) == 0) + { + options.check = parse_int_arg(value_of("check"), "check") != 0; + } + else if(arg.rfind("--ref=", 0) == 0) + { + options.ref = parse_int_arg(value_of("ref"), "ref") != 0; + } + else if(arg.rfind("--group-lengths=", 0) == 0 || + arg.rfind("--grouped-length=", 0) == 0) + { + options.group_lengths.clear(); + const std::string list = + value_of(arg.rfind("--group-lengths=", 0) == 0 ? "group-lengths" + : "grouped-length"); + std::size_t start = 0; + while(start <= list.size()) + { + const std::size_t comma = list.find(',', start); + const std::string item = list.substr(start, comma - start); + if(!item.empty()) + { + options.group_lengths.push_back(parse_int_arg(item.c_str(), "group-lengths")); + } + if(comma == std::string::npos) break; + start = comma + 1; + } + } + else if(arg.rfind("--cpu-ref=", 0) == 0) + { + options.cpu_ref = parse_int_arg(value_of("cpu-ref"), "cpu-ref") != 0; + } + else if(arg.rfind("--dump=", 0) == 0) + { + options.dump = parse_int_arg(value_of("dump"), "dump") != 0; + } + else if(arg.rfind("--num-cu=", 0) == 0) + { + options.num_cu = parse_int_arg(value_of("num-cu"), "num-cu"); + } + else if(arg.rfind("--hipblaslt-ref=", 0) == 0) + { + options.hipblaslt_ref = parse_int_arg(value_of("hipblaslt-ref"), "hipblaslt-ref") != 0; + } + else if(arg.rfind("--hipblaslt-perf=", 0) == 0) + { + options.hipblaslt_perf = parse_int_arg(value_of("hipblaslt-perf"), "hipblaslt-perf") != 0; + } + else + { + std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); + std::exit(2); + } + } + if(options.layout != "NN" && options.layout != "TN") + { + std::fprintf(stderr, "--layout must be NN or TN\n"); + std::exit(2); + } +} + +void build_case(const Options& options, Case& c) +{ + c.tn = options.layout == "TN"; + c.n = options.n; + c.k = options.k; + c.group_m = options.group_lengths.empty() + ? std::vector(static_cast(options.groups), options.m) + : options.group_lengths; + c.g = static_cast(c.group_m.size()); + c.m = c.group_m.front(); + c.total_m = 0; + for(int g = 0; g < c.g; ++g) c.total_m += c.group_m[g]; + // TN's A panel is [K, sum(M_g)]: it is addressed across the whole group set, + // so its pitch is the summed M (padded up to a whole tile, which is also the + // 256-element slack the panel fetch needs past the last group's columns). + const int a_pitch_tn = (c.total_m + 255) / 256 * 256; + + c.strides_a.resize(c.g); + c.strides_b.resize(c.g); + c.a_elements.resize(c.g); + c.b_elements.resize(c.g); + c.host_a.resize(c.g); + c.host_b.resize(c.g); + c.host_d_ref.resize(c.g); + + for(int g = 0; g < c.g; ++g) + { + // NN: A row-major [M_g, K]. TN: A is the [K, M_g] panel of the + // column-major [M_g, K] view, so its pitch is the M extent (the same + // choice the device-args smoke makes for its TN cases). + const int stride_a = c.tn ? a_pitch_tn : c.k; + const int stride_b = c.n; + c.strides_a[g] = stride_a; + c.strides_b[g] = stride_b; + c.a_elements[g] = + (c.tn ? c.k : c.group_m[g]) * stride_a + (c.tn ? static_cast(kAtailPad) : 0); + c.b_elements[g] = c.k * stride_b; + + c.host_a[g].assign(static_cast(c.a_elements[g]), + ck_tile::type_convert(0.0f)); + c.host_b[g].assign(static_cast(c.b_elements[g]), + ck_tile::type_convert(0.0f)); + c.host_d_ref[g].assign(static_cast(c.group_m[g]) * c.n, 0.0f); + + for(int row = 0; row < c.group_m[g]; ++row) + { + for(int kk = 0; kk < c.k; ++kk) + { + const std::size_t at = + c.tn ? static_cast(kk) * stride_a + row + : static_cast(row) * stride_a + kk; + c.host_a[g][at] = ck_tile::type_convert(a_value(row, kk)); + } + } + for(int kk = 0; kk < c.k; ++kk) + { + for(int col = 0; col < c.n; ++col) + { + c.host_b[g][static_cast(kk) * stride_b + col] = + ck_tile::type_convert(b_value(kk, col)); + } + } + } +} + +void cpu_reference(Case& c) +{ + for(int g = 0; g < c.g; ++g) + { + const int stride_a = c.strides_a[g]; + const int stride_b = c.strides_b[g]; + for(int col = 0; col < c.n; ++col) + { + for(int row = 0; row < c.m; ++row) + { + float acc = 0.0f; + for(int kk = 0; kk < c.k; ++kk) + { + const std::size_t at = + c.tn ? static_cast(kk) * stride_a + row + : static_cast(row) * stride_a + kk; + const float a = ck_tile::type_convert(c.host_a[g][at]); + const float b = ck_tile::type_convert( + c.host_b[g][static_cast(kk) * stride_b + col]); + acc += a * b; + } + c.host_d_ref[g][static_cast(row) * c.n + col] = acc; + } + } + } +} + +#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) +// --------------------------------------------------------------------------- +// Device-side functional reference: one plain hipBLASLt GEMM per group, then a +// device-side element compare -- the output never touches the CPU, so +// production shapes verify in seconds. +// +// The col-major formulation (same construction bf16_gemm_hipblaslt.cpp uses): +// D_row[M_g, N] = A . B <=> C_col[N, M_g] = B_view[N, K] . A_view[K, M_g] +// with B_view the [N, K] col-major view of the [K, N] row-major B (ld = N), +// and A_view the operand buffer seen as +// NN: [K, M_g] col-major, ld = K (the [M_g, K] row-major A) -> opB = N +// TN: [M_g, K] col-major, ld = pitch (the [K, pitch] A panel) -> opB = T +// The panel leading dimension is the one thing the bf16_gemm_hipblaslt wrapper +// cannot express for the grouped TN operand (it hardcodes ldb = k); here it is +// passed explicitly. +// --------------------------------------------------------------------------- +void hipblaslt_check_status(hipblasStatus_t status, const char* what) +{ + if(status != HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "hipBLASLt error %d at %s\n", + static_cast(status), + what); + std::exit(1); + } +} + +// --------------------------------------------------------------------------- +// Tuned-solution table -- same file and semantics as bf16_gemm_hipblaslt.cpp +// (BF16GEMM_HIPBLASLT_SOLUTION_FILE selects it, BF16GEMM_NO_TUNING=1 disables; +// lines "m,n,k,transA,transB,solution_index", '*' = wildcard, '#' comments). +// The keys live in THIS test's hipblasLt parameter space (m = caller N, +// n = M_g, k = K; transA is always N here, transB = T for the TN panel form), +// which is the same space the wrapper tunes in. Resolution: best-scoring +// entry -> hipblaslt_ext getAlgosFromIndex + matmulIsAlgoSupported -> else the +// stock heuristic, warning once per rejected key. +// --------------------------------------------------------------------------- +struct TunedSolution +{ + long long m = -1, n = -1, k = -1; // -1 = wildcard + char ta = '?', tb = '?'; + int index = -1; +}; + +static long long tuned_num(const std::string& s) +{ + return (s == "*") ? -1 : std::strtoll(s.c_str(), nullptr, 10); +} + +static bool tuned_hit(long long key, long long want) +{ + return key < 0 || key == want; +} + +static bool tuned_hit(char key, char want) +{ + return key == '*' || key == '?' || key == want; +} + +static const std::vector& tuned_table() +{ + static const std::vector table = [] { + std::vector t; + const char* disable = std::getenv("BF16GEMM_NO_TUNING"); + if(disable != nullptr && std::atoi(disable) != 0) + { + return t; + } + const char* path = std::getenv("BF16GEMM_HIPBLASLT_SOLUTION_FILE"); + if(path == nullptr || *path == '\0') + { + return t; + } + FILE* f = std::fopen(path, "r"); + if(f == nullptr) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: cannot open '%s', using the stock heuristic\n", + path); + return t; + } + char line[512]; + while(std::fgets(line, sizeof(line), f) != nullptr) + { + char* p = line; + while(*p == ' ' || *p == '\t') + { + ++p; + } + if(*p == '#' || *p == '\n' || *p == '\r' || *p == '\0') + { + continue; + } + std::string fields[6]; + int nf = 0; + for(char* tok = std::strtok(p, ", \t\r\n"); tok != nullptr && nf < 6; + tok = std::strtok(nullptr, ", \t\r\n")) + { + fields[nf++] = tok; + } + if(nf < 6) + { + std::fprintf(stderr, "[hipblasLt] WARNING: malformed line ignored: %s", line); + continue; + } + TunedSolution e; + e.m = tuned_num(fields[0]); + e.n = tuned_num(fields[1]); + e.k = tuned_num(fields[2]); + e.ta = static_cast(std::toupper(static_cast(fields[3][0]))); + e.tb = static_cast(std::toupper(static_cast(fields[4][0]))); + e.index = static_cast(std::strtol(fields[5].c_str(), nullptr, 10)); + t.push_back(e); + } + std::fclose(f); + std::fprintf(stderr, + "[hipblasLt] tuned-solution table: %zu entries from %s\n", + t.size(), + path); + return t; + }(); + return table; +} + +static bool tuned_pick(long long m, long long n, long long k, char ta, char tb, int& index) +{ + int best_score = 6; + bool found = false; + for(const TunedSolution& e : tuned_table()) + { + if(!(tuned_hit(e.m, m) && tuned_hit(e.n, n) && tuned_hit(e.k, k) && + tuned_hit(e.ta, ta) && tuned_hit(e.tb, tb))) + { + continue; + } + const int score = (e.m < 0 ? 1 : 0) + (e.n < 0 ? 1 : 0) + (e.k < 0 ? 1 : 0) + + ((e.ta == '*' || e.ta == '?') ? 1 : 0) + + ((e.tb == '*' || e.tb == '?') ? 1 : 0); + if(score < best_score) + { + best_score = score; + index = e.index; + found = true; + if(score == 0) + { + break; // exact key: cannot do better + } + } + } + return found; +} + +// D = 1.0 * op(B) . op(A) + 0.0, all bf16 in / bf16 out, fp32 compute. +// Descriptor + algorithm selection is split from the execution: the perf +// reference builds each group's plan once and times bare hipblasLtMatmul +// calls, exactly what a production consumer would do. +struct HipblasltPlan +{ + hipblasLtMatmulDesc_t desc = nullptr; + hipblasLtMatrixLayout_t a_layout = nullptr; // the B operand, col-major [N, K], ld = N + hipblasLtMatrixLayout_t b_layout = nullptr; // the A operand + hipblasLtMatrixLayout_t c_layout = nullptr; // D, col-major [N, M_g], ld = N + hipblasLtMatmulAlgo_t algo{}; +}; + +bool hipblaslt_make_plan(HipblasltPlan& plan, + hipblasLtHandle_t handle, + void* workspace, + std::size_t workspace_size, + int m_g, + int n, + int k, + int ld_a, + bool tn) +{ + (void)workspace; // the workspace pointer only matters at Matmul time + hipblasLtMatmulPreference_t pref = nullptr; + + hipblasOperation_t op_a = HIPBLAS_OP_N; + hipblasOperation_t op_b = tn ? HIPBLAS_OP_T : HIPBLAS_OP_N; + + hipblaslt_check_status( + hipblasLtMatmulDescCreate(&plan.desc, HIPBLAS_COMPUTE_32F, HIP_R_32F), + "hipblasLtMatmulDescCreate"); + hipblaslt_check_status( + hipblasLtMatmulDescSetAttribute(plan.desc, + HIPBLASLT_MATMUL_DESC_TRANSA, + &op_a, + sizeof(op_a)), + "SetAttribute(TRANSA)"); + hipblaslt_check_status( + hipblasLtMatmulDescSetAttribute(plan.desc, + HIPBLASLT_MATMUL_DESC_TRANSB, + &op_b, + sizeof(op_b)), + "SetAttribute(TRANSB)"); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.a_layout, + HIP_R_16BF, + static_cast(n), + static_cast(k), + static_cast(n)), + "LayoutCreate(B view)"); + // opT stores the [K, M_g] logical matrix through the [M_g, K] panel: ld is + // the panel pitch, not K. + const std::size_t b_ld = tn ? static_cast(ld_a) + : static_cast(k); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.b_layout, + HIP_R_16BF, + static_cast(k), + static_cast(m_g), + b_ld), + "LayoutCreate(A view)"); + hipblaslt_check_status( + hipblasLtMatrixLayoutCreate(&plan.c_layout, + HIP_R_16BF, + static_cast(n), + static_cast(m_g), + static_cast(n)), + "LayoutCreate(D)"); + + hipblaslt_check_status( + hipblasLtMatmulPreferenceCreate(&pref), + "hipblasLtMatmulPreferenceCreate"); + hipblaslt_check_status( + hipblasLtMatmulPreferenceSetAttribute( + pref, + HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, + sizeof(workspace_size)), + "PreferenceSetAttribute(WORKSPACE)"); + + // 1) tuned table (best-scoring entry, resolved through the ext API); + // 2) stock heuristic. A table entry the library rejects only warns. + const char ta_ch = (op_a == HIPBLAS_OP_N) ? 'N' : 'T'; + const char tb_ch = (op_b == HIPBLAS_OP_N) ? 'N' : 'T'; + int tuned_index = -1; + if(tuned_pick(n, m_g, k, ta_ch, tb_ch, tuned_index)) + { + std::vector resolved(1); + std::vector algo_index(1, tuned_index); + const float probe_alpha = 1.0f; + const float probe_beta = 0.0f; + std::size_t required = 0; + if(hipblaslt_ext::getAlgosFromIndex(handle, algo_index, resolved) != + HIPBLAS_STATUS_SUCCESS || + resolved[0].state != HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d does not exist; " + "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", + tuned_index, + n, + m_g, + k, + ta_ch, + tb_ch); + } + else if(hipblaslt_ext::matmulIsAlgoSupported(handle, + plan.desc, + &probe_alpha, + plan.a_layout, + plan.b_layout, + &probe_beta, + plan.c_layout, + plan.c_layout, + resolved[0].algo, + required) != + HIPBLAS_STATUS_SUCCESS) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d unsupported here; " + "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", + tuned_index, + n, + m_g, + k, + ta_ch, + tb_ch); + } + else if(required > workspace_size) + { + std::fprintf(stderr, + "[hipblasLt] WARNING: solution index %d needs %zu bytes " + "workspace (have %zu); stock heuristic\n", + tuned_index, + required, + workspace_size); + } + else + { + plan.algo = resolved[0].algo; + std::fprintf(stderr, + "[hipblasLt] tuned: m=%d n=%d k=%d trans=%c%c -> index=%d\n", + n, + m_g, + k, + ta_ch, + tb_ch, + tuned_index); + hipblasLtMatmulPreferenceDestroy(pref); + return true; + } + } + + hipblasLtMatmulHeuristicResult_t algo{}; + int algo_count = 0; + hipblaslt_check_status( + hipblasLtMatmulAlgoGetHeuristic(handle, + plan.desc, + plan.a_layout, + plan.b_layout, + plan.c_layout, + plan.c_layout, + pref, + 1, + &algo, + &algo_count), + "hipblasLtMatmulAlgoGetHeuristic"); + const bool ok = algo_count > 0; + if(ok) + { + plan.algo = algo.algo; + } + hipblasLtMatmulPreferenceDestroy(pref); + return ok; +} + +void hipblaslt_run_plan(const HipblasltPlan& plan, + hipblasLtHandle_t handle, + void* workspace, + std::size_t workspace_size, + hipStream_t stream, + const void* b_dev, + const void* a_dev, + void* d_dev) +{ + const float alpha = 1.0f; + const float beta = 0.0f; + hipblaslt_check_status( + hipblasLtMatmul(handle, + plan.desc, + &alpha, + b_dev, + plan.a_layout, + a_dev, + plan.b_layout, + &beta, + d_dev, + plan.c_layout, + d_dev, + plan.c_layout, + &plan.algo, + workspace, + workspace_size, + stream), + "hipblasLtMatmul"); +} + +void hipblaslt_destroy_plan(HipblasltPlan& plan) +{ + hipblasLtMatrixLayoutDestroy(plan.c_layout); + hipblasLtMatrixLayoutDestroy(plan.b_layout); + hipblasLtMatrixLayoutDestroy(plan.a_layout); + hipblasLtMatmulDescDestroy(plan.desc); + plan = HipblasltPlan{}; +} + +// Element compare on the device: bad = |got - ref| > atol + rtol*|ref|, +// max_err tracked via the non-negative-float bit pattern (uint order == float +// order for non-negative values). +__global__ void blas_compare_kernel(const Bf16* __restrict__ got, + const Bf16* __restrict__ ref, + long long count, + float atol, + float rtol, + unsigned long long* bad, + unsigned int* max_bits) +{ + const long long i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if(i >= count) + { + return; + } + const float g = ck_tile::type_convert(got[i]); + const float r = ck_tile::type_convert(ref[i]); + const float d = fabsf(g - r); + if(d > atol + rtol * fabsf(r)) + { + atomicAdd(bad, 1ull); + } + atomicMax(max_bits, __float_as_uint(d)); +} + +struct BlasCheckResult +{ + long long bad = 0; + double max_err = 0.0; +}; + +BlasCheckResult blas_compare(hipStream_t stream, + const void* got, + const void* ref, + long long count) +{ + unsigned long long* bad_dev = nullptr; + unsigned int* max_bits_dev = nullptr; + hip_check(hipMalloc(&bad_dev, sizeof(*bad_dev)), "hipMalloc(bad)"); + hip_check(hipMalloc(&max_bits_dev, sizeof(*max_bits_dev)), "hipMalloc(max)"); + hip_check(hipMemset(bad_dev, 0, sizeof(*bad_dev)), "hipMemset(bad)"); + hip_check(hipMemset(max_bits_dev, 0, sizeof(*max_bits_dev)), "hipMemset(max)"); + + const int block = 256; + const long long grid = (count + block - 1) / block; + blas_compare_kernel<<(grid), block, 0, stream>>>( + static_cast(got), + static_cast(ref), + count, + 1.0f, + 0.01f, + bad_dev, + max_bits_dev); + hip_check(hipGetLastError(), "blas_compare_kernel launch"); + + unsigned long long bad = 0; + unsigned int max_bits = 0; + hip_check(hipMemcpyAsync(&bad, bad_dev, sizeof(bad), hipMemcpyDeviceToHost, stream), + "hipMemcpyAsync(bad)"); + hip_check(hipMemcpyAsync(&max_bits, max_bits_dev, sizeof(max_bits), hipMemcpyDeviceToHost, stream), + "hipMemcpyAsync(max)"); + hip_check(hipStreamSynchronize(stream), "blas_compare sync"); + hip_check(hipFree(bad_dev), "hipFree(bad)"); + hip_check(hipFree(max_bits_dev), "hipFree(max)"); + + BlasCheckResult result; + result.bad = static_cast(bad); + // The bit pattern came back over memcpy: rebuild the non-negative float. + static_assert(sizeof(max_bits) == sizeof(float), "bit transfer size"); + std::memcpy(&result.max_err, &max_bits, sizeof(result.max_err)); + return result; +} +#endif // CK_TILE_V7_TEST_HAS_HIPBLASLT + +struct DeviceCase +{ + std::vector a_ptrs; + std::vector b_ptrs; + std::vector c_ptrs; + std::vector a_blocks; + std::vector b_blocks; + std::vector c_blocks; + void* kargs = nullptr; + std::size_t kargs_sz = 0; + bool allocated = false; + + ~DeviceCase() { release(); } + + void release() + { + if(!allocated) return; + for(auto* p : a_blocks) hip_check(hipFree(p), "hipFree"); + for(auto* p : b_blocks) hip_check(hipFree(p), "hipFree"); + for(auto* p : c_blocks) hip_check(hipFree(p), "hipFree"); + if(kargs != nullptr) hip_check(hipFree(kargs), "hipFree"); + a_blocks.clear(); b_blocks.clear(); c_blocks.clear(); + kargs = nullptr; + allocated = false; + } + + void allocate(const Case& c, const Options& options) + { + a_ptrs.resize(c.g); + b_ptrs.resize(c.g); + c_ptrs.resize(c.g); + a_blocks.resize(c.g); + b_blocks.resize(c.g); + c_blocks.resize(c.g); + + for(int g = 0; g < c.g; ++g) + { + const std::size_t a_bytes = + static_cast(c.a_elements[g]) * sizeof(Bf16); + const std::size_t b_bytes = + static_cast(c.b_elements[g]) * sizeof(Bf16); + const std::size_t c_bytes = + static_cast(c.m) * c.n * sizeof(Bf16); + hip_check(hipMalloc(&a_blocks[g], align_to(a_bytes, 256)), "hipMalloc(A)"); + hip_check(hipMalloc(&b_blocks[g], align_to(b_bytes, 256)), "hipMalloc(B)"); + hip_check(hipMalloc(&c_blocks[g], align_to(c_bytes, 256)), "hipMalloc(C)"); + hip_check(hipMemcpy(a_blocks[g], + c.host_a[g].data(), + a_bytes, + hipMemcpyHostToDevice), + "hipMemcpy(A)"); + hip_check(hipMemcpy(b_blocks[g], + c.host_b[g].data(), + b_bytes, + hipMemcpyHostToDevice), + "hipMemcpy(B)"); + hip_check(hipMemset(c_blocks[g], 0, c_bytes), "hipMemset(C)"); + a_ptrs[g] = a_blocks[g]; + b_ptrs[g] = b_blocks[g]; + c_ptrs[g] = c_blocks[g]; + } + + // block_start/block_end must be filled: the framework's non-persistent + // kernel path maps a 1D block id to (group, m tile, n tile) through them, + // which is how the v7 uniform instances resolve their tile. + const ck_tile::index_t n_tiles = (c.n + 255) / 256; + ck_tile::index_t block_cursor = 0; + std::vector host_kargs; + host_kargs.reserve(c.g); + for(int g = 0; g < c.g; ++g) + { + host_kargs.emplace_back(ck_tile::UniversalGemmKernelArgs<>{{a_ptrs[g]}, + {b_ptrs[g]}, + {}, + c_ptrs[g], + c.group_m[g], + c.n, + c.k, + {c.strides_a[g]}, + {c.strides_b[g]}, + {}, + c.n, + 1}, + block_cursor, + block_cursor + + ((c.group_m[g] + 255) / 256) * n_tiles); + block_cursor += ((c.group_m[g] + 255) / 256) * n_tiles; + } + (void)options; + kargs_sz = host_kargs.size() * sizeof(ck_tile::GemmTransKernelArg); + hip_check(hipMalloc(&kargs, kargs_sz), "hipMalloc(kargs)"); + hip_check(hipMemcpy(kargs, host_kargs.data(), kargs_sz, hipMemcpyHostToDevice), + "hipMemcpy(kargs)"); + allocated = true; + } +}; + +// --dump=1: the first few host operands, the CPU reference and the device +// output. Values, not aggregates: an aggregate cannot tell a broken mapping +// from a broken reference. +void dump_case(const Case& c, const std::vector& c_blocks) +{ + std::printf("dump host_a[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(ck_tile::type_convert(c.host_a[0][i]))); + std::printf("\ndump host_b[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(ck_tile::type_convert(c.host_b[0][i]))); + std::printf("\ndump host_d[0][0..7]:"); + for(int i = 0; i < 8; ++i) + std::printf(" %g", static_cast(c.host_d_ref[0][i])); + std::vector raw(16); + hip_check(hipMemcpy(raw.data(), + c_blocks[0], + raw.size() * sizeof(unsigned short), + hipMemcpyDeviceToHost), + "hipMemcpy(dump)"); + std::printf("\ndump device[0][0..7]:"); + for(int i = 0; i < 8; ++i) + { + const Bf16 value = *reinterpret_cast(&raw[i]); + std::printf(" %g(0x%04x)", + static_cast(ck_tile::type_convert(value)), + raw[i]); + } + std::printf("\n"); +} + +void build_problem(const Case& c, ck_tile_hcu_grouped_gemm_problem_v1& problem) +{ + problem = {}; + problem.struct_size = sizeof(problem); + problem.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + problem.data_type = CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_BF16_V1; + problem.layout = c.tn ? CK_TILE_HCU_GROUPED_GEMM_LAYOUT_TN_V1 + : CK_TILE_HCU_GROUPED_GEMM_LAYOUT_NN_V1; + problem.group_count = c.g; + problem.problem_flags = 0; + problem.policy_disable_flags = 0; + // Ragged per-group M: declared the way the ABI has it -- a device-lengths + // dimension with a host-visible capacity (the summed rows, plus one tile of + // rounding per group). A common M stays expressible too. + const bool common_m = (c.total_m == c.m * c.g); + problem.m = {common_m ? CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1 + : CK_TILE_HCU_GROUPED_GEMM_DIMENSION_DEVICE_LENGTHS_WITH_CAPACITY_V1, + 0, + common_m ? static_cast(c.m) + : static_cast(c.total_m) + 256 * c.g}; + problem.n = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.n}; + problem.k = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.k}; +} + +int select_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + std::uint32_t& instance_id, + std::string& name) +{ + ck_tile_hcu_grouped_gemm_selection_v1 selection{}; + selection.struct_size = sizeof(selection); + selection.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + const int status = ck_tile_hcu_grouped_gemm_select_device_args_v1(&problem, &selection); + if(status != 0) return status; + instance_id = selection.instance_id; + const char* instance_name = + ck_tile_hcu_grouped_gemm_instance_name_v1(selection.instance_id); + name = instance_name == nullptr ? std::string("?") : std::string(instance_name); + return 0; +} + +int run_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + void* kargs, + std::uint32_t instance_id, + std::uint32_t num_cu, + hipStream_t stream) +{ + ck_tile_hcu_grouped_gemm_launch_v1 launch{}; + launch.struct_size = sizeof(launch); + launch.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + launch.device_args = kargs; + launch.num_cu = num_cu; + launch.stream = reinterpret_cast(stream); + launch.expected_instance_id = instance_id; + ck_tile_hcu_grouped_gemm_selection_v1 launched{}; + launched.struct_size = sizeof(launched); + launched.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + return ck_tile_hcu_grouped_gemm_run_device_args_v1(&problem, &launch, &launched); +} + +double time_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, + void* kargs, + std::uint32_t instance_id, + std::uint32_t num_cu, + hipStream_t stream, + int warmup, + int iters) +{ + for(int i = 0; i < warmup; ++i) + { + run_instance(problem, kargs, instance_id, num_cu, stream); + } + hipEvent_t start{}; + hipEvent_t stop{}; + hip_check(hipEventCreate(&start), "hipEventCreate"); + hip_check(hipEventCreate(&stop), "hipEventCreate"); + hip_check(hipEventRecord(start, stream), "hipEventRecord"); + for(int i = 0; i < iters; ++i) + { + run_instance(problem, kargs, instance_id, num_cu, stream); + } + hip_check(hipEventRecord(stop, stream), "hipEventRecord"); + hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); + float ms = 0.0f; + hip_check(hipEventElapsedTime(&ms, start, stop), "hipEventElapsedTime"); + hip_check(hipEventDestroy(start), "hipEventDestroy"); + hip_check(hipEventDestroy(stop), "hipEventDestroy"); + return static_cast(ms) / iters; +} + +struct CheckResult +{ + double max_err = 0.0; + std::size_t bad = 0; + std::size_t total = 0; + // First few mismatches, as (group, row, col, got, expected): the pattern is + // what identifies a broken tile/row/column mapping, and aggregate counts do + // not carry it. + struct Sample + { + int g; + int row; + int col; + float got; + float ref; + }; + std::vector samples; +}; + +CheckResult check_against_reference(const Case& c, + const std::vector& blocks, + int stride_c) +{ + CheckResult result; + for(int g = 0; g < c.g; ++g) + { + std::vector host_d(static_cast(c.m) * stride_c); + hip_check(hipMemcpy(host_d.data(), + blocks[g], + host_d.size() * sizeof(Bf16), + hipMemcpyDeviceToHost), + "hipMemcpy(D)"); + for(int row = 0; row < c.m; ++row) + { + for(int col = 0; col < c.n; ++col) + { + const std::size_t at = static_cast(row) * stride_c + col; + const double got = static_cast(ck_tile::type_convert(host_d[at])); + const double ref = + static_cast(c.host_d_ref[g][static_cast(row) * c.n + col]); + const double err = std::fabs(got - ref); + if(err > result.max_err) result.max_err = err; + ++result.total; + if(err > 0.5) + { + ++result.bad; + if(result.samples.size() < 8) + { + result.samples.push_back({g, + row, + col, + static_cast(got), + static_cast(ref)}); + } + } + } + } + } + return result; +} + +void print_check(const char* tag, const CheckResult& checked) +{ + std::printf("check %s max_err=%.6g bad=%zu/%zu\n", + tag, + checked.max_err, + checked.bad, + checked.total); + for(const auto& s : checked.samples) + { + std::printf(" sample g=%d row=%d col=%d got=%.6g ref=%.6g\n", + s.g, + s.row, + s.col, + static_cast(s.got), + static_cast(s.ref)); + } +} + +} // namespace + +namespace v7_backend { +int run(int argc, char* argv[]) +{ + Options options; + parse_args(argc, argv, options); + + Case c; + build_case(options, c); + if(options.check && options.cpu_ref) + { + cpu_reference(c); + } + + DeviceCase device; + device.allocate(c, options); + + ck_tile_hcu_grouped_gemm_problem_v1 problem{}; + build_problem(c, problem); + + const std::uint32_t num_cu = options.num_cu > 0 + ? static_cast(options.num_cu) + : static_cast(ck_tile::get_num_cu()); + hipStream_t stream = nullptr; + std::printf("grid num_cu=%u (device=%d)\n", + num_cu, + static_cast(ck_tile::get_num_cu())); + + bool ok = true; + + const char* enable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V8_TN" + : "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN"; + const char* disable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V8_TN" + : "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V7_NN"; + + // --- v7 instance ----------------------------------------------------- + setenv(enable_switch, "1", 1); + unsetenv(disable_switch); + + std::uint32_t v7_id = 0; + std::string v7_name; + if(select_instance(problem, v7_id, v7_name) != 0) + { + std::printf("layout=%s result=FAIL reason=select_status\n", options.layout.c_str()); + return 1; + } + const std::uint32_t expected = c.tn ? kV8TnInstance : kV7NnInstance; + std::printf("select layout=%s instance_id=%u expected=%u name=%s\n", + options.layout.c_str(), + v7_id, + expected, + v7_name.c_str()); + ok = ok && v7_id == expected; + + const int run_status = + run_instance(problem, device.kargs, v7_id, num_cu, stream); + if(run_status != 0) + { + std::printf("layout=%s result=FAIL reason=run_status=%d\n", options.layout.c_str(), run_status); + return 1; + } + const double v7_ms = + time_instance(problem, device.kargs, v7_id, num_cu, stream, options.warmup, options.iters); + + if(options.dump) + { + dump_case(c, device.c_ptrs); + } + + if(options.check) + { + if(options.cpu_ref) + { + const CheckResult checked = check_against_reference(c, device.c_ptrs, c.n); + print_check("port_vs_cpu", checked); + ok = ok && checked.bad == 0; + } + } + + // --- hipBLASLt functional reference (device-side, no CPU pass) ---------- + if(options.hipblaslt_ref) + { +#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT + hipblasLtHandle_t handle = nullptr; + hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); + void* workspace = nullptr; + const std::size_t workspace_size = 256ull << 20; + hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); + + std::vector ref_blocks(static_cast(c.g), nullptr); + std::vector plans(static_cast(c.g)); + for(int g = 0; g < c.g; ++g) + { + const std::size_t bytes = + static_cast(c.group_m[g]) * c.n * sizeof(Bf16); + hip_check(hipMalloc(&ref_blocks[g], bytes), "hipMalloc(blas ref)"); + hip_check(hipMemset(ref_blocks[g], 0, bytes), "hipMemset(blas ref)"); + if(!hipblaslt_make_plan(plans[g], + handle, + workspace, + workspace_size, + c.group_m[g], + c.n, + c.k, + c.strides_a[g], + c.tn)) + { + std::fprintf(stderr, "hipBLASLt: no algorithm for group %d\n", g); + std::exit(1); + } + // One plain GEMM per group on the same stream the instance used. + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + ref_blocks[g]); + } + hip_check(hipStreamSynchronize(stream), "hipblaslt ref sync"); + + long long blas_bad = 0; + double blas_max_err = 0.0; + for(int g = 0; g < c.g; ++g) + { + const BlasCheckResult r = + blas_compare(stream, + device.c_ptrs[g], + ref_blocks[g], + static_cast(c.group_m[g]) * c.n); + blas_bad += r.bad; + blas_max_err = blas_max_err > r.max_err ? blas_max_err : r.max_err; + hip_check(hipFree(ref_blocks[g]), "hipFree(blas ref)"); + } + std::printf("blas_check layout=%s bad=%lld max_err=%.4f tolerance=1.0+0.01*|ref|\n", + options.layout.c_str(), + blas_bad, + blas_max_err); + ok = ok && blas_bad == 0; + for(auto& plan : plans) + { + hipblaslt_destroy_plan(plan); + } + hip_check(hipFree(workspace), "hipFree(blas workspace)"); + hipblasLtDestroy(handle); +#else + std::printf("blas_check result=FAIL reason=built_without_hipblaslt\n"); + ok = false; +#endif + } + + // --- hipBLASLt performance reference -------------------------------------- + // Plans are built once per group; the timed loop issues bare + // hipblasLtMatmul calls for all groups per iteration -- the launch pattern a + // production consumer would use against the same buffers. + if(options.hipblaslt_perf) + { +#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT + hipblasLtHandle_t handle = nullptr; + hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); + void* workspace = nullptr; + const std::size_t workspace_size = 256ull << 20; + hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); + + std::vector plans(static_cast(c.g)); + bool plans_ok = true; + for(int g = 0; g < c.g && plans_ok; ++g) + { + plans_ok = hipblaslt_make_plan(plans[g], + handle, + workspace, + workspace_size, + c.group_m[g], + c.n, + c.k, + c.strides_a[g], + c.tn); + } + if(!plans_ok) + { + std::printf("blas_perf layout=%s result=FAIL reason=no_algo\n", + options.layout.c_str()); + ok = false; + } + else + { + for(int i = 0; i < options.warmup; ++i) + { + for(int g = 0; g < c.g; ++g) + { + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + device.c_ptrs[g]); + } + } + hipEvent_t start{}; + hipEvent_t stop{}; + hip_check(hipEventCreate(&start), "hipEventCreate"); + hip_check(hipEventCreate(&stop), "hipEventCreate"); + hip_check(hipEventRecord(start, stream), "hipEventRecord"); + for(int i = 0; i < options.iters; ++i) + { + for(int g = 0; g < c.g; ++g) + { + hipblaslt_run_plan(plans[g], + handle, + workspace, + workspace_size, + stream, + device.b_blocks[g], + device.a_blocks[g], + device.c_ptrs[g]); + } + } + hip_check(hipEventRecord(stop, stream), "hipEventRecord"); + hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); + float elapsed_ms = 0.0f; + hip_check(hipEventElapsedTime(&elapsed_ms, start, stop), "hipEventElapsedTime"); + const double blas_ms = + static_cast(elapsed_ms) / static_cast(options.iters); + const double blas_tf = + 2.0 * static_cast(c.total_m) * c.n * c.k / + (blas_ms * 1.0e-3) / 1.0e12; + std::printf("blas_perf layout=%s m=%d n=%d k=%d groups=%d " + "blas_ms=%.4f blas_tf=%.1f\n", + options.layout.c_str(), + c.m, + c.n, + c.k, + c.g, + blas_ms, + blas_tf); + hip_check(hipEventDestroy(start), "hipEventDestroy"); + hip_check(hipEventDestroy(stop), "hipEventDestroy"); + } + for(auto& plan : plans) + { + hipblaslt_destroy_plan(plan); + } + hip_check(hipFree(workspace), "hipFree(blas workspace)"); + hipblasLtDestroy(handle); +#else + std::printf("blas_perf result=FAIL reason=built_without_hipblaslt\n"); + ok = false; +#endif + } + + const double flops = 2.0 * static_cast(c.total_m) * c.n * c.k; + std::printf("perf layout=%s m=%d n=%d k=%d groups=%d v7_ms=%.4f v7_tf=%.1f\n", + options.layout.c_str(), + c.m, + c.n, + c.k, + c.g, + v7_ms, + flops / (v7_ms * 1.0e-3) / 1.0e12); + + // --- default CK instance for the same kargs (same-process A/B) ---------- + if(options.ref) + { + setenv(disable_switch, "1", 1); + std::uint32_t ref_id = 0; + std::string ref_name; + if(select_instance(problem, ref_id, ref_name) != 0) + { + std::printf("ref layout=%s result=SKIP reason=select_status\n", options.layout.c_str()); + } + else if(ref_id == v7_id) + { + std::printf("ref layout=%s result=SKIP reason=same_instance\n", options.layout.c_str()); + } + else + { + // Contract check: the TN candidates in the registry consume the + // TRANSPOSED view of the problem (C^T = B^T A), i.e. a different + // kargs meaning than the direct view this driver builds for the v7 + // entry. Running them on these kargs would compare two different + // problems, so skip instead of printing a meaningless number. + ck_tile_hcu_grouped_gemm_candidate_info_v1 ref_info{}; + ref_info.struct_size = sizeof(ref_info); + ref_info.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; + const bool ref_transposed_view = + ck_tile_hcu_grouped_gemm_query_instance_v1(ref_id, &ref_info) == 0 && + (ref_info.contract_flags & + CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_CONTRACT_C_TRANSPOSED_VIEW_V1) != 0; + + const int ref_status = ref_transposed_view + ? 0 + : run_instance(problem, device.kargs, ref_id, num_cu, stream); + if(ref_transposed_view) + { + std::printf("ref layout=%s instance_id=%u name=%s result=SKIP " + "reason=transposed_view_contract\n", + options.layout.c_str(), + ref_id, + ref_name.c_str()); + } + else if(ref_status != 0) + { + std::printf("ref layout=%s instance_id=%u result=SKIP reason=run_status=%d\n", + options.layout.c_str(), + ref_id, + ref_status); + } + else + { + const double ref_ms = time_instance( + problem, device.kargs, ref_id, num_cu, stream, options.warmup, options.iters); + std::printf("ref layout=%s instance_id=%u name=%s ref_ms=%.4f ref_tf=%.1f\n", + options.layout.c_str(), + ref_id, + ref_name.c_str(), + ref_ms, + flops / (ref_ms * 1.0e-3) / 1.0e12); + } + } + unsetenv(disable_switch); + } + + std::printf("layout=%s result=%s\n", options.layout.c_str(), ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} + +} // namespace v7_backend + +namespace wasp_backend { using Bf16 = ck_tile::bf16_t; using Row = ck_tile::tensor_layout::gemm::RowMajor; @@ -192,10 +1545,6 @@ struct Options // computed from, so a shifted/stale consumer stripe shows up as a step at // the exact column where the mapping breaks. int dump_row = -1; - // Ablation mask for the timing decomposition (see gemm_pipeline_ag_bg_cr_wasp.hpp): - // 1 = no B loads, 2 = no MMACs, 4 = every A read from row 0, 8 = no barriers. - // Needs CK_TILE_GROUPED_GEMM_WASP_ABLATION. Results become garbage by design. - int ablate = 0; }; // Round-to-nearest-even fp32 -> bf16, returned as the exact float value the @@ -383,18 +1732,17 @@ bool parse_group_lens(const char* text, std::vector& out) void print_usage() { std::printf( - "usage: grouped_gemm_wasp_test [--m=N] [--n=N] [--k=N] [--grouped-length=a,b,c]\n" - " [--layout=NT] [--output-dtype=bf16]\n" - " [--bench=0|1] [--persistent-kernel=1]\n" - " [--warmup-iters=N] [--bench-iters=N]\n" - " [--ref=0|1] [--rel-tol=F] [--abs-tol=F]\n" - " [--hipblaslt-ref=0|1] [--hipblaslt-perf=0|1]\n" - " [--allow-unaligned] [--quiet]\n" - " extras not present in the reference harness:\n" - " [--groups=N] equal groups of --m rows\n" - " [--warmup=N] [--iters=N] aliases\n" - " [--pattern=random|ones|bindex|aindex]\n" - " [--dump-row=N] [--debug-ctas]\n"); + "usage: grouped_gemm_fused_test --layout=NN|TN|NT [test options]\n" + " --layout=NN selects v7; --layout=TN selects v8; --layout=NT selects WASP\n" + " common shape: [--m=N] [--n=N] [--k=N] [--groups=N]\n" + " [--group-lengths=a,b,c] (alias: --grouped-length=a,b,c)\n" + " common timing: [--warmup=N] [--iters=N] (aliases: --warmup-iters, --bench-iters)\n" + " NN/TN options: [--check=0|1] [--ref=0|1] [--cpu-ref=0|1]\n" + " [--num-cu=N] [--hipblaslt-ref=0|1] [--hipblaslt-perf=0|1]\n" + " NT options: [--bench=0|1]\n" + " [--allow-unaligned] [--quiet] [--output-dtype=bf16]\n" + " [--pattern=random|ones|bindex|aindex] [--dump-row=N]\n" + " [--rel-tol=F] [--abs-tol=F] [--debug-ctas]\n"); } // --------------------------------------------------------------------------- @@ -1052,13 +2400,19 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre hipblasLtHandle_t handle = nullptr; hipblasLtMatmulDesc_t desc = nullptr; hipblasLtMatrixLayout_t b_layout = nullptr; - hipblasLtMatrixLayout_t a_layout = nullptr; - hipblasLtMatrixLayout_t d_layout = nullptr; hipblasLtMatmulPreference_t preference = nullptr; - hipblasLtMatmulAlgo_t algo{}; constexpr std::size_t workspace_size = 256ull << 20; void* workspace = nullptr; DevBuf blas_out; + struct BlasPlan + { + hipblasLtMatrixLayout_t a_layout = nullptr; + hipblasLtMatrixLayout_t d_layout = nullptr; + hipblasLtMatmulAlgo_t algo{}; + }; + const std::size_t blas_group_count = + opt.hipblaslt_ref ? sh.lens.size() : std::size_t{1}; + std::vector plans(blas_group_count); auto blas_check = [](hipblasStatus_t status, const char* operation) { if(status != HIPBLAS_STATUS_SUCCESS) @@ -1097,12 +2451,6 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre blas_check(hipblasLtMatrixLayoutCreate( &b_layout, HIP_R_16BF, k, n, k), "hipblasLtMatrixLayoutCreate(B)"); - blas_check(hipblasLtMatrixLayoutCreate( - &a_layout, HIP_R_16BF, k, sh.lens[0], k), - "hipblasLtMatrixLayoutCreate(A)"); - blas_check(hipblasLtMatrixLayoutCreate( - &d_layout, HIP_R_16BF, n, sh.lens[0], n), - "hipblasLtMatrixLayoutCreate(D)"); blas_check(hipblasLtMatmulPreferenceCreate(&preference), "hipblasLtMatmulPreferenceCreate"); blas_check(hipblasLtMatmulPreferenceSetAttribute( @@ -1111,87 +2459,105 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre &workspace_size, sizeof(workspace_size)), "hipblasLtMatmulPreferenceSetAttribute"); - bool selected_tuned = false; - int tuned_index = -1; - if(pick_tuned_solution(n, sh.lens[0], k, 'T', 'N', tuned_index)) - { - std::vector resolved(1); - std::vector indices{tuned_index}; - if(hipblaslt_ext::getAlgosFromIndex(handle, indices, resolved) == - HIPBLAS_STATUS_SUCCESS && - resolved[0].state == HIPBLAS_STATUS_SUCCESS) + for(std::size_t g = 0; g < plans.size(); ++g) + { + auto& plan = plans[g]; + const int m = sh.lens[g]; + blas_check(hipblasLtMatrixLayoutCreate( + &plan.a_layout, HIP_R_16BF, k, m, k), + "hipblasLtMatrixLayoutCreate(A)"); + blas_check(hipblasLtMatrixLayoutCreate( + &plan.d_layout, HIP_R_16BF, n, m, n), + "hipblasLtMatrixLayoutCreate(D)"); + bool selected_tuned = false; + int tuned_index = -1; + if(pick_tuned_solution(n, m, k, 'T', 'N', tuned_index)) { - const float probe_alpha = 1.0F; - const float probe_beta = 0.0F; - std::size_t required_workspace = 0; - const hipblasStatus_t supported = - hipblaslt_ext::matmulIsAlgoSupported(handle, - desc, - &probe_alpha, - b_layout, - a_layout, - &probe_beta, - d_layout, - d_layout, - resolved[0].algo, - required_workspace); - if(supported == HIPBLAS_STATUS_SUCCESS && - required_workspace <= workspace_size) + std::vector resolved(1); + std::vector indices{tuned_index}; + if(hipblaslt_ext::getAlgosFromIndex(handle, indices, resolved) == + HIPBLAS_STATUS_SUCCESS && + resolved[0].state == HIPBLAS_STATUS_SUCCESS) { - algo = resolved[0].algo; - selected_tuned = true; - std::fprintf(stderr, - "[hipBLASLt] tuned NT shape m=%d n=%d k=%d " - "transA=T transB=N -> solution_index=%d " - "(workspace=%zu B)\n", - n, - sh.lens[0], - k, - tuned_index, - required_workspace); + const float probe_alpha = 1.0F; + const float probe_beta = 0.0F; + std::size_t required_workspace = 0; + const hipblasStatus_t supported = + hipblaslt_ext::matmulIsAlgoSupported(handle, + desc, + &probe_alpha, + b_layout, + plan.a_layout, + &probe_beta, + plan.d_layout, + plan.d_layout, + resolved[0].algo, + required_workspace); + if(supported == HIPBLAS_STATUS_SUCCESS && + required_workspace <= workspace_size) + { + plan.algo = resolved[0].algo; + selected_tuned = true; + std::fprintf(stderr, + "[hipBLASLt] tuned NT group=%zu shape m=%d n=%d k=%d " + "transA=T transB=N -> solution_index=%d " + "(workspace=%zu B)\n", + g, + n, + m, + k, + tuned_index, + required_workspace); + } + else + { + std::fprintf(stderr, + "[hipBLASLt] WARNING: tuned solution %d for group %zu " + "is unsupported or needs %zu bytes workspace " + "(available %zu); using heuristic\n", + tuned_index, + g, + required_workspace, + workspace_size); + } } else { std::fprintf(stderr, - "[hipBLASLt] WARNING: tuned solution %d is unsupported " - "or needs %zu bytes workspace (available %zu); " - "using heuristic\n", + "[hipBLASLt] WARNING: tuned solution %d for group %zu " + "is unavailable; using heuristic\n", tuned_index, - required_workspace, - workspace_size); + g); } } - else + if(!selected_tuned) { + hipblasLtMatmulHeuristicResult_t heuristic{}; + int algo_count = 0; + blas_check(hipblasLtMatmulAlgoGetHeuristic(handle, + desc, + b_layout, + plan.a_layout, + plan.d_layout, + plan.d_layout, + preference, + 1, + &heuristic, + &algo_count), + "hipblasLtMatmulAlgoGetHeuristic"); + if(algo_count == 0) + { + std::fprintf(stderr, + "hipBLASLt did not return an algorithm for NT group %zu\n", + g); + std::exit(1); + } + plan.algo = heuristic.algo; std::fprintf(stderr, - "[hipBLASLt] WARNING: tuned solution %d is unavailable; " - "using heuristic\n", - tuned_index); + "[hipBLASLt] group=%zu using heuristic algorithm\n", + g); } } - if(!selected_tuned) - { - hipblasLtMatmulHeuristicResult_t heuristic{}; - int algo_count = 0; - blas_check(hipblasLtMatmulAlgoGetHeuristic(handle, - desc, - b_layout, - a_layout, - d_layout, - d_layout, - preference, - 1, - &heuristic, - &algo_count), - "hipblasLtMatmulAlgoGetHeuristic"); - if(algo_count == 0) - { - std::fprintf(stderr, "hipBLASLt did not return an algorithm for NT shape\n"); - std::exit(1); - } - algo = heuristic.algo; - std::fprintf(stderr, "[hipBLASLt] using heuristic algorithm\n"); - } blas_check(hipblasLtMatmulPreferenceDestroy(preference), "hipblasLtMatmulPreferenceDestroy"); if(hipMalloc(&workspace, workspace_size) != hipSuccess || @@ -1203,29 +2569,38 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre const float alpha = 1.0F; const float beta = 0.0F; - auto run_blas = [&](void* output) { - blas_check(hipblasLtMatmul(handle, - desc, - &alpha, - b_dev.Data(), - b_layout, - a_dev.Data(), - a_layout, - &beta, - output, - d_layout, - output, - d_layout, - &algo, - workspace, - workspace_size, - stream), - "hipblasLtMatmul"); + auto run_blas = [&](void* output, std::size_t groups_to_run) { + for(std::size_t g = 0; g < groups_to_run; ++g) + { + const std::size_t row_offset = static_cast(sh.prefix[g]); + const std::size_t b_offset = g * static_cast(n) * + static_cast(k); + const std::size_t a_offset = row_offset * static_cast(k); + const std::size_t c_offset = row_offset * static_cast(n); + auto* output_group = static_cast(output) + c_offset; + blas_check(hipblasLtMatmul(handle, + desc, + &alpha, + b_dev.Data() + b_offset, + b_layout, + a_dev.Data() + a_offset, + plans[g].a_layout, + &beta, + output_group, + plans[g].d_layout, + output_group, + plans[g].d_layout, + &plans[g].algo, + workspace, + workspace_size, + stream), + "hipblasLtMatmul"); + } }; if(opt.hipblaslt_ref) { - run_blas(blas_out.Data()); + run_blas(blas_out.Data(), plans.size()); if(hipStreamSynchronize(stream) != hipSuccess) { std::fprintf(stderr, "hipBLASLt reference synchronization failed\n"); @@ -1255,9 +2630,15 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre } if(opt.hipblaslt_perf) { + std::fprintf(stderr, + "[hipBLASLt] profiling one matmul for group 0 " + "(m=%d, n=%d, k=%d)\n", + sh.lens[0], + n, + k); for(int i = 0; i < opt.warmup; ++i) { - run_blas(blas_out.Data()); + run_blas(blas_out.Data(), 1); } hipEvent_t start = nullptr; hipEvent_t stop = nullptr; @@ -1269,7 +2650,7 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre hip_runtime_check(hipEventRecord(start, stream), "hipEventRecord(start)"); for(int i = 0; i < opt.iters; ++i) { - run_blas(blas_out.Data()); + run_blas(blas_out.Data(), 1); } hip_runtime_check(hipEventRecord(stop, stream), "hipEventRecord(stop)"); hip_runtime_check(hipEventSynchronize(stop), "hipEventSynchronize(stop)"); @@ -1282,8 +2663,11 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre } hip_runtime_check(hipFree(workspace), "hipFree(hipBLASLt workspace)"); - hipblasLtMatrixLayoutDestroy(d_layout); - hipblasLtMatrixLayoutDestroy(a_layout); + for(auto& plan : plans) + { + hipblasLtMatrixLayoutDestroy(plan.d_layout); + hipblasLtMatrixLayoutDestroy(plan.a_layout); + } hipblasLtMatrixLayoutDestroy(b_layout); hipblasLtMatmulDescDestroy(desc); hipblasLtDestroy(handle); @@ -1389,7 +2773,8 @@ CaseResult run_case(const Options& opt, const ShapeInfo& sh, int num_cu, hipStre } // namespace -int main(int argc, char* argv[]) +namespace wasp_backend { +int run(int argc, char* argv[]) { Options opt{}; for(int i = 1; i < argc; ++i) @@ -1464,16 +2849,16 @@ int main(int argc, char* argv[]) { opt.dump_row = parse_int(arg.c_str() + 11, opt.dump_row); } - else if(arg.rfind("--ablate=", 0) == 0) - { - opt.ablate = parse_int(arg.c_str() + 9, opt.ablate); - } // ---- reference-harness vocabulary ---- - else if(arg.rfind("--grouped-length=", 0) == 0) + else if(arg.rfind("--grouped-length=", 0) == 0 || + arg.rfind("--group-lengths=", 0) == 0) { - if(!parse_group_lens(arg.c_str() + 17, opt.group_lens)) + const std::size_t prefix = + arg.rfind("--grouped-length=", 0) == 0 ? std::strlen("--grouped-length=") + : std::strlen("--group-lengths="); + if(!parse_group_lens(arg.c_str() + prefix, opt.group_lens)) { - std::printf("bad --grouped-length: %s\n", arg.c_str()); + std::printf("bad group lengths: %s\n", arg.c_str()); return 2; } } @@ -1532,12 +2917,6 @@ int main(int argc, char* argv[]) opt.out_dtype.c_str()); return 2; } - if((opt.hipblaslt_ref || opt.hipblaslt_perf) && - (opt.groups != 1 || !opt.group_lens.empty())) - { - std::printf("wasp test: hipBLASLt comparison/timing currently requires one group\n"); - return 2; - } if(!opt.persistent_kernel) { // The CK grouped kernel walks a grid-stride tile list inside the device @@ -1609,25 +2988,6 @@ int main(int argc, char* argv[]) opt.iters, opt.ref ? "on" : "off"); -#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) - if(opt.ablate != 0) - { - const int mask = opt.ablate; - if(hipMemcpyToSymbol(ck_tile::wasp_detail::g_wasp_ablate, - &mask, - sizeof(mask), - 0, - hipMemcpyHostToDevice) != hipSuccess) - { - std::printf("wasp test: ERROR: could not set the ablation mask\n"); - return 2; - } - std::printf("wasp test: ABLATION mask=%d (1=no B loads, 2=no MMACs, " - "4=A read row0, 8=no barriers) -- results are garbage by design\n", - mask); - } -#endif - const CaseResult r = run_case(opt, sh, num_cu, stream); if(r.skipped) @@ -2066,3 +3426,67 @@ int main(int argc, char* argv[]) std::printf("wasp test: PASS\n"); return 0; } + +} // namespace wasp_backend + +void print_fused_usage() +{ + std::printf( + "usage: grouped_gemm_fused_test --layout=NN|TN|NT [test options]\n" + " --layout=NN selects v7; --layout=TN selects v8; --layout=NT selects WASP\n" + " common shape: [--m=N] [--n=N] [--k=N] [--groups=N]\n" + " [--group-lengths=a,b,c] (alias: --grouped-length=a,b,c)\n" + " common timing: [--warmup=N] [--iters=N] " + "(aliases: --warmup-iters, --bench-iters)\n" + " NN/TN options: [--check=0|1] [--ref=0|1] [--cpu-ref=0|1]\n" + " [--num-cu=N] [--hipblaslt-ref=0|1] [--hipblaslt-perf=0|1]\n" + " NT options: [--bench=0|1] [--allow-unaligned] [--quiet]\n" + " [--output-dtype=bf16] [--pattern=random|ones|bindex|aindex]\n" + " [--dump-row=N] [--rel-tol=F] [--abs-tol=F] [--debug-ctas]\n"); +} + +int main(int argc, char* argv[]) +{ + std::string layout; + bool layout_seen = false; + std::vector forwarded_argv; + forwarded_argv.reserve(static_cast(argc) + 1); + forwarded_argv.push_back(argv[0]); + for(int i = 1; i < argc; ++i) + { + const std::string arg = argv[i]; + if(arg == "--help") + { + print_fused_usage(); + return 0; + } + if(arg.rfind("--layout=", 0) == 0) + { + if(layout_seen) + { + std::fprintf(stderr, "--layout may be specified only once\n"); + return 2; + } + layout_seen = true; + layout = arg.substr(std::strlen("--layout=")); + } + forwarded_argv.push_back(argv[i]); + } + if(!layout_seen || layout.empty()) + { + std::fprintf(stderr, "missing required --layout=NN|TN|NT\n"); + return 2; + } + forwarded_argv.push_back(nullptr); + const int forwarded_argc = static_cast(forwarded_argv.size()) - 1; + if(layout == "NN" || layout == "TN") + { + return v7_backend::run(forwarded_argc, forwarded_argv.data()); + } + if(layout == "NT") + { + return wasp_backend::run(forwarded_argc, forwarded_argv.data()); + } + std::fprintf(stderr, "unsupported layout '%s' (expected NN, TN or NT)\n", layout.c_str()); + return 2; +} diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp deleted file mode 100644 index ccdc343..0000000 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_v7_instance_test.cpp +++ /dev/null @@ -1,1434 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2026 Hygon Information Technologies Co., Ltd. -// -// Correctness + performance driver for the v7 BF16 NN / TN grouped-GEMM -// instances (InstanceId 5000 / 5001). -// -// WHAT THIS DRIVES -// The instances themselves, through the versioned device-args ABI Primus uses: -// ck_tile_hcu_grouped_gemm_select_device_args_v1 -> must return the v7 id -// ck_tile_hcu_grouped_gemm_run_device_args_v1 -> launches the kernel -// Nothing is called behind the ABI's back: the kargs array is built here exactly -// as the ABI's builders build it (one ck_tile::GemmTransKernelArg per group, -// uploaded once, consumed by the device-side tile walk). -// -// The v7 routes are OPT-IN, so the driver sets the same switch production -// would set: -// CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN / _TN -// and --ref=1 re-runs the same kargs through the DEFAULT instance for the same -// problem (switch disabled) to give a same-process, same-descriptor A/B. -// -// WHAT IT CHECKS -// * --check=1 (default): an independent CPU reference (fp32 accumulate over -// bf16 inputs, which are exact in fp32 -- only the summation order differs) -// and, with --ref=1, the default CK instance as a triangle. -// * always: selection identity, launch status, and the achieved time/TFLOPS. -// -// Layout contract (the one the v7 NN / TN entries consume): -// NN : A [M_g, K] row-major (stride_a = K) -// B [K, N] row-major (stride_b = N) -// TN : A [K, M_g] row-major panel, per-group a_ptr, stride_a = the panel pitch -// B [K, N] row-major (stride_b = N) -// D [M_g, N] row-major (stride_e = N) -// TN's A panel load covers a whole 256-wide column tile, so the A allocation -// carries slack past the last group (kAtailPad) -- the caller-side pad -// obligation the v7 TN entry's header documents. -// -// Usage: -// grouped_gemm_v7_instance_test [--layout=NN|TN] [--m=512] [--n=512] -// [--k=256] [--groups=2] [--iters=20] -// [--warmup=3] [--check=1] [--ref=1] -// [--hipblaslt-ref=0] [--hipblaslt-perf=0] -// --hipblaslt-ref=1: verify against one hipBLASLt GEMM per group, compared -// on the device (no CPU pass over the output) -// --hipblaslt-perf=1: time hipBLASLt (plans built once, bare Matmul in the -// loop) and print blas_ms / blas_tf next to the v7 numbers - -#include - -#include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm/grouped_gemm_device_args.h" -#include "ck_tile/host/device_prop.hpp" -#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" - -#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) -#include -#include -#endif - -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using Bf16 = ck_tile::bf16_t; - -constexpr std::uint32_t kV7NnInstance = 5000; -constexpr std::uint32_t kV8TnInstance = 5001; - -// Slack the TN A panel needs past the last group's columns (see the header). -constexpr std::size_t kAtailPad = 1024; - -struct Options -{ - std::string layout = "NN"; - int m = 512; - int n = 512; - int k = 256; - int groups = 2; - int iters = 20; - int warmup = 3; - bool check = true; - bool ref = true; - // The independent fp32 CPU reference the checked runs compare against (bf16 - // inputs are exact in fp32 -- only the summation order differs). --check=0 - // skips it for perf runs on production shapes. - bool cpu_ref = true; - // Ragged group lengths: the shapes the v7 harness benchmarks - // (--grouped-length=8100,8100,8091,8091). Empty => `groups` copies of `m`. - std::vector group_lengths; - bool dump = false; - // Grid override. 0 => the device's CU count (the production value). - int num_cu = 0; - // Device-side functional reference: run the same kargs through hipBLASLt - // (one plain GEMM per group) and compare on the GPU -- no CPU pass over the - // output, so production shapes check in seconds. Built only when the - // hipBLASLt install is present (see CMakeLists CK_TILE_V7_TEST_HAS_HIPBLASLT). - bool hipblaslt_ref = false; - // Performance reference: time hipBLASLt (plans built once, bare Matmul in - // the loop) with the same warmup/iters as the v7 timing, and print - // blas_ms / blas_tf next to v7_ms / v7_tf. - bool hipblaslt_perf = false; -}; - -std::size_t align_to(std::size_t value, std::size_t alignment) -{ - return (value + alignment - 1) / alignment * alignment; -} - -// Every hip* call below whose result carries [[nodiscard]] goes through this, so -// the driver both stays -Wunused-result clean and reports the first failure -// instead of running on with garbage. -void hip_check(hipError_t status, const char* what) -{ - if(status != hipSuccess) - { - std::fprintf(stderr, "HIP error %d (%s) at %s\n", static_cast(status), - hipGetErrorName(status), what); - std::exit(1); - } -} - -// Small multiples of 0.125: exact in bf16 (8 mantissa bits), so the only -// difference between the CPU reference and the kernel is the summation order. -// Both indices contribute, so no lane pattern of the kernels cancels a term. -float a_value(int row, int kk) -{ - return static_cast((row % 7) - 3) * 0.5f + - static_cast(kk % 3) * 0.25f; -} -float b_value(int kk, int col) -{ - return static_cast((col % 5) - 2) * 0.25f + - static_cast(kk % 4) * 0.125f; -} - -struct Case -{ - // Per-group row counts (ragged) plus their sum. `m` stays the first group's - // length for the single-group paths; every buffer is sized from group_m. - std::vector group_m; - int total_m = 0; - bool tn = false; - int m = 0; - int n = 0; - int k = 0; - int g = 0; - std::vector strides_a; - std::vector strides_b; - std::vector a_elements; - std::vector b_elements; - std::vector> host_a; - std::vector> host_b; - std::vector> host_d_ref; -}; - -int parse_int_arg(const char* text, const char* name) -{ - char* end = nullptr; - const long value = std::strtol(text, &end, 10); - if(end == text || value < 0) - { - std::fprintf(stderr, "bad --%s=%s\n", name, text); - std::exit(2); - } - return static_cast(value); -} - -void parse_args(int argc, char* argv[], Options& options) -{ - for(int i = 1; i < argc; ++i) - { - const std::string arg = argv[i]; - const auto value_of = [&](const char* name) -> const char* { - const std::size_t at = arg.find('='); - if(at == std::string::npos) - { - std::fprintf(stderr, "%s expects --%s=VALUE\n", arg.c_str(), name); - std::exit(2); - } - return argv[i] + at + 1; - }; - if(arg.rfind("--layout=", 0) == 0) - { - options.layout = value_of("layout"); - } - else if(arg.rfind("--m=", 0) == 0) - { - options.m = parse_int_arg(value_of("m"), "m"); - } - else if(arg.rfind("--n=", 0) == 0) - { - options.n = parse_int_arg(value_of("n"), "n"); - } - else if(arg.rfind("--k=", 0) == 0) - { - options.k = parse_int_arg(value_of("k"), "k"); - } - else if(arg.rfind("--groups=", 0) == 0) - { - options.groups = parse_int_arg(value_of("groups"), "groups"); - } - else if(arg.rfind("--iters=", 0) == 0) - { - options.iters = parse_int_arg(value_of("iters"), "iters"); - } - else if(arg.rfind("--warmup=", 0) == 0) - { - options.warmup = parse_int_arg(value_of("warmup"), "warmup"); - } - else if(arg.rfind("--check=", 0) == 0) - { - options.check = parse_int_arg(value_of("check"), "check") != 0; - } - else if(arg.rfind("--ref=", 0) == 0) - { - options.ref = parse_int_arg(value_of("ref"), "ref") != 0; - } - else if(arg.rfind("--group-lengths=", 0) == 0) - { - options.group_lengths.clear(); - const std::string list = value_of("group-lengths"); - std::size_t start = 0; - while(start <= list.size()) - { - const std::size_t comma = list.find(',', start); - const std::string item = list.substr(start, comma - start); - if(!item.empty()) - { - options.group_lengths.push_back(parse_int_arg(item.c_str(), "group-lengths")); - } - if(comma == std::string::npos) break; - start = comma + 1; - } - } - else if(arg.rfind("--cpu-ref=", 0) == 0) - { - options.cpu_ref = parse_int_arg(value_of("cpu-ref"), "cpu-ref") != 0; - } - else if(arg.rfind("--dump=", 0) == 0) - { - options.dump = parse_int_arg(value_of("dump"), "dump") != 0; - } - else if(arg.rfind("--num-cu=", 0) == 0) - { - options.num_cu = parse_int_arg(value_of("num-cu"), "num-cu"); - } - else if(arg.rfind("--hipblaslt-ref=", 0) == 0) - { - options.hipblaslt_ref = parse_int_arg(value_of("hipblaslt-ref"), "hipblaslt-ref") != 0; - } - else if(arg.rfind("--hipblaslt-perf=", 0) == 0) - { - options.hipblaslt_perf = parse_int_arg(value_of("hipblaslt-perf"), "hipblaslt-perf") != 0; - } - else - { - std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); - std::exit(2); - } - } - if(options.layout != "NN" && options.layout != "TN") - { - std::fprintf(stderr, "--layout must be NN or TN\n"); - std::exit(2); - } -} - -void build_case(const Options& options, Case& c) -{ - c.tn = options.layout == "TN"; - c.n = options.n; - c.k = options.k; - c.group_m = options.group_lengths.empty() - ? std::vector(static_cast(options.groups), options.m) - : options.group_lengths; - c.g = static_cast(c.group_m.size()); - c.m = c.group_m.front(); - c.total_m = 0; - for(int g = 0; g < c.g; ++g) c.total_m += c.group_m[g]; - // TN's A panel is [K, sum(M_g)]: it is addressed across the whole group set, - // so its pitch is the summed M (padded up to a whole tile, which is also the - // 256-element slack the panel fetch needs past the last group's columns). - const int a_pitch_tn = (c.total_m + 255) / 256 * 256; - - c.strides_a.resize(c.g); - c.strides_b.resize(c.g); - c.a_elements.resize(c.g); - c.b_elements.resize(c.g); - c.host_a.resize(c.g); - c.host_b.resize(c.g); - c.host_d_ref.resize(c.g); - - for(int g = 0; g < c.g; ++g) - { - // NN: A row-major [M_g, K]. TN: A is the [K, M_g] panel of the - // column-major [M_g, K] view, so its pitch is the M extent (the same - // choice the device-args smoke makes for its TN cases). - const int stride_a = c.tn ? a_pitch_tn : c.k; - const int stride_b = c.n; - c.strides_a[g] = stride_a; - c.strides_b[g] = stride_b; - c.a_elements[g] = - (c.tn ? c.k : c.group_m[g]) * stride_a + (c.tn ? static_cast(kAtailPad) : 0); - c.b_elements[g] = c.k * stride_b; - - c.host_a[g].assign(static_cast(c.a_elements[g]), - ck_tile::type_convert(0.0f)); - c.host_b[g].assign(static_cast(c.b_elements[g]), - ck_tile::type_convert(0.0f)); - c.host_d_ref[g].assign(static_cast(c.group_m[g]) * c.n, 0.0f); - - for(int row = 0; row < c.group_m[g]; ++row) - { - for(int kk = 0; kk < c.k; ++kk) - { - const std::size_t at = - c.tn ? static_cast(kk) * stride_a + row - : static_cast(row) * stride_a + kk; - c.host_a[g][at] = ck_tile::type_convert(a_value(row, kk)); - } - } - for(int kk = 0; kk < c.k; ++kk) - { - for(int col = 0; col < c.n; ++col) - { - c.host_b[g][static_cast(kk) * stride_b + col] = - ck_tile::type_convert(b_value(kk, col)); - } - } - } -} - -void cpu_reference(Case& c) -{ - for(int g = 0; g < c.g; ++g) - { - const int stride_a = c.strides_a[g]; - const int stride_b = c.strides_b[g]; - for(int col = 0; col < c.n; ++col) - { - for(int row = 0; row < c.m; ++row) - { - float acc = 0.0f; - for(int kk = 0; kk < c.k; ++kk) - { - const std::size_t at = - c.tn ? static_cast(kk) * stride_a + row - : static_cast(row) * stride_a + kk; - const float a = ck_tile::type_convert(c.host_a[g][at]); - const float b = ck_tile::type_convert( - c.host_b[g][static_cast(kk) * stride_b + col]); - acc += a * b; - } - c.host_d_ref[g][static_cast(row) * c.n + col] = acc; - } - } - } -} - -#if defined(CK_TILE_V7_TEST_HAS_HIPBLASLT) -// --------------------------------------------------------------------------- -// Device-side functional reference: one plain hipBLASLt GEMM per group, then a -// device-side element compare -- the output never touches the CPU, so -// production shapes verify in seconds. -// -// The col-major formulation (same construction bf16_gemm_hipblaslt.cpp uses): -// D_row[M_g, N] = A . B <=> C_col[N, M_g] = B_view[N, K] . A_view[K, M_g] -// with B_view the [N, K] col-major view of the [K, N] row-major B (ld = N), -// and A_view the operand buffer seen as -// NN: [K, M_g] col-major, ld = K (the [M_g, K] row-major A) -> opB = N -// TN: [M_g, K] col-major, ld = pitch (the [K, pitch] A panel) -> opB = T -// The panel leading dimension is the one thing the bf16_gemm_hipblaslt wrapper -// cannot express for the grouped TN operand (it hardcodes ldb = k); here it is -// passed explicitly. -// --------------------------------------------------------------------------- -void hipblaslt_check_status(hipblasStatus_t status, const char* what) -{ - if(status != HIPBLAS_STATUS_SUCCESS) - { - std::fprintf(stderr, - "hipBLASLt error %d at %s\n", - static_cast(status), - what); - std::exit(1); - } -} - -// --------------------------------------------------------------------------- -// Tuned-solution table -- same file and semantics as bf16_gemm_hipblaslt.cpp -// (BF16GEMM_HIPBLASLT_SOLUTION_FILE selects it, BF16GEMM_NO_TUNING=1 disables; -// lines "m,n,k,transA,transB,solution_index", '*' = wildcard, '#' comments). -// The keys live in THIS test's hipblasLt parameter space (m = caller N, -// n = M_g, k = K; transA is always N here, transB = T for the TN panel form), -// which is the same space the wrapper tunes in. Resolution: best-scoring -// entry -> hipblaslt_ext getAlgosFromIndex + matmulIsAlgoSupported -> else the -// stock heuristic, warning once per rejected key. -// --------------------------------------------------------------------------- -struct TunedSolution -{ - long long m = -1, n = -1, k = -1; // -1 = wildcard - char ta = '?', tb = '?'; - int index = -1; -}; - -static long long tuned_num(const std::string& s) -{ - return (s == "*") ? -1 : std::strtoll(s.c_str(), nullptr, 10); -} - -static bool tuned_hit(long long key, long long want) -{ - return key < 0 || key == want; -} - -static bool tuned_hit(char key, char want) -{ - return key == '*' || key == '?' || key == want; -} - -static const std::vector& tuned_table() -{ - static const std::vector table = [] { - std::vector t; - const char* disable = std::getenv("BF16GEMM_NO_TUNING"); - if(disable != nullptr && std::atoi(disable) != 0) - { - return t; - } - const char* path = std::getenv("BF16GEMM_HIPBLASLT_SOLUTION_FILE"); - if(path == nullptr || *path == '\0') - { - return t; - } - FILE* f = std::fopen(path, "r"); - if(f == nullptr) - { - std::fprintf(stderr, - "[hipblasLt] WARNING: cannot open '%s', using the stock heuristic\n", - path); - return t; - } - char line[512]; - while(std::fgets(line, sizeof(line), f) != nullptr) - { - char* p = line; - while(*p == ' ' || *p == '\t') - { - ++p; - } - if(*p == '#' || *p == '\n' || *p == '\r' || *p == '\0') - { - continue; - } - std::string fields[6]; - int nf = 0; - for(char* tok = std::strtok(p, ", \t\r\n"); tok != nullptr && nf < 6; - tok = std::strtok(nullptr, ", \t\r\n")) - { - fields[nf++] = tok; - } - if(nf < 6) - { - std::fprintf(stderr, "[hipblasLt] WARNING: malformed line ignored: %s", line); - continue; - } - TunedSolution e; - e.m = tuned_num(fields[0]); - e.n = tuned_num(fields[1]); - e.k = tuned_num(fields[2]); - e.ta = static_cast(std::toupper(static_cast(fields[3][0]))); - e.tb = static_cast(std::toupper(static_cast(fields[4][0]))); - e.index = static_cast(std::strtol(fields[5].c_str(), nullptr, 10)); - t.push_back(e); - } - std::fclose(f); - std::fprintf(stderr, - "[hipblasLt] tuned-solution table: %zu entries from %s\n", - t.size(), - path); - return t; - }(); - return table; -} - -static bool tuned_pick(long long m, long long n, long long k, char ta, char tb, int& index) -{ - int best_score = 6; - bool found = false; - for(const TunedSolution& e : tuned_table()) - { - if(!(tuned_hit(e.m, m) && tuned_hit(e.n, n) && tuned_hit(e.k, k) && - tuned_hit(e.ta, ta) && tuned_hit(e.tb, tb))) - { - continue; - } - const int score = (e.m < 0 ? 1 : 0) + (e.n < 0 ? 1 : 0) + (e.k < 0 ? 1 : 0) + - ((e.ta == '*' || e.ta == '?') ? 1 : 0) + - ((e.tb == '*' || e.tb == '?') ? 1 : 0); - if(score < best_score) - { - best_score = score; - index = e.index; - found = true; - if(score == 0) - { - break; // exact key: cannot do better - } - } - } - return found; -} - -// D = 1.0 * op(B) . op(A) + 0.0, all bf16 in / bf16 out, fp32 compute. -// Descriptor + algorithm selection is split from the execution: the perf -// reference builds each group's plan once and times bare hipblasLtMatmul -// calls, exactly what a production consumer would do. -struct HipblasltPlan -{ - hipblasLtMatmulDesc_t desc = nullptr; - hipblasLtMatrixLayout_t a_layout = nullptr; // the B operand, col-major [N, K], ld = N - hipblasLtMatrixLayout_t b_layout = nullptr; // the A operand - hipblasLtMatrixLayout_t c_layout = nullptr; // D, col-major [N, M_g], ld = N - hipblasLtMatmulAlgo_t algo{}; -}; - -bool hipblaslt_make_plan(HipblasltPlan& plan, - hipblasLtHandle_t handle, - void* workspace, - std::size_t workspace_size, - int m_g, - int n, - int k, - int ld_a, - bool tn) -{ - (void)workspace; // the workspace pointer only matters at Matmul time - hipblasLtMatmulPreference_t pref = nullptr; - - hipblasOperation_t op_a = HIPBLAS_OP_N; - hipblasOperation_t op_b = tn ? HIPBLAS_OP_T : HIPBLAS_OP_N; - - hipblaslt_check_status( - hipblasLtMatmulDescCreate(&plan.desc, HIPBLAS_COMPUTE_32F, HIP_R_32F), - "hipblasLtMatmulDescCreate"); - hipblaslt_check_status( - hipblasLtMatmulDescSetAttribute(plan.desc, - HIPBLASLT_MATMUL_DESC_TRANSA, - &op_a, - sizeof(op_a)), - "SetAttribute(TRANSA)"); - hipblaslt_check_status( - hipblasLtMatmulDescSetAttribute(plan.desc, - HIPBLASLT_MATMUL_DESC_TRANSB, - &op_b, - sizeof(op_b)), - "SetAttribute(TRANSB)"); - hipblaslt_check_status( - hipblasLtMatrixLayoutCreate(&plan.a_layout, - HIP_R_16BF, - static_cast(n), - static_cast(k), - static_cast(n)), - "LayoutCreate(B view)"); - // opT stores the [K, M_g] logical matrix through the [M_g, K] panel: ld is - // the panel pitch, not K. - const std::size_t b_ld = tn ? static_cast(ld_a) - : static_cast(k); - hipblaslt_check_status( - hipblasLtMatrixLayoutCreate(&plan.b_layout, - HIP_R_16BF, - static_cast(k), - static_cast(m_g), - b_ld), - "LayoutCreate(A view)"); - hipblaslt_check_status( - hipblasLtMatrixLayoutCreate(&plan.c_layout, - HIP_R_16BF, - static_cast(n), - static_cast(m_g), - static_cast(n)), - "LayoutCreate(D)"); - - hipblaslt_check_status( - hipblasLtMatmulPreferenceCreate(&pref), - "hipblasLtMatmulPreferenceCreate"); - hipblaslt_check_status( - hipblasLtMatmulPreferenceSetAttribute( - pref, - HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, - &workspace_size, - sizeof(workspace_size)), - "PreferenceSetAttribute(WORKSPACE)"); - - // 1) tuned table (best-scoring entry, resolved through the ext API); - // 2) stock heuristic. A table entry the library rejects only warns. - const char ta_ch = (op_a == HIPBLAS_OP_N) ? 'N' : 'T'; - const char tb_ch = (op_b == HIPBLAS_OP_N) ? 'N' : 'T'; - int tuned_index = -1; - if(tuned_pick(n, m_g, k, ta_ch, tb_ch, tuned_index)) - { - std::vector resolved(1); - std::vector algo_index(1, tuned_index); - const float probe_alpha = 1.0f; - const float probe_beta = 0.0f; - std::size_t required = 0; - if(hipblaslt_ext::getAlgosFromIndex(handle, algo_index, resolved) != - HIPBLAS_STATUS_SUCCESS || - resolved[0].state != HIPBLAS_STATUS_SUCCESS) - { - std::fprintf(stderr, - "[hipblasLt] WARNING: solution index %d does not exist; " - "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", - tuned_index, - n, - m_g, - k, - ta_ch, - tb_ch); - } - else if(hipblaslt_ext::matmulIsAlgoSupported(handle, - plan.desc, - &probe_alpha, - plan.a_layout, - plan.b_layout, - &probe_beta, - plan.c_layout, - plan.c_layout, - resolved[0].algo, - required) != - HIPBLAS_STATUS_SUCCESS) - { - std::fprintf(stderr, - "[hipblasLt] WARNING: solution index %d unsupported here; " - "stock heuristic for m=%d n=%d k=%d trans=%c%c\n", - tuned_index, - n, - m_g, - k, - ta_ch, - tb_ch); - } - else if(required > workspace_size) - { - std::fprintf(stderr, - "[hipblasLt] WARNING: solution index %d needs %zu bytes " - "workspace (have %zu); stock heuristic\n", - tuned_index, - required, - workspace_size); - } - else - { - plan.algo = resolved[0].algo; - std::fprintf(stderr, - "[hipblasLt] tuned: m=%d n=%d k=%d trans=%c%c -> index=%d\n", - n, - m_g, - k, - ta_ch, - tb_ch, - tuned_index); - hipblasLtMatmulPreferenceDestroy(pref); - return true; - } - } - - hipblasLtMatmulHeuristicResult_t algo{}; - int algo_count = 0; - hipblaslt_check_status( - hipblasLtMatmulAlgoGetHeuristic(handle, - plan.desc, - plan.a_layout, - plan.b_layout, - plan.c_layout, - plan.c_layout, - pref, - 1, - &algo, - &algo_count), - "hipblasLtMatmulAlgoGetHeuristic"); - const bool ok = algo_count > 0; - if(ok) - { - plan.algo = algo.algo; - } - hipblasLtMatmulPreferenceDestroy(pref); - return ok; -} - -void hipblaslt_run_plan(const HipblasltPlan& plan, - hipblasLtHandle_t handle, - void* workspace, - std::size_t workspace_size, - hipStream_t stream, - const void* b_dev, - const void* a_dev, - void* d_dev) -{ - const float alpha = 1.0f; - const float beta = 0.0f; - hipblaslt_check_status( - hipblasLtMatmul(handle, - plan.desc, - &alpha, - b_dev, - plan.a_layout, - a_dev, - plan.b_layout, - &beta, - d_dev, - plan.c_layout, - d_dev, - plan.c_layout, - &plan.algo, - workspace, - workspace_size, - stream), - "hipblasLtMatmul"); -} - -void hipblaslt_destroy_plan(HipblasltPlan& plan) -{ - hipblasLtMatrixLayoutDestroy(plan.c_layout); - hipblasLtMatrixLayoutDestroy(plan.b_layout); - hipblasLtMatrixLayoutDestroy(plan.a_layout); - hipblasLtMatmulDescDestroy(plan.desc); - plan = HipblasltPlan{}; -} - -// Element compare on the device: bad = |got - ref| > atol + rtol*|ref|, -// max_err tracked via the non-negative-float bit pattern (uint order == float -// order for non-negative values). -__global__ void blas_compare_kernel(const Bf16* __restrict__ got, - const Bf16* __restrict__ ref, - long long count, - float atol, - float rtol, - unsigned long long* bad, - unsigned int* max_bits) -{ - const long long i = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if(i >= count) - { - return; - } - const float g = ck_tile::type_convert(got[i]); - const float r = ck_tile::type_convert(ref[i]); - const float d = fabsf(g - r); - if(d > atol + rtol * fabsf(r)) - { - atomicAdd(bad, 1ull); - } - atomicMax(max_bits, __float_as_uint(d)); -} - -struct BlasCheckResult -{ - long long bad = 0; - double max_err = 0.0; -}; - -BlasCheckResult blas_compare(hipStream_t stream, - const void* got, - const void* ref, - long long count) -{ - unsigned long long* bad_dev = nullptr; - unsigned int* max_bits_dev = nullptr; - hip_check(hipMalloc(&bad_dev, sizeof(*bad_dev)), "hipMalloc(bad)"); - hip_check(hipMalloc(&max_bits_dev, sizeof(*max_bits_dev)), "hipMalloc(max)"); - hip_check(hipMemset(bad_dev, 0, sizeof(*bad_dev)), "hipMemset(bad)"); - hip_check(hipMemset(max_bits_dev, 0, sizeof(*max_bits_dev)), "hipMemset(max)"); - - const int block = 256; - const long long grid = (count + block - 1) / block; - blas_compare_kernel<<(grid), block, 0, stream>>>( - static_cast(got), - static_cast(ref), - count, - 1.0f, - 0.01f, - bad_dev, - max_bits_dev); - hip_check(hipGetLastError(), "blas_compare_kernel launch"); - - unsigned long long bad = 0; - unsigned int max_bits = 0; - hip_check(hipMemcpyAsync(&bad, bad_dev, sizeof(bad), hipMemcpyDeviceToHost, stream), - "hipMemcpyAsync(bad)"); - hip_check(hipMemcpyAsync(&max_bits, max_bits_dev, sizeof(max_bits), hipMemcpyDeviceToHost, stream), - "hipMemcpyAsync(max)"); - hip_check(hipStreamSynchronize(stream), "blas_compare sync"); - hip_check(hipFree(bad_dev), "hipFree(bad)"); - hip_check(hipFree(max_bits_dev), "hipFree(max)"); - - BlasCheckResult result; - result.bad = static_cast(bad); - // The bit pattern came back over memcpy: rebuild the non-negative float. - static_assert(sizeof(max_bits) == sizeof(float), "bit transfer size"); - std::memcpy(&result.max_err, &max_bits, sizeof(result.max_err)); - return result; -} -#endif // CK_TILE_V7_TEST_HAS_HIPBLASLT - -struct DeviceCase -{ - std::vector a_ptrs; - std::vector b_ptrs; - std::vector c_ptrs; - std::vector a_blocks; - std::vector b_blocks; - std::vector c_blocks; - void* kargs = nullptr; - std::size_t kargs_sz = 0; - bool allocated = false; - - ~DeviceCase() { release(); } - - void release() - { - if(!allocated) return; - for(auto* p : a_blocks) hip_check(hipFree(p), "hipFree"); - for(auto* p : b_blocks) hip_check(hipFree(p), "hipFree"); - for(auto* p : c_blocks) hip_check(hipFree(p), "hipFree"); - if(kargs != nullptr) hip_check(hipFree(kargs), "hipFree"); - a_blocks.clear(); b_blocks.clear(); c_blocks.clear(); - kargs = nullptr; - allocated = false; - } - - void allocate(const Case& c, const Options& options) - { - a_ptrs.resize(c.g); - b_ptrs.resize(c.g); - c_ptrs.resize(c.g); - a_blocks.resize(c.g); - b_blocks.resize(c.g); - c_blocks.resize(c.g); - - for(int g = 0; g < c.g; ++g) - { - const std::size_t a_bytes = - static_cast(c.a_elements[g]) * sizeof(Bf16); - const std::size_t b_bytes = - static_cast(c.b_elements[g]) * sizeof(Bf16); - const std::size_t c_bytes = - static_cast(c.m) * c.n * sizeof(Bf16); - hip_check(hipMalloc(&a_blocks[g], align_to(a_bytes, 256)), "hipMalloc(A)"); - hip_check(hipMalloc(&b_blocks[g], align_to(b_bytes, 256)), "hipMalloc(B)"); - hip_check(hipMalloc(&c_blocks[g], align_to(c_bytes, 256)), "hipMalloc(C)"); - hip_check(hipMemcpy(a_blocks[g], - c.host_a[g].data(), - a_bytes, - hipMemcpyHostToDevice), - "hipMemcpy(A)"); - hip_check(hipMemcpy(b_blocks[g], - c.host_b[g].data(), - b_bytes, - hipMemcpyHostToDevice), - "hipMemcpy(B)"); - hip_check(hipMemset(c_blocks[g], 0, c_bytes), "hipMemset(C)"); - a_ptrs[g] = a_blocks[g]; - b_ptrs[g] = b_blocks[g]; - c_ptrs[g] = c_blocks[g]; - } - - // block_start/block_end must be filled: the framework's non-persistent - // kernel path maps a 1D block id to (group, m tile, n tile) through them, - // which is how the v7 uniform instances resolve their tile. - const ck_tile::index_t n_tiles = (c.n + 255) / 256; - ck_tile::index_t block_cursor = 0; - std::vector host_kargs; - host_kargs.reserve(c.g); - for(int g = 0; g < c.g; ++g) - { - host_kargs.emplace_back(ck_tile::UniversalGemmKernelArgs<>{{a_ptrs[g]}, - {b_ptrs[g]}, - {}, - c_ptrs[g], - c.group_m[g], - c.n, - c.k, - {c.strides_a[g]}, - {c.strides_b[g]}, - {}, - c.n, - 1}, - block_cursor, - block_cursor + - ((c.group_m[g] + 255) / 256) * n_tiles); - block_cursor += ((c.group_m[g] + 255) / 256) * n_tiles; - } - (void)options; - kargs_sz = host_kargs.size() * sizeof(ck_tile::GemmTransKernelArg); - hip_check(hipMalloc(&kargs, kargs_sz), "hipMalloc(kargs)"); - hip_check(hipMemcpy(kargs, host_kargs.data(), kargs_sz, hipMemcpyHostToDevice), - "hipMemcpy(kargs)"); - allocated = true; - } -}; - -// --dump=1: the first few host operands, the CPU reference and the device -// output. Values, not aggregates: an aggregate cannot tell a broken mapping -// from a broken reference. -void dump_case(const Case& c, const std::vector& c_blocks) -{ - std::printf("dump host_a[0][0..7]:"); - for(int i = 0; i < 8; ++i) - std::printf(" %g", static_cast(ck_tile::type_convert(c.host_a[0][i]))); - std::printf("\ndump host_b[0][0..7]:"); - for(int i = 0; i < 8; ++i) - std::printf(" %g", static_cast(ck_tile::type_convert(c.host_b[0][i]))); - std::printf("\ndump host_d[0][0..7]:"); - for(int i = 0; i < 8; ++i) - std::printf(" %g", static_cast(c.host_d_ref[0][i])); - std::vector raw(16); - hip_check(hipMemcpy(raw.data(), - c_blocks[0], - raw.size() * sizeof(unsigned short), - hipMemcpyDeviceToHost), - "hipMemcpy(dump)"); - std::printf("\ndump device[0][0..7]:"); - for(int i = 0; i < 8; ++i) - { - const Bf16 value = *reinterpret_cast(&raw[i]); - std::printf(" %g(0x%04x)", - static_cast(ck_tile::type_convert(value)), - raw[i]); - } - std::printf("\n"); -} - -void build_problem(const Case& c, ck_tile_hcu_grouped_gemm_problem_v1& problem) -{ - problem = {}; - problem.struct_size = sizeof(problem); - problem.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; - problem.data_type = CK_TILE_HCU_GROUPED_GEMM_DATA_TYPE_BF16_V1; - problem.layout = c.tn ? CK_TILE_HCU_GROUPED_GEMM_LAYOUT_TN_V1 - : CK_TILE_HCU_GROUPED_GEMM_LAYOUT_NN_V1; - problem.group_count = c.g; - problem.problem_flags = 0; - problem.policy_disable_flags = 0; - // Ragged per-group M: declared the way the ABI has it -- a device-lengths - // dimension with a host-visible capacity (the summed rows, plus one tile of - // rounding per group). A common M stays expressible too. - const bool common_m = (c.total_m == c.m * c.g); - problem.m = {common_m ? CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1 - : CK_TILE_HCU_GROUPED_GEMM_DIMENSION_DEVICE_LENGTHS_WITH_CAPACITY_V1, - 0, - common_m ? static_cast(c.m) - : static_cast(c.total_m) + 256 * c.g}; - problem.n = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.n}; - problem.k = {CK_TILE_HCU_GROUPED_GEMM_DIMENSION_COMMON_V1, 0, c.k}; -} - -int select_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, - std::uint32_t& instance_id, - std::string& name) -{ - ck_tile_hcu_grouped_gemm_selection_v1 selection{}; - selection.struct_size = sizeof(selection); - selection.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; - const int status = ck_tile_hcu_grouped_gemm_select_device_args_v1(&problem, &selection); - if(status != 0) return status; - instance_id = selection.instance_id; - const char* instance_name = - ck_tile_hcu_grouped_gemm_instance_name_v1(selection.instance_id); - name = instance_name == nullptr ? std::string("?") : std::string(instance_name); - return 0; -} - -int run_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, - void* kargs, - std::uint32_t instance_id, - std::uint32_t num_cu, - hipStream_t stream) -{ - ck_tile_hcu_grouped_gemm_launch_v1 launch{}; - launch.struct_size = sizeof(launch); - launch.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; - launch.device_args = kargs; - launch.num_cu = num_cu; - launch.stream = reinterpret_cast(stream); - launch.expected_instance_id = instance_id; - ck_tile_hcu_grouped_gemm_selection_v1 launched{}; - launched.struct_size = sizeof(launched); - launched.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; - return ck_tile_hcu_grouped_gemm_run_device_args_v1(&problem, &launch, &launched); -} - -double time_instance(const ck_tile_hcu_grouped_gemm_problem_v1& problem, - void* kargs, - std::uint32_t instance_id, - std::uint32_t num_cu, - hipStream_t stream, - int warmup, - int iters) -{ - for(int i = 0; i < warmup; ++i) - { - run_instance(problem, kargs, instance_id, num_cu, stream); - } - hipEvent_t start{}; - hipEvent_t stop{}; - hip_check(hipEventCreate(&start), "hipEventCreate"); - hip_check(hipEventCreate(&stop), "hipEventCreate"); - hip_check(hipEventRecord(start, stream), "hipEventRecord"); - for(int i = 0; i < iters; ++i) - { - run_instance(problem, kargs, instance_id, num_cu, stream); - } - hip_check(hipEventRecord(stop, stream), "hipEventRecord"); - hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); - float ms = 0.0f; - hip_check(hipEventElapsedTime(&ms, start, stop), "hipEventElapsedTime"); - hip_check(hipEventDestroy(start), "hipEventDestroy"); - hip_check(hipEventDestroy(stop), "hipEventDestroy"); - return static_cast(ms) / iters; -} - -struct CheckResult -{ - double max_err = 0.0; - std::size_t bad = 0; - std::size_t total = 0; - // First few mismatches, as (group, row, col, got, expected): the pattern is - // what identifies a broken tile/row/column mapping, and aggregate counts do - // not carry it. - struct Sample - { - int g; - int row; - int col; - float got; - float ref; - }; - std::vector samples; -}; - -CheckResult check_against_reference(const Case& c, - const std::vector& blocks, - int stride_c) -{ - CheckResult result; - for(int g = 0; g < c.g; ++g) - { - std::vector host_d(static_cast(c.m) * stride_c); - hip_check(hipMemcpy(host_d.data(), - blocks[g], - host_d.size() * sizeof(Bf16), - hipMemcpyDeviceToHost), - "hipMemcpy(D)"); - for(int row = 0; row < c.m; ++row) - { - for(int col = 0; col < c.n; ++col) - { - const std::size_t at = static_cast(row) * stride_c + col; - const double got = static_cast(ck_tile::type_convert(host_d[at])); - const double ref = - static_cast(c.host_d_ref[g][static_cast(row) * c.n + col]); - const double err = std::fabs(got - ref); - if(err > result.max_err) result.max_err = err; - ++result.total; - if(err > 0.5) - { - ++result.bad; - if(result.samples.size() < 8) - { - result.samples.push_back({g, - row, - col, - static_cast(got), - static_cast(ref)}); - } - } - } - } - } - return result; -} - -void print_check(const char* tag, const CheckResult& checked) -{ - std::printf("check %s max_err=%.6g bad=%zu/%zu\n", - tag, - checked.max_err, - checked.bad, - checked.total); - for(const auto& s : checked.samples) - { - std::printf(" sample g=%d row=%d col=%d got=%.6g ref=%.6g\n", - s.g, - s.row, - s.col, - static_cast(s.got), - static_cast(s.ref)); - } -} - -} // namespace - -int main(int argc, char* argv[]) -{ - Options options; - parse_args(argc, argv, options); - - Case c; - build_case(options, c); - if(options.check && options.cpu_ref) - { - cpu_reference(c); - } - - DeviceCase device; - device.allocate(c, options); - - ck_tile_hcu_grouped_gemm_problem_v1 problem{}; - build_problem(c, problem); - - const std::uint32_t num_cu = options.num_cu > 0 - ? static_cast(options.num_cu) - : static_cast(ck_tile::get_num_cu()); - hipStream_t stream = nullptr; - std::printf("grid num_cu=%u (device=%d)\n", - num_cu, - static_cast(ck_tile::get_num_cu())); - - bool ok = true; - - const char* enable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V8_TN" - : "CK_TILE_GROUPED_GEMM_ENABLE_BF16_V7_NN"; - const char* disable_switch = c.tn ? "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V8_TN" - : "CK_TILE_GROUPED_GEMM_DISABLE_BF16_V7_NN"; - - // --- v7 instance ----------------------------------------------------- - setenv(enable_switch, "1", 1); - unsetenv(disable_switch); - - std::uint32_t v7_id = 0; - std::string v7_name; - if(select_instance(problem, v7_id, v7_name) != 0) - { - std::printf("layout=%s result=FAIL reason=select_status\n", options.layout.c_str()); - return 1; - } - const std::uint32_t expected = c.tn ? kV8TnInstance : kV7NnInstance; - std::printf("select layout=%s instance_id=%u expected=%u name=%s\n", - options.layout.c_str(), - v7_id, - expected, - v7_name.c_str()); - ok = ok && v7_id == expected; - - const int run_status = - run_instance(problem, device.kargs, v7_id, num_cu, stream); - if(run_status != 0) - { - std::printf("layout=%s result=FAIL reason=run_status=%d\n", options.layout.c_str(), run_status); - return 1; - } - const double v7_ms = - time_instance(problem, device.kargs, v7_id, num_cu, stream, options.warmup, options.iters); - - if(options.dump) - { - dump_case(c, device.c_ptrs); - } - - if(options.check) - { - if(options.cpu_ref) - { - const CheckResult checked = check_against_reference(c, device.c_ptrs, c.n); - print_check("port_vs_cpu", checked); - ok = ok && checked.bad == 0; - } - } - - // --- hipBLASLt functional reference (device-side, no CPU pass) ---------- - if(options.hipblaslt_ref) - { -#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT - hipblasLtHandle_t handle = nullptr; - hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); - void* workspace = nullptr; - const std::size_t workspace_size = 256ull << 20; - hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); - - std::vector ref_blocks(static_cast(c.g), nullptr); - std::vector plans(static_cast(c.g)); - for(int g = 0; g < c.g; ++g) - { - const std::size_t bytes = - static_cast(c.group_m[g]) * c.n * sizeof(Bf16); - hip_check(hipMalloc(&ref_blocks[g], bytes), "hipMalloc(blas ref)"); - hip_check(hipMemset(ref_blocks[g], 0, bytes), "hipMemset(blas ref)"); - if(!hipblaslt_make_plan(plans[g], - handle, - workspace, - workspace_size, - c.group_m[g], - c.n, - c.k, - c.strides_a[g], - c.tn)) - { - std::fprintf(stderr, "hipBLASLt: no algorithm for group %d\n", g); - std::exit(1); - } - // One plain GEMM per group on the same stream the instance used. - hipblaslt_run_plan(plans[g], - handle, - workspace, - workspace_size, - stream, - device.b_blocks[g], - device.a_blocks[g], - ref_blocks[g]); - } - hip_check(hipStreamSynchronize(stream), "hipblaslt ref sync"); - - long long blas_bad = 0; - double blas_max_err = 0.0; - for(int g = 0; g < c.g; ++g) - { - const BlasCheckResult r = - blas_compare(stream, - device.c_ptrs[g], - ref_blocks[g], - static_cast(c.group_m[g]) * c.n); - blas_bad += r.bad; - blas_max_err = blas_max_err > r.max_err ? blas_max_err : r.max_err; - hip_check(hipFree(ref_blocks[g]), "hipFree(blas ref)"); - } - std::printf("blas_check layout=%s bad=%lld max_err=%.4f tolerance=1.0+0.01*|ref|\n", - options.layout.c_str(), - blas_bad, - blas_max_err); - ok = ok && blas_bad == 0; - for(auto& plan : plans) - { - hipblaslt_destroy_plan(plan); - } - hip_check(hipFree(workspace), "hipFree(blas workspace)"); - hipblasLtDestroy(handle); -#else - std::printf("blas_check result=FAIL reason=built_without_hipblaslt\n"); - ok = false; -#endif - } - - // --- hipBLASLt performance reference -------------------------------------- - // Plans are built once per group; the timed loop issues bare - // hipblasLtMatmul calls for all groups per iteration -- the launch pattern a - // production consumer would use against the same buffers. - if(options.hipblaslt_perf) - { -#ifdef CK_TILE_V7_TEST_HAS_HIPBLASLT - hipblasLtHandle_t handle = nullptr; - hipblaslt_check_status(hipblasLtCreate(&handle), "hipblasLtCreate"); - void* workspace = nullptr; - const std::size_t workspace_size = 256ull << 20; - hip_check(hipMalloc(&workspace, workspace_size), "hipMalloc(blas workspace)"); - - std::vector plans(static_cast(c.g)); - bool plans_ok = true; - for(int g = 0; g < c.g && plans_ok; ++g) - { - plans_ok = hipblaslt_make_plan(plans[g], - handle, - workspace, - workspace_size, - c.group_m[g], - c.n, - c.k, - c.strides_a[g], - c.tn); - } - if(!plans_ok) - { - std::printf("blas_perf layout=%s result=FAIL reason=no_algo\n", - options.layout.c_str()); - ok = false; - } - else - { - for(int i = 0; i < options.warmup; ++i) - { - for(int g = 0; g < c.g; ++g) - { - hipblaslt_run_plan(plans[g], - handle, - workspace, - workspace_size, - stream, - device.b_blocks[g], - device.a_blocks[g], - device.c_ptrs[g]); - } - } - hipEvent_t start{}; - hipEvent_t stop{}; - hip_check(hipEventCreate(&start), "hipEventCreate"); - hip_check(hipEventCreate(&stop), "hipEventCreate"); - hip_check(hipEventRecord(start, stream), "hipEventRecord"); - for(int i = 0; i < options.iters; ++i) - { - for(int g = 0; g < c.g; ++g) - { - hipblaslt_run_plan(plans[g], - handle, - workspace, - workspace_size, - stream, - device.b_blocks[g], - device.a_blocks[g], - device.c_ptrs[g]); - } - } - hip_check(hipEventRecord(stop, stream), "hipEventRecord"); - hip_check(hipEventSynchronize(stop), "hipEventSynchronize"); - float elapsed_ms = 0.0f; - hip_check(hipEventElapsedTime(&elapsed_ms, start, stop), "hipEventElapsedTime"); - const double blas_ms = - static_cast(elapsed_ms) / static_cast(options.iters); - const double blas_tf = - 2.0 * static_cast(c.total_m) * c.n * c.k / - (blas_ms * 1.0e-3) / 1.0e12; - std::printf("blas_perf layout=%s m=%d n=%d k=%d groups=%d " - "blas_ms=%.4f blas_tf=%.1f\n", - options.layout.c_str(), - c.m, - c.n, - c.k, - c.g, - blas_ms, - blas_tf); - hip_check(hipEventDestroy(start), "hipEventDestroy"); - hip_check(hipEventDestroy(stop), "hipEventDestroy"); - } - for(auto& plan : plans) - { - hipblaslt_destroy_plan(plan); - } - hip_check(hipFree(workspace), "hipFree(blas workspace)"); - hipblasLtDestroy(handle); -#else - std::printf("blas_perf result=FAIL reason=built_without_hipblaslt\n"); - ok = false; -#endif - } - - const double flops = 2.0 * static_cast(c.total_m) * c.n * c.k; - std::printf("perf layout=%s m=%d n=%d k=%d groups=%d v7_ms=%.4f v7_tf=%.1f\n", - options.layout.c_str(), - c.m, - c.n, - c.k, - c.g, - v7_ms, - flops / (v7_ms * 1.0e-3) / 1.0e12); - - // --- default CK instance for the same kargs (same-process A/B) ---------- - if(options.ref) - { - setenv(disable_switch, "1", 1); - std::uint32_t ref_id = 0; - std::string ref_name; - if(select_instance(problem, ref_id, ref_name) != 0) - { - std::printf("ref layout=%s result=SKIP reason=select_status\n", options.layout.c_str()); - } - else if(ref_id == v7_id) - { - std::printf("ref layout=%s result=SKIP reason=same_instance\n", options.layout.c_str()); - } - else - { - // Contract check: the TN candidates in the registry consume the - // TRANSPOSED view of the problem (C^T = B^T A), i.e. a different - // kargs meaning than the direct view this driver builds for the v7 - // entry. Running them on these kargs would compare two different - // problems, so skip instead of printing a meaningless number. - ck_tile_hcu_grouped_gemm_candidate_info_v1 ref_info{}; - ref_info.struct_size = sizeof(ref_info); - ref_info.abi_version = CK_TILE_HCU_GROUPED_GEMM_DEVICE_ARGS_ABI_VERSION; - const bool ref_transposed_view = - ck_tile_hcu_grouped_gemm_query_instance_v1(ref_id, &ref_info) == 0 && - (ref_info.contract_flags & - CK_TILE_HCU_GROUPED_GEMM_CANDIDATE_CONTRACT_C_TRANSPOSED_VIEW_V1) != 0; - - const int ref_status = ref_transposed_view - ? 0 - : run_instance(problem, device.kargs, ref_id, num_cu, stream); - if(ref_transposed_view) - { - std::printf("ref layout=%s instance_id=%u name=%s result=SKIP " - "reason=transposed_view_contract\n", - options.layout.c_str(), - ref_id, - ref_name.c_str()); - } - else if(ref_status != 0) - { - std::printf("ref layout=%s instance_id=%u result=SKIP reason=run_status=%d\n", - options.layout.c_str(), - ref_id, - ref_status); - } - else - { - const double ref_ms = time_instance( - problem, device.kargs, ref_id, num_cu, stream, options.warmup, options.iters); - std::printf("ref layout=%s instance_id=%u name=%s ref_ms=%.4f ref_tf=%.1f\n", - options.layout.c_str(), - ref_id, - ref_name.c_str(), - ref_ms, - flops / (ref_ms * 1.0e-3) / 1.0e12); - } - } - unsetenv(disable_switch); - } - - std::printf("layout=%s result=%s\n", options.layout.c_str(), ok ? "PASS" : "FAIL"); - return ok ? 0 : 1; -} diff --git a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp index cff2147..18361de 100644 --- a/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp +++ b/example_hcu/ck_tile/19_grouped_gemm/grouped_gemm_wasp_epilogue.hpp @@ -379,13 +379,6 @@ struct GroupedGemmWaspStoreEpilogue { return; // producer waves own no output } - if((ck_tile::wasp_detail::wasp_ablate() & 16) != 0) - { - // Ablation (see gemm_pipeline_ag_bg_cr_wasp.hpp): run the whole GEMM - // but throw the writeback away, so the epilogue's cost can be timed - // on its own. - return; - } // Only used by the DEBUG_RECORD block below; store_raw recomputes its // own. Keep for that diagnostic without tripping -Wunused otherwise. [[maybe_unused]] const ck_tile::index_t consumer_id = @@ -449,7 +442,7 @@ struct GroupedGemmWaspStoreEpilogue // RunGemmWasp path calls this directly with kargs-derived scalars, so no // output tile-window object is built or kept live. // - // The role filter and the ablation gate live HERE (not only in the window + // The role filter lives HERE (not only in the window // operator() above): RunGemmWasp calls store_raw directly on EVERY wave of // the workgroup, so producer waves (wid < kNumProducerWaves) must be dropped // here, exactly as the window path drops them before extracting the window. @@ -469,10 +462,6 @@ struct GroupedGemmWaspStoreEpilogue { return; // producer waves own no output } - if((ck_tile::wasp_detail::wasp_ablate() & 16) != 0) - { - return; // ablation: run the GEMM but throw the writeback away - } const ck_tile::index_t consumer_id = wid - Pipeline::kNumProducerWaves; const ck_tile::index_t lane = ck_tile::get_lane_id(); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp index bb818b0..ff5ea69 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_wasp.hpp @@ -71,20 +71,16 @@ // ----------------------------------------------------------------------------- // THE 1/4-PHASE STAGGER (the standalone kernel's C / C' personas). // -// WHAT IT IS. Both consumer bodies execute the same items per k-stage, but C' -// issues them rotated by ONE ITEM: C pairs each read with its MMAC inside one -// interval, C' lets one read cross each barrier and consumes it at the top of -// the next interval. In items (one item = one 8-unit m-chunk = 32 MMACs): +// Each k-stage has four quarters. Consumer waves 0..3 execute +// Q0(wait B, read A.h0, drain LDS), Q1(MMAC A.h0), BAR, +// Q2(read half A.h1, prefetch B(k+1), read half A.h1, drain LDS), +// Q3(MMAC A.h1), BAR. Consumer waves 4..7 rotate Q3 from the previous stage +// ahead of Q0, then execute Q0, BAR, Q1, Q2, BAR. The first stage has no +// carried Q3; the last carried Q3 is drained after the loop. The producer issues +// A(k+1) before the first barrier and waits for its DMA between the two barriers, +// publishing the completed stage at the second barrier. // -// C I1 = [R c0 M c0 R c1 M c1] C I2 = [R c2 M c2 R c3 M c3] -// C' I1 = [M c3(prev) R c0 M c0 R c1] C' I2 = [M c1 R c2 M c2 R c3] -// -// with a k-stage split into two intervals by BAR(k).mid, and c0/c1 the h0 chunks, -// c2/c3 the h1 chunks. C leads an interval with its ds_reads (LDS busy) while C' -// leads with its MMACs (matrix busy), and they swap at the tail. This is the -// same shape as the reference kernel's stage_body_C / stage_body_Cp; the -// granularity is one m-chunk instead of one K-half because the fragment has to -// fit (see below). +// Older chunk-level cadence measurements below predate this Q0-Q3 ordering. // // MEASURED, this shape, two cadences: // @@ -103,13 +99,8 @@ // configuration that loses; every uniform body is better than any mix, and the // loss is not codegen (the spill is 16 / 52 B in all three (b) rows). // -// The mechanism is the rendezvous: a barrier releases when the SLOWEST wave -// arrives, and the two personas have different interval costs (C drains in -// place twice per interval, C' gets one cheap drain and one free one), so the -// workgroup runs at C's pace while C' idles -- and the LDS queue interleaving of -// two out-of-phase read streams makes both of them slower than either uniform -// stream. There is no phase offset to be had once a hard sync re-times every -// wave twice a stage. +// Those earlier measurements used a different schedule and are retained as +// historical results; they do not describe the Q0-Q3 pipeline below. // // THE REGISTER SIDE of the same experiment: the reference kernel keeps a whole // K-half fragment (64 VGPR) live across the barrier. Here that lands on the @@ -195,32 +186,6 @@ static_assert(sizeof(stage16_t) == 16, "16 B staging unit"); // VALU and an MMAC from different waves on the same SIMD cannot issue in the // same cycle. Use the pipeline's a_read_half (one base VGPR + immediates). -#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) -// --------------------------------------------------------------------------- -// Ablation switch, set from the host with hipMemcpyToSymbol. Results are WRONG -// by construction -- the only thing being measured is the timing decomposition, -// because three separate hypotheses about where this kernel's time goes were -// each falsified by measurement and modelling is clearly not working. -// -// bit 0 : skip the B global loads -// bit 1 : skip the MMACs -// bit 2 : read every A fragment from row 0's address (same LDS traffic count, -// degenerate address pattern -- isolates "16 distinct addresses" from -// "the reads themselves") -// bit 3 : skip the workgroup barriers -// --------------------------------------------------------------------------- -__device__ static int g_wasp_ablate = 0; -#endif - -CK_TILE_DEVICE int wasp_ablate() -{ -#if defined(CK_TILE_GROUPED_GEMM_WASP_ABLATION) - return g_wasp_ablate; -#else - return 0; -#endif -} - CK_TILE_DEVICE void ds_write_b128(uint32_t byte_addr, const dword4_t& v) { asm volatile("ds_write_b128 %0, %1" ::"v"(byte_addr), "v"(v) : "memory"); @@ -400,14 +365,9 @@ CK_TILE_DEVICE dword4_t wasp_buffer_load_b128_idxen(buffer_rsrc_t rsrc, // Bare workgroup barrier -- the reference kernel's wait_all_warp_arrived() // verbatim: sched_barrier(0), s_barrier, sched_barrier(0). // -// NO LDS DRAIN. This is load bearing for the k-stage cadence: the C' persona -// issues its A(k,1) read at the very end of an interval and consumes it on the -// far side of the barrier. A draining barrier (s_waitcnt lgkmcnt(0) + s_barrier) -// forces that read to complete before ANY wave may pass, which is exactly the -// serialisation the 1/4-phase offset exists to remove -- and it also puts every -// wave's LDS latency on the critical path to the barrier. Each consumer waits -// for its own A data explicitly (the waits inside Q1/Q3), so nothing is lost by -// not draining here. +// NO LDS DRAIN. Q0 and Q2 drain their own A reads before the corresponding +// MMACs, so the barrier need not force every wave's LDS queue to drain before +// any wave may pass. // // The sched_barrier pair is the reference kernel's, not decoration: it stops the // machine scheduler from hoisting a read across the barrier (which would break @@ -421,10 +381,7 @@ CK_TILE_DEVICE void wasp_barrier() CK_TILE_DEVICE void wasp_barrier_maybe() { - if((wasp_ablate() & 8) == 0) - { - wasp_barrier(); - } + wasp_barrier(); } CK_TILE_DEVICE void wasp_sync_lds() @@ -433,8 +390,8 @@ CK_TILE_DEVICE void wasp_sync_lds() // // An earlier revision also drained vmcnt(0) on the theory that the A stage // (written with ds_write, an LDS-unit store) might be ordered by the global - // counter. The ablation disproves the cost model behind that: the barriers - // are ~20% of the runtime, and vmcnt(0) makes every wave -- including the 8 + // counter. Timing decomposition showed that the barriers are ~20% of the + // runtime, and vmcnt(0) makes every wave -- including the 8 // consumers, whose global loads are its own B operands and have nothing to // do with publishing A -- stall on its own outstanding global traffic at // each of the num_loop barriers. The producer's A writes are LDS stores and @@ -442,14 +399,9 @@ CK_TILE_DEVICE void wasp_sync_lds() asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier\n" ::: "memory"); } -// Barrier with the ablation bit honored (see wasp_ablate). Uniform branch, so -// it costs one scalar compare when the ablation is off. CK_TILE_DEVICE void wasp_sync_lds_maybe() { - if((wasp_ablate() & 8) == 0) - { - wasp_sync_lds(); - } + wasp_sync_lds(); } // Retire outstanding direct-to-LDS DMAs (vmcnt). The producer calls this before @@ -667,8 +619,8 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 // a_base and lda are per-TILE, not per k-stage, so they are passed in rather // than re-derived here. The first version called calculate_offset on every - // one of the num_loop calls -- 64 times per tile per wave -- and the ablation - // says this kernel's time is dominated by exactly that kind of scalar + // one of the num_loop calls -- 64 times per tile per wave -- and timing + // decomposition showed this kernel's time is dominated by exactly that kind of scalar // bookkeeping, not by the MMACs (removing every MMAC saves ~10%). template CK_TILE_DEVICE AProdSrc make_a_prod_src(const ADataTypePtr* a_base, @@ -739,11 +691,6 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 { using namespace wasp_detail; - if((wasp_ablate() & 64) != 0) - { - return; // ablation: idle producer (see wasp_ablate) - } - const index_t slot = (k_stage % kNumStages) * kAStageBytes; const uint32_t soffset = static_cast(k_stage * kStageK) * sizeof(uint16_t); @@ -818,14 +765,6 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 { using namespace wasp_detail; const uint32_t s_col_base = static_cast(k_stage * kStageK) * sizeof(uint16_t); - if((wasp_ablate() & 1) != 0) - { - buf[0][0].d4 = 0; - buf[1][0].d4 = 0; - buf[0][1].d4 = 0; - buf[1][1].d4 = 0; - return; - } // Destination order is the reference kernel's: (n0,kh0) (n0,kh1) // (n1,kh0) (n1,kh1) -- buf[kh][n], the immediate 0/64 choosing the half. wasp_buffer_load_b128_x4_idxen_soff2(b_rsrc, bp.p0, bp.p1, s_col_base, @@ -854,13 +793,12 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 // The A fragment is ONE WHOLE K-HALF: 16 x 16 B per lane = 64 VGPR, read with // ONE base register plus compile-time immediates, and consumed by ONE // 64-MMAC burst. This is the standalone kernel's shape AND grouping - // (ds_read_a_half + mmac_half), and the k-stage body around it is unchanged: - // two intervals per stage, one 16-unit read burst and one 64-MMAC burst in - // each, in opposite order for the two consumer personas. See c_I1/c_I2 (C) - // and cp_I1/cp_I2 (C'). + // (ds_read_a_half + mmac_half). Each stage still has two 16-read / 64-MMAC + // pairs; the second read is split into two 8-read issues around B prefetch. + // The C and C' functions below place those pairs in their staggered quarters. // ------------------------------------------------------------------ - // 16 consecutive ds_read_b128 off ONE base VGPR, with the m-iteration row + // 8 consecutive ds_read_b128 off ONE base VGPR, with the m-iteration row // offsets (16 rows x 128 B = 0x800) AND the 2-deep A ring slot as // compile-time immediates. 32768 + 30720 = 63488 fits the 16-bit unsigned // offset field. @@ -899,68 +837,67 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 // bf16_gemm_nt_hipc.cpp:1247 and :1281-1300). template CK_TILE_DEVICE static void - ds_read_b128_x16_imm(uint32_t base, - wasp_detail::dword4_t& o0, wasp_detail::dword4_t& o1, wasp_detail::dword4_t& o2, wasp_detail::dword4_t& o3, - wasp_detail::dword4_t& o4, wasp_detail::dword4_t& o5, wasp_detail::dword4_t& o6, wasp_detail::dword4_t& o7, - wasp_detail::dword4_t& o8, wasp_detail::dword4_t& o9, wasp_detail::dword4_t& o10, wasp_detail::dword4_t& o11, - wasp_detail::dword4_t& o12, wasp_detail::dword4_t& o13, wasp_detail::dword4_t& o14, wasp_detail::dword4_t& o15) + ds_read_b128_x8_imm(uint32_t base, + wasp_detail::dword4_t& o0, + wasp_detail::dword4_t& o1, + wasp_detail::dword4_t& o2, + wasp_detail::dword4_t& o3, + wasp_detail::dword4_t& o4, + wasp_detail::dword4_t& o5, + wasp_detail::dword4_t& o6, + wasp_detail::dword4_t& o7) { - static_assert(SLOT + 30720 <= 65535, + static_assert(SLOT + 14336 <= 65535, "ds_read_b128 imm offset overflows the 16-bit unsigned field"); - asm volatile("ds_read_b128 %0, %16 offset:%c17\n\t" - "ds_read_b128 %1, %16 offset:%c18\n\t" - "ds_read_b128 %2, %16 offset:%c19\n\t" - "ds_read_b128 %3, %16 offset:%c20\n\t" - "ds_read_b128 %4, %16 offset:%c21\n\t" - "ds_read_b128 %5, %16 offset:%c22\n\t" - "ds_read_b128 %6, %16 offset:%c23\n\t" - "ds_read_b128 %7, %16 offset:%c24\n\t" - "ds_read_b128 %8, %16 offset:%c25\n\t" - "ds_read_b128 %9, %16 offset:%c26\n\t" - "ds_read_b128 %10, %16 offset:%c27\n\t" - "ds_read_b128 %11, %16 offset:%c28\n\t" - "ds_read_b128 %12, %16 offset:%c29\n\t" - "ds_read_b128 %13, %16 offset:%c30\n\t" - "ds_read_b128 %14, %16 offset:%c31\n\t" - "ds_read_b128 %15, %16 offset:%c32\n" - : "=&v"(o0), "=&v"(o1), "=&v"(o2), "=&v"(o3), - "=&v"(o4), "=&v"(o5), "=&v"(o6), "=&v"(o7), - "=&v"(o8), "=&v"(o9), "=&v"(o10), "=&v"(o11), - "=&v"(o12), "=&v"(o13), "=&v"(o14), "=&v"(o15) + asm volatile("ds_read_b128 %0, %8 offset:%c9\n\t" + "ds_read_b128 %1, %8 offset:%c10\n\t" + "ds_read_b128 %2, %8 offset:%c11\n\t" + "ds_read_b128 %3, %8 offset:%c12\n\t" + "ds_read_b128 %4, %8 offset:%c13\n\t" + "ds_read_b128 %5, %8 offset:%c14\n\t" + "ds_read_b128 %6, %8 offset:%c15\n\t" + "ds_read_b128 %7, %8 offset:%c16\n" + : "=&v"(o0), "=&v"(o1), "=&v"(o2), "=&v"(o3), + "=&v"(o4), "=&v"(o5), "=&v"(o6), "=&v"(o7) : "v"(base), - "n"(SLOT + 0), "n"(SLOT + 2048), "n"(SLOT + 4096), "n"(SLOT + 6144), - "n"(SLOT + 8192), "n"(SLOT + 10240), "n"(SLOT + 12288), "n"(SLOT + 14336), - "n"(SLOT + 16384), "n"(SLOT + 18432), "n"(SLOT + 20480), "n"(SLOT + 22528), - "n"(SLOT + 24576), "n"(SLOT + 26624), "n"(SLOT + 28672), "n"(SLOT + 30720) + "n"(SLOT + 0), "n"(SLOT + 2048), "n"(SLOT + 4096), + "n"(SLOT + 6144), "n"(SLOT + 8192), "n"(SLOT + 10240), + "n"(SLOT + 12288), "n"(SLOT + 14336) : "memory"); } - // One K-half of the 256M x K16 A operand for the whole wave tile. SLOT is - // the A-ring slot ((k_stage % 2) * kAStageBytes), which the k loop unrolled - // by parity makes compile-time. + // Read one eight-register half of a K-half. Splitting at the B prefetch + // lets Q2 overlap the load with LDS reads without changing their addresses. template - CK_TILE_DEVICE void - a_read_half(const AReadBases& b, - wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters]) const + CK_TILE_DEVICE void a_read_half_low( + const AReadBases& b, + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters]) const { - ds_read_b128_x16_imm(H1 ? b.h1 : b.h0, - reinterpret_cast(a_regs[0]), - reinterpret_cast(a_regs[1]), - reinterpret_cast(a_regs[2]), - reinterpret_cast(a_regs[3]), - reinterpret_cast(a_regs[4]), - reinterpret_cast(a_regs[5]), - reinterpret_cast(a_regs[6]), - reinterpret_cast(a_regs[7]), - reinterpret_cast(a_regs[8]), - reinterpret_cast(a_regs[9]), - reinterpret_cast(a_regs[10]), - reinterpret_cast(a_regs[11]), - reinterpret_cast(a_regs[12]), - reinterpret_cast(a_regs[13]), - reinterpret_cast(a_regs[14]), - reinterpret_cast(a_regs[15])); + ds_read_b128_x8_imm(H1 ? b.h1 : b.h0, + reinterpret_cast(a_regs[0]), + reinterpret_cast(a_regs[1]), + reinterpret_cast(a_regs[2]), + reinterpret_cast(a_regs[3]), + reinterpret_cast(a_regs[4]), + reinterpret_cast(a_regs[5]), + reinterpret_cast(a_regs[6]), + reinterpret_cast(a_regs[7])); + } + template + CK_TILE_DEVICE void a_read_half_high( + const AReadBases& b, + wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters]) const + { + ds_read_b128_x8_imm(H1 ? b.h1 : b.h0, + reinterpret_cast(a_regs[8]), + reinterpret_cast(a_regs[9]), + reinterpret_cast(a_regs[10]), + reinterpret_cast(a_regs[11]), + reinterpret_cast(a_regs[12]), + reinterpret_cast(a_regs[13]), + reinterpret_cast(a_regs[14]), + reinterpret_cast(a_regs[15])); } // 64 MMACs for one K-half: kk(2) x m(16) x n(2). The standalone kernel's @@ -981,24 +918,11 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 // the single drain. The LDS queue is shared with the other seven consumers, // so "wait for my first four reads" does not return any sooner; it just // takes two stalls instead of one. - // The Q1 drain -- the reference kernel's `wait lgkmcnt(0); wait vmcnt(0)`: - // lgkmcnt for A's LDS read, vmcnt for the B prefetch that Q1's MMACs consume. - CK_TILE_DEVICE void wasp_wait_q1(wasp_detail::WaspVmemDiag& vmem_diag, - index_t stage) const - { - asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); - wasp_detail::wasp_wait_vmcnt0(vmem_diag, stage); - } - CK_TILE_DEVICE void mmac_half(const wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], const wasp_detail::stage16_t (&b_half)[WaspAccTile::kNIters], AccTile& acc) const { using namespace wasp_detail; - if((wasp_ablate() & 2) != 0) - { - return; // ablation: skip MMACs - } // The reference kernel peels the first MMAC out of the chain so that // set_prio<3> can be raised AFTER one has issued and set_prio<0> dropped // after the very last, with no per-iteration branch in the chain. @@ -1040,9 +964,7 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 } // ------------------------------------------------------------------ - // PERSONA C's k-stage: (Q0,Q1) | BAR k.mid | (Q2,Q3) | BAR k.end. - // Every read is drained inside the interval that issued it, so every A slot - // is provably dead a full stage before the producer refills it. + // PERSONA C: Q0, Q1 | BAR | Q2, Q3 | BAR. // ------------------------------------------------------------------ template CK_TILE_DEVICE void c_I1(index_t consumer_id, @@ -1056,22 +978,15 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 { using namespace wasp_detail; (void)consumer_id; - (void)k_stage; (void)num_loop; constexpr index_t SLOT = CUR * kAStageBytes; - // This interval issues no B prefetch and no stage-dependent addressing, - // but keeps the same signature as its I2 sibling by design. - (void)k_stage; - (void)num_loop; - if((wasp_ablate() & 32) != 0) - { - return; // ablation: idle consumer - } - // Q0: A(k, h0). - a_read_half(a_rb, a_regs); - // Q1: A(h0) landed (lgkmcnt) + B(k) landed (vmcnt) -- B(k) was prefetched - // in the previous stage and B(k+1) has not been issued yet. - wasp_wait_q1(vmem_diag, k_stage); + // Q0: B(k) is ready before the A reads; drain those LDS reads before Q1. + wasp_detail::wasp_wait_vmcnt0(vmem_diag, k_stage); + a_read_half_low(a_rb, a_regs); + a_read_half_high(a_rb, a_regs); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + + // Q1: 64 MMACs. mmac_half(a_regs, b_regs[CUR][0], acc); } @@ -1090,35 +1005,22 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 using namespace wasp_detail; constexpr index_t SLOT = CUR * kAStageBytes; (void)consumer_id; - if((wasp_ablate() & 32) != 0) - { - return; // ablation: idle consumer - } - // Q2: issue the B prefetch -- 4 x buffer_load_dwordx4 in ONE asm block, - // the reference kernel's form (two IDXEN pairs, k-half on the immediate, - // the stage's K column on the uniform SGPR soffset) -- and read A(k, h1). - // - // The MMACs of the next interval consume these bytes, so the Q1 drain - // below is the reference kernel's `lgkmcnt(0); vmcnt(0)`. + // Q2: overlap the next B prefetch with the two halves of A(k, h1). + a_read_half_low(a_rb, a_regs); if(k_stage + 1 < num_loop) { load_B_stage(b_rsrc, b_pairs, k_stage + 1, b_regs[CUR ^ 1], vmem_diag); } - a_read_half(a_rb, a_regs); - // Q3: A(h1) landed. B(k)'s second half is resident, so NO vmcnt here -- - // draining it would stall on the B(k+1) prefetch just issued. + a_read_half_high(a_rb, a_regs); asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + + // Q3: 64 MMACs. mmac_half(a_regs, b_regs[CUR][1], acc); } // ------------------------------------------------------------------ - // PERSONA C's k-stage, rotated by ONE QUARTER: - // (Q3 of k-1, Q0) | BAR k.mid | (Q1, Q2) | BAR k.end. - // - // What crosses a barrier is the MMAC and the A fragment registers it - // consumes. Each interval carries the same work as C's (one 16-unit read, one - // 64-MMAC burst, one real drain): a rendezvous releases when the SLOWEST wave - // arrives, so a heavier persona would just be paid for by the other half. + // PERSONA C' starts each stage with the previous stage's Q3, then executes + // Q0 | BAR | Q1, Q2 | BAR. // ------------------------------------------------------------------ template CK_TILE_DEVICE void cp_I1(index_t consumer_id, @@ -1127,31 +1029,23 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 const AReadBases& a_rb, wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], + wasp_detail::WaspVmemDiag& vmem_diag, AccTile& acc) const { using namespace wasp_detail; (void)consumer_id; - (void)k_stage; (void)num_loop; constexpr index_t SLOT = CUR * kAStageBytes; - (void)k_stage; - (void)num_loop; - if((wasp_ablate() & 32) != 0) - { - return; // ablation: idle consumer - } if constexpr(!IS_FIRST) { - // Q3 of the PREVIOUS stage: the quarter that crosses the barrier. - // Its data was read at the end of the previous stage and is drained - // HERE, so the wait is for a read issued a whole barrier ago -- that - // is what the deferral buys. B(k-1) is still resident (its buffer is - // CUR^1 here), which is why cp_I2's prefetch must not move earlier. - asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + // Q3 of the previous stage; its A reads were drained in that Q2. mmac_half(a_regs, b_regs[CUR ^ 1][1], acc); } - // Q0: A(k, h0). - a_read_half(a_rb, a_regs); + // Q0: B(k) is ready before the A reads; drain those reads before Q1. + wasp_detail::wasp_wait_vmcnt0(vmem_diag, k_stage); + a_read_half_low(a_rb, a_regs); + a_read_half_high(a_rb, a_regs); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); } template @@ -1159,35 +1053,27 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 index_t k_stage, index_t num_loop, const AReadBases& a_rb, - wasp_detail::buffer_rsrc_t b_rsrc, + wasp_detail::buffer_rsrc_t b_rsrc, const BPairs& b_pairs, wasp_detail::stage16_t (&b_regs)[2][2][WaspAccTile::kNIters], wasp_detail::stage16_t (&a_regs)[WaspAccTile::kMIters], - wasp_detail::WaspVmemDiag& vmem_diag, - AccTile& acc) const + wasp_detail::WaspVmemDiag& vmem_diag, + AccTile& acc) const { using namespace wasp_detail; constexpr index_t SLOT = CUR * kAStageBytes; (void)consumer_id; - if((wasp_ablate() & 32) != 0) - { - return; // ablation: idle consumer - } - // Q1: A(k, h0) landed -- it was read before the mid barrier, so this wait - // is cheap by construction -- and B(k) landed (prefetched one stage ago). - wasp_wait_q1(vmem_diag, k_stage); + // Q1: 64 MMACs. mmac_half(a_regs, b_regs[CUR][0], acc); - // Q2: issue the B prefetch -- 4 x buffer_load_dwordx4 in ONE asm block, - // the reference kernel's form (two IDXEN pairs, k-half on the immediate, - // the stage's K column on the uniform SGPR soffset) -- and read A(k, h1). - // The bytes land in b_regs[CUR ^ 1] and are consumed by the next - // interval's Q1 MMACs; that is what wasp_wait_q1()'s vmcnt half waits for. + // Q2: overlap the next B prefetch with the two halves of A(k, h1). + a_read_half_low(a_rb, a_regs); if(k_stage + 1 < num_loop) { load_B_stage(b_rsrc, b_pairs, k_stage + 1, b_regs[CUR ^ 1], vmem_diag); } - a_read_half(a_rb, a_regs); + a_read_half_high(a_rb, a_regs); + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); } // Persona C driver: two intervals per k-stage, mid barrier between them. @@ -1257,14 +1143,16 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 } if((num_loop & 1) != 0) { - cp_I1<0, /*IS_FIRST=*/true>(consumer_id, 0, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<0, /*IS_FIRST=*/true>( + consumer_id, 0, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(0).mid cp_I2<0>( consumer_id, 0, num_loop, a_rb, b_rsrc, b_pairs, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(0).end for(index_t k = 1; k < num_loop; k += 2) { - cp_I1<1, false>(consumer_id, k, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<1, false>( + consumer_id, k, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).mid cp_I2<1>(consumer_id, k, @@ -1277,7 +1165,8 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).end - cp_I1<0, false>(consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<0, false>( + consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k+1).mid cp_I2<0>(consumer_id, k + 1, @@ -1291,16 +1180,14 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 acc); wasp_detail::wasp_barrier_maybe(); // BAR(k+1).end } - if((wasp_detail::wasp_ablate() & 32) == 0) - { - // Odd num_loop leaves the final stage at CUR=0. - asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); - mmac_half(a_regs, b_regs[0][1], acc); - } + // Odd num_loop leaves the final stage at CUR=0. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[0][1], acc); } else { - cp_I1<0, /*IS_FIRST=*/true>(consumer_id, 0, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<0, /*IS_FIRST=*/true>( + consumer_id, 0, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(0).mid cp_I2<0>( consumer_id, 0, num_loop, a_rb, b_rsrc, b_pairs, b_regs, a_regs, vmem_diag, acc); @@ -1308,7 +1195,8 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 for(index_t k = 1; k < num_loop - 1; k += 2) { - cp_I1<1, false>(consumer_id, k, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<1, false>( + consumer_id, k, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).mid cp_I2<1>(consumer_id, k, @@ -1321,7 +1209,8 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).end - cp_I1<0, false>(consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<0, false>( + consumer_id, k + 1, num_loop, a_rb, b_regs, a_regs, vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k+1).mid cp_I2<0>(consumer_id, k + 1, @@ -1335,7 +1224,14 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 acc); wasp_detail::wasp_barrier_maybe(); // BAR(k+1).end } - cp_I1<1, false>(consumer_id, num_loop - 1, num_loop, a_rb, b_regs, a_regs, acc); + cp_I1<1, false>(consumer_id, + num_loop - 1, + num_loop, + a_rb, + b_regs, + a_regs, + vmem_diag, + acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).mid cp_I2<1>(consumer_id, num_loop - 1, @@ -1348,12 +1244,9 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 vmem_diag, acc); wasp_detail::wasp_barrier_maybe(); // BAR(k).end - if((wasp_detail::wasp_ablate() & 32) == 0) - { - // Even num_loop leaves the final stage at CUR=1. - asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); - mmac_half(a_regs, b_regs[1][1], acc); - } + // Even num_loop leaves the final stage at CUR=1. + asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory"); + mmac_half(a_regs, b_regs[1][1], acc); } } // ------------------------------------------------------------------ @@ -1524,27 +1417,33 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 wasp_detail::wasp_wait_vmcnt0(); // A(0) landed wasp_detail::wasp_barrier_maybe(); // BAR prologue.end: publish A(0) - for(index_t k = 0; k < num_loop; ++k) + for(index_t k = 0; k + 1 < num_loop; ++k) { // Issuing k+1 here instead of after the barrier buys its DMA a // whole extra interval to land (it is confirmed below), and the // slot it writes -- (k+1)%2 -- was last read by the consumers in // stage k-1, i.e. one full interval before this issue. - if(k + 1 < num_loop) - { - producer_fill(a_src, k + 1, wid); // issue k+1 -> slot (k+1)%2 + producer_fill(a_src, k + 1, wid); // issue k+1 -> slot (k+1)%2 #if defined(CK_TILE_GROUPED_GEMM_WASP_VMEM_WAIT_COUNT) - if(k + 1 == wasp_detail::kVmemSampleStage) - { - vmem_diag.issued_at = wasp_detail::wasp_memtime(); - vmem_diag.issued = true; - } -#endif + if(k + 1 == wasp_detail::kVmemSampleStage) + { + vmem_diag.issued_at = wasp_detail::wasp_memtime(); + vmem_diag.issued = true; } +#endif wasp_detail::wasp_barrier_maybe(); // BAR(k).mid wasp_detail::wasp_wait_vmcnt0_producer(vmem_diag, k + 1); // A(k+1) landed wasp_detail::wasp_barrier_maybe(); // BAR(k).end: publish A(k+1) } + + // The final stage has no successor to issue, but still completes + // the same barrier/wait sequence as every other stage. + if(num_loop > 0) + { + wasp_detail::wasp_barrier_maybe(); // BAR(k).mid + wasp_detail::wasp_wait_vmcnt0_producer(vmem_diag, num_loop); + wasp_detail::wasp_barrier_maybe(); // BAR(k).end + } } else { @@ -1577,16 +1476,12 @@ struct GemmPipelineAgBgCrWasp : public GemmPipelineAgBgCrCompV3 (static_cast(inrow) << 32) | (nrow + 16u)}; // Destination of the prefetch, double buffered [parity][k-half][n]. wasp_detail::stage16_t b_regs[2][2][WaspAccTile::kNIters]; - // Prologue, mirroring the producer's: prime B(0), barrier, wait for - // it, barrier. The second barrier is the one that publishes A(0) -- - // the producer reaches it only after waiting for that DMA. - // Prime B(0); the first stage's MMACs consume it (see the wait below). + // Prime B(0); Q0 waits for it after the matching prologue barriers. if(num_loop > 0) { load_B_stage(b_rsrc, b_pr, 0, b_regs[0], vmem_diag); } wasp_detail::wasp_barrier_maybe(); // BAR prologue.mid - wasp_detail::wasp_wait_vmcnt0(); // B(0) landed wasp_detail::wasp_barrier_maybe(); // BAR prologue.end // waves 0..3 -> C, 4..7 -> C' From 34416ce5285da497d22b52a214aa9b1f2e918cfa Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Thu, 15 Oct 2026 16:18:16 +0800 Subject: [PATCH 10/11] feat(conv): add fp32 conv example --- example_hcu/ck_tile/14_conv/CMakeLists.txt | 1 + example_hcu/ck_tile/14_conv/conv_fwd_fp32.cpp | 303 ++++++++++++++++++ .../ops/conv/warp/warp_gemm_dispatcher.hpp | 52 +++ .../warp/warp_gemm_attribute_mmac_impl.hpp | 111 +++++++ 4 files changed, 467 insertions(+) create mode 100644 example_hcu/ck_tile/14_conv/conv_fwd_fp32.cpp diff --git a/example_hcu/ck_tile/14_conv/CMakeLists.txt b/example_hcu/ck_tile/14_conv/CMakeLists.txt index 62fc8f1..a699ad1 100644 --- a/example_hcu/ck_tile/14_conv/CMakeLists.txt +++ b/example_hcu/ck_tile/14_conv/CMakeLists.txt @@ -6,3 +6,4 @@ add_example_executable(tile_example_conv_fwd_dsreadm conv_fwd_dsreadm.cpp) set_tests_properties(tile_example_conv_fwd_dsreadm PROPERTIES SKIP_RETURN_CODE 77) add_example_executable(tile_example_conv_fwd_lds_direct conv_fwd_lds_direct.cpp) add_example_executable(tile_example_conv_fwd_bf16_f32 conv_fwd_bf16_f32.cpp) +add_example_executable(tile_example_conv_fwd_fp32 conv_fwd_fp32.cpp) \ No newline at end of file diff --git a/example_hcu/ck_tile/14_conv/conv_fwd_fp32.cpp b/example_hcu/ck_tile/14_conv/conv_fwd_fp32.cpp new file mode 100644 index 0000000..0e6ef09 --- /dev/null +++ b/example_hcu/ck_tile/14_conv/conv_fwd_fp32.cpp @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 Hygon Information Technology Co., Ltd. + +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/host/reference/reference_conv_fwd.hpp" +#include "ck_tile/ops/conv.hpp" +#include "ck_tile/ops/elementwise.hpp" +#include "ck_tile/ops/epilogue.hpp" + +// fp32 inputs + fp32 output, fp32 accumulation (MMAC 16x16x8 f32) +using ADataType = float; +using BDataType = float; +using CDataType = float; +using AccDataType = float; + +using ALayout = ck_tile::tensor_layout::convolution::NHWGC; +using BLayout = ck_tile::tensor_layout::convolution::GKYXC; +using CLayout = ck_tile::tensor_layout::convolution::NHWGK; + +using AElementOp = ck_tile::element_wise::PassThrough; +using BElementOp = ck_tile::element_wise::PassThrough; +using CElementOp = ck_tile::element_wise::PassThrough; + +static constexpr ck_tile::index_t MBlockTile = 64; +static constexpr ck_tile::index_t NBlockTile = 64; +static constexpr ck_tile::index_t KBlockTile = 16; + +static constexpr ck_tile::index_t MWarp = 2; +static constexpr ck_tile::index_t NWarp = 2; + +static constexpr ck_tile::index_t MWarpIter = 2; +static constexpr ck_tile::index_t NWarpIter = 2; + +static constexpr ck_tile::index_t MmmacIter = 1; +static constexpr ck_tile::index_t NmmacIter = 1; + +static constexpr ck_tile::index_t MmmacInterleave = 1; +static constexpr ck_tile::index_t NmmacInterleave = 1; + +static constexpr ck_tile::index_t NumPrefetch = 2; + +static constexpr ck_tile::index_t BlockSize = MWarp * NWarp * ck_tile::get_warp_size(); +static constexpr ck_tile::index_t MinBlockPerCU = 2; + +using BlockTile = ck_tile::sequence; +using BlockWarps = ck_tile::sequence; +using WarpTile = ck_tile::sequence; + +using TileShape = ck_tile::ConvIgemmTileShape; + +// gmem load / smem store: 4 x fp32 (16B); smem load: 2 x fp32 per mmac (kABKPerLane); +// c store: 2 x fp32 (LdsCShuffleEpilogue stores 8B per vector) +using TileTraits = + ck_tile::ConvIgemmUniversalTileTraits, + ck_tile::tuple, + ck_tile::sequence<4, 4>, + ck_tile::sequence<4, 4>, + ck_tile::sequence<2, 2>, + 2>; + +using Policy = ck_tile::ConvIgemmFwdPipelineDefaultPolicy; + +using TilePartitioner = ck_tile::ConvIgemmTilePartitioner; + +static constexpr auto ConvSpec = ck_tile::ConvFwdSpec::Default; + +template +bool run_grouped_conv_fwd(const ck_tile::conv::ConvParam& conv_param, + bool time_kernel, + double rtol, + double atol) +{ + static constexpr ck_tile::index_t NDim = NDimSpatial + 3; + + const auto in_g_n_c_wis_desc = + ck_tile::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed(conv_param); + + const auto wei_g_k_c_xs_desc = + ck_tile::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed(conv_param); + + const auto out_g_n_k_wos_desc = + ck_tile::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed(conv_param); + + ck_tile::HostTensor in(in_g_n_c_wis_desc); + ck_tile::HostTensor wei(wei_g_k_c_xs_desc); + ck_tile::HostTensor out_host(out_g_n_k_wos_desc); + ck_tile::HostTensor out_device(out_g_n_k_wos_desc); + + std::cout << "in: " << in.mDesc << std::endl; + std::cout << "wei: " << wei.mDesc << std::endl; + std::cout << "out: " << out_host.mDesc << std::endl; + + ck_tile::FillUniformDistribution{0.0f, 1.0f}(in); + ck_tile::FillUniformDistribution{-0.05f, 0.05f}(wei); + + ck_tile::DeviceMem in_device_buf(in.get_element_space_size_in_bytes()); + ck_tile::DeviceMem wei_device_buf(wei.get_element_space_size_in_bytes()); + ck_tile::DeviceMem out_device_buf(out_device.get_element_space_size_in_bytes()); + + in_device_buf.ToDevice(in.data()); + wei_device_buf.ToDevice(wei.data()); + out_host.SetZero(); + out_device.SetZero(); + + std::array a_g_n_c_wis_lengths{}; + std::array a_g_n_c_wis_strides{}; + std::array b_g_k_c_xs_lengths{}; + std::array b_g_k_c_xs_strides{}; + std::array c_g_n_k_wos_lengths{}; + std::array c_g_n_k_wos_strides{}; + std::array conv_filter_strides{}; + std::array conv_filter_dilations{}; + std::array input_left_pads{}; + std::array input_right_pads{}; + + auto copy = [](const auto& x, auto& y) { std::copy(x.begin(), x.end(), y.begin()); }; + + copy(in_g_n_c_wis_desc.get_lengths(), a_g_n_c_wis_lengths); + copy(in_g_n_c_wis_desc.get_strides(), a_g_n_c_wis_strides); + copy(wei_g_k_c_xs_desc.get_lengths(), b_g_k_c_xs_lengths); + copy(wei_g_k_c_xs_desc.get_strides(), b_g_k_c_xs_strides); + copy(out_g_n_k_wos_desc.get_lengths(), c_g_n_k_wos_lengths); + copy(out_g_n_k_wos_desc.get_strides(), c_g_n_k_wos_strides); + copy(conv_param.conv_filter_strides_, conv_filter_strides); + copy(conv_param.conv_filter_dilations_, conv_filter_dilations); + copy(conv_param.input_left_pads_, input_left_pads); + copy(conv_param.input_right_pads_, input_right_pads); + + // gpu conv + using Problem = ck_tile::ConvIgemmPipelineProblem< + NDimSpatial, + ck_tile::tuple, + TileShape, + TileTraits, + NumPrefetch, + true, + true>; + + using EpilogueProblem = ck_tile::LdsCShuffleEpilogueProblem< + AccDataType, + CDataType, + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence, + true>; + + using WG = typename ck_tile::conv::WarpGemmMmacDispatcher; + + using Epilogue = ck_tile::LdsCShuffleEpilogue; + + using Pipeline = ck_tile::ConvIgemmPipelineAGmemBGmemCRegV1; + + using Kernel = ck_tile::ConvIgemmFwdKernel; + + auto host_args = Kernel::MakeHostArgs(in_device_buf.GetDeviceBuffer(), + wei_device_buf.GetDeviceBuffer(), + out_device_buf.GetDeviceBuffer(), + a_g_n_c_wis_lengths, + a_g_n_c_wis_strides, + b_g_k_c_xs_lengths, + b_g_k_c_xs_strides, + c_g_n_k_wos_lengths, + c_g_n_k_wos_strides, + conv_filter_strides, + conv_filter_dilations, + input_left_pads, + input_right_pads); + + auto kernel_args = Kernel::MakeKernelArgs(host_args); + + const auto grid_size = + kernel_args.tile_partitioner.CalculateGridSize() * conv_param.G_; + + float avg_time = ck_tile::launch_kernel( + {nullptr, time_kernel}, + ck_tile::make_kernel( + Kernel{}, grid_size, BlockSize, 0, kernel_args)); + + if(time_kernel) + { + std::size_t flop = conv_param.GetFlops(); + std::size_t num_btype = conv_param.GetByte(); + + float tflops = static_cast(flop) / 1.E9 / avg_time; + float gb_per_sec = num_btype / 1.E6 / avg_time; + std::cout << "Perf: " << avg_time << " ms, " << tflops << " TFlops, " << gb_per_sec + << " GB/s" << std::endl; + } + else + { + std::cout << "timing:off" << std::endl; + } + + auto ref_conv = ck_tile::ReferenceConvFwd(); + + auto ref_invoker = ref_conv.MakeInvoker(); + auto ref_argument = ref_conv.MakeArgument(in, + wei, + out_host, + conv_param.conv_filter_strides_, + conv_param.conv_filter_dilations_, + conv_param.input_left_pads_, + conv_param.input_right_pads_, + AElementOp{}, + BElementOp{}, + CElementOp{}); + + ref_invoker.Run(ref_argument); + + out_device_buf.FromDevice(out_device.data()); + + return ck_tile::check_err(out_device, out_host, "Error: incorrect results!", rtol, atol); +} + +bool run_grouped_conv_fwd_example(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "2", "spatial dim") + .insert("g", "1", "group") + .insert("n", "64", "batch size") + .insert("k", "64", "out channels") + .insert("c", "64", "in channels") + .insert("r", "3", "filter x") + .insert("s", "3", "filter y") + .insert("h", "4", "input x") + .insert("w", "4", "input y") + .insert("u", "1", "stride h") + .insert("v", "1", "stride w") + .insert("l", "1", "dilation h") + .insert("j", "1", "dilation w") + .insert("p", "1", "pad h") + .insert("q", "1", "pad w") + .insert("time", "1", "time kernel (0=no, 1=yes)") + .insert("rtol", "1e-4", "relative validation tolerance") + .insert("atol", "1e-4", "absolute validation tolerance"); + + bool result = arg_parser.parse(argc, argv); + + if(!result) + return false; + + ck_tile::conv::ConvParam conv_param{arg_parser.get_int("m"), + arg_parser.get_int("g"), + arg_parser.get_int("n"), + arg_parser.get_int("k"), + arg_parser.get_int("c"), + {arg_parser.get_int("r"), arg_parser.get_int("s")}, + {arg_parser.get_int("h"), arg_parser.get_int("w")}, + {arg_parser.get_int("u"), arg_parser.get_int("v")}, + {arg_parser.get_int("l"), arg_parser.get_int("j")}, + {arg_parser.get_int("p"), arg_parser.get_int("p")}, + {arg_parser.get_int("q"), arg_parser.get_int("q")}}; + + if(arg_parser.get_int("m") == 2) + { + return run_grouped_conv_fwd<2>( + conv_param, + arg_parser.get_bool("time"), + arg_parser.get_double("rtol"), + arg_parser.get_double("atol")); + } + else + { + printf("unsupport spatial dim\n"); + return false; + } +} + +int main(int argc, char* argv[]) +{ + const auto target = ck_tile::get_hcu_target_enum(); + if(target != ck_tile::hcu_target_enum::gfx928 && target != ck_tile::hcu_target_enum::gfx92a && + target != ck_tile::hcu_target_enum::gfx936 && target != ck_tile::hcu_target_enum::gfx938) + { + std::cerr << "fp32 MMAC conv is not supported on " << ck_tile::get_device_name() + << std::endl; + return 77; + } + + return run_grouped_conv_fwd_example(argc, argv) ? 0 : 1; +} diff --git a/include/ck_tile/ops/conv/warp/warp_gemm_dispatcher.hpp b/include/ck_tile/ops/conv/warp/warp_gemm_dispatcher.hpp index 2972a7d..4dd00c8 100644 --- a/include/ck_tile/ops/conv/warp/warp_gemm_dispatcher.hpp +++ b/include/ck_tile/ops/conv/warp/warp_gemm_dispatcher.hpp @@ -130,6 +130,58 @@ struct WarpGemmMmacDispatcher>; }; +template +struct WarpGemmMmacDispatcher +{ + using Type = + WarpGemmImpl>; +}; + +template +struct WarpGemmMmacDispatcher +{ + using Type = WarpGemmImpl< + WarpGemmAttributeMmacIterateKTransC>; +}; + } // namespace impl template ; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + // mmac inst size + static constexpr index_t kM = 16; + static constexpr index_t kN = 16; + static constexpr index_t kK = 8; + + // thread layout + static constexpr index_t kAMLane = 16; + static constexpr index_t kBNLane = 16; + static constexpr index_t kABKLane = 4; + static constexpr index_t kABKPerLane = 2; // 2 element per thread vec + + // mmac C layout + static constexpr index_t kCMLane = 16; + static constexpr index_t kCNLane = 4; + + static constexpr index_t kCMPerLane = 1; + static constexpr index_t kCNPerLane = 4; + static constexpr index_t kCN0PerLane = 4; // jump layout + static constexpr index_t kCN1PerLane = 1; + static constexpr index_t kCM1PerLane = kCMPerLane; + static constexpr index_t kCM0PerLane = kCMPerLane; + + // c_vec += a_vec * b_vec + CK_TILE_DEVICE void + operator()(CVecType& c_vec, const AVecType& a_vec, const BVecType& b_vec) const + { +#if defined(__gfx928__) || defined(__gfx936__) || defined(__gfx938__) + c_vec = __builtin_hcu_mmac_16x16x8_f32(a_vec, b_vec, c_vec); +#elif defined(__gfx926__) || defined(__gfx92a__) + c_vec = __builtin_hcu_mmac_16x16x4_f32(a_vec[0], b_vec[0], c_vec); + c_vec = __builtin_hcu_mmac_16x16x4_f32(a_vec[1], b_vec[1], c_vec); +#else + ignore = c_vec; + ignore = a_vec; + ignore = b_vec; +#endif + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { + CVecType c_vec{0.f}; + operator()(c_vec, a_vec, b_vec); + return c_vec; + } +}; + +// FP32 +// C transposed (operand swap, matching the conv TransC iterators like F16 TransCLegacy) +struct WarpGemmAttributeMmacImplF32F32F32M16N16K8TransC + : public WarpGemmAttributeMmacImplF32F32F32M16N16K8 +{ + using Base = WarpGemmAttributeMmacImplF32F32F32M16N16K8; + + using ADataType = typename Base::ADataType; + using BDataType = typename Base::BDataType; + using CDataType = typename Base::CDataType; + + using AVecType = typename Base::AVecType; + using BVecType = typename Base::BVecType; + using CVecType = typename Base::CVecType; + + // mmac inst size + static constexpr index_t kM = Base::kM; + static constexpr index_t kN = Base::kN; + static constexpr index_t kK = Base::kK; + + // thread layout + static constexpr index_t kAMLane = Base::kAMLane; + static constexpr index_t kBNLane = Base::kBNLane; + static constexpr index_t kABKLane = Base::kABKLane; + static constexpr index_t kABKPerLane = Base::kABKPerLane; // 2 element per thread vec + + // tranposed mmac C layout + static constexpr index_t kCMLane = Base::kCNLane; + static constexpr index_t kCNLane = Base::kCMLane; + + static constexpr index_t kCNPerLane = Base::kCMPerLane; + static constexpr index_t kCM0PerLane = Base::kCN0PerLane; + static constexpr index_t kCM1PerLane = Base::kCN1PerLane; + + // FIXME: workaournd + static constexpr index_t kCMPerLane = -1; + static constexpr index_t kCN0PerLane = -1; + static constexpr index_t kCN1PerLane = -1; + + // c_vec += b_vec * a_vec + CK_TILE_DEVICE void + operator()(CVecType& c_vec, const AVecType& a_vec, const BVecType& b_vec) const + { + Base::operator()(c_vec, b_vec, a_vec); + } + + // c_vec = b_vec * a_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { + return Base::operator()(b_vec, a_vec); + } +}; + //int8 struct WarpGemmAttributeMmacImplI8I8I32M16N16K32 { From f0d5b11c448b3dac2e088d32b60270f10008474e Mon Sep 17 00:00:00 2001 From: wangwanqi Date: Fri, 16 Oct 2026 17:38:47 +0800 Subject: [PATCH 11/11] feat(gemm): add IsSupportedArgument for gemm --- example_hcu/ck_tile/03_gemm/gemm_basic.cpp | 5 + .../ck_tile/ops/gemm/kernel/gemm_kernel.hpp | 107 ++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/example_hcu/ck_tile/03_gemm/gemm_basic.cpp b/example_hcu/ck_tile/03_gemm/gemm_basic.cpp index 540ad21..8f8861c 100644 --- a/example_hcu/ck_tile/03_gemm/gemm_basic.cpp +++ b/example_hcu/ck_tile/03_gemm/gemm_basic.cpp @@ -101,6 +101,11 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s) const dim3 grids = Kernel::GridSize(args.M, args.N, args.kbatch); 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:" diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index d9bb750..9318b86 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -67,6 +67,113 @@ struct GemmKernel { return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); } + + CK_TILE_HOST static bool IsSupportedArgument(const GemmCommonKargs& kargs) + { + // Vector-load divisibility is only required on the paths this kernel actually vectorizes + // (A RowMajor along K, B ColumnMajor along K, C RowMajor along N); the transposed layouts + // are loaded/stored with vector length 1 and impose no such constraint. + if constexpr(std::is_same_v) + { + if(kargs.K % TilePartitioner::kK != 0 && GemmPipeline::kPadK == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support K that is not a multiple of KPerBlock without padding!"); + } + return false; + } + if(kargs.K % GemmPipeline::VectorSizeA != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("K is not a multiple of vector load size for A tensor!"); + } + return false; + } + } + else + { + if(kargs.M % TilePartitioner::kM != 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!"); + } + return false; + } + } + + if constexpr(std::is_same_v) + { + if(kargs.N % TilePartitioner::kN != 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!"); + } + return false; + } + } + else + { + if(kargs.K % TilePartitioner::kK != 0 && GemmPipeline::kPadK == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support K that is not a multiple of KPerBlock without padding!"); + } + return false; + } + if(kargs.K % GemmPipeline::VectorSizeB != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("K is not a multiple of vector load size for B tensor!"); + } + return false; + } + } + + if constexpr(std::is_same_v) + { + if(kargs.N % TilePartitioner::kN != 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!"); + } + return false; + } + if(kargs.N % GemmPipeline::VectorSizeC != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("N is not a multiple of vector load size for C tensor!"); + } + return false; + } + } + else + { + if(kargs.M % TilePartitioner::kM != 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!"); + } + return false; + } + } + + return true; + } CK_TILE_DEVICE void operator()(GemmCommonKargs kargs) const {