diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc index 1cd5024e6b..9abae017de 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm.cc @@ -2,8 +2,9 @@ // 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 "python/tools/kernel_explorer/kernels/gemm_rocblas.h" +#include "python/tools/kernel_explorer/kernels/gemm_tunable.h" #include #include @@ -22,6 +23,7 @@ void InitGemm(py::module mod) { InitRocBlasGemm(mod); InitComposableKernelGemm(mod); + InitTunableGemm(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 index 5d0f4310b3..39f9bdbc8b 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc @@ -5,74 +5,16 @@ #include -#include #include #include #include #include #include -#include "ck/ck.hpp" -#include "ck/library/tensor_operation_instance/gpu/gemm.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 "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; -}; - -} // namespace - -using Row = ck::tensor_layout::gemm::RowMajor; -using Col = ck::tensor_layout::gemm::ColumnMajor; - -using Nop = ck::tensor_operation::element_wise::PassThrough; - -// to be moved to onnxruntime once we have a monolithicly tunable gemm wrapper and it is enabled for onnxruntime -template -auto GetCKGemmTypeStringAndOps() { - using CKDataType = typename DataTypeAdaptor::type; - using DeviceGemm = ck::tensor_operation::device::DeviceGemm< - ALayout, BLayout, Row, - CKDataType, CKDataType, CKDataType, - Nop, Nop, Nop>; - using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory; - - std::vector>>> ret; - for (auto&& impl : InstanceFactory::GetInstances()) { - auto type_string = impl->GetTypeString(); - auto invoker = impl->MakeInvokerPointer(); - auto ck_gemm_op = [impl = std::move(impl), invoker = std::move(invoker)](const GemmParams* params) -> Status { - auto nop = Nop{}; - auto arg = impl->MakeArgumentPointer(params->a, params->b, params->c, - params->m, params->n, params->k, - params->lda, params->ldb, params->ldc, - nop, nop, nop); - TUNABLE_OP_RETURN_UNSUPPOTED_ARGUMENT_IF(!impl->IsSupportedArgument(arg.get()), - impl->GetTypeString(), " does not support ", params->Signature()); - invoker->Run(arg.get(), StreamConfig{params->stream}); - return Status::OK(); - }; - ret.emplace_back(std::make_pair(std::move(type_string), std::move(ck_gemm_op))); - } - return ret; -} - template class CKGemm : public IKernelExplorer { public: diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h index 782de08556..7058df89b0 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h @@ -5,10 +5,67 @@ #include +#include +#include +#include + +#include "ck/ck.hpp" +#include "ck/library/tensor_operation_instance/gpu/gemm.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 "python/tools/kernel_explorer/kernels/gemm.h" + namespace py = pybind11; namespace onnxruntime { +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; + +using Nop = ck::tensor_operation::element_wise::PassThrough; + +// to be moved to onnxruntime once we have a monolithicly tunable gemm wrapper and it is enabled for onnxruntime +template +auto GetCKGemmTypeStringAndOps() { + using CKDataType = typename DataTypeAdaptor::type; + using DeviceGemm = ck::tensor_operation::device::DeviceGemm< + ALayout, BLayout, Row, + CKDataType, CKDataType, CKDataType, + Nop, Nop, Nop>; + using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory; + + std::vector>>> ret; + for (auto&& impl : InstanceFactory::GetInstances()) { + auto type_string = impl->GetTypeString(); + auto invoker = impl->MakeInvokerPointer(); + auto ck_gemm_op = [impl = std::move(impl), invoker = std::move(invoker)](const GemmParams* params) -> Status { + auto nop = Nop{}; + auto arg = impl->MakeArgumentPointer(params->a, params->b, params->c, + params->m, params->n, params->k, + params->lda, params->ldb, params->ldc, + nop, nop, nop); + TUNABLE_OP_RETURN_UNSUPPOTED_ARGUMENT_IF(!impl->IsSupportedArgument(arg.get()), + impl->GetTypeString(), " does not support ", params->Signature()); + invoker->Run(arg.get(), StreamConfig{params->stream}); + return Status::OK(); + }; + ret.emplace_back(std::make_pair(std::move(type_string), std::move(ck_gemm_op))); + } + return ret; +} + void InitComposableKernelGemm(py::module mod); -} +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc index 6577eb3d4d..1a23a401ca 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc @@ -10,34 +10,12 @@ #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 { -// to be moved to onnxruntime once we have a monolithicly tunable gemm wrapper and it is enabled for onnxruntime -template -Status RocBlasGemmOp(const GemmParams* params) { - // 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. - auto status = rocblasGemmHelper( - params->handle, - params->opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose, - params->opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose, - params->n, params->m, params->k, - &(params->alpha), - params->b, params->ldb, - params->a, params->lda, - &(params->beta), - params->c, params->ldc); - ORT_RETURN_IF(status != rocblas_status_success, rocblas_status_to_string(status)); - return Status::OK(); -} - template class RocBlasGemm : public IKernelExplorer { public: diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h index 046177ed45..f6c5c733e6 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h @@ -5,10 +5,34 @@ #include +#include "core/common/common.h" +#include "core/providers/rocm/shared_inc/fpgeneric.h" +#include "python/tools/kernel_explorer/kernels/gemm.h" + namespace py = pybind11; namespace onnxruntime { +// to be moved to onnxruntime once we have a monolithicly tunable gemm wrapper and it is enabled for onnxruntime +template +Status RocBlasGemmOp(const GemmParams* params) { + // 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. + auto status = rocblasGemmHelper( + params->handle, + params->opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose, + params->opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose, + params->n, params->m, params->k, + &(params->alpha), + params->b, params->ldb, + params->a, params->lda, + &(params->beta), + params->c, params->ldc); + ORT_RETURN_IF(status != rocblas_status_success, rocblas_status_to_string(status)); + return Status::OK(); +} + void InitRocBlasGemm(py::module mod); -} +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py index d5ff5a697d..4b11932265 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py @@ -135,6 +135,15 @@ def test_ck_gemm_bert_cases(dtype, size, transab): _test_gemm(getattr(ke, wrapper_name), dtype, *size, *transab) +# Tunable is basically wrapped around of rocblas and ck gemm, so no need for full tests +@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_gemm_tunable_bert_cases(dtype, size, transab): + wrapper_name = "GemmTunable_{}_{}".format(dtype_to_suffix(dtype), transab_to_suffix(transab)) + _test_gemm(getattr(ke, wrapper_name), dtype, *size, *transab) + + def profile_gemm_func(f, dtype, m, n, k): a_shape = (m, k) b_shape = (k, n) @@ -172,6 +181,8 @@ def profile(): profile_gemm_func(getattr(ke, "RocblasGemm" + dtype_suffix), dtype, m, n, k) transab_suffix = "_" + transab_to_suffix((False, False)) profile_gemm_func(getattr(ke, "CKGemm" + dtype_suffix + transab_suffix), dtype, m, n, k) + profile_gemm_func(getattr(ke, "GemmTunable" + dtype_suffix + transab_suffix), dtype, m, n, k) + print() print() diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc new file mode 100644 index 0000000000..e33f3e0e29 --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.cc @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "python/tools/kernel_explorer/kernels/gemm_tunable.h" + +#include + +#include +#include +#include + +#include "core/providers/rocm/rocm_common.h" +#include "contrib_ops/rocm/bert/tunable_op.h" +#include "python/tools/kernel_explorer/kernels/gemm.h" +#include "python/tools/kernel_explorer/kernels/gemm_ck.h" +#include "python/tools/kernel_explorer/kernels/gemm_rocblas.h" + +namespace onnxruntime { + +template +class GemmTunableOp : public contrib::rocm::TunableOp> { + public: + GemmTunableOp() { + this->ops_.emplace_back(RocBlasGemmOp); + for (auto&& [_, op] : GetCKGemmTypeStringAndOps()) { + ORT_UNUSED_PARAMETER(_); + this->ops_.emplace_back(std::move(op)); + } + } +}; + +template +class GemmTunable : public IKernelExplorer { + public: + GemmTunable(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) { + ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_)); + params_.handle = rocblas_handle_; + params_.opa = opa; + params_.opb = opb; + params_.m = m; + params_.n = n; + params_.k = k; + params_.alpha = alpha; + params_.a = static_cast(a.ptr()); + params_.lda = lda; + params_.b = static_cast(b.ptr()); + params_.ldb = ldb; + params_.beta = beta; + params_.c = static_cast(c.ptr()); + params_.ldc = ldc; + + op_.EnableTuning(); + } + + ~GemmTunable() { + ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_)); + rocblas_handle_ = nullptr; + } + + void Run() override { + ORT_THROW_IF_ERROR(op_(¶ms_)); + } + + std::vector ListOps() const { + return {"Tunable"}; + } + + bool SelectOp(const std::string& name) { + return name == "Tunable"; + } + + private: + using ParamsT = GemmParams; + ParamsT params_; + + // tunable is stateful, store it as an instance + GemmTunableOp op_{}; + rocblas_handle rocblas_handle_; +}; + +#define REGISTER_OP(type, alayout, blayout, layout_string) \ + py::class_>(m, "GemmTunable_" #type "_" layout_string) \ + .def(py::init()) \ + .def("SetRepeats", &GemmTunable::SetRepeats) \ + .def("Profile", &GemmTunable::Profile) \ + .def("Run", &GemmTunable::Run) \ + .def("ListOps", &GemmTunable::ListOps) \ + .def("SelectOp", &GemmTunable::SelectOp); + +#define REGISTER_OP_FOR_ALL_TRANSAB(type) \ + REGISTER_OP(type, Row, Row, "NN"); \ + REGISTER_OP(type, Row, Col, "NT"); \ + REGISTER_OP(type, Col, Row, "TN"); \ + REGISTER_OP(type, Col, Col, "TT"); + +void InitTunableGemm(py::module m) { + REGISTER_OP_FOR_ALL_TRANSAB(float); + REGISTER_OP_FOR_ALL_TRANSAB(half); +} + +} // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.h b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.h new file mode 100644 index 0000000000..48aa2b792b --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_tunable.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 InitTunableGemm(py::module mod); + +}