Remove pre-main CUDART usage and setup env in main.

NVIDIA · Apr 5, 2024 · c5a40d8 · c5a40d8
1 parent d77f335
commit c5a40d8
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 13 deletions.
diff --git a/nvbench/benchmark_base.cuh b/nvbench/benchmark_base.cuh
@@ -20,7 +20,6 @@
 
 #include <nvbench/axes_metadata.cuh>
 #include <nvbench/device_info.cuh>
-#include <nvbench/device_manager.cuh>
 #include <nvbench/state.cuh>
 #include <nvbench/stopping_criterion.cuh>
 
@@ -53,7 +52,6 @@ struct benchmark_base
   template <typename TypeAxes>
   explicit benchmark_base(TypeAxes type_axes)
       : m_axes(type_axes)
-      , m_devices(nvbench::device_manager::get().get_devices())
   {}
 
   virtual ~benchmark_base();

diff --git a/nvbench/main.cuh b/nvbench/main.cuh
@@ -28,17 +28,6 @@
 #include <cstdlib>
 #include <iostream>
 
-static int _nvbench_env [[maybe_unused]] = []() -> int
-{
-  // See NVIDIA/NVBench#136
-#ifdef _MSC_VER
-  int retval = _putenv_s("CUDA_MODULE_LOADING", "EAGER");
-#else
-  int retval = setenv("CUDA_MODULE_LOADING", "EAGER", 1);
-#endif
-  return retval;
-}();
-
 #define NVBENCH_MAIN                                                                               \
   int main(int argc, char const *const *argv)                                                      \
   try                                                                                              \
@@ -70,9 +59,17 @@ static int _nvbench_env [[maybe_unused]] = []() -> int
   nvbench::option_parser parser;                                                                   \
   parser.parse(argc, argv)
 
+// See NVIDIA/NVBench#136 for CUDA_MODULE_LOADING
+#ifdef _MSC_VER
+#define NVBENCH_INITIALIZE_CUDA_ENV _putenv_s("CUDA_MODULE_LOADING", "EAGER")
+#else
+#define NVBENCH_INITIALIZE_CUDA_ENV setenv("CUDA_MODULE_LOADING", "EAGER", 1)
+#endif
+
 #define NVBENCH_MAIN_BODY(argc, argv)                                                              \
   do                                                                                               \
   {                                                                                                \
+    NVBENCH_INITIALIZE_CUDA_ENV;                                                                   \
     NVBENCH_INITIALIZE_DRIVER_API;                                                                 \
     NVBENCH_MAIN_PARSE(argc, argv);                                                                \
     auto &printer = parser.get_printer();                                                          \

diff --git a/nvbench/option_parser.cu b/nvbench/option_parser.cu
@@ -22,6 +22,7 @@
 #include <nvbench/benchmark_manager.cuh>
 #include <nvbench/csv_printer.cuh>
 #include <nvbench/criterion_manager.cuh>
+#include <nvbench/device_manager.cuh>
 #include <nvbench/git_revision.cuh>
 #include <nvbench/json_printer.cuh>
 #include <nvbench/markdown_printer.cuh>
@@ -781,6 +782,8 @@ catch (std::exception &e)
 
 void option_parser::replay_global_args()
 {
+  // Initialize benchmark with all devices:
+  m_benchmarks.back()->set_devices(nvbench::device_manager::get().get_devices());
   this->parse_range(m_global_benchmark_args.cbegin(), m_global_benchmark_args.cend());
 }