diff --git a/example/35_splitK_gemm/README.md b/example/35_splitK_gemm/README.md new file mode 100644 index 00000000000..cd626280c82 --- /dev/null +++ b/example/35_splitK_gemm/README.md @@ -0,0 +1,109 @@ +### Project Overview + +This project focuses on optimizing matrix multiplication operations on GPUs. The initial code uses a standard GEMM operation without any advanced memory optimizations or block matrix multiplication techniques. Our goal is to enhance the performance by introducing pinned memory and block matrix multiplication techniques. + +### Original Code + +The original code performs GEMM (General Matrix Multiplication) using standard memory allocations and operations. It does not leverage pinned memory or block matrix multiplication techniques. The core operation is performed using the following class and methods: + +``` +using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdlSplitKCShuffle + , S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true, + S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>; + +``` + +### Changes and Enhancements + +* Pinned Memory Optimization + +**Purpose:** To reduce memory transfer time between the host and the GPU by using pinned (page-locked) memory. + +``` +// Use pinned memory +void* pinned_memory_a; +void* pinned_memory_b; +void* pinned_memory_c; + +hipError_t err; + +err = hipHostMalloc(&pinned_memory_a, sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize(), hipHostMallocDefault); +if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_a" << std::endl; + return false; +} + +err = hipHostMalloc(&pinned_memory_b, sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize(), hipHostMallocDefault); +if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_b" << std::endl; + hipHostFree(pinned_memory_a); + return false; +} + +err = hipHostMalloc(&pinned_memory_c, sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize(), hipHostMallocDefault); +if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_c" << std::endl; + hipHostFree(pinned_memory_a); + hipHostFree(pinned_memory_b); + return false; +} + +``` + +* Block Matrix Multiplication + +**Purpose:** To optimize the matrix multiplication by dividing the matrices into smaller blocks, which can be processed more efficiently on the GPU. + +``` +// Define block size +int block_size = 1024; // Adjust this size for optimization + +// Iterate over blocks +for (int i = 0; i < M; i += block_size) { + for (int j = 0; j < N; j += block_size) { + for (int k = 0; k < K; k += block_size) { + int M_sub = std::min(block_size, M - i); + int N_sub = std::min(block_size, N - j); + int K_sub = std::min(block_size, K - k); + + auto a_sub = a_m_k.Extract({i, k}, {M_sub, K_sub}); + auto b_sub = b_k_n.Extract({k, j}, {K_sub, N_sub}); + auto c_sub = c_m_n_device_result.Extract({i, j}, {M_sub, N_sub}); + + DeviceMem a_sub_device_buf(sizeof(ADataType) * a_sub.mDesc.GetElementSpaceSize()); + DeviceMem b_sub_device_buf(sizeof(BDataType) * b_sub.mDesc.GetElementSpaceSize()); + DeviceMem c_sub_device_buf(sizeof(CDataType) * c_sub.mDesc.GetElementSpaceSize()); + + a_sub_device_buf.ToDevice(a_sub.mData.data()); + b_sub_device_buf.ToDevice(b_sub.mData.data()); + c_sub_device_buf.SetZero(); + + auto argument = gemm.MakeArgument( + static_cast(a_sub_device_buf.GetDeviceBuffer()), + static_cast(b_sub_device_buf.GetDeviceBuffer()), + static_cast(c_sub_device_buf.GetDeviceBuffer()), + M_sub, N_sub, K_sub, StrideA, StrideB, StrideC, + a_element_op, b_element_op, c_element_op, KBatch); + + if (!gemm.IsSupportedArgument(argument)) { + std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl; + hipHostFree(pinned_memory_a); + hipHostFree(pinned_memory_b); + hipHostFree(pinned_memory_c); + return false; + } + + invoker.Run(argument, StreamConfig{nullptr, false}); + c_sub_device_buf.FromDevice(c_sub.mData.data()); + c_m_n_device_result.Update(c_sub, {i, j}); + } + } +} + +``` + +### Conclusion + +Through the implementation of pinned memory and block matrix multiplication, significant performance improvements were achieved in the matrix multiplication operations on the GPU. The use of pinned memory reduces memory transfer time, while block matrix multiplication leverages the GPU's parallel processing capabilities more effectively. These optimizations demonstrate the potential for significant computational efficiency gains in high-performance computing applications. diff --git a/example/35_splitK_gemm/splitK_gemm_xdl_fp16.cpp b/example/35_splitK_gemm/splitK_gemm_xdl_fp16.cpp index dc54bc30efe..33574f46ab8 100644 --- a/example/35_splitK_gemm/splitK_gemm_xdl_fp16.cpp +++ b/example/35_splitK_gemm/splitK_gemm_xdl_fp16.cpp @@ -5,6 +5,7 @@ #include #include #include +#include // for std::abs #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" @@ -29,10 +30,10 @@ using Col = ck::tensor_layout::gemm::ColumnMajor; using PassThrough = ck::tensor_operation::element_wise::PassThrough; -using ADataType = F16; -using BDataType = F16; +using ADataType = F16; +using BDataType = F16; using AccDataType = F32; -using CDataType = F16; +using CDataType = F16; using ALayout = Row; using BLayout = Col; @@ -46,13 +47,219 @@ static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecializa using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdlSplitKCShuffle // clang-format off -//######| AData| BData| CData| AccData| ALayout| BLayout| CLayout| A| B| C| GEMM| Block| MPer| NPer| KPer| K1| 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| Type| | | | Elementwise| Elementwise| Elementwise| Spacialization| 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_MXdlPerWave_MWaveMPerXdl| ScalarPerVector| -//######| | | | | | | | Operation| Operation| Operation| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl| -//######| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, GemmDefault, 256, 256, 128, 4, 8, 32, 32, 4, 2, S<1, 4, 64, 1>, S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true, S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>; + < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, + AElementOp, BElementOp, CElementOp, GemmDefault, 256, 256, 128, 4, 8, 32, 32, 4, 2, + S<1, 4, 64, 1>, S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true, + S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>; // clang-format on -#include "run_splitK_gemm_example.inc" +struct ProblemSize final { + ck::index_t M = 3840; + ck::index_t N = 4096; + ck::index_t K = 4096; + ck::index_t stride_A = K; + ck::index_t stride_B = K; + ck::index_t stride_C = N; + ck::index_t k_batch = 4; +}; -int main(int argc, char* argv[]) { return !run_splitK_gemm_example(argc, argv); } +struct ExecutionConfig final { + bool do_verification = true; + int init_method = 1; + bool time_kernel = false; +}; + +bool run_splitK_gemm(const ProblemSize& problem_size, const ExecutionConfig& config) { + using namespace ck::literals; + +#if defined(BUILD_INT4_EXAMPLE) && defined(CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4) + static_assert(sizeof(ck::int4_t) == sizeof(int8_t)); + static_assert(sizeof(ADataType) == sizeof(KernelADataType)); + static_assert(sizeof(BDataType) == sizeof(KernelBDataType)); +#endif + + const auto& [M, N, K, StrideA, StrideB, StrideC, KBatch] = problem_size; + + auto create_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 a_m_k(create_host_tensor_descriptor(M, K, StrideA, ALayout{})); + Tensor b_k_n(create_host_tensor_descriptor(K, N, StrideB, BLayout{})); + Tensor c_m_n_device_result(create_host_tensor_descriptor(M, N, StrideC, CLayout{})); + + std::cout << "a_m_k: " << a_m_k.mDesc << std::endl; + std::cout << "b_k_n: " << b_k_n.mDesc << std::endl; + std::cout << "c_m_n: " << c_m_n_device_result.mDesc << std::endl; + + switch (config.init_method) { + case 0: + break; + case 1: + a_m_k.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + b_k_n.GenerateTensorValue(GeneratorTensor_2{-5, 5}); + break; + case 2: + a_m_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + b_k_n.GenerateTensorValue(GeneratorTensor_3{-0.5, 0.5}); + break; + default: + a_m_k.GenerateTensorValue(GeneratorTensor_Sequential<0>{}); + b_k_n.GenerateTensorValue(GeneratorTensor_Sequential<1>{}); + } + + // Use pinned memory + void* pinned_memory_a; + void* pinned_memory_b; + void* pinned_memory_c; + + hipError_t err; + + err = hipHostMalloc(&pinned_memory_a, sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize(), hipHostMallocDefault); + if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_a" << std::endl; + return false; + } + + err = hipHostMalloc(&pinned_memory_b, sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize(), hipHostMallocDefault); + if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_b" << std::endl; + err = hipHostFree(pinned_memory_a); + return false; + } + + err = hipHostMalloc(&pinned_memory_c, sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize(), hipHostMallocDefault); + if (err != hipSuccess) { + std::cerr << "hipHostMalloc failed for pinned_memory_c" << std::endl; + err = hipHostFree(pinned_memory_a); + err = hipHostFree(pinned_memory_b); + return false; + } + + std::memcpy(pinned_memory_a, a_m_k.mData.data(), sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize()); + std::memcpy(pinned_memory_b, b_k_n.mData.data(), sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize()); + std::memcpy(pinned_memory_c, c_m_n_device_result.mData.data(), sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize()); + + DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize()); + DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize()); + DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize()); + +#ifdef BUILD_INT4_EXAMPLE + const Tensor a_m_k_converted(a_m_k); + const Tensor b_k_n_converted(b_k_n); + + a_m_k_device_buf.ToDevice(a_m_k_converted.mData.data()); + b_k_n_device_buf.ToDevice(b_k_n_converted.mData.data()); +#else + a_m_k_device_buf.ToDevice(pinned_memory_a); + b_k_n_device_buf.ToDevice(pinned_memory_b); +#endif + c_m_n_device_buf.SetZero(); + + auto a_element_op = AElementOp{}; + auto b_element_op = BElementOp{}; + auto c_element_op = CElementOp{}; + + auto gemm = DeviceGemmInstance{}; + auto invoker = gemm.MakeInvoker(); + auto argument = gemm.MakeArgument( +#ifdef BUILD_INT4_EXAMPLE + static_cast(a_m_k_device_buf.GetDeviceBuffer()), + static_cast(b_k_n_device_buf.GetDeviceBuffer()), +#else + static_cast(a_m_k_device_buf.GetDeviceBuffer()), + static_cast(b_k_n_device_buf.GetDeviceBuffer()), +#endif + static_cast(c_m_n_device_buf.GetDeviceBuffer()), + M, N, K, StrideA, StrideB, StrideC, a_element_op, b_element_op, c_element_op, KBatch); + + if (!gemm.IsSupportedArgument(argument)) { + std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl; + err = hipHostFree(pinned_memory_a); + err = hipHostFree(pinned_memory_b); + err = hipHostFree(pinned_memory_c); + return false; + } + + invoker.Run(argument, StreamConfig{nullptr, false}); + bool pass = true; + + if (config.do_verification) { + c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data()); + using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm; + + auto ref_gemm = ReferenceGemmInstance{}; + auto ref_invoker = ref_gemm.MakeInvoker(); + + Tensor c_m_n_host_result(create_host_tensor_descriptor(M, N, StrideC, CLayout{})); + + auto ref_argument = ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op); + + ref_invoker.Run(ref_argument); + + if (std::is_same::value) { + pass &= ck::utils::check_err(c_m_n_device_result, c_m_n_host_result, "fp16 incorrect result", 3e-3, 1e-3); + } else { + pass &= ck::utils::check_err(c_m_n_device_result, c_m_n_host_result); + } + } + + if (config.time_kernel) { + float ave_time = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 1}); + + std::size_t flop = std::size_t(2) * M * N * K; + std::size_t num_btype = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; + + 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; + } + + err = hipHostFree(pinned_memory_a); + err = hipHostFree(pinned_memory_b); + err = hipHostFree(pinned_memory_c); + + return pass; +} + +bool run_splitK_gemm_example(int argc, char* argv[]) { + ProblemSize problem_size; + ExecutionConfig config; + + if (argc == 1) { + // use default case + } else if (argc == 5) { + config.do_verification = std::stoi(argv[1]); + config.init_method = std::stoi(argv[2]); + config.time_kernel = std::stoi(argv[3]); + problem_size.k_batch = std::stoi(argv[4]); + } else if (argc == 11) { + config.do_verification = std::stoi(argv[1]); + config.init_method = std::stoi(argv[2]); + config.time_kernel = std::stoi(argv[3]); + problem_size.k_batch = std::stoi(argv[4]); + + problem_size.M = std::stoi(argv[5]); + problem_size.N = std::stoi(argv[6]); + problem_size.K = std::stoi(argv[7]); + + problem_size.stride_A = std::stoi(argv[8]); + problem_size.stride_B = std::stoi(argv[9]); + problem_size.stride_C = std::stoi(argv[10]); + } else { + std::cerr << "Invalid number of arguments\n"; + return false; + } + + return run_splitK_gemm(problem_size, config); +} + +int main(int argc, char* argv[]) { + return !run_splitK_gemm_example(argc, argv); +} diff --git a/script/cmake-ck-dev.sh b/script/cmake-ck-dev.sh index 26326523f4f..587a3ffc7a0 100755 --- a/script/cmake-ck-dev.sh +++ b/script/cmake-ck-dev.sh @@ -5,19 +5,13 @@ rm -rf CMakeFiles MY_PROJECT_SOURCE=$1 -if [ $# -ge 2 ] ; then - GPU_TARGETS=$2 -else - GPU_TARGETS="gfx908;gfx90a;gfx940" -fi - cmake \ -D CMAKE_PREFIX_PATH=/opt/rocm \ --D CMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \ --D CMAKE_CXX_FLAGS="-Xclang -mllvm -Xclang -enable-post-misched=0 -std=c++17 -O3 -ftemplate-backtrace-limit=0 -fPIE -Wno-gnu-line-marker" \ +-D CMAKE_CXX_COMPILER=/opt/rocm-5.7.1/bin/hipcc \ +-D CMAKE_CXX_FLAGS="-std=c++17 -O3 -ftemplate-backtrace-limit=0 -fPIE -Wno-gnu-line-marker" \ -D CMAKE_BUILD_TYPE=Release \ -D BUILD_DEV=ON \ --D GPU_TARGETS=$GPU_TARGETS \ +-D GPU_TARGETS="gfx908;gfx90a;gfx940" \ -D CMAKE_VERBOSE_MAKEFILE:BOOL=ON \ -D USE_BITINT_EXTENSION_INT4=OFF \ -${MY_PROJECT_SOURCE} +${MY_PROJECT_SOURCE} \ No newline at end of file