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 <abhishek.udupa@microsoft.com>

onnxruntime/core/framework/tunable.h

@@ -15,6 +15,7 @@
 #include <thread>
 #include <type_traits>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -74,6 +75,7 @@ class Op {
  public:
   template <typename T>
   explicit Op(T&& c) : callable_{std::make_unique<CallableImpl<T>>(std::forward<T>(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 <typename T>
   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 <typename ParamsT, typename TimerT>
 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<ParamsT, TimerT>* 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<ParamsT>& 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<Op<ParamsT>>& 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<double>::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<Op<ParamsT>&>(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<int>(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<Op<ParamsT>> ops_;
 
  private:
   // mapping from Signature to best impl
   std::unordered_map<std::string, int> 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<TunableOp<ParamsT, TimerT>*> nested_tunable_ops_;
 };
 
 }  // namespace tunable

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 <typename T>
+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<float>(const float*) {
+  return rocblas_datatype_f32_r;
+}
+
+template <>
+constexpr rocblas_datatype RocBlasDataTypeFor<half>(const half*) {
+  return rocblas_datatype_f16_r;
+}
+
+template <>
+constexpr rocblas_datatype RocBlasDataTypeFor<double>(const double*) {
+  return rocblas_datatype_f64_r;
+}
+
+template <>
+constexpr rocblas_datatype RocBlasDataTypeFor<BFloat16>(const BFloat16*) {
+  return rocblas_datatype_bf16_r;
+}
+
+template <typename T>
+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<float>(const float*) {
+  return rocblas_datatype_f32_r;
+}
+
+template <>
+constexpr rocblas_datatype RocBlasComputeTypeFor<half>(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<double>(const double*) {
+  return rocblas_datatype_f64_r;
+}
+
+template <>
+constexpr rocblas_datatype RocBlasComputeTypeFor<BFloat16>(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 <typename T>
+class IndexedRocBlasGemmOp {
+ public:
+  IndexedRocBlasGemmOp()
+      : index_(0) {}
+  IndexedRocBlasGemmOp(int index)
+      : index_(index) {}
+
+  Status operator()(const GemmParams<T>* 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<T>*) {
+    return Status::OK();
+  }
+
+ private:
+  int index_;
+};
+
+template <typename T>
+class RocBlasGemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
+ public:
+  RocBlasGemmTunableOp() {
+    // Ensure that the default implementation is always present
+    this->ops_.emplace_back(IndexedRocBlasGemmOp<T>{0});
+  }
+
+  Status IsSupported(const GemmParams<T>* params) {
+    ORT_UNUSED_PARAMETER(params);
+    return Status::OK();
+  }
+
+ protected:
+  virtual int FindFastest(const GemmParams<T>* params) override {
+    auto solution_indices = this->GetSolutions(params);
+    std::vector<Op<GemmParams<T>>> candidates;
+    for (int solution_idx : solution_indices) {
+      candidates.emplace_back(IndexedRocBlasGemmOp<T>{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<int> GetSolutions(const GemmParams<T>* 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<int> 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 <typename T>
 Status RocBlasGemmOp(const GemmParams<T>* params) {
   RocblasHandleStreamGuard guard(params->handle, params->stream);

onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh

@@ -31,12 +31,16 @@ bool IsZero(half v) {
   return __half2float(v) == 0.0f;
 }
 
-
 template <typename T, typename ALayout, typename BLayout>
 class GemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
  public:
   GemmTunableOp() {
     this->ops_.emplace_back(RocBlasGemmOp<T>);
+
+#ifdef USE_ROCBLAS_EXTENSION_API
+    this->RegisterNestedTunableOp(&rocblas_gemm_tunable_op_);
+#endif /* #ifdef USE_ROCBLAS_EXTENSION_API */
+
     for (auto&& [_, op] : GetCKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
       ORT_UNUSED_PARAMETER(_);
       this->ops_.emplace_back(std::move(op));
@@ -67,6 +71,11 @@ class GemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
       delete params;
     }
   }
+
+ private:
+#ifdef USE_ROCBLAS_EXTENSION_API
+  RocBlasGemmTunableOp<T> rocblas_gemm_tunable_op_;
+#endif
 };
 
 }  // namespace internal