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