Add composable kernel GEMM baseline for kernel explorer (#12364)

* Split GemmBase RocBlasGemm

* Add composable kernel GEMM baseline

* Make linter happy

* Address review comment

* Update bert cases with batchsize

* Adjust includes to fix IWYU lint

* Only builds and links used ck kernels to improve building time

* Remove warmup run on SelectImpl

* Add comment to utility function

* Mute cpplint

* Make RocBlasGemm<T>::SelectImpl semantically correct

* Add reduced basic test cases for ck gemm

* More robust gemm testing

* Fix warnings

* Fix grammar

cmake/external/composable_kernel.cmake

@@ -0,0 +1,19 @@
+set(composable_kernel_URL https://github.com/ROCmSoftwarePlatform/composable_kernel.git)
+set(composable_kernel_TAG 85978e0201bb94bf6e59b325e1f5f19266845d08) # 2022-07-22 11:52:10 -0700
+
+set(BUILD_DEV OFF)
+set(PATCH ${PROJECT_SOURCE_DIR}/patches/composable_kernel/Fix_Clang_Build.patch)
+
+include(FetchContent)
+FetchContent_Declare(composable_kernel
+  GIT_REPOSITORY ${composable_kernel_URL}
+  GIT_TAG        ${composable_kernel_TAG}
+  PATCH_COMMAND  git apply --reverse --check ${PATCH} || git apply --ignore-space-change --ignore-whitespace ${PATCH}
+)
+
+FetchContent_MakeAvailable(composable_kernel)
+
+add_library(onnxruntime_composable_kernel_includes INTERFACE)
+target_include_directories(onnxruntime_composable_kernel_includes INTERFACE
+  ${composable_kernel_SOURCE_DIR}/include
+  ${composable_kernel_SOURCE_DIR}/library/include)

cmake/onnxruntime_kernel_explorer.cmake

@@ -13,6 +13,8 @@ if(NOT HIP_FOUND)
   message(FATAL_ERROR "hip is required but is not found")
 endif()
 
+include(composable_kernel)
+
 set(KERNEL_EXPLORER_ROOT ${ONNXRUNTIME_ROOT}/python/tools/kernel_explorer)
 set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/rocm/bert)
 
@@ -30,6 +32,11 @@ target_include_directories(kernel_explorer PUBLIC
 target_link_libraries(kernel_explorer
   PRIVATE
     $<TARGET_PROPERTY:onnxruntime_pybind11_state,LINK_LIBRARIES>
+    onnxruntime_composable_kernel_includes
+    # Currently we shall not use composablekernels::device_operations, the target includes all conv dependencies, which
+    # are extremely slow to compile. Instead, we only link all gemm related objects. See the following link on updating.
+    # https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/85978e0201/library/src/tensor_operation_instance/gpu/CMakeLists.txt#L33-L54
+    device_gemm_instance
     ${HIP_LIB})
 target_compile_definitions(kernel_explorer
   PUBLIC ROCM_USE_FLOAT16
@@ -38,7 +45,7 @@ target_compile_definitions(kernel_explorer
 # handle kernel_explorer sources as hip language
 target_compile_options(kernel_explorer PRIVATE "-xhip")
 # TODO: use predefined AMDGPU_TARGETS
-target_compile_options(kernel_explorer PRIVATE "--offload-arch=gfx906" "--offload-arch=gfx908" "--offload-arch=gfx90a")
+target_compile_options(kernel_explorer PRIVATE "--offload-arch=gfx908" "--offload-arch=gfx90a")
 # https://github.com/ROCm-Developer-Tools/HIP/blob/4514f350849b1090954295f8f87a5f8d78bd781b/hip-lang-config.cmake.in
 target_link_libraries(kernel_explorer PRIVATE ${CLANGRT_BUILTINS})
 

cmake/patches/composable_kernel/Fix_Clang_Build.patch

@@ -0,0 +1,53 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 9f706207..fb444227 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -8,7 +8,7 @@ list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
+ enable_testing()
+ 
+ set(ROCM_SYMLINK_LIBS OFF)
+-find_package(ROCM 0.8 REQUIRED PATHS /opt/rocm)
++find_package(ROCM 0.7 REQUIRED PATHS /opt/rocm)
+ 
+ include(ROCMInstallTargets)
+ include(ROCMPackageConfigHelpers)
+@@ -28,6 +28,17 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
+ set(CMAKE_CXX_EXTENSIONS OFF)
+ message("CMAKE_CXX_COMPILER_ID: ${CMAKE_CXX_COMPILER_ID}")
+ 
++add_compile_options(
++    -xhip
++    --offload-arch=gfx908
++    --offload-arch=gfx90a
++    -O3 # otherwise, "Illegal instruction detected" for gfx908
++    "SHELL:-mllvm -amdgpu-early-inline-all=true"
++    "SHELL:-mllvm -amdgpu-function-calls=false" # otherwise, "local memory (65920) exceeds limit (65536) in function"
++    -fhip-new-launch-api
++)
++
++
+ ## OpenMP
+ if(CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+ 	# workaround issue hipcc in rocm3.5 cannot find openmp
+@@ -237,9 +248,6 @@ rocm_package_setup_component(tests
+ )
+ 
+ add_subdirectory(library)
+-add_subdirectory(example)
+-add_subdirectory(test)
+-add_subdirectory(profiler)
+ 
+ #Create an interface target for the include only files and call it "composablekernels"
+ include(CMakePackageConfigHelpers)
+@@ -265,11 +273,3 @@ rocm_install(FILES
+ 
+ set(CPACK_RESOURCE_FILE_LICENSE "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE")
+ set(CPACK_RPM_PACKAGE_LICENSE "MIT")
+-
+-rocm_create_package(
+-    NAME composablekernel
+-    DESCRIPTION "High Performance Composable Kernel for AMD GPUs"
+-    MAINTAINER "MIOpen Kernels Dev Team <dl.MIOpen@amd.com>"
+-    LDCONFIG
+-    HEADER_ONLY
+-)

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

@@ -2,109 +2,16 @@
 // 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 <type_traits>
 #include <pybind11/pybind11.h>
 
-#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/operator.h"
-
 namespace py = pybind11;
 
 namespace onnxruntime {
 
-enum class BlasOp {
-  N,
-  T,
-};
-
-// We don't assume the implementation is row-majored or column-majored. But for testing convenience, we assume all
-// our wrappers have row-majored convention, since it is the native layout to numpy and pytorch.
-template <typename T>
-class GemmBase : public Operator {
- public:
-  GemmBase(
-      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)
-      : Operator(),
-        opa_(opa),
-        opb_(opb),
-        m_(m),
-        n_(n),
-        k_(k),
-        alpha_(alpha),
-        a_(reinterpret_cast<T*>(a.ptr())),
-        lda_(lda),
-        b_(reinterpret_cast<T*>(b.ptr())),
-        ldb_(ldb),
-        beta_(beta),
-        c_(reinterpret_cast<T*>(c.ptr())),
-        ldc_(ldc) {}
-
- protected:
-  BlasOp opa_;
-  BlasOp opb_;
-  int64_t m_;
-  int64_t n_;
-  int64_t k_;
-  T alpha_;
-  T* a_;
-  int64_t lda_;
-  T* b_;
-  int64_t ldb_;
-  T beta_;
-  T* c_;
-  int64_t ldc_;
-};
-
-template <typename T>
-class RocBlasGemm : public GemmBase<T> {
- public:
-  RocBlasGemm(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)
-      : GemmBase<T>(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc),
-        opa_(opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose),
-        opb_(opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose) {
-    ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_));
-  }
-
-  ~RocBlasGemm() {
-    ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_));
-    rocblas_handle_ = nullptr;
-  }
-
-  void Run() {
-    // NOTE: rocblas assume 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.
-    ROCBLAS_CALL_THROW(
-        rocblasGemmHelper(this->rocblas_handle_, this->opb_, this->opa_,
-                          this->n_, this->m_, this->k_,
-                          &(this->alpha_),
-                          this->b_, this->ldb_,
-                          this->a_, this->lda_,
-                          &(this->beta_),
-                          this->c_, this->ldc_));
-  }
-
- private:
-  rocblas_handle rocblas_handle_;
-  rocblas_operation opa_;
-  rocblas_operation opb_;
-};
-
 void InitGemm(py::module mod) {
   auto blas_op = mod.def_submodule("blas_op");
 
@@ -113,19 +20,8 @@ void InitGemm(py::module mod) {
       .value("T", BlasOp::T, "Transpose")
       .export_values();
 
-  // float
-  py::class_<RocBlasGemm<float>>(mod, "RocblasGemm_float")
-      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double, DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
-      .def("SetRepeats", &RocBlasGemm<float>::SetRepeats)
-      .def("Profile", &RocBlasGemm<float>::Profile)
-      .def("Run", &RocBlasGemm<float>::Run);
-
-  // half
-  py::class_<RocBlasGemm<half>>(mod, "RocblasGemm_half")
-      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double, DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
-      .def("SetRepeats", &RocBlasGemm<half>::SetRepeats)
-      .def("Profile", &RocBlasGemm<half>::Profile)
-      .def("Run", &RocBlasGemm<half>::Run);
+  InitRocBlasGemm(mod);
+  InitComposableKernelGemm(mod);
 }
 
 }  // namespace onnxruntime

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

@@ -3,12 +3,70 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
+
 #include <pybind11/pybind11.h>
 
+#include "python/tools/kernel_explorer/device_array.h"
+#include "python/tools/kernel_explorer/operator.h"
+
 namespace py = pybind11;
 
 namespace onnxruntime {
 
+enum class BlasOp {
+  N,
+  T,
+};
+
+// We don't assume the implementation is row-majored or column-majored. But for testing convenience, we assume all
+// our wrappers have row-majored convention, since it is the native layout to numpy and pytorch.
+template <typename T>
+class GemmBase : public Operator {
+ public:
+  GemmBase(
+      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)
+      : Operator(),
+        opa_(opa),
+        opb_(opb),
+        m_(m),
+        n_(n),
+        k_(k),
+        alpha_(alpha),
+        a_(reinterpret_cast<T*>(a.ptr())),
+        lda_(lda),
+        b_(reinterpret_cast<T*>(b.ptr())),
+        ldb_(ldb),
+        beta_(beta),
+        c_(reinterpret_cast<T*>(c.ptr())),
+        ldc_(ldc) {}
+
+  virtual std::vector<std::string> ListImpls() const = 0;
+  virtual bool SelectImpl(const std::string& name) = 0;
+
+ protected:
+  BlasOp opa_;
+  BlasOp opb_;
+  int64_t m_;
+  int64_t n_;
+  int64_t k_;
+  T alpha_;
+  T* a_;
+  int64_t lda_;
+  T* b_;
+  int64_t ldb_;
+  T beta_;
+  T* c_;
+  int64_t ldc_;
+};
+
 void InitGemm(py::module mod);
 
-}
+}  // namespace onnxruntime

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

@@ -0,0 +1,187 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "python/tools/kernel_explorer/kernels/gemm_ck.h"
+
+#include <pybind11/stl.h>
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "ck/ck.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 "ck/library/tensor_operation_instance/gpu/gemm.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;
+};
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+}  // namespace
+
+template <typename T>
+class CKGemm : public GemmBase<T> {
+  using CKDataType = typename DataTypeAdaptor<T>::type;
+
+  template <typename ALayout, typename BLayout>
+  using DeviceGemm = ck::tensor_operation::device::DeviceGemm<
+      ALayout, BLayout, Row,
+      CKDataType, CKDataType, CKDataType,
+      ck::tensor_operation::element_wise::PassThrough,
+      ck::tensor_operation::element_wise::PassThrough,
+      ck::tensor_operation::element_wise::PassThrough>;
+
+  template <typename Op>
+  using CKGemmInstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<Op>;
+
+  // composable kernel's GetInstances() returns a collection of <std::unique_ptr<DeviceGemm<...>>>. We need to manually
+  // cast them to std::unique_ptr<BaseOperator>> to avoid type problem.
+  template <typename Op>
+  static std::vector<std::unique_ptr<ck::tensor_operation::device::BaseOperator>> GetInstances() {
+    auto ptrs = CKGemmInstanceFactory<Op>::GetInstances();
+    std::vector<std::unique_ptr<ck::tensor_operation::device::BaseOperator>> results;
+    for (auto&& ptr : ptrs) {
+      results.emplace_back(std::move(ptr));
+    }
+    return results;
+  }
+
+  bool UpdateArgumentAndInvoker() {
+    auto nop = ck::tensor_operation::element_wise::PassThrough{};
+    if (this->opa_ == BlasOp::N && this->opb_ == BlasOp::N) {
+      using DeviceOp = DeviceGemm<Row, Row>;
+      auto typed_impl = static_cast<DeviceOp*>(impls_[selected_impl_].get());
+      arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_,
+                                             this->m_, this->n_, this->k_,
+                                             this->lda_, this->ldb_, this->ldc_,
+                                             nop, nop, nop);
+      invoker_ = typed_impl->MakeInvokerPointer();
+    } else if (this->opa_ == BlasOp::T && this->opb_ == BlasOp::N) {
+      using DeviceOp = DeviceGemm<Col, Row>;
+      auto typed_impl = static_cast<DeviceOp*>(impls_[selected_impl_].get());
+      arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_,
+                                             this->m_, this->n_, this->k_,
+                                             this->lda_, this->ldb_, this->ldc_,
+                                             nop, nop, nop);
+      invoker_ = typed_impl->MakeInvokerPointer();
+    } else if (this->opa_ == BlasOp::N && this->opb_ == BlasOp::T) {
+      using DeviceOp = DeviceGemm<Row, Col>;
+      auto typed_impl = static_cast<DeviceOp*>(impls_[selected_impl_].get());
+      arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_,
+                                             this->m_, this->n_, this->k_,
+                                             this->lda_, this->ldb_, this->ldc_,
+                                             nop, nop, nop);
+      invoker_ = typed_impl->MakeInvokerPointer();
+    } else if (this->opa_ == BlasOp::T && this->opb_ == BlasOp::T) {
+      using DeviceOp = DeviceGemm<Col, Col>;
+      auto typed_impl = static_cast<DeviceOp*>(impls_[selected_impl_].get());
+      arg_ = typed_impl->MakeArgumentPointer(this->a_, this->b_, this->c_,
+                                             this->m_, this->n_, this->k_,
+                                             this->lda_, this->ldb_, this->ldc_,
+                                             nop, nop, nop);
+      invoker_ = typed_impl->MakeInvokerPointer();
+    }
+
+    return impls_[selected_impl_]->IsSupportedArgument(arg_.get());
+  }
+
+ public:
+  CKGemm(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)
+      : GemmBase<T>(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc) {
+    ORT_ENFORCE(alpha == 1.0);
+    ORT_ENFORCE(beta == 0.0);
+
+    if (opa == BlasOp::N && opb == BlasOp::N) {
+      impls_ = this->GetInstances<DeviceGemm<Row, Row>>();
+    } else if (opa == BlasOp::T && opb == BlasOp::N) {
+      impls_ = this->GetInstances<DeviceGemm<Col, Row>>();
+    } else if (opa == BlasOp::N && opb == BlasOp::T) {
+      impls_ = this->GetInstances<DeviceGemm<Row, Col>>();
+    } else if (opa == BlasOp::T && opb == BlasOp::T) {
+      impls_ = this->GetInstances<DeviceGemm<Col, Col>>();
+    }
+
+    ORT_ENFORCE(!impls_.empty());
+    selected_impl_ = 0;
+    UpdateArgumentAndInvoker();
+  }
+
+  std::vector<std::string> ListImpls() const override {
+    std::vector<std::string> results;
+    std::transform(impls_.cbegin(), impls_.cend(), std::back_inserter(results),
+                   [](const auto& it) { return it->GetTypeString(); });
+    return results;
+  }
+
+  bool SelectImpl(const std::string& name) override {
+    for (size_t i = 0; i < impls_.size(); i++) {
+      if (impls_[i]->GetTypeString() == name) {
+        selected_impl_ = i;
+        return UpdateArgumentAndInvoker();
+      }
+    }
+
+    ORT_THROW("Cannot find implementation ", name);
+  }
+
+  void Run() override {
+    invoker_->Run(arg_.get());
+  }
+
+ private:
+  std::vector<std::unique_ptr<ck::tensor_operation::device::BaseOperator>> impls_;
+  size_t selected_impl_{};
+  std::unique_ptr<ck::tensor_operation::device::BaseArgument> arg_;
+  std::unique_ptr<ck::tensor_operation::device::BaseInvoker> invoker_;
+};
+
+void InitComposableKernelGemm(py::module mod) {
+  // float
+  py::class_<CKGemm<float>>(mod, "CKGemm_float")
+      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double,
+                    DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
+      .def("SetRepeats", &CKGemm<float>::SetRepeats)
+      .def("Profile", &CKGemm<float>::Profile)
+      .def("Run", &CKGemm<float>::Run)
+      .def("ListImpls", &CKGemm<float>::ListImpls)
+      .def("SelectImpl", &CKGemm<float>::SelectImpl);
+
+  // half
+  py::class_<CKGemm<half>>(mod, "CKGemm_half")
+      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double,
+                    DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
+      .def("SetRepeats", &CKGemm<half>::SetRepeats)
+      .def("Profile", &CKGemm<half>::Profile)
+      .def("Run", &CKGemm<half>::Run)
+      .def("ListImpls", &CKGemm<half>::ListImpls)
+      .def("SelectImpl", &CKGemm<half>::SelectImpl);
+}
+
+}  // namespace onnxruntime

