From f39354d7cb16d5540dc099b0f50417ce0ddea21f Mon Sep 17 00:00:00 2001 From: cloudhan Date: Fri, 5 Aug 2022 08:32:20 +0800 Subject: [PATCH] Add composable kernel GEMM baseline for kernel explorer (#12364) * Split GemmBase RocBlasGemm * Add composable kernel GEMM baseline * Make linter happy * Address review comment * Update bert cases with batchsize * Adjust includes to fix IWYU lint * Only builds and links used ck kernels to improve building time * Remove warmup run on SelectImpl * Add comment to utility function * Mute cpplint * Make RocBlasGemm::SelectImpl semantically correct * Add reduced basic test cases for ck gemm * More robust gemm testing * Fix warnings * Fix grammar --- cmake/external/composable_kernel.cmake | 19 ++ cmake/onnxruntime_kernel_explorer.cmake | 9 +- .../composable_kernel/Fix_Clang_Build.patch | 53 +++++ .../tools/kernel_explorer/kernels/gemm.cc | 112 +---------- .../tools/kernel_explorer/kernels/gemm.h | 60 +++++- .../tools/kernel_explorer/kernels/gemm_ck.cc | 187 ++++++++++++++++++ .../tools/kernel_explorer/kernels/gemm_ck.h | 14 ++ .../kernel_explorer/kernels/gemm_rocblas.cc | 92 +++++++++ .../kernel_explorer/kernels/gemm_rocblas.h | 14 ++ .../kernel_explorer/kernels/gemm_test.py | 157 +++++++++------ 10 files changed, 550 insertions(+), 167 deletions(-) create mode 100644 cmake/external/composable_kernel.cmake create mode 100644 cmake/patches/composable_kernel/Fix_Clang_Build.patch create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc create mode 100644 onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h diff --git a/cmake/external/composable_kernel.cmake b/cmake/external/composable_kernel.cmake new file mode 100644 index 0000000000..0803ffb16d --- /dev/null +++ b/cmake/external/composable_kernel.cmake @@ -0,0 +1,19 @@ +set(composable_kernel_URL https://github.com/ROCmSoftwarePlatform/composable_kernel.git) +set(composable_kernel_TAG 85978e0201bb94bf6e59b325e1f5f19266845d08) # 2022-07-22 11:52:10 -0700 + +set(BUILD_DEV OFF) +set(PATCH ${PROJECT_SOURCE_DIR}/patches/composable_kernel/Fix_Clang_Build.patch) + +include(FetchContent) +FetchContent_Declare(composable_kernel + GIT_REPOSITORY ${composable_kernel_URL} + GIT_TAG ${composable_kernel_TAG} + PATCH_COMMAND git apply --reverse --check ${PATCH} || git apply --ignore-space-change --ignore-whitespace ${PATCH} +) + +FetchContent_MakeAvailable(composable_kernel) + +add_library(onnxruntime_composable_kernel_includes INTERFACE) +target_include_directories(onnxruntime_composable_kernel_includes INTERFACE + ${composable_kernel_SOURCE_DIR}/include + ${composable_kernel_SOURCE_DIR}/library/include) diff --git a/cmake/onnxruntime_kernel_explorer.cmake b/cmake/onnxruntime_kernel_explorer.cmake index d16b844cb6..ad91139297 100644 --- a/cmake/onnxruntime_kernel_explorer.cmake +++ b/cmake/onnxruntime_kernel_explorer.cmake @@ -13,6 +13,8 @@ if(NOT HIP_FOUND) message(FATAL_ERROR "hip is required but is not found") endif() +include(composable_kernel) + set(KERNEL_EXPLORER_ROOT ${ONNXRUNTIME_ROOT}/python/tools/kernel_explorer) set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/rocm/bert) @@ -30,6 +32,11 @@ target_include_directories(kernel_explorer PUBLIC target_link_libraries(kernel_explorer PRIVATE $ + onnxruntime_composable_kernel_includes + # Currently we shall not use composablekernels::device_operations, the target includes all conv dependencies, which + # are extremely slow to compile. Instead, we only link all gemm related objects. See the following link on updating. + # https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/85978e0201/library/src/tensor_operation_instance/gpu/CMakeLists.txt#L33-L54 + device_gemm_instance ${HIP_LIB}) target_compile_definitions(kernel_explorer PUBLIC ROCM_USE_FLOAT16 @@ -38,7 +45,7 @@ target_compile_definitions(kernel_explorer # handle kernel_explorer sources as hip language target_compile_options(kernel_explorer PRIVATE "-xhip") # TODO: use predefined AMDGPU_TARGETS -target_compile_options(kernel_explorer PRIVATE "--offload-arch=gfx906" "--offload-arch=gfx908" "--offload-arch=gfx90a") +target_compile_options(kernel_explorer PRIVATE "--offload-arch=gfx908" "--offload-arch=gfx90a") # https://github.com/ROCm-Developer-Tools/HIP/blob/4514f350849b1090954295f8f87a5f8d78bd781b/hip-lang-config.cmake.in target_link_libraries(kernel_explorer PRIVATE ${CLANGRT_BUILTINS}) diff --git a/cmake/patches/composable_kernel/Fix_Clang_Build.patch b/cmake/patches/composable_kernel/Fix_Clang_Build.patch new file mode 100644 index 0000000000..931e736074 --- /dev/null +++ b/cmake/patches/composable_kernel/Fix_Clang_Build.patch @@ -0,0 +1,53 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 9f706207..fb444227 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -8,7 +8,7 @@ list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake") + enable_testing() + + set(ROCM_SYMLINK_LIBS OFF) +-find_package(ROCM 0.8 REQUIRED PATHS /opt/rocm) ++find_package(ROCM 0.7 REQUIRED PATHS /opt/rocm) + + include(ROCMInstallTargets) + include(ROCMPackageConfigHelpers) +@@ -28,6 +28,17 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) + set(CMAKE_CXX_EXTENSIONS OFF) + message("CMAKE_CXX_COMPILER_ID: ${CMAKE_CXX_COMPILER_ID}") + ++add_compile_options( ++ -xhip ++ --offload-arch=gfx908 ++ --offload-arch=gfx90a ++ -O3 # otherwise, "Illegal instruction detected" for gfx908 ++ "SHELL:-mllvm -amdgpu-early-inline-all=true" ++ "SHELL:-mllvm -amdgpu-function-calls=false" # otherwise, "local memory (65920) exceeds limit (65536) in function" ++ -fhip-new-launch-api ++) ++ ++ + ## OpenMP + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + # workaround issue hipcc in rocm3.5 cannot find openmp +@@ -237,9 +248,6 @@ rocm_package_setup_component(tests + ) + + add_subdirectory(library) +-add_subdirectory(example) +-add_subdirectory(test) +-add_subdirectory(profiler) + + #Create an interface target for the include only files and call it "composablekernels" + include(CMakePackageConfigHelpers) +@@ -265,11 +273,3 @@ rocm_install(FILES + + set(CPACK_RESOURCE_FILE_LICENSE "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE") + set(CPACK_RPM_PACKAGE_LICENSE "MIT") +- +-rocm_create_package( +- NAME composablekernel +- DESCRIPTION "High Performance Composable Kernel for AMD GPUs" +- MAINTAINER "MIOpen Kernels Dev Team " +- LDCONFIG +- HEADER_ONLY +-) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc index bcf674c911..1cd5024e6b 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc @@ -2,109 +2,16 @@ // Licensed under the MIT License. #include "python/tools/kernel_explorer/kernels/gemm.h" +#include "python/tools/kernel_explorer/kernels/gemm_rocblas.h" +#include "python/tools/kernel_explorer/kernels/gemm_ck.h" #include #include -#include "core/providers/rocm/rocm_common.h" -#include "core/providers/rocm/shared_inc/fpgeneric.h" -#include "python/tools/kernel_explorer/device_array.h" -#include "python/tools/kernel_explorer/operator.h" - namespace py = pybind11; namespace onnxruntime { -enum class BlasOp { - N, - T, -}; - -// We don't assume the implementation is row-majored or column-majored. But for testing convenience, we assume all -// our wrappers have row-majored convention, since it is the native layout to numpy and pytorch. -template -class GemmBase : public Operator { - public: - GemmBase( - BlasOp opa, BlasOp opb, - int64_t m, int64_t n, int64_t k, - double alpha, - DeviceArray& a, int64_t lda, - DeviceArray& b, int64_t ldb, - double beta, - DeviceArray& c, int64_t ldc) - : Operator(), - opa_(opa), - opb_(opb), - m_(m), - n_(n), - k_(k), - alpha_(alpha), - a_(reinterpret_cast(a.ptr())), - lda_(lda), - b_(reinterpret_cast(b.ptr())), - ldb_(ldb), - beta_(beta), - c_(reinterpret_cast(c.ptr())), - ldc_(ldc) {} - - protected: - BlasOp opa_; - BlasOp opb_; - int64_t m_; - int64_t n_; - int64_t k_; - T alpha_; - T* a_; - int64_t lda_; - T* b_; - int64_t ldb_; - T beta_; - T* c_; - int64_t ldc_; -}; - -template -class RocBlasGemm : public GemmBase { - public: - RocBlasGemm(BlasOp opa, BlasOp opb, - int64_t m, int64_t n, int64_t k, - double alpha, - DeviceArray& a, int64_t lda, - DeviceArray& b, int64_t ldb, - double beta, - DeviceArray& c, int64_t ldc) - : GemmBase(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc), - opa_(opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose), - opb_(opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose) { - ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_)); - } - - ~RocBlasGemm() { - ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_)); - rocblas_handle_ = nullptr; - } - - void Run() { - // NOTE: rocblas assume the storage is column-majored, swapping A and B makes it have the same interface - // as those with row-majored convention. That is, if you treat the storage as row-majored but view the matrices as - // transposed, then by using the property Transpose(A*B) = Tranpose(B)*Transpose(A), the correctness is obvious. - ROCBLAS_CALL_THROW( - rocblasGemmHelper(this->rocblas_handle_, this->opb_, this->opa_, - this->n_, this->m_, this->k_, - &(this->alpha_), - this->b_, this->ldb_, - this->a_, this->lda_, - &(this->beta_), - this->c_, this->ldc_)); - } - - private: - rocblas_handle rocblas_handle_; - rocblas_operation opa_; - rocblas_operation opb_; -}; - void InitGemm(py::module mod) { auto blas_op = mod.def_submodule("blas_op"); @@ -113,19 +20,8 @@ void InitGemm(py::module mod) { .value("T", BlasOp::T, "Transpose") .export_values(); - // float - py::class_>(mod, "RocblasGemm_float") - .def(py::init()) - .def("SetRepeats", &RocBlasGemm::SetRepeats) - .def("Profile", &RocBlasGemm::Profile) - .def("Run", &RocBlasGemm::Run); - - // half - py::class_>(mod, "RocblasGemm_half") - .def(py::init()) - .def("SetRepeats", &RocBlasGemm::SetRepeats) - .def("Profile", &RocBlasGemm::Profile) - .def("Run", &RocBlasGemm::Run); + InitRocBlasGemm(mod); + InitComposableKernelGemm(mod); } } // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.h index 736030d517..006b98e55d 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.h +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.h @@ -3,12 +3,70 @@ #pragma once +#include +#include + #include +#include "python/tools/kernel_explorer/device_array.h" +#include "python/tools/kernel_explorer/operator.h" + namespace py = pybind11; namespace onnxruntime { +enum class BlasOp { + N, + T, +}; + +// We don't assume the implementation is row-majored or column-majored. But for testing convenience, we assume all +// our wrappers have row-majored convention, since it is the native layout to numpy and pytorch. +template +class GemmBase : public Operator { + public: + GemmBase( + BlasOp opa, BlasOp opb, + int64_t m, int64_t n, int64_t k, + double alpha, + DeviceArray& a, int64_t lda, + DeviceArray& b, int64_t ldb, + double beta, + DeviceArray& c, int64_t ldc) + : Operator(), + opa_(opa), + opb_(opb), + m_(m), + n_(n), + k_(k), + alpha_(alpha), + a_(reinterpret_cast(a.ptr())), + lda_(lda), + b_(reinterpret_cast(b.ptr())), + ldb_(ldb), + beta_(beta), + c_(reinterpret_cast(c.ptr())), + ldc_(ldc) {} + + virtual std::vector ListImpls() const = 0; + virtual bool SelectImpl(const std::string& name) = 0; + + protected: + BlasOp opa_; + BlasOp opb_; + int64_t m_; + int64_t n_; + int64_t k_; + T alpha_; + T* a_; + int64_t lda_; + T* b_; + int64_t ldb_; + T beta_; + T* c_; + int64_t ldc_; +}; + void InitGemm(py::module mod); -} +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc new file mode 100644 index 0000000000..ebfc92fe1e --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc @@ -0,0 +1,187 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "python/tools/kernel_explorer/kernels/gemm_ck.h" + +#include + +#include +#include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/device_gemm.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#include "ck/library/tensor_operation_instance/gpu/gemm.hpp" + +#include "python/tools/kernel_explorer/kernels/gemm.h" + +namespace py = pybind11; + +namespace onnxruntime { + +namespace { + +template +struct DataTypeAdaptor { + using type = T; +}; + +template <> +struct DataTypeAdaptor { + using type = ck::half_t; +}; + +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; + +} // namespace + +template +class CKGemm : public GemmBase { + using CKDataType = typename DataTypeAdaptor::type; + + template + using DeviceGemm = ck::tensor_operation::device::DeviceGemm< + ALayout, BLayout, Row, + CKDataType, CKDataType, CKDataType, + ck::tensor_operation::element_wise::PassThrough, + ck::tensor_operation::element_wise::PassThrough, + ck::tensor_operation::element_wise::PassThrough>; + + template + using CKGemmInstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory; + + // composable kernel's GetInstances() returns a collection of >>. We need to manually + // cast them to std::unique_ptr> to avoid type problem. + template + static std::vector> GetInstances() { + auto ptrs = CKGemmInstanceFactory::GetInstances(); + std::vector> results; + for (auto&& ptr : ptrs) { + results.emplace_back(std::move(ptr)); + } + return results; + } + + bool UpdateArgumentAndInvoker() { + auto nop = ck::tensor_operation::element_wise::PassThrough{}; + if (this->opa_ == BlasOp::N && this->opb_ == BlasOp::N) { + using DeviceOp = DeviceGemm; + auto typed_impl = static_cast(impls_[selected_impl_].get()); + arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_, + this->m_, this->n_, this->k_, + this->lda_, this->ldb_, this->ldc_, + nop, nop, nop); + invoker_ = typed_impl->MakeInvokerPointer(); + } else if (this->opa_ == BlasOp::T && this->opb_ == BlasOp::N) { + using DeviceOp = DeviceGemm; + auto typed_impl = static_cast(impls_[selected_impl_].get()); + arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_, + this->m_, this->n_, this->k_, + this->lda_, this->ldb_, this->ldc_, + nop, nop, nop); + invoker_ = typed_impl->MakeInvokerPointer(); + } else if (this->opa_ == BlasOp::N && this->opb_ == BlasOp::T) { + using DeviceOp = DeviceGemm; + auto typed_impl = static_cast(impls_[selected_impl_].get()); + arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_, + this->m_, this->n_, this->k_, + this->lda_, this->ldb_, this->ldc_, + nop, nop, nop); + invoker_ = typed_impl->MakeInvokerPointer(); + } else if (this->opa_ == BlasOp::T && this->opb_ == BlasOp::T) { + using DeviceOp = DeviceGemm; + auto typed_impl = static_cast(impls_[selected_impl_].get()); + arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_, + this->m_, this->n_, this->k_, + this->lda_, this->ldb_, this->ldc_, + nop, nop, nop); + invoker_ = typed_impl->MakeInvokerPointer(); + } + + return impls_[selected_impl_]->IsSupportedArgument(arg_.get()); + } + + public: + CKGemm(BlasOp opa, BlasOp opb, + int64_t m, int64_t n, int64_t k, + double alpha, + DeviceArray& a, int64_t lda, + DeviceArray& b, int64_t ldb, + double beta, + DeviceArray& c, int64_t ldc) + : GemmBase(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc) { + ORT_ENFORCE(alpha == 1.0); + ORT_ENFORCE(beta == 0.0); + + if (opa == BlasOp::N && opb == BlasOp::N) { + impls_ = this->GetInstances>(); + } else if (opa == BlasOp::T && opb == BlasOp::N) { + impls_ = this->GetInstances>(); + } else if (opa == BlasOp::N && opb == BlasOp::T) { + impls_ = this->GetInstances>(); + } else if (opa == BlasOp::T && opb == BlasOp::T) { + impls_ = this->GetInstances>(); + } + + ORT_ENFORCE(!impls_.empty()); + selected_impl_ = 0; + UpdateArgumentAndInvoker(); + } + + std::vector ListImpls() const override { + std::vector results; + std::transform(impls_.cbegin(), impls_.cend(), std::back_inserter(results), + [](const auto& it) { return it->GetTypeString(); }); + return results; + } + + bool SelectImpl(const std::string& name) override { + for (size_t i = 0; i < impls_.size(); i++) { + if (impls_[i]->GetTypeString() == name) { + selected_impl_ = i; + return UpdateArgumentAndInvoker(); + } + } + + ORT_THROW("Cannot find implementation ", name); + } + + void Run() override { + invoker_->Run(arg_.get()); + } + + private: + std::vector> impls_; + size_t selected_impl_{}; + std::unique_ptr arg_; + std::unique_ptr invoker_; +}; + +void InitComposableKernelGemm(py::module mod) { + // float + py::class_>(mod, "CKGemm_float") + .def(py::init()) + .def("SetRepeats", &CKGemm::SetRepeats) + .def("Profile", &CKGemm::Profile) + .def("Run", &CKGemm::Run) + .def("ListImpls", &CKGemm::ListImpls) + .def("SelectImpl", &CKGemm::SelectImpl); + + // half + py::class_>(mod, "CKGemm_half") + .def(py::init()) + .def("SetRepeats", &CKGemm::SetRepeats) + .def("Profile", &CKGemm::Profile) + .def("Run", &CKGemm::Run) + .def("ListImpls", &CKGemm::ListImpls) + .def("SelectImpl", &CKGemm::SelectImpl); +} + +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h new file mode 100644 index 0000000000..782de08556 --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +namespace py = pybind11; + +namespace onnxruntime { + +void InitComposableKernelGemm(py::module mod); + +} diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc new file mode 100644 index 0000000000..684fb8c9b9 --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "python/tools/kernel_explorer/kernels/gemm_rocblas.h" + +#include +#include + +#include +#include + +#include "core/providers/rocm/rocm_common.h" +#include "core/providers/rocm/shared_inc/fpgeneric.h" +#include "python/tools/kernel_explorer/device_array.h" +#include "python/tools/kernel_explorer/kernels/gemm.h" + +namespace py = pybind11; + +namespace onnxruntime { + +template +class RocBlasGemm : public GemmBase { + public: + RocBlasGemm(BlasOp opa, BlasOp opb, + int64_t m, int64_t n, int64_t k, + double alpha, + DeviceArray& a, int64_t lda, + DeviceArray& b, int64_t ldb, + double beta, + DeviceArray& c, int64_t ldc) + : GemmBase(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc), + opa_(opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose), + opb_(opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose) { + ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_)); + } + + ~RocBlasGemm() { + ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_)); + rocblas_handle_ = nullptr; + } + + void Run() override { + // NOTE: rocblas assumes the storage is column-majored, swapping A and B makes it have the same interface + // as those with row-majored convention. That is, if you treat the storage as row-majored but view the matrices as + // transposed, then by using the property Transpose(A*B) = Tranpose(B)*Transpose(A), the correctness is obvious. + ROCBLAS_CALL_THROW( + rocblasGemmHelper(this->rocblas_handle_, this->opb_, this->opa_, + this->n_, this->m_, this->k_, + &(this->alpha_), + this->b_, this->ldb_, + this->a_, this->lda_, + &(this->beta_), + this->c_, this->ldc_)); + } + + std::vector ListImpls() const override { + return {"Rocblas"}; + } + + bool SelectImpl(const std::string& name) override { + return name == "Rocblas"; + } + + private: + rocblas_handle rocblas_handle_; + rocblas_operation opa_; + rocblas_operation opb_; +}; + +void InitRocBlasGemm(py::module mod) { + // float + py::class_>(mod, "RocblasGemm_float") + .def(py::init()) + .def("SetRepeats", &RocBlasGemm::SetRepeats) + .def("Profile", &RocBlasGemm::Profile) + .def("Run", &RocBlasGemm::Run) + .def("ListImpls", &RocBlasGemm::ListImpls) + .def("SelectImpl", &RocBlasGemm::SelectImpl); + + // half + py::class_>(mod, "RocblasGemm_half") + .def(py::init()) + .def("SetRepeats", &RocBlasGemm::SetRepeats) + .def("Profile", &RocBlasGemm::Profile) + .def("Run", &RocBlasGemm::Run) + .def("ListImpls", &RocBlasGemm::ListImpls) + .def("SelectImpl", &RocBlasGemm::SelectImpl); +} + +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h new file mode 100644 index 0000000000..046177ed45 --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +namespace py = pybind11; + +namespace onnxruntime { + +void InitRocBlasGemm(py::module mod); + +} diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py index c29cef90a6..f3d5705303 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py @@ -10,70 +10,122 @@ import numpy as np import pytest -def dtype_to_bytes(dtype): - type_map = { - "float16": 2, - "float32": 4, - } - return type_map[dtype] - - -def dtype_to_funcs(dtype): - type_map = { - "float16": list(filter(lambda x: "RocblasGemm_half" in x, dir(ke))), - "float32": list(filter(lambda x: "RocblasGemm_float" in x, dir(ke))), - } - return type_map[dtype] +def dtype_to_suffix(dtype): + return { + "float32": "float", + "float16": "half", + }[dtype] def _test_gemm(func, dtype: str, m: int, n: int, k: int, transa=False, transb=False): + assert dtype in ["float32", "float16"] + a_shape = (k, m) if transa else (m, k) b_shape = (n, k) if transb else (k, n) np.random.seed(0) - a = (np.random.rand(*a_shape) * 2 - 1).astype(dtype) - b = (np.random.rand(*b_shape) * 2 - 1).astype(dtype) + a = (np.random.rand(*a_shape) + 0.5).astype(dtype).astype("float64") + b = (np.random.rand(*b_shape) + 0.5).astype(dtype).astype("float64") + ref_c = (a.T if transa else a) @ (b.T if transb else b) + + # The machine epsilon, unit roundoff, the smallest positive floating point number n such that the floating point + # number that represents 1 + n is greater than 1. + machine_eps = 2.0 ** -(24 if dtype == "float32" else 11) + + # The following implements error bound 5.7 in paper I. C. Ipsen and H. Zhou, “Probabilistic error analysis for + # Inner Products,” SIAM Journal on Matrix Analysis and Applications, vol. 41, no. 4, pp. 1726–1741, 2020. + # NOTE: the bound is not tight for float16 when k is large + absa_mul_absb = np.abs(a.T if transa else a) @ np.abs(b.T if transb else b) + coeff = np.max(absa_mul_absb / np.abs(ref_c)) + gamma_2k = (1.0 + machine_eps) ** (2 * k) - 1.0 + bound_5_7 = coeff * np.sqrt(np.log(2 / 1e-10) * machine_eps * gamma_2k / 2) + bound = bound_5_7 + + a = a.astype(dtype) + b = b.astype(dtype) my_c = np.zeros((m, n), dtype=dtype) dev_a = ke.DeviceArray(a) dev_b = ke.DeviceArray(b) dev_c = ke.DeviceArray(my_c) - f = getattr(ke, func) - opa = ke.blas_op.T if transa else ke.blas_op.N opb = ke.blas_op.T if transb else ke.blas_op.N lda = a_shape[1] ldb = b_shape[1] alpha = 1.0 beta = 0.0 - my_gemm = f(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n) - my_gemm.Run() - dev_c.UpdateHostNumpyArray() + my_gemm = func(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n) - rtol = 1e-3 if dtype == "float16" else 1e-5 - atol = 1e-3 if dtype == "float16" else 1e-5 + failures = {} + print(f"m={m:<5} n={n:<5} k={k:<5} dtype={dtype} bound: {bound}") - print(a, b) - ref_c = (a.T if transa else a) @ (b.T if transb else b) - np.testing.assert_allclose(ref_c, my_c, rtol=rtol, atol=atol) + for impl in my_gemm.ListImpls(): + if not my_gemm.SelectImpl(impl): + continue + + my_gemm.Run() + dev_c.UpdateHostNumpyArray() + + try: + np.testing.assert_allclose(my_c, ref_c, rtol=bound) + except Exception as err: + failures[impl] = str(err) + + if failures: + raise Exception(failures) dtypes = ["float32", "float16"] -transabs = product([True, False], repeat=2) -basic_sizes = product([1, 3, 4, 16, 127, 128, 129, 133, 1024], repeat=3) +all_transabs = list(product([True, False], repeat=2)) +all_basic_sizes = list(product([1, 3, 4, 16, 127, 128, 129, 133, 1024], repeat=3)) -@pytest.mark.parametrize( - "dtype, size, transab", - product(dtypes, basic_sizes, transabs), -) -def test_gemm_all_cases(dtype, size, transab): - for f in dtype_to_funcs(dtype): - _test_gemm(f, dtype, *size, *transab) +def get_bert_sizes(full=True): + bert_base_sizes = [ + # m, n, k + (384, 768, 768), + (384, 768, 768 * 3), + (384, 768, 768 * 4), + (384, 768 * 4, 768), + (384, 1024, 1024), + (384, 1024, 1024 * 3), + (384, 1024, 1024 * 4), + (384, 1024 * 4, 1024), + ] + + # we then multiply m with the batch size + if full: + batch_sizes = [1, 64] + else: + batch_sizes = [1] + bert_sizes = [] + for bsz in batch_sizes: + bert_sizes.extend([(m * bsz, n, k) for m, n, k in bert_base_sizes]) + return bert_sizes -def profile_gemm_func(func, dtype, m, n, k): +@pytest.mark.parametrize("dtype", dtypes) +@pytest.mark.parametrize("size", all_basic_sizes + get_bert_sizes(full=False)) +@pytest.mark.parametrize("transab", all_transabs) +def test_rocblas_gemm_all_cases(dtype, size, transab): + _test_gemm(getattr(ke, "RocblasGemm_" + dtype_to_suffix(dtype)), dtype, *size, *transab) + + +# ck has various impls to be tested, use the full basic cases will result too many cases to test. +# So we use a reduced combination here. +reduced_basic_sizes = list(product([1, 4, 127, 133], [3, 16, 128], [3, 129, 1024])) +no_transabs = [[False, False]] + + +@pytest.mark.parametrize("dtype", dtypes) +@pytest.mark.parametrize("size", reduced_basic_sizes + get_bert_sizes(full=False)) +@pytest.mark.parametrize("transab", no_transabs) +def test_ck_gemm_bert_cases(dtype, size, transab): + _test_gemm(getattr(ke, "CKGemm_" + dtype_to_suffix(dtype)), dtype, *size, *transab) + + +def profile_gemm_func(f, dtype, m, n, k): a_shape = (m, k) b_shape = (k, n) @@ -86,7 +138,6 @@ def profile_gemm_func(func, dtype, m, n, k): dev_b = ke.DeviceArray(b) dev_c = ke.DeviceArray(my_c) - f = getattr(ke, func) opa = ke.blas_op.N opb = ke.blas_op.N lda = a_shape[1] @@ -94,30 +145,22 @@ def profile_gemm_func(func, dtype, m, n, k): alpha = 1.0 beta = 0.0 my_gemm = f(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n) - time_ms = my_gemm.Profile() - - time_us = time_ms * 1000 - tflops = (m * k * n * 2) / (time_ms * 1e-3) / 1e12 - - print(f"RocBLAS GEMM {dtype} m={m:<4} k={k:<4} n={n:<4}, {time_us:>8.4f} us, {tflops:>5.2f} tflops") + for impl in my_gemm.ListImpls(): + if not my_gemm.SelectImpl(impl): + print(f"{impl:<50} {dtype} m={m:<4} k={k:<4} n={n:<4}, not supported") + continue + time_ms = my_gemm.Profile() + time_us = time_ms * 1000 + tflops = (m * k * n * 2) / (time_ms * 1e-3) / 1e12 + print(f"{impl:<50} {dtype} m={m:<4} k={k:<4} n={n:<4}, {time_us:>8.4f} us, {tflops:>5.2f} tflops") def profile(): - dtypes = ["float32", "float16"] - bert_sizes = [ - # m, k, n - (384, 384, 64), - (384, 768, 768), - (384, 768, 3072), - (384, 1024, 1024), - (384, 1024, 4096), - (384, 3072, 768), - (384, 4096, 1024), - ] for dtype in dtypes: - for f in dtype_to_funcs(dtype): - for m, k, n in bert_sizes: - profile_gemm_func(f, dtype, m, k, n) + for m, n, k in get_bert_sizes(full=True): + suffix = dtype_to_suffix(dtype) + profile_gemm_func(getattr(ke, "RocblasGemm_" + suffix), dtype, m, n, k) + profile_gemm_func(getattr(ke, "CKGemm_" + suffix), dtype, m, n, k) print() print()