Create stream in state.get_cuda_stream if needed.

This simplifies usage in the common case where the caller just want an expected stream.

The const method `get_cuda_stream_optional` still returns the optional if this behavior is undesired.

Author: alliepiper

nvbench/detail/measure_cold.cu

@@ -37,13 +37,7 @@ namespace nvbench::detail
 
 measure_cold_base::measure_cold_base(state &exec_state)
     : m_state{exec_state}
-    , m_launch{nvbench::launch([this]() -> decltype(auto) {
-      if (!m_state.get_cuda_stream().has_value())
-      {
-        m_state.set_cuda_stream(nvbench::cuda_stream{m_state.get_device()});
-      }
-      return m_state.get_cuda_stream().value();
-    }())}
+    , m_launch{exec_state.get_cuda_stream()}
     , m_criterion_params{exec_state.get_criterion_params()}
     , m_stopping_criterion{nvbench::criterion_manager::get().get_criterion(
         exec_state.get_stopping_criterion())}

nvbench/detail/measure_cupti.cu

@@ -165,13 +165,7 @@ measure_cupti_base::measure_cupti_base(state &exec_state)
 // (formatter doesn't handle `try :` very well...)
 try
   : m_state{exec_state}
-  , m_launch{[this]() -> decltype(auto) {
-                   if (!m_state.get_cuda_stream().has_value())
-                   {
-                    m_state.set_cuda_stream(nvbench::cuda_stream{m_state.get_device()});
-                   }
-                   return m_state.get_cuda_stream().value();
-                 }()}
+  , m_launch{exec_state.get_cuda_stream()}
   , m_cupti{*m_state.get_device(), add_metrics(m_state)}
 {}
 // clang-format on

nvbench/detail/measure_hot.cu

@@ -36,13 +36,7 @@ namespace nvbench::detail
 
 measure_hot_base::measure_hot_base(state &exec_state)
     : m_state{exec_state}
-    , m_launch{nvbench::launch([this]() -> decltype(auto) {
-      if (!m_state.get_cuda_stream().has_value())
-      {
-        m_state.set_cuda_stream(nvbench::cuda_stream{m_state.get_device()});
-      }
-      return m_state.get_cuda_stream().value();
-    }())}
+    , m_launch{exec_state.get_cuda_stream()}
     , m_min_samples{exec_state.get_min_samples()}
     , m_min_time{exec_state.get_min_time()}
     , m_skip_time{exec_state.get_skip_time()}

nvbench/state.cuh

@@ -64,7 +64,18 @@ struct state
   state &operator=(const state &) = delete;
   state &operator=(state &&)      = default;
 
-  [[nodiscard]] const std::optional<nvbench::cuda_stream> &get_cuda_stream() const
+  /// If a stream exists, return that. Otherwise, create a new stream using the current
+  /// device (or the current device if none is set), save it, and return it.
+  /// @sa get_cuda_stream_optional
+  [[nodiscard]] nvbench::cuda_stream &get_cuda_stream()
+  {
+    if (!m_cuda_stream.has_value())
+    {
+      m_cuda_stream = nvbench::cuda_stream{m_device};
+    }
+    return m_cuda_stream.value();
+  }
+  [[nodiscard]] const std::optional<nvbench::cuda_stream> &get_cuda_stream_optional() const
   {
     return m_cuda_stream;
   }

testing/state.cu

@@ -56,11 +56,12 @@ void test_streams()
   state_tester state{bench};
 
   // Confirm that the stream hasn't been initialized yet
-  ASSERT(!state.get_cuda_stream().has_value());
+  ASSERT(!state.get_cuda_stream_optional().has_value());
 
   // Test non-owning stream
   cudaStream_t default_stream = 0;
   state.set_cuda_stream(nvbench::cuda_stream{default_stream, false});
+  ASSERT(state.get_cuda_stream_optional() == default_stream);
   ASSERT(state.get_cuda_stream() == default_stream);
 
   // Test owning stream