onnxruntime/python/tools/kernel_explorer/kernels/gemm_ck.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 InitComposableKernelGemm(py::module mod);
+
+}

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

@@ -0,0 +1,92 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "python/tools/kernel_explorer/kernels/gemm_rocblas.h"
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <string>
+#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 {
+
+template <typename T>
+class RocBlasGemm : public GemmBase<T> {
+ public:
+  RocBlasGemm(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)
+      : GemmBase<T>(opa, opb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc),
+        opa_(opa == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose),
+        opb_(opb == BlasOp::N ? rocblas_operation_none : rocblas_operation_transpose) {
+    ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_));
+  }
+
+  ~RocBlasGemm() {
+    ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_));
+    rocblas_handle_ = nullptr;
+  }
+
+  void Run() override {
+    // 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.
+    ROCBLAS_CALL_THROW(
+        rocblasGemmHelper(this->rocblas_handle_, this->opb_, this->opa_,
+                          this->n_, this->m_, this->k_,
+                          &(this->alpha_),
+                          this->b_, this->ldb_,
+                          this->a_, this->lda_,
+                          &(this->beta_),
+                          this->c_, this->ldc_));
+  }
+
+  std::vector<std::string> ListImpls() const override {
+    return {"Rocblas"};
+  }
+
+  bool SelectImpl(const std::string& name) override {
+    return name == "Rocblas";
+  }
+
+ private:
+  rocblas_handle rocblas_handle_;
+  rocblas_operation opa_;
+  rocblas_operation opb_;
+};
+
+void InitRocBlasGemm(py::module mod) {
+  // float
+  py::class_<RocBlasGemm<float>>(mod, "RocblasGemm_float")
+      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double,
+                    DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
+      .def("SetRepeats", &RocBlasGemm<float>::SetRepeats)
+      .def("Profile", &RocBlasGemm<float>::Profile)
+      .def("Run", &RocBlasGemm<float>::Run)
+      .def("ListImpls", &RocBlasGemm<float>::ListImpls)
+      .def("SelectImpl", &RocBlasGemm<float>::SelectImpl);
+
+  // half
+  py::class_<RocBlasGemm<half>>(mod, "RocblasGemm_half")
+      .def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, double,
+                    DeviceArray&, int64_t, DeviceArray&, int64_t, double, DeviceArray&, int64_t>())
+      .def("SetRepeats", &RocBlasGemm<half>::SetRepeats)
+      .def("Profile", &RocBlasGemm<half>::Profile)
+      .def("Run", &RocBlasGemm<half>::Run)
+      .def("ListImpls", &RocBlasGemm<half>::ListImpls)
+      .def("SelectImpl", &RocBlasGemm<half>::SelectImpl);
+}
+
+}  // namespace onnxruntime

