Enable/Disbale tunable GEMM by using tunable switch in provider options and env var (#13116)

Related PRs #12853

This allows the user enable/disbale tunable GEMM on demand.

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -428,7 +428,16 @@ typedef struct OrtCUDAProviderOptions {
 */
 typedef struct OrtROCMProviderOptions {
 #ifdef __cplusplus
-  OrtROCMProviderOptions() : device_id{}, miopen_conv_exhaustive_search{0}, gpu_mem_limit{SIZE_MAX}, arena_extend_strategy{}, do_copy_in_default_stream{1}, has_user_compute_stream{}, user_compute_stream{}, default_memory_arena_cfg{} {}
+  OrtROCMProviderOptions()
+      : device_id{},
+        miopen_conv_exhaustive_search{0},
+        gpu_mem_limit{SIZE_MAX},
+        arena_extend_strategy{},
+        do_copy_in_default_stream{1},
+        has_user_compute_stream{},
+        user_compute_stream{},
+        tunable_op_enabled{false},
+        default_memory_arena_cfg{} {}
 #endif
 
   /** \brief ROCM device Id
@@ -474,6 +483,12 @@ typedef struct OrtROCMProviderOptions {
   */
   void* user_compute_stream;
 
+  /** \brief Enable TunableOp.
+  *   Set it to 1 to enable TunableOp. Otherwise, it is disabled by default.
+  *   This option can be superseded by environment variable ORT_ROCM_TUNABLE_OP_ENABLED.
+  */
+  int tunable_op_enabled;
+
   /** \brief ROCM memory arena configuration parameters
   */
   OrtArenaCfg* default_memory_arena_cfg;
@@ -3568,13 +3583,13 @@ struct OrtApi {
   * \since Version 1.13.
   */
   void(ORT_API_CALL* ReleaseCANNProviderOptions)(_Frees_ptr_opt_ OrtCANNProviderOptions* input);
-  
+
  /*  \brief Get OrtDevice type from MemoryInfo
-  * 
+  *
   *  \since Version 1.14
   */
   void(ORT_API_CALL* MemoryInfoGetDeviceType)(_In_ const OrtMemoryInfo* ptr, _Out_ OrtMemoryInfoDeviceType* out);
-  
+
 
 #ifdef __cplusplus
   OrtApi(const OrtApi&)=delete; // Prevent users from accidentally copying the API structure, it should always be passed as a pointer

onnxruntime/core/providers/rocm/math/gemm.cc

@@ -128,7 +128,7 @@ Status Gemm<T>::ComputeInternal(OpKernelContext* ctx) const {
   }
 
   return tunable::blas::column_major::Gemm(
-      false, Stream(),
+      IsTunableOpEnabled(), Stream(),
       RocblasHandle(),
       trans_B_ ? BlasOp::Trans : BlasOp::NonTrans,
       trans_A_ ? BlasOp::Trans : BlasOp::NonTrans,

onnxruntime/core/providers/rocm/math/matmul_impl.cc

@@ -88,7 +88,7 @@ Status MatMulImpl(const RocmKernel* op, MatMulComputeHelper& helper,
     BlasOp transA = transa ? BlasOp::Trans : BlasOp::NonTrans;
     BlasOp transB = transb ? BlasOp::Trans : BlasOp::NonTrans;
     return tunable::blas::column_major::Gemm(
-        false, op->Stream(),
+        op->IsTunableOpEnabled(), op->Stream(),
         op->RocblasHandle(), transB, transA, static_cast<int64_t>(helper.N()),
         static_cast<int64_t>(helper.M()), static_cast<int64_t>(helper.K()), t_alpha,
         reinterpret_cast<const HipT*>(right_x_data), ldb,

onnxruntime/core/providers/rocm/rocm_execution_provider.cc

@@ -154,6 +154,40 @@ ROCMExecutionProvider::PerThreadContext::~PerThreadContext() {
   }
 }
 
+std::optional<std::string> LoadEnv(const std::string& name, const std::unordered_set<std::string>& valid_values) {
+  ORT_ENFORCE(!valid_values.empty());
+  auto env = onnxruntime::GetEnvironmentVar(name);
+  if (env.empty()) {
+    return std::nullopt;
+  }
+
+  LOGS_DEFAULT(WARNING) << "Environment variable "<< name << " is used. It is reserved for internal testing prupose. "
+                           "End users should opt for provider options or session options and must not rely on it.";
+
+  if (valid_values.find(env) == valid_values.cend()) {
+    std::ostringstream oss;
+    auto it = valid_values.cbegin();
+    oss << *it++;
+    while(it != valid_values.cend()) {
+      oss << ", " << *it++;
+    }
+    ORT_THROW("Value of environment variable ", name," must be ",oss.str(), ", but got ", env);
+  }
+
+  return env;
+}
+
+void OverrideTunableOpInfoByEnv(ROCMExecutionProviderInfo& info) {
+  std::optional<std::string> env;
+
+  env = LoadEnv("ORT_ROCM_TUNABLE_OP_ENABLED", {"0", "1"});
+  const bool env_tunable_op_enabled = env == "1";
+  if (env.has_value() && env_tunable_op_enabled != info.tunable_op.enabled) {
+    LOGS_DEFAULT(INFO) << "ORT_ROCM_TUNABLE_OP_ENABLED is set to " << env_tunable_op_enabled;
+    info.tunable_op.enabled = env_tunable_op_enabled;
+  }
+}
+
 ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& info)
     : IExecutionProvider{onnxruntime::kRocmExecutionProvider},
       info_{info} {
@@ -180,6 +214,8 @@ ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& in
   size_t free = 0;
   size_t total = 0;
   HIP_CALL_THROW(hipMemGetInfo(&free, &total));
+
+  OverrideTunableOpInfoByEnv(info_);
 }
 
 ROCMExecutionProvider::~ROCMExecutionProvider() {
@@ -202,6 +238,20 @@ ROCMExecutionProvider::~ROCMExecutionProvider() {
   }
 }
 
+void ROCMExecutionProvider::EnableTunableOp() {
+  LOGS_DEFAULT(INFO) << "Enable TunableOp for ROCm Execution Provider";
+  info_.tunable_op.enabled = true;
+}
+
+void ROCMExecutionProvider::DisableTunableOp() {
+  LOGS_DEFAULT(INFO) << "Disable TunableOp for ROCm Execution Provider";
+  info_.tunable_op.enabled = false;
+}
+
+bool ROCMExecutionProvider::IsTunableOpEnabled() const {
+  return info_.tunable_op.enabled;
+}
+
 std::unique_ptr<profiling::EpProfiler> ROCMExecutionProvider::GetProfiler() {
   return std::make_unique<profiling::RocmProfiler>();
 }

onnxruntime/core/providers/rocm/rocm_execution_provider.h

@@ -119,6 +119,10 @@ class ROCMExecutionProvider : public IExecutionProvider {
   static AllocatorPtr CreateRocmAllocator(OrtDevice::DeviceId device_id, size_t rocm_mem_limit, ArenaExtendStrategy arena_extend_strategy,
                                           ROCMExecutionProviderExternalAllocatorInfo external_alloc_info, OrtArenaCfg* arena_cfg);
 
+  void EnableTunableOp();
+  void DisableTunableOp();
+  bool IsTunableOpEnabled() const;
+
   std::unique_ptr<profiling::EpProfiler> GetProfiler() override;
 
  private:

onnxruntime/core/providers/rocm/rocm_execution_provider_info.cc

@@ -21,6 +21,7 @@ constexpr const char* kGpuExternalAlloc = "gpu_external_alloc";
 constexpr const char* kGpuExternalFree = "gpu_external_free";
 constexpr const char* kGpuExternalEmptyCache = "gpu_external_empty_cache";
 constexpr const char* kMiopenConvUseMaxWorkspace = "miopen_conv_use_max_workspace";
+constexpr const char* kTunableOpEnabled = "tunable_op_enabled";
 }  // namespace provider_option_names
 }  // namespace rocm
 
@@ -81,6 +82,12 @@ ROCMExecutionProviderInfo ROCMExecutionProviderInfo::FromProviderOptions(const P
           .AddAssignmentToReference(rocm::provider_option_names::kMiopenConvExhaustiveSearch, info.miopen_conv_exhaustive_search)
           .AddAssignmentToReference(rocm::provider_option_names::kDoCopyInDefaultStream, info.do_copy_in_default_stream)
           .AddAssignmentToReference(rocm::provider_option_names::kMiopenConvUseMaxWorkspace, info.miopen_conv_use_max_workspace)
+          .AddValueParser(
+              rocm::provider_option_names::kTunableOpEnabled,
+              [&info](const std::string& value_str) -> Status {
+                ORT_RETURN_IF_ERROR(ParseStringWithClassicLocale(value_str, info.tunable_op.enabled));
+                return Status::OK();
+              })
           .Parse(options));
 
   ROCMExecutionProviderExternalAllocatorInfo alloc_info{alloc, free, empty_cache};
@@ -100,6 +107,7 @@ ProviderOptions ROCMExecutionProviderInfo::ToProviderOptions(const ROCMExecution
       {rocm::provider_option_names::kMiopenConvExhaustiveSearch, MakeStringWithClassicLocale(info.miopen_conv_exhaustive_search)},
       {rocm::provider_option_names::kDoCopyInDefaultStream, MakeStringWithClassicLocale(info.do_copy_in_default_stream)},
       {rocm::provider_option_names::kMiopenConvUseMaxWorkspace, MakeStringWithClassicLocale(info.miopen_conv_use_max_workspace)},
+      {rocm::provider_option_names::kTunableOpEnabled, MakeStringWithClassicLocale(info.tunable_op.enabled)},
   };
 
   return options;

onnxruntime/core/providers/rocm/rocm_execution_provider_info.h

@@ -3,8 +3,10 @@
 
 #pragma once
 
+#include <functional>
 #include <limits>
 
+#include "core/common/hash_combine.h"
 #include "core/framework/arena_extend_strategy.h"
 #include "core/framework/ortdevice.h"
 #include "core/framework/provider_options.h"
@@ -34,6 +36,12 @@ struct ROCMExecutionProviderExternalAllocatorInfo {
   }
 };
 
+namespace rocm {
+struct TunableOpInfo {
+  bool enabled{false};
+};
+}  // namespace rocm
+
 struct ROCMExecutionProviderInfo {
   OrtDevice::DeviceId device_id{0};
   size_t gpu_mem_limit{std::numeric_limits<size_t>::max()};                         // Will be over-ridden by contents of `default_memory_arena_cfg` (if specified)
@@ -42,6 +50,7 @@ struct ROCMExecutionProviderInfo {
   bool do_copy_in_default_stream{true};
   bool has_user_compute_stream{false};
   void* user_compute_stream{nullptr};
+  rocm::TunableOpInfo tunable_op{};
   // The following OrtArenaCfg instance only characterizes the behavior of the default memory
   // arena allocator and not any other auxiliary allocator that may also be part of the ROCM EP.
   // For example, auxiliary allocators `HIP_PINNED` and `HIP_CPU` will not be configured using this
@@ -54,3 +63,12 @@ struct ROCMExecutionProviderInfo {
   static ProviderOptions ToProviderOptions(const ROCMExecutionProviderInfo& info);
 };
 }  // namespace onnxruntime
+
+template<>
+struct std::hash<::onnxruntime::rocm::TunableOpInfo> {
+  size_t operator()(const ::onnxruntime::rocm::TunableOpInfo& info) const {
+    size_t seed_and_value{0xbc9f1d34};
+    onnxruntime::HashCombine(info.enabled, seed_and_value);
+    return seed_and_value;
+  }
+};

onnxruntime/core/providers/rocm/rocm_kernel.h

@@ -66,6 +66,8 @@ class RocmKernel : public OpKernel {
 
   inline hipStream_t Stream() const { return static_cast<hipStream_t>(provider_->GetComputeStream()); }
 
+  bool IsTunableOpEnabled() const { return provider_->IsTunableOpEnabled(); }
+
   // To support hipMemcpyAsync, the cpu memory should be allocated in pinned memory
   // and it can only be released after the copy has finished
   template <typename T>

onnxruntime/core/providers/rocm/rocm_provider_factory.cc

@@ -173,6 +173,7 @@ struct ROCM_Provider : Provider {
     info.has_user_compute_stream = params->has_user_compute_stream;
     info.user_compute_stream = params->user_compute_stream;
     info.default_memory_arena_cfg = params->default_memory_arena_cfg;
+    info.tunable_op.enabled = params->tunable_op_enabled;
 
     return std::make_shared<ROCMProviderFactory>(info);
   }

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -371,6 +371,7 @@ const ROCMExecutionProviderInfo GetRocmExecutionProviderInfo(ProviderInfo_ROCM*
     info.miopen_conv_exhaustive_search = miopen_conv_exhaustive_search;
     info.do_copy_in_default_stream = do_copy_in_default_stream;
     info.external_allocator_info = external_allocator_info;
+    info.tunable_op = tunable_op;
   }
   return info;
 }

onnxruntime/python/onnxruntime_pybind_state_common.cc

@@ -34,6 +34,8 @@ onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExten
 bool miopen_conv_exhaustive_search = false;
 // TODO remove deprecated global config
 bool do_copy_in_default_stream = true;
+// TODO remove deprecated global config
+onnxruntime::rocm::TunableOpInfo tunable_op{};
 onnxruntime::ROCMExecutionProviderExternalAllocatorInfo external_allocator_info{};
 // TODO remove deprecated global config
 onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo;

onnxruntime/python/onnxruntime_pybind_state_common.h

@@ -199,6 +199,8 @@ namespace python {
 extern bool miopen_conv_exhaustive_search;
 // TODO remove deprecated global config
 extern bool do_copy_in_default_stream;
+// TODO remove deprecated global config
+extern onnxruntime::rocm::TunableOpInfo tunable_op;
 extern onnxruntime::ROCMExecutionProviderExternalAllocatorInfo external_allocator_info;
 extern onnxruntime::ArenaExtendStrategy arena_extend_strategy;
 }  // namespace python

onnxruntime/test/python/onnxruntime_test_python.py

@@ -159,7 +159,7 @@ class TestInferenceSession(unittest.TestCase):
 
             """
             int8_use_native_calibration_table = "false"
-            option['trt_int8_use_native_calibration_table'] = int8_use_native_calibration_table 
+            option['trt_int8_use_native_calibration_table'] = int8_use_native_calibration_table
             int8_enable = "true"
             option['trt_int8_enable'] = int8_enable
             calib_table_name = '/home/onnxruntime/table.flatbuffers' # this file is not existed
@@ -348,6 +348,31 @@ class TestInferenceSession(unittest.TestCase):
                 runBaseTest2()
                 # raise OSError("could not load any of: " + ' '.join(libnames))
 
+        if "ROCMExecutionProvider" in onnxrt.get_available_providers():
+
+            def runRocmOptionsTest():
+                sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), providers=["ROCMExecutionProvider"])
+                self.assertIn("ROCMExecutionProvider", sess.get_providers())
+                options = sess.get_provider_options()
+
+                def test_get_and_set_option_with_values(option_name, option_values):
+                    provider_options = sess.get_provider_options()
+                    self.assertIn("ROCMExecutionProvider", provider_options)
+                    rocm_options = options["ROCMExecutionProvider"]
+                    self.assertIn(option_name, rocm_options)
+                    for option_value in option_values:
+                        rocm_options[option_name] = option_value
+                        sess.set_providers(["ROCMExecutionProvider"], [rocm_options])
+                        new_provider_options = sess.get_provider_options()
+                        self.assertEqual(
+                            new_provider_options.get("ROCMExecutionProvider", {}).get(option_name),
+                            str(option_value),
+                        )
+
+                test_get_and_set_option_with_values("tunable_op_enabled", ["1", "0"])
+
+            runRocmOptionsTest()
+
     def testInvalidSetProviders(self):
         with self.assertRaises(RuntimeError) as context:
             sess = onnxrt.InferenceSession(get_name("mul_1.onnx"), providers=["CPUExecutionProvider"])

orttraining/orttraining/python/orttraining_python_module.cc

@@ -101,7 +101,8 @@ bool GetProviderInstanceHash(const std::string& type,
              (static_cast<size_t>(info.arena_extend_strategy) << 16) ^
              (static_cast<size_t>(info.miopen_conv_exhaustive_search) << 18) ^
              (static_cast<size_t>(info.do_copy_in_default_stream) << 20) ^
-             (static_cast<size_t>(info.has_user_compute_stream) << 22);
+             (static_cast<size_t>(info.has_user_compute_stream) << 22) ^
+             std::hash<rocm::TunableOpInfo>{}(info.tunable_op);
       return true;
     }
 #endif