Allow CUDA EP enable or disable TunableOp via session options and environment variable (#13601)

This ports #13116 from ROCm EP to CUDA EP

include/onnxruntime/core/providers/cuda/cuda_provider_options.h

@@ -27,4 +27,5 @@ struct OrtCUDAProviderOptionsV2 {
   int cudnn_conv_use_max_workspace;                        // flag specifying if maximum workspace can be used in cudnn conv algo search.
   int enable_cuda_graph;                                   // flag specifying if the CUDA graph is to be captured for the model.
   int cudnn_conv1d_pad_to_nc1d;                            // flag specifying if pad Conv1D's input [N,C,D] to [N,C,1,D] or [N,C,D,1].
+  int tunable_op_enabled;                                  // flag specifying if TunableOp is enabled.
 };

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -368,7 +368,16 @@ typedef enum OrtCudnnConvAlgoSearch {
 */
 typedef struct OrtCUDAProviderOptions {
 #ifdef __cplusplus
-  OrtCUDAProviderOptions() : device_id{}, cudnn_conv_algo_search{OrtCudnnConvAlgoSearchExhaustive}, 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{} {}
+  OrtCUDAProviderOptions()
+      : device_id{},
+        cudnn_conv_algo_search{OrtCudnnConvAlgoSearchExhaustive},
+        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{},
+        tunable_op_enabled{false} {}
 #endif
 
   /** \brief CUDA device Id
@@ -419,6 +428,12 @@ typedef struct OrtCUDAProviderOptions {
   */
   OrtArenaCfg* default_memory_arena_cfg;
 
+  /** \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_CUDA_TUNABLE_OP_ENABLED.
+  */
+  int tunable_op_enabled;
+
 } OrtCUDAProviderOptions;
 
 /** \brief ROCM Provider Options
@@ -435,8 +450,8 @@ typedef struct OrtROCMProviderOptions {
         do_copy_in_default_stream{1},
         has_user_compute_stream{},
         user_compute_stream{},
-        tunable_op_enabled{false},
-        default_memory_arena_cfg{} {}
+        default_memory_arena_cfg{},
+        tunable_op_enabled{false} {}
 #endif
 
   /** \brief ROCM device Id
@@ -482,16 +497,16 @@ typedef struct OrtROCMProviderOptions {
   */
   void* user_compute_stream;
 
+  /** \brief ROCM memory arena configuration parameters
+  */
+  OrtArenaCfg* default_memory_arena_cfg;
+
   /** \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;
-
 } OrtROCMProviderOptions;
 
 /** \brief TensorRT Provider Options

onnxruntime/core/platform/env_var_utils.h

@@ -3,7 +3,12 @@
 
 #pragma once
 
+#include <unordered_set>
+
 #include "core/common/common.h"
+#ifndef SHARED_PROVIDER
+#include "core/common/logging/logging.h"
+#endif
 #include "core/common/optional.h"
 #include "core/common/parse_string.h"
 #include "core/platform/env.h"
@@ -43,4 +48,44 @@ T ParseEnvironmentVariableWithDefault(const std::string& name, const T& default_
 
   return default_value;
 }
+
+/**
+ * Parses an environment variable value for testing convenience.
+ *
+ * This function ensures the value is valid and also produces a warning on its usage.
+ */
+template <typename T>
+std::optional<T> ParseTestOnlyEnvironmentVariable(const std::string& name,
+                                                  const std::unordered_set<std::string>& valid_values,
+                                                  const std::string& hint = "") {
+  ORT_ENFORCE(!valid_values.empty());
+
+#ifndef SHARED_PROVIDER
+  const std::string raw_env = Env::Default().GetEnvironmentVar(name);
+#else
+  const std::string raw_env = GetEnvironmentVar(name);
+#endif
+  if (raw_env.empty()) {
+    return std::nullopt;
+  }
+  if (valid_values.find(raw_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 ", raw_env);
+  }
+
+  auto env = onnxruntime::ParseEnvironmentVariable<T>(name);
+
+  std::string default_hint = "End users should opt for provider options or session options.";
+  const std::string& logged_hint = hint.empty() ? default_hint : hint;
+
+  LOGS_DEFAULT(WARNING) << "Environment variable " << name << " is used. It is reserved for internal testing prupose. "
+                        << logged_hint;
+
+  return env;
+}
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

@@ -3,6 +3,7 @@
 
 #include "core/common/inlined_containers.h"
 #include "core/providers/shared_library/provider_api.h"
+#include "core/platform/env_var_utils.h"
 #include "core/providers/cuda/cuda_execution_provider.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cuda_allocator.h"
@@ -198,6 +199,15 @@ void CUDAExecutionProvider::PerThreadContext::IncrementRegularRunCountBeforeGrap
 }
 #endif
 
+void OverrideTunableOpInfoByEnv(CUDAExecutionProviderInfo& info) {
+  auto env_tunable_op_enabled = onnxruntime::ParseTestOnlyEnvironmentVariable<bool>(
+      "ORT_CUDA_TUNABLE_OP_ENABLED", {"0", "1"}, "Use provider_options \"tunable_op_enabled\" instead.");
+  if (env_tunable_op_enabled.has_value() && env_tunable_op_enabled != info.tunable_op.enabled) {
+    LOGS_DEFAULT(INFO) << "ORT_CUDA_TUNABLE_OP_ENABLED is set to " << *env_tunable_op_enabled;
+    info.tunable_op.enabled = *env_tunable_op_enabled;
+  }
+}
+
 CUDAExecutionProvider::CUDAExecutionProvider(const CUDAExecutionProviderInfo& info)
     : IExecutionProvider{onnxruntime::kCudaExecutionProvider},
       info_{info} {
@@ -224,6 +234,8 @@ CUDAExecutionProvider::CUDAExecutionProvider(const CUDAExecutionProviderInfo& in
   size_t free = 0;
   size_t total = 0;
   CUDA_CALL_THROW(cudaMemGetInfo(&free, &total));
+
+  OverrideTunableOpInfoByEnv(info_);
 }
 
 CUDAExecutionProvider::~CUDAExecutionProvider() {
@@ -246,6 +258,20 @@ CUDAExecutionProvider::~CUDAExecutionProvider() {
   }
 }
 
+void CUDAExecutionProvider::EnableTunableOp() {
+  LOGS_DEFAULT(INFO) << "Enable TunableOp for CUDA Execution Provider";
+  info_.tunable_op.enabled = true;
+}
+
+void CUDAExecutionProvider::DisableTunableOp() {
+  LOGS_DEFAULT(INFO) << "Disable TunableOp for CUDA Execution Provider";
+  info_.tunable_op.enabled = false;
+}
+
+bool CUDAExecutionProvider::IsTunableOpEnabled() const {
+  return info_.tunable_op.enabled;
+}
+
 std::unique_ptr<profiling::EpProfiler> CUDAExecutionProvider::GetProfiler() {
   return std::make_unique<profiling::CudaProfiler>();
 }

onnxruntime/core/providers/cuda/cuda_execution_provider.h

@@ -126,6 +126,10 @@ class CUDAExecutionProvider : public IExecutionProvider {
   static AllocatorPtr CreateCudaAllocator(OrtDevice::DeviceId device_id, size_t cuda_mem_limit, ArenaExtendStrategy arena_extend_strategy,
                                           CUDAExecutionProviderExternalAllocatorInfo external_alloc_info, OrtArenaCfg* arena_cfg);
 
+  void EnableTunableOp();
+  void DisableTunableOp();
+  bool IsTunableOpEnabled() const;
+
   std::unique_ptr<profiling::EpProfiler> GetProfiler() override;
 
 #if defined(CUDA_VERSION) && CUDA_VERSION >= 10000

onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc

@@ -24,6 +24,7 @@ constexpr const char* kGpuExternalEmptyCache = "gpu_external_empty_cache";
 constexpr const char* kCudnnConvUseMaxWorkspace = "cudnn_conv_use_max_workspace";
 constexpr const char* kEnableCudaGraph = "enable_cuda_graph";
 constexpr const char* kCudnnConv1dPadToNc1d = "cudnn_conv1d_pad_to_nc1d";
+constexpr const char* kTunableOpEnabled = "tunable_op_enabled";
 }  // namespace provider_option_names
 }  // namespace cuda
 
@@ -92,6 +93,12 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P
           .AddAssignmentToReference(cuda::provider_option_names::kCudnnConvUseMaxWorkspace, info.cudnn_conv_use_max_workspace)
           .AddAssignmentToReference(cuda::provider_option_names::kEnableCudaGraph, info.enable_cuda_graph)
           .AddAssignmentToReference(cuda::provider_option_names::kCudnnConv1dPadToNc1d, info.cudnn_conv1d_pad_to_nc1d)
+          .AddValueParser(
+              cuda::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));
 
   CUDAExecutionProviderExternalAllocatorInfo alloc_info{alloc, free, empty_cache};
@@ -113,7 +120,8 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const CUDAExecution
       {cuda::provider_option_names::kDoCopyInDefaultStream, MakeStringWithClassicLocale(info.do_copy_in_default_stream)},
       {cuda::provider_option_names::kCudnnConvUseMaxWorkspace, MakeStringWithClassicLocale(info.cudnn_conv_use_max_workspace)},
       {cuda::provider_option_names::kEnableCudaGraph, MakeStringWithClassicLocale(info.enable_cuda_graph)},
-      {cuda::provider_option_names::kCudnnConv1dPadToNc1d, MakeStringWithClassicLocale(info.cudnn_conv1d_pad_to_nc1d)}
+      {cuda::provider_option_names::kCudnnConv1dPadToNc1d, MakeStringWithClassicLocale(info.cudnn_conv1d_pad_to_nc1d)},
+      {cuda::provider_option_names::kTunableOpEnabled, MakeStringWithClassicLocale(info.tunable_op.enabled)},
   };
 
   return options;

onnxruntime/core/providers/cuda/cuda_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 CUDAExecutionProviderExternalAllocatorInfo {
   }
 };
 
+namespace cuda {
+struct TunableOpInfo {
+  bool enabled{false};
+};
+}  // namespace cuda
+
 struct CUDAExecutionProviderInfo {
   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)
@@ -57,8 +65,19 @@ struct CUDAExecutionProviderInfo {
   // By default, for Conv1D, will pad [N,C,D] to [N,C,D,1], if turn on, will pad to [N,C,1,D].
   bool cudnn_conv1d_pad_to_nc1d{false};
 
+  cuda::TunableOpInfo tunable_op{};
+
   static CUDAExecutionProviderInfo FromProviderOptions(const ProviderOptions& options);
   static ProviderOptions ToProviderOptions(const CUDAExecutionProviderInfo& info);
   static ProviderOptions ToProviderOptions(const OrtCUDAProviderOptionsV2& info);
 };
 }  // namespace onnxruntime
+
+template<>
+struct std::hash<::onnxruntime::cuda::TunableOpInfo> {
+  size_t operator()(const ::onnxruntime::cuda::TunableOpInfo& info) const {
+    size_t seed_and_value{0xbc9f1d34};
+    onnxruntime::HashCombine(info.enabled, seed_and_value);
+    return seed_and_value;
+  }
+};

onnxruntime/core/providers/cuda/cuda_kernel.h

@@ -66,6 +66,8 @@ class CudaKernel : public OpKernel {
 
   inline cudaStream_t Stream() const { return static_cast<cudaStream_t>(provider_->GetComputeStream()); }
 
+  bool IsTunableOpEnabled() const { return provider_->IsTunableOpEnabled(); }
+
   // To support cudaMemcpyAsync, 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/cuda/cuda_provider_factory.cc

@@ -242,6 +242,7 @@ struct CUDA_Provider : Provider {
     info.cudnn_conv_use_max_workspace = params->cudnn_conv_use_max_workspace != 0;
     info.enable_cuda_graph = params->enable_cuda_graph != 0;
     info.cudnn_conv1d_pad_to_nc1d = params->cudnn_conv1d_pad_to_nc1d != 0;
+    info.tunable_op.enabled = params->tunable_op_enabled;
 
     return std::make_shared<CUDAProviderFactory>(info);
   }

onnxruntime/core/providers/rocm/rocm_execution_provider.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "core/providers/shared_library/provider_api.h"
+#include "core/platform/env_var_utils.h"
 #include "core/providers/rocm/rocm_execution_provider.h"
 #include "core/providers/rocm/rocm_common.h"
 #include "core/providers/rocm/rocm_allocator.h"
@@ -154,37 +155,12 @@ 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;
+  auto env_tunable_op_enabled = onnxruntime::ParseTestOnlyEnvironmentVariable<bool>(
+      "ORT_ROCM_TUNABLE_OP_ENABLED", {"0", "1"}, "Use provider_options \"tunable_op_enabled\" instead.");
+  if (env_tunable_op_enabled.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;
   }
 }
 

onnxruntime/core/providers/rocm/rocm_execution_provider_info.h

@@ -50,7 +50,6 @@ 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
@@ -59,6 +58,8 @@ struct ROCMExecutionProviderInfo {
   ROCMExecutionProviderExternalAllocatorInfo external_allocator_info{};
   bool miopen_conv_use_max_workspace{false};
 
+  rocm::TunableOpInfo tunable_op{};
+
   static ROCMExecutionProviderInfo FromProviderOptions(const ProviderOptions& options);
   static ProviderOptions ToProviderOptions(const ROCMExecutionProviderInfo& info);
 };

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -339,6 +339,7 @@ const CUDAExecutionProviderInfo GetCudaExecutionProviderInfo(ProviderInfo_CUDA*
     info.cudnn_conv_algo_search = cudnn_conv_algo_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

@@ -24,6 +24,8 @@ size_t gpu_mem_limit = std::numeric_limits<size_t>::max();
 OrtCudnnConvAlgoSearch cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive;
 // TODO remove deprecated global config
 bool do_copy_in_default_stream = true;
+// TODO remove deprecated global config
+onnxruntime::cuda::TunableOpInfo tunable_op{};
 onnxruntime::CUDAExecutionProviderExternalAllocatorInfo external_allocator_info{};
 // TODO remove deprecated global config
 onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo;

onnxruntime/python/onnxruntime_pybind_state_common.h

@@ -177,6 +177,8 @@ namespace python {
 extern OrtCudnnConvAlgoSearch cudnn_conv_algo_search;
 // TODO remove deprecated global config
 extern bool do_copy_in_default_stream;
+// TODO remove deprecated global config
+extern onnxruntime::cuda::TunableOpInfo tunable_op;
 extern onnxruntime::CUDAExecutionProviderExternalAllocatorInfo external_allocator_info;
 extern onnxruntime::ArenaExtendStrategy arena_extend_strategy;
 }  // namespace python

onnxruntime/test/python/onnxruntime_test_python.py

@@ -237,6 +237,8 @@ class TestInferenceSession(unittest.TestCase):
 
                 test_get_and_set_option_with_values("do_copy_in_default_stream", [0, 1])
 
+                test_get_and_set_option_with_values("tunable_op_enabled", ["1", "0"])
+
                 option["gpu_external_alloc"] = "0"
                 option["gpu_external_free"] = "0"
                 option["gpu_external_empty_cache"] = "0"

orttraining/orttraining/python/orttraining_python_module.cc

@@ -87,7 +87,8 @@ bool GetProviderInstanceHash(const std::string& type,
              (static_cast<size_t>(info.arena_extend_strategy) << 16) ^
              (static_cast<size_t>(info.cudnn_conv_algo_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<cuda::TunableOpInfo>{}(info.tunable_op);
       return true;
     }
 #endif