onnxruntime/python/tools/kernel_explorer/kernels/gemm_rocblas.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 InitRocBlasGemm(py::module mod);
+
+}

onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py

@@ -10,70 +10,122 @@ import numpy as np
 import pytest
 
 
-def dtype_to_bytes(dtype):
-    type_map = {
-        "float16": 2,
-        "float32": 4,
-    }
-    return type_map[dtype]
-
-
-def dtype_to_funcs(dtype):
-    type_map = {
-        "float16": list(filter(lambda x: "RocblasGemm_half" in x, dir(ke))),
-        "float32": list(filter(lambda x: "RocblasGemm_float" in x, dir(ke))),
-    }
-    return type_map[dtype]
+def dtype_to_suffix(dtype):
+    return {
+        "float32": "float",
+        "float16": "half",
+    }[dtype]
 
 
 def _test_gemm(func, dtype: str, m: int, n: int, k: int, transa=False, transb=False):
+    assert dtype in ["float32", "float16"]
+
     a_shape = (k, m) if transa else (m, k)
     b_shape = (n, k) if transb else (k, n)
 
     np.random.seed(0)
-    a = (np.random.rand(*a_shape) * 2 - 1).astype(dtype)
-    b = (np.random.rand(*b_shape) * 2 - 1).astype(dtype)
+    a = (np.random.rand(*a_shape) + 0.5).astype(dtype).astype("float64")
+    b = (np.random.rand(*b_shape) + 0.5).astype(dtype).astype("float64")
+    ref_c = (a.T if transa else a) @ (b.T if transb else b)
+
+    # The machine epsilon, unit roundoff, the smallest positive floating point number n such that the floating point
+    # number that represents 1 + n is greater than 1.
+    machine_eps = 2.0 ** -(24 if dtype == "float32" else 11)
+
+    # The following implements error bound 5.7 in paper I. C. Ipsen and H. Zhou, “Probabilistic error analysis for
+    # Inner Products,” SIAM Journal on Matrix Analysis and Applications, vol. 41, no. 4, pp. 1726–1741, 2020.
+    # NOTE: the bound is not tight for float16 when k is large
+    absa_mul_absb = np.abs(a.T if transa else a) @ np.abs(b.T if transb else b)
+    coeff = np.max(absa_mul_absb / np.abs(ref_c))
+    gamma_2k = (1.0 + machine_eps) ** (2 * k) - 1.0
+    bound_5_7 = coeff * np.sqrt(np.log(2 / 1e-10) * machine_eps * gamma_2k / 2)
+    bound = bound_5_7
+
+    a = a.astype(dtype)
+    b = b.astype(dtype)
 
     my_c = np.zeros((m, n), dtype=dtype)
     dev_a = ke.DeviceArray(a)
     dev_b = ke.DeviceArray(b)
     dev_c = ke.DeviceArray(my_c)
 
