Add op tuning functionality and an example for vector add. (#12060)

* Add op tuning functionality and example for vector add.

* Add namespace.

* Various improvements.

* use unique pointer

* fix lint errors

* Check return error.

onnxruntime/contrib_ops/rocm/bert/timer.cc

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/rocm/bert/timer.h"
+
+#define HIP_CHECKED_CALL(expr)                                                             \
+  do {                                                                                     \
+    auto status = expr;                                                                    \
+    if (status != hipSuccess) {                                                            \
+      std::printf("HIP Error at %s:%d\n    Error name  : %s\n    Error string: %s\n",      \
+                  __FILE__, __LINE__, hipGetErrorName(status), hipGetErrorString(status)); \
+      std::abort();                                                                        \
+    }                                                                                      \
+  } while (0)
+
+namespace onnxruntime {
+namespace contrib {
+namespace rocm {
+
+Timer::Timer() {
+  HIP_CHECKED_CALL(hipEventCreate(&start_));
+  HIP_CHECKED_CALL(hipEventCreate(&end_));
+}
+
+void Timer::Start() {
+  HIP_CHECKED_CALL(hipDeviceSynchronize());
+  HIP_CHECKED_CALL(hipEventRecord(start_, nullptr));
+}
+
+void Timer::End() {
+  HIP_CHECKED_CALL(hipEventRecord(end_, nullptr));
+  HIP_CHECKED_CALL(hipEventSynchronize(end_));
+}
+
+float Timer::time() {
+  float time;
+  // time is in ms with a resolution of 1 us
+  HIP_CHECKED_CALL(hipEventElapsedTime(&time, start_, end_));
+  return time;
+}
+
+Timer::~Timer() {
+  HIP_CHECKED_CALL(hipEventDestroy(start_));
+  HIP_CHECKED_CALL(hipEventDestroy(end_));
+}
+
+}  // namespace rocm
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/rocm/bert/timer.h

@@ -5,6 +5,10 @@
 
 #include "hip/hip_runtime.h"
 
+namespace onnxruntime {
+namespace contrib {
+namespace rocm {
+
 class Timer {
  public:
   Timer();
@@ -16,3 +20,7 @@ class Timer {
  private:
   hipEvent_t start_, end_; 
 };
+
+}  // namespace rocm
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/rocm/bert/tunable_op.h

@@ -0,0 +1,91 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <iostream>
+#include <memory>
+#include <map>
+#include <string>
+#include <vector>
+#include "contrib_ops/rocm/bert/timer.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace rocm {
+
+struct OpParams {
+  virtual std::string signature() const = 0;
+};
+
+class Op {
+ public:
+  Op() : repeats_(100) {}
+
+  virtual void Run(const OpParams*) = 0;
+
+  void SetRepeats(int n) {
+    repeats_ = n;
+  }
+
+  float Profile(const OpParams* op_params) {
+    // warm up
+    for (int i = 0; i < 5; i++) {
+      Run(op_params);
+    }
+    timer_.Start();
+    for (int i = 0; i < repeats_; i++) {
+      Run(op_params);
+    }
+    timer_.End();
+    return timer_.time()/repeats_;
+  }
+
+  virtual ~Op() {}
+
+ private:
+  Timer timer_;
+  int repeats_;
+};
+
+class TunableOp {
+ public:
+  TunableOp() {}
+
+  void Run(const OpParams* op_params_) {
+    int id;
+    if (kernel_map_.find(op_params_->signature()) == kernel_map_.end()) {
+      id = FindFastest(op_params_);
+      kernel_map_.insert({op_params_->signature(), id});
+    } else {
+      id = kernel_map_[op_params_->signature()];
+    }
+    ops_[id]->Run(op_params_);
+  }
+
+  virtual ~TunableOp() {}
+
+ protected:
+  std::vector<std::unique_ptr<Op>> ops_;
+
+ private:
+  int FindFastest(const OpParams* op_params_) {
+    assert(ops_.size() > 0);
+    float min_time = ops_[0]->Profile(op_params_);
+    int id = 0;
+    for (int i = 1; i < ops_.size(); i++) {
+      float time = ops_[i]->Profile(op_params_);
+      if (time < min_time) {
+        min_time = time;
+        id = i;
+      }
+    }
+    return id;
+  }
+
+  std::map<std::string, int> kernel_map_;
+};
+
+}  // namespace rocm
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/python/tools/kernel_explorer/CMakeLists.txt

@@ -18,13 +18,15 @@ list(APPEND CMAKE_PREFIX_PATH ${onnxruntime_ROCM_HOME} ${onnxruntime_ROCM_HOME}/
 find_package(hip)
 find_package(PythonLibs 3.8 EXACT REQUIRED)
 
+set(onnxruntime_DIR ../../../../onnxruntime/)
+set(onnxruntime_BERT_DIR ${onnxruntime_DIR}/contrib_ops/rocm/bert/)
 include_directories(${PYTHON_INCLUDE_DIRS})
 include_directories(${pybind11_INCLUDE_DIRS})
-include_directories(../../../../onnxruntime/contrib_ops/rocm/bert/)
+include_directories(${onnxruntime_DIR})
 include_directories(.)
 
-FILE(GLOB kernel_srcs kernels/*.cpp)
-add_library(kernel_explorer SHARED kernel_explorer.cpp timer.cpp ${kernel_srcs})
+FILE(GLOB kernel_srcs kernels/*.cc)
+add_library(kernel_explorer SHARED kernel_explorer.cc ${onnxruntime_BERT_DIR}/timer.cc ${kernel_srcs})
 target_link_libraries(kernel_explorer ${PYTHON_LIBRARIES})
 set_target_properties(kernel_explorer PROPERTIES PREFIX "")
 

onnxruntime/python/tools/kernel_explorer/README.md

@@ -1,5 +1,5 @@
 # Kernel Explorer 
-Kernel Explorer is a tool to help develop, test, and profile GPU kernels.
+Kernel Explorer hooks up GPU kernel code with a Python frontend to help develop, test, profile, and auto-tune GPU kernels.
 
 ## Example Usage
 ```

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

@@ -5,7 +5,7 @@
 
 #include "device_array.h"
 #include "operator.h"
-#include "fast_gelu_impl_kernel.h"
+#include "contrib_ops/rocm/bert/fast_gelu_impl_kernel.h"
 
 template <typename T, int ThreadsPerBlock, int VecSize>
 void LaunchFastGelu(const T* input, const T* bias, T* output, int input_length, int bias_length) {

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

@@ -3,7 +3,8 @@
 
 #include <pybind11/pybind11.h>
 #include "hip/hip_fp16.h"
-#include "vector_add_kernel.h"
+#include "kernels/vector_add_kernel.h"
+#include "kernels/vector_add_tunable_op.h"
 
 namespace py = pybind11;
 
@@ -28,6 +29,29 @@ class VectorAdd: public Operator<T> {
   int n_;
 };
 
+template <typename T>
+class VectorAddTunable: public Operator<T> {
+ public:
+  VectorAddTunable(DeviceArray& x, DeviceArray& y, DeviceArray& z, int n) :
+    x_(reinterpret_cast<T*>(x.ptr())),
+    y_(reinterpret_cast<T*>(y.ptr())),
+    z_(reinterpret_cast<T*>(z.ptr())),
+    n_(n),
+    Operator<T>() {}
+
+  void Run() {
+    VectorAddParams<T> op_params(x_, y_, z_, n_);
+    op_.Run(&op_params);
+  }
+
+ private:
+  T* x_;
+  T* y_;
+  T* z_;
+  int n_;
+  VectorAddTunableOp<T> op_;
+};
+
 #define REGISTER_OP(name, type, threads_per_block, vec_size)                                              \
   py::class_<name<type, threads_per_block, vec_size>>(m, #name"_"#type"_"#threads_per_block"_"#vec_size)  \
     .def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, int>())                                       \
@@ -55,4 +79,9 @@ class VectorAdd: public Operator<T> {
 void InitVectorAdd(py::module m) {
   REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK(VectorAdd, half);
   REGISTER_OP_FOR_ALL_THREADS_PER_BLOCK(VectorAdd, float);
+  py::class_<VectorAddTunable<half>>(m, "VectorAdd_half_Tunable")
+    .def(py::init<DeviceArray&, DeviceArray&, DeviceArray&, int>())
+    .def("SetRepeats", &VectorAddTunable<half>::SetRepeats)
+    .def("Profile", &VectorAddTunable<half>::Profile)
+    .def("Run", &VectorAddTunable<half>::Run);
 }

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

@@ -0,0 +1,63 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+#include "contrib_ops/rocm/bert/tunable_op.h"
+#include "kernels/vector_add_kernel.h"
+
+using onnxruntime::contrib::rocm::OpParams;
+using onnxruntime::contrib::rocm::Op;
+using onnxruntime::contrib::rocm::TunableOp;
+
+template<typename T>
+struct VectorAddParams : OpParams {
+  VectorAddParams(const T* x, const T* y, T* z, int n) : x(x), y(y), z(z), n(n) {}
+
+  std::string signature() const {
+    return std::to_string(n);
+  }
+
+  const T* x;
+  const T* y;
+  T* z;
+  int n;
+};
+
+template <typename T, int ThreadsPerBlock, int VecSize>
+class VectorAddOp : public Op {
+ public:
+  VectorAddOp() : Op() {}
+
+  void Run(const OpParams* op_params) {
+    const VectorAddParams<T>* vector_add_params = static_cast<const VectorAddParams<T>*>(op_params);
+    LaunchVectorAdd<T, ThreadsPerBlock, VecSize>(vector_add_params->x,
+                                                 vector_add_params->y,
+                                                 vector_add_params->z,
+                                                 vector_add_params->n);
+  }
+};
+
+#define ADD_OP(threads_per_block)                                              \
+  ops_.push_back(std::make_unique<VectorAddOp<half, threads_per_block, 1>>()); \
+  ops_.push_back(std::make_unique<VectorAddOp<half, threads_per_block, 2>>()); \
+  ops_.push_back(std::make_unique<VectorAddOp<half, threads_per_block, 4>>()); \
+  ops_.push_back(std::make_unique<VectorAddOp<half, threads_per_block, 8>>());
+
+template <typename T>
+class VectorAddTunableOp : public TunableOp {
+ public:
+  VectorAddTunableOp() {
+    ADD_OP(64);
+    ADD_OP(128);
+    ADD_OP(192);
+    ADD_OP(256);
+    ADD_OP(320);
+    ADD_OP(384);
+    ADD_OP(448);
+    ADD_OP(512);
+  }
+};

onnxruntime/python/tools/kernel_explorer/operator.h

@@ -3,7 +3,9 @@
 
 #pragma once
 
-#include "timer.h"
+#include "contrib_ops/rocm/bert/timer.h"
+
+using onnxruntime::contrib::rocm::Timer;
 
 namespace onnxruntime {
 namespace rocm {

onnxruntime/python/tools/kernel_explorer/timer.cpp

@@ -1,31 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-
-#include "timer.h"
-
-Timer::Timer() {
-  hipEventCreate(&start_);
-  hipEventCreate(&end_);
-}
-
-void Timer::Start() {
-  hipDeviceSynchronize();
-  hipEventRecord(start_, nullptr);
-}
-
-void Timer::End() {
-  hipEventRecord(end_, nullptr);
-  hipEventSynchronize(end_);
-}
-
-float Timer::time() {
-  float time;
-  // time is in ms with a resolution of 1 us
-  hipEventElapsedTime(&time, start_, end_);
-  return time;
-}
-
-Timer::~Timer() {
-  hipEventDestroy(start_);
-  hipEventDestroy(end_);
-}