add TuningContext

microsoft · Feb 10, 2023 · 0b6c029 · 0b6c029
1 parent 10a7350
commit 0b6c029
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/onnxruntime/core/providers/rocm/math/softmax_common.h b/onnxruntime/core/providers/rocm/math/softmax_common.h
@@ -11,10 +11,10 @@ namespace onnxruntime {
 namespace rocm {
 
 template <typename InputT, typename OutputT>
-struct SoftmaxParams : onnxruntime::rocm::tunable::OpParams {
-  SoftmaxParams(hipStream_t stream, OutputT* output, const InputT* input, int softmax_elements,
-                int input_stride, int output_stride, int batch_count, bool is_log_softmax)
-      : OpParams(stream), output(output), input(input), softmax_elements(softmax_elements), input_stride(input_stride), output_stride(output_stride), batch_count(batch_count), is_log_softmax(is_log_softmax) {}
+struct SoftmaxParams : tunable::OpParams {
+  SoftmaxParams(tunable::RocmTuningContext* tuning_ctx, hipStream_t stream, OutputT* output, const InputT* input,
+                int softmax_elements, int input_stride, int output_stride, int batch_count, bool is_log_softmax)
+      : OpParams(tuning_ctx, stream), output(output), input(input), softmax_elements(softmax_elements), input_stride(input_stride), output_stride(output_stride), batch_count(batch_count), is_log_softmax(is_log_softmax) {}
 
   std::string Signature() const override {
     std::string sig = std::to_string(batch_count) + "_" + std::to_string(softmax_elements);

diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu
@@ -26,7 +26,7 @@ class SoftmaxBlockwise : public IKernelExplorer {
  public:
   SoftmaxBlockwise(DeviceArray& output, DeviceArray& input, int softmax_elements,
                    int input_stride, int output_stride, int batch_count, bool is_log_softmax)
-      : params_(this->Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
+      : params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
                 softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) {
     type_string_ = "SoftmaxBlockwise_" + std::to_string(VecSize);
   }
@@ -55,7 +55,7 @@ class SoftmaxBlockwiseStaticSelection : public IKernelExplorer {
  public:
   SoftmaxBlockwiseStaticSelection(DeviceArray& output, DeviceArray& input, int softmax_elements,
                                   int input_stride, int output_stride, int batch_count, bool is_log_softmax)
-      : params_(this->Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
+      : params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
                 softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) {}
 
   void Run() override {
@@ -80,9 +80,9 @@ class SoftmaxTunable : public IKernelExplorer {
  public:
   SoftmaxTunable(DeviceArray& output, DeviceArray& input, int softmax_elements,
                  int input_stride, int output_stride, int batch_count, bool is_log_softmax)
-      : params_(this->Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
+      : params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
                 softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) {
-    op_.EnableTuning();
+    params_.TuningContext()->EnableTunableOp();
   }
 
   void Run() override {
@@ -109,7 +109,7 @@ class CKSoftmax : public IKernelExplorer {
  public:
   CKSoftmax(DeviceArray& output, DeviceArray& input, int softmax_elements,
             int input_stride, int output_stride, int batch_count, bool is_log_softmax)
-      : params_(this->Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
+      : params_(TuningContext(), Stream(), static_cast<T*>(output.ptr()), static_cast<T*>(input.ptr()),
                 softmax_elements, input_stride, output_stride, batch_count, is_log_softmax) {
     for (auto&& [type_string, op] : rocm::GetCKSoftmaxTypeStringAndOps<T, T, rocm::AccumulationType_t<T>>()) {
       type_strings_.emplace_back(std::move(type_string));