-    f = getattr(ke, func)
-
     opa = ke.blas_op.T if transa else ke.blas_op.N
     opb = ke.blas_op.T if transb else ke.blas_op.N
     lda = a_shape[1]
     ldb = b_shape[1]
     alpha = 1.0
     beta = 0.0
-    my_gemm = f(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n)
-    my_gemm.Run()
-    dev_c.UpdateHostNumpyArray()
+    my_gemm = func(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n)
 
-    rtol = 1e-3 if dtype == "float16" else 1e-5
-    atol = 1e-3 if dtype == "float16" else 1e-5
+    failures = {}
+    print(f"m={m:<5} n={n:<5} k={k:<5} dtype={dtype} bound: {bound}")
 
-    print(a, b)
-    ref_c = (a.T if transa else a) @ (b.T if transb else b)
-    np.testing.assert_allclose(ref_c, my_c, rtol=rtol, atol=atol)
+    for impl in my_gemm.ListImpls():
+        if not my_gemm.SelectImpl(impl):
+            continue
+
+        my_gemm.Run()
+        dev_c.UpdateHostNumpyArray()
+
+        try:
+            np.testing.assert_allclose(my_c, ref_c, rtol=bound)
+        except Exception as err:
+            failures[impl] = str(err)
+
+    if failures:
+        raise Exception(failures)
 
 
 dtypes = ["float32", "float16"]
