Ensure we only synchronize if the streams differ in set_stream

NVIDIA · Sep 16, 2024 · a92af5d · a92af5d
1 parent 84f67e2
commit a92af5d
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh b/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh
@@ -221,11 +221,14 @@ public:
 
   //! @brief Replaces the stored stream
   //! @param __stream the new stream
-  //! @note Synchronizes with the old stream
-  _CCCL_NODISCARD constexpr void get_stream(::cuda::stream_ref __stream)
+  //! @note Always synchronizes with the old stream
+  constexpr void set_stream(::cuda::stream_ref __stream)
   {
-    __stream_.wait();
-    __stream_ = __stream;
+    if (__stream != __stream_)
+    {
+      __stream_.wait();
+      __stream_ = __stream;
+    }
   }
 
   //! @brief Swaps the contents with those of another \c uninitialized_async_buffer