Add Tunable GEMM composed from rocblas and composable kernels (#12599)

* Add tunable gemm

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 <type_traits>
 #include <pybind11/pybind11.h>
@@ -22,6 +23,7 @@ void InitGemm(py::module mod) {
 
   InitRocBlasGemm(mod);
   InitComposableKernelGemm(mod);
+  InitTunableGemm(mod);
 }
 
 }  // namespace onnxruntime

onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.cc

@@ -5,74 +5,16 @@
 
 #include <pybind11/stl.h>
 
-#include <algorithm>
 #include <memory>
 #include <string>
 #include <type_traits>
 #include <utility>
 #include <vector>
 
-#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 <typename T>
-struct DataTypeAdaptor {
-  using type = T;
-};
-
-template <>
-struct DataTypeAdaptor<half> {
-  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 <typename T, typename ALayout, typename BLayout>
-auto GetCKGemmTypeStringAndOps() {
-  using CKDataType = typename DataTypeAdaptor<T>::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<DeviceGemm>;
-
-  std::vector<std::pair<std::string, contrib::rocm::Op<GemmParams<T>>>> 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<T>* 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 <typename T, typename ALayout, typename BLayout>
 class CKGemm : public IKernelExplorer {
  public:

onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.h

@@ -5,10 +5,67 @@
 
 #include <pybind11/pybind11.h>
 
+#include <string>
+#include <utility>
+#include <vector>
+
+#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 <typename T>
+struct DataTypeAdaptor {
+  using type = T;
+};
+
+template <>
+struct DataTypeAdaptor<half> {
+  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 <typename T, typename ALayout, typename BLayout>
+auto GetCKGemmTypeStringAndOps() {
+  using CKDataType = typename DataTypeAdaptor<T>::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<DeviceGemm>;
+
+  std::vector<std::pair<std::string, contrib::rocm::Op<GemmParams<T>>>> 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<T>* 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

onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.cc

@@ -10,34 +10,12 @@
 #include <vector>
 
 #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 <typename T>
-Status RocBlasGemmOp(const GemmParams<T>* 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 <typename T>
 class RocBlasGemm : public IKernelExplorer {
  public:

onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.h

@@ -5,10 +5,34 @@
 
 #include <pybind11/pybind11.h>
 
+#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 <typename T>
+Status RocBlasGemmOp(const GemmParams<T>* 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

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()
 

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 <pybind11/stl.h>
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#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 <typename T, typename ALayout, typename BLayout>
+class GemmTunableOp : public contrib::rocm::TunableOp<GemmParams<T>> {
+ public:
+  GemmTunableOp() {
+    this->ops_.emplace_back(RocBlasGemmOp<T>);
+    for (auto&& [_, op] : GetCKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
+      ORT_UNUSED_PARAMETER(_);
+      this->ops_.emplace_back(std::move(op));
+    }
+  }
+};
+
+template <typename T, typename ALayout, typename BLayout>
+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<T*>(a.ptr());
+    params_.lda = lda;
+    params_.b = static_cast<T*>(b.ptr());
+    params_.ldb = ldb;
+    params_.beta = beta;
+    params_.c = static_cast<T*>(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_(&params_));
+  }
+
+  std::vector<std::string> ListOps() const {
+    return {"Tunable"};
+  }
+
+  bool SelectOp(const std::string& name) {
+    return name == "Tunable";
+  }
+
+ private:
+  using ParamsT = GemmParams<T>;
+  ParamsT params_;
+
+  // tunable is stateful, store it as an instance
+  GemmTunableOp<T, ALayout, BLayout> op_{};
+  rocblas_handle rocblas_handle_;
+};
+
+#define REGISTER_OP(type, alayout, blayout, layout_string)                                   \
+  py::class_<GemmTunable<type, alayout, blayout>>(m, "GemmTunable_" #type "_" layout_string) \
+      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t,                               \
+                    double,                                                                  \
+                    DeviceArray&, int64_t,                                                   \
+                    DeviceArray&, int64_t,                                                   \
+                    double,                                                                  \
+                    DeviceArray&, int64_t>())                                                \
+      .def("SetRepeats", &GemmTunable<type, alayout, blayout>::SetRepeats)                   \
+      .def("Profile", &GemmTunable<type, alayout, blayout>::Profile)                         \
+      .def("Run", &GemmTunable<type, alayout, blayout>::Run)                                 \
+      .def("ListOps", &GemmTunable<type, alayout, blayout>::ListOps)                         \
+      .def("SelectOp", &GemmTunable<type, alayout, blayout>::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

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 <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace onnxruntime {
+
+void InitTunableGemm(py::module mod);
+
+}