-transabs = product([True, False], repeat=2)
-basic_sizes = product([1, 3, 4, 16, 127, 128, 129, 133, 1024], repeat=3)
+all_transabs = list(product([True, False], repeat=2))
+all_basic_sizes = list(product([1, 3, 4, 16, 127, 128, 129, 133, 1024], repeat=3))
 
 
-@pytest.mark.parametrize(
-    "dtype, size, transab",
-    product(dtypes, basic_sizes, transabs),
-)
-def test_gemm_all_cases(dtype, size, transab):
-    for f in dtype_to_funcs(dtype):
-        _test_gemm(f, dtype, *size, *transab)
+def get_bert_sizes(full=True):
+    bert_base_sizes = [
+        # m, n, k
+        (384, 768, 768),
+        (384, 768, 768 * 3),
+        (384, 768, 768 * 4),
+        (384, 768 * 4, 768),
+        (384, 1024, 1024),
+        (384, 1024, 1024 * 3),
+        (384, 1024, 1024 * 4),
+        (384, 1024 * 4, 1024),
+    ]
+
+    # we then multiply m with the batch size
+    if full:
+        batch_sizes = [1, 64]
+    else:
+        batch_sizes = [1]
+    bert_sizes = []
+    for bsz in batch_sizes:
+        bert_sizes.extend([(m * bsz, n, k) for m, n, k in bert_base_sizes])
+    return bert_sizes
 
 
-def profile_gemm_func(func, dtype, m, n, k):
+@pytest.mark.parametrize("dtype", dtypes)
+@pytest.mark.parametrize("size", all_basic_sizes + get_bert_sizes(full=False))
+@pytest.mark.parametrize("transab", all_transabs)
+def test_rocblas_gemm_all_cases(dtype, size, transab):
+    _test_gemm(getattr(ke, "RocblasGemm_" + dtype_to_suffix(dtype)), dtype, *size, *transab)
+
+
+# ck has various impls to be tested, use the full basic cases will result too many cases to test.
+# So we use a reduced combination here.
+reduced_basic_sizes = list(product([1, 4, 127, 133], [3, 16, 128], [3, 129, 1024]))
+no_transabs = [[False, False]]
+
+
+@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_ck_gemm_bert_cases(dtype, size, transab):
+    _test_gemm(getattr(ke, "CKGemm_" + dtype_to_suffix(dtype)), dtype, *size, *transab)
+
+
+def profile_gemm_func(f, dtype, m, n, k):
     a_shape = (m, k)
     b_shape = (k, n)
 
@@ -86,7 +138,6 @@ def profile_gemm_func(func, dtype, m, n, k):
     dev_b = ke.DeviceArray(b)
     dev_c = ke.DeviceArray(my_c)
 
-    f = getattr(ke, func)
     opa = ke.blas_op.N
     opb = ke.blas_op.N
     lda = a_shape[1]
@@ -94,30 +145,22 @@ def profile_gemm_func(func, dtype, m, n, k):
     alpha = 1.0
     beta = 0.0
     my_gemm = f(opa, opb, m, n, k, alpha, dev_a, lda, dev_b, ldb, beta, dev_c, n)
-    time_ms = my_gemm.Profile()
-
-    time_us = time_ms * 1000
-    tflops = (m * k * n * 2) / (time_ms * 1e-3) / 1e12
-
-    print(f"RocBLAS GEMM {dtype} m={m:<4} k={k:<4} n={n:<4}, {time_us:>8.4f} us, {tflops:>5.2f} tflops")
+    for impl in my_gemm.ListImpls():
+        if not my_gemm.SelectImpl(impl):
+            print(f"{impl:<50} {dtype} m={m:<4} k={k:<4} n={n:<4}, not supported")
+            continue
+        time_ms = my_gemm.Profile()
+        time_us = time_ms * 1000
+        tflops = (m * k * n * 2) / (time_ms * 1e-3) / 1e12
+        print(f"{impl:<50} {dtype} m={m:<4} k={k:<4} n={n:<4}, {time_us:>8.4f} us, {tflops:>5.2f} tflops")
 
 
 def profile():
-    dtypes = ["float32", "float16"]
-    bert_sizes = [
-        # m, k, n
-        (384, 384, 64),
-        (384, 768, 768),
-        (384, 768, 3072),
-        (384, 1024, 1024),
-        (384, 1024, 4096),
-        (384, 3072, 768),
-        (384, 4096, 1024),
-    ]
     for dtype in dtypes:
-        for f in dtype_to_funcs(dtype):
-            for m, k, n in bert_sizes:
-                profile_gemm_func(f, dtype, m, k, n)
+        for m, n, k in get_bert_sizes(full=True):
+            suffix = dtype_to_suffix(dtype)
+            profile_gemm_func(getattr(ke, "RocblasGemm_" + suffix), dtype, m, n, k)
+            profile_gemm_func(getattr(ke, "CKGemm_" + suffix), dtype, m, n, k)
             print()
         print()