From 7d684d125559685bc4a6c32248cb5e60f04d84af Mon Sep 17 00:00:00 2001 From: Abhishek Udupa Date: Thu, 8 Dec 2022 14:21:17 -0800 Subject: [PATCH] Include algorithm selection exposed by ROCBLAS extensions API in GEMM autotuning (#13831) ### Description Extend GEMM autotuning by including algorithms exposed by a ROCBLAS extension API. ### Motivation and Context Based on our request, the ROCm team has implemented extension APIs in ROCBLAS that provides a list of application GEMM algorithms/implementations for a given input size, along with an API that actually performs the GEMM using the specified implementation/algorithm. We have observed that the ROCBLAS algorithm/implementation selection logic does not always pick the optimal. This PR uses the extension APIs to integrate the exposed ROCBLAS algorithms/implementations into the autotuning framework. The feature is disabled by default (the ROCBlas extension APIs are slated to be released with ROCm 5.5, and are not yet generally available). To enable: build with `--cmake-extra-defines USE_ROCBLAS_EXTENSION_API=1 CMAKE_HIP_FLAGS=-DUSE_ROCBLAS_EXTENSION_API` and then enable tuning in the provider options. Co-authored-by: Abhishek Udupa --- onnxruntime/core/framework/tunable.h | 55 ++++-- .../providers/rocm/tunable/gemm_rocblas.h | 180 ++++++++++++++++++ .../providers/rocm/tunable/gemm_tunable.cuh | 11 +- 3 files changed, 234 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/framework/tunable.h b/onnxruntime/core/framework/tunable.h index 883135b2b6..aee72331d5 100644 --- a/onnxruntime/core/framework/tunable.h +++ b/onnxruntime/core/framework/tunable.h @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -74,6 +75,7 @@ class Op { public: template explicit Op(T&& c) : callable_{std::make_unique>(std::forward(c))} {} + Op(Op&&) = default; Status operator()(const ParamsT* param) { return (*callable_)(param); } Status IsSupported(const ParamsT* param) { return (*callable_).IsSupported(param); } @@ -87,6 +89,7 @@ class Op { template struct CallableImpl : ICallable { explicit CallableImpl(T&& c) : c_{std::move(c)} {} + CallableImpl(CallableImpl&&) = default; Status operator()(const ParamsT* param) override { return c_(param); } Status IsSupported(const ParamsT* param) override { @@ -117,6 +120,9 @@ class Op { template class TunableOp { public: + TunableOp() = default; + TunableOp(TunableOp&&) = default; + Status operator()(const ParamsT* params) { int id; if (tuning_) { @@ -137,10 +143,16 @@ class TunableOp { void EnableTuning() { tuning_ = true; + for (auto nested_op_ptr : nested_tunable_ops_) { + nested_op_ptr->EnableTuning(); + } } void DisableTuning() { tuning_ = false; + for (auto nested_op_ptr : nested_tunable_ops_) { + nested_op_ptr->DisableTuning(); + } } // We might want to do some tricks to the `params`, e.g., some op will use a buffer for input and output at the same @@ -164,6 +176,20 @@ class TunableOp { default_id_ = id; } + void RegisterNestedTunableOp(TunableOp* op_ptr) { + nested_tunable_ops_.insert(op_ptr); + if (tuning_) { + op_ptr->EnableTuning(); + } else { + op_ptr->DisableTuning(); + } + + // Add an op for this tunable op as well. + ops_.emplace_back([op_ptr](const ParamsT* params) { + return op_ptr->operator()(params); + }); + } + private: static void WarmUp(Op& op, const ParamsT* param) { constexpr const int num_iter = 4; @@ -209,42 +235,49 @@ class TunableOp { #endif } - int FindFastest(const ParamsT* params) { + protected: + virtual int FindFastest(const ParamsT* params) { + return FindFastestImpl(params, ops_); + } + + int FindFastestImpl(const ParamsT* params, const std::vector>& candidates) { auto op_sig = OpSignature(); auto param_sig = params->Signature(); - LOGS_DEFAULT(VERBOSE) << "FindFastest for " << op_sig << '(' << param_sig << ')'; + LOGS_DEFAULT(VERBOSE) << "FindFastestImpl for " << op_sig << '(' << param_sig << ')'; auto min_time = std::numeric_limits::infinity(); int id = -1; - for (size_t i = 0; i < this->ops_.size(); i++) { - if (!IsSupported(ops_[i], params)) { - LOGS_DEFAULT(VERBOSE) << "FindFastest found unsupported " << op_sig << '(' << param_sig << ") id=" << i; + + for (size_t i = 0; i < candidates.size(); i++) { + auto& candidate = const_cast&>(candidates[i]); + if (!IsSupported(candidate, params)) { + LOGS_DEFAULT(VERBOSE) << "FindFastestImpl found unsupported " << op_sig + << '(' << param_sig << ") id=" << i; continue; } - WarmUp(ops_[i], params); - auto time = Profile(ops_[i], params); + WarmUp(candidate, params); + auto time = Profile(candidate, params); if (time < min_time) { min_time = time; id = static_cast(i); } } ORT_ENFORCE(id >= 0, "Cannot found viable op"); - LOGS_DEFAULT(VERBOSE) << "FindFastest for " << op_sig << '(' << param_sig << ") found fastest with id=" << id; + LOGS_DEFAULT(VERBOSE) << "FindFastestImpl for " << op_sig << '(' << param_sig << ") found fastest with id=" << id; std::this_thread::sleep_for(std::chrono::milliseconds(50)); return id; } - protected: std::vector> ops_; private: // mapping from Signature to best impl std::unordered_map kernel_map_; - // the default impl to use when tuning is disabled int default_id_{0}; - bool tuning_{false}; + // Registered tunable sub-ops for nested tuning + std::unordered_set*> nested_tunable_ops_; }; } // namespace tunable diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_rocblas.h b/onnxruntime/core/providers/rocm/tunable/gemm_rocblas.h index 99f0916f8a..079287c58c 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_rocblas.h +++ b/onnxruntime/core/providers/rocm/tunable/gemm_rocblas.h @@ -6,6 +6,7 @@ #include "core/common/common.h" #include "core/providers/rocm/shared_inc/fpgeneric.h" #include "core/providers/rocm/tunable/gemm_common.h" +#include "core/providers/rocm/tunable/rocm_tunable.h" namespace onnxruntime { namespace rocm { @@ -32,6 +33,185 @@ class RocblasHandleStreamGuard { hipStream_t original_stream_; }; +#ifdef USE_ROCBLAS_EXTENSION_API + +template +constexpr rocblas_datatype RocBlasDataTypeFor(const T*) { + static_assert(sizeof(T) == -1, "Unsupported type for rocBLAS operation."); + // The code below should be unreachable due to the static_assert above. + // But the compiler doesn't like not having a return statement, so we + // return something sensible. + return rocblas_datatype_f32_r; +} + +template <> +constexpr rocblas_datatype RocBlasDataTypeFor(const float*) { + return rocblas_datatype_f32_r; +} + +template <> +constexpr rocblas_datatype RocBlasDataTypeFor(const half*) { + return rocblas_datatype_f16_r; +} + +template <> +constexpr rocblas_datatype RocBlasDataTypeFor(const double*) { + return rocblas_datatype_f64_r; +} + +template <> +constexpr rocblas_datatype RocBlasDataTypeFor(const BFloat16*) { + return rocblas_datatype_bf16_r; +} + +template +constexpr rocblas_datatype RocBlasComputeTypeFor(const T*) { + static_assert(sizeof(T) == -1, "Unsupported type for rocBLAS operation."); + // The code below should be unreachable due to the static_assert above. + // But the compiler doesn't like not having a return statement, so we + // return something sensible. + return rocblas_datatype_f32_r; +} + +template <> +constexpr rocblas_datatype RocBlasComputeTypeFor(const float*) { + return rocblas_datatype_f32_r; +} + +template <> +constexpr rocblas_datatype RocBlasComputeTypeFor(const half*) { + // Note that we're returning the _compute_ type for a given datatype. + // As of 12/2022, using compute type FP16 for 16-bit floats was much + // slower than using compute type FP32. So we use FP32 compute even for + // FP16 datatypes. This is how GEMM is implemented even in the function + // rocblasGemmHelper (see fpgeneric.h) + return rocblas_datatype_f32_r; +} + +template <> +constexpr rocblas_datatype RocBlasComputeTypeFor(const double*) { + return rocblas_datatype_f64_r; +} + +template <> +constexpr rocblas_datatype RocBlasComputeTypeFor(const BFloat16*) { + // Note that we're returning the _compute_ type for a given datatype. + // As of 12/2022, using compute type FP16 for 16-bit floats was much + // slower than using compute type FP32. So we use FP32 compute even for + // BF16 datatypes. This is how GEMM is implemented even in the function + // rocblasGemmHelper (see fpgeneric.h) + return rocblas_datatype_f32_r; +} + +template +class IndexedRocBlasGemmOp { + public: + IndexedRocBlasGemmOp() + : index_(0) {} + IndexedRocBlasGemmOp(int index) + : index_(index) {} + + Status operator()(const GemmParams* params) { + RocblasHandleStreamGuard guard(params->handle, params->stream); + return ROCBLAS_CALL( + rocblas_gemm_ex( + 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, RocBlasDataTypeFor(params->b), params->ldb, + params->a, RocBlasDataTypeFor(params->a), params->lda, + &(params->beta), + params->c, RocBlasDataTypeFor(params->c), params->ldc, + params->c, RocBlasDataTypeFor(params->c), params->ldc, + RocBlasComputeTypeFor(params->a), + rocblas_gemm_algo_standard, + index_, + rocblas_gemm_flags_none)); + } + + Status IsSupported(const GemmParams*) { + return Status::OK(); + } + + private: + int index_; +}; + +template +class RocBlasGemmTunableOp : public tunable::TunableOp> { + public: + RocBlasGemmTunableOp() { + // Ensure that the default implementation is always present + this->ops_.emplace_back(IndexedRocBlasGemmOp{0}); + } + + Status IsSupported(const GemmParams* params) { + ORT_UNUSED_PARAMETER(params); + return Status::OK(); + } + + protected: + virtual int FindFastest(const GemmParams* params) override { + auto solution_indices = this->GetSolutions(params); + std::vector>> candidates; + for (int solution_idx : solution_indices) { + candidates.emplace_back(IndexedRocBlasGemmOp{solution_idx}); + } + + auto id = this->FindFastestImpl(params, candidates); + // memoize the result + this->ops_.emplace_back(std::move(candidates[id])); + return this->ops_.size() - 1; + } + + private: + std::vector GetSolutions(const GemmParams* params) { + int num_solutions = 0; + // Get the number of candidate solutions + ROCBLAS_CALL_THROW(rocblas_gemm_ex_get_solutions( + 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, RocBlasDataTypeFor(params->b), params->ldb, + params->a, RocBlasDataTypeFor(params->a), params->lda, + &(params->beta), + params->c, RocBlasDataTypeFor(params->c), params->ldc, + params->c, RocBlasDataTypeFor(params->c), params->ldc, + RocBlasComputeTypeFor(params->a), + rocblas_gemm_algo_standard, + rocblas_gemm_flags_none, + NULL, + &num_solutions)); + + // Get the actual candidate solutions + std::vector solutions(num_solutions); + ROCBLAS_CALL_THROW(rocblas_gemm_ex_get_solutions( + 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, RocBlasDataTypeFor(params->b), params->ldb, + params->a, RocBlasDataTypeFor(params->a), params->lda, + &(params->beta), + params->c, RocBlasDataTypeFor(params->c), params->ldc, + params->c, RocBlasDataTypeFor(params->c), params->ldc, + RocBlasComputeTypeFor(params->a), + rocblas_gemm_algo_standard, + rocblas_gemm_flags_none, + solutions.data(), + &num_solutions)); + + return solutions; + } +}; + +#endif /* #ifdef USE_ROCBLAS_EXTENSION_API */ + template Status RocBlasGemmOp(const GemmParams* params) { RocblasHandleStreamGuard guard(params->handle, params->stream); diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh index 7ac95f4322..6880cf0356 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh +++ b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh @@ -31,12 +31,16 @@ bool IsZero(half v) { return __half2float(v) == 0.0f; } - template class GemmTunableOp : public tunable::TunableOp> { public: GemmTunableOp() { this->ops_.emplace_back(RocBlasGemmOp); + +#ifdef USE_ROCBLAS_EXTENSION_API + this->RegisterNestedTunableOp(&rocblas_gemm_tunable_op_); +#endif /* #ifdef USE_ROCBLAS_EXTENSION_API */ + for (auto&& [_, op] : GetCKGemmTypeStringAndOps()) { ORT_UNUSED_PARAMETER(_); this->ops_.emplace_back(std::move(op)); @@ -67,6 +71,11 @@ class GemmTunableOp : public tunable::TunableOp> { delete params; } } + + private: +#ifdef USE_ROCBLAS_EXTENSION_API + RocBlasGemmTunableOp rocblas_gemm_tunable_op_; +#endif }; } // namespace internal