[CUDAX] Remove synchronization from set_stream and add a stream argument to destroy in async_buffer (#5697)

* Remove synchronization from set_stream and add a stream argument to destroy

* Fix format

* Adjust comments

Author: pciolkosz

cudax/include/cuda/experimental/__container/async_buffer.cuh

@@ -550,6 +550,7 @@ public:
   }
 
   //! @brief Returns the stored stream
+  //! @note Stream used to allocate the buffer is initially stored in the buffer, but can be changed with `set_stream`
   [[nodiscard]] _CCCL_HIDE_FROM_ABI constexpr stream_ref stream() const noexcept
   {
     return __buf_.stream();
@@ -559,15 +560,6 @@ public:
   //! @param __new_stream the new stream
   //! @note Always synchronizes with the old stream
   _CCCL_HIDE_FROM_ABI constexpr void set_stream(stream_ref __new_stream)
-  {
-    __buf_.set_stream(__new_stream);
-  }
-
-  //! @brief Replaces the stored stream
-  //! @param __new_stream the new stream
-  //! @warning This does not synchronize between \p __new_stream and the current stream. It is the user's responsibility
-  //! to ensure proper stream order going forward
-  _CCCL_HIDE_FROM_ABI constexpr void set_stream_unsynchronized(stream_ref __new_stream) noexcept
   {
     __buf_.set_stream_unsynchronized(__new_stream);
   }
@@ -596,6 +588,15 @@ public:
   }
 
   //! @brief Destroys the async_buffer, deallocates the buffer and destroys the memory resource
+  //! @param __stream The stream to deallocate the buffer on.
+  //! @warning After this explicit destroy call, the buffer can only be assigned to or destroyed.
+  _CCCL_HIDE_FROM_ABI void destroy(::cuda::stream_ref __stream)
+  {
+    __buf_.destroy(__stream);
+  }
+
+  //! @brief Destroys the async_buffer, deallocates the buffer and destroys the memory resource
+  //! @note Uses the stored stream to deallocate the buffer, equivalent to calling buffer.destroy(buffer.stream())
   //! @warning After this explicit destroy call, the buffer can only be assigned to or destroyed.
   _CCCL_HIDE_FROM_ABI void destroy()
   {

cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh

@@ -249,20 +249,32 @@ public:
   }
 
   //! @brief Destroys an \c uninitialized_async_buffer, deallocates the buffer in stream order on the stream that
-  //! was used to create the buffer and destroys the memory resource.
+  //! is stored in the buffer and destroys the memory resource.
+  //! @param __stream The stream to deallocate the buffer on.
   //! @warning destroy does not destroy any objects that may or may not reside within the buffer. It is the
   //! user's responsibility to ensure that all objects within the buffer have been properly destroyed.
-  _CCCL_HIDE_FROM_ABI void destroy()
+  _CCCL_HIDE_FROM_ABI void destroy(::cuda::stream_ref __stream)
   {
     if (__buf_)
     {
-      __mr_.deallocate(__stream_, __buf_, __get_allocation_size(__count_));
+      __mr_.deallocate(__stream, __buf_, __get_allocation_size(__count_));
       __buf_   = nullptr;
       __count_ = 0;
     }
+    // TODO should we make sure we move the mr only once by moving it to the if above?
+    // It won't work for 0 count buffers, so we would probably need a separate bool to track it
     auto __tmp_mr = ::cuda::std::move(__mr_);
   }
 
+  //! @brief Destroys an \c uninitialized_async_buffer, deallocates the buffer in stream order on the stream that
+  //! is stored in the buffer and destroys the memory resource.
+  //! @warning destroy does not destroy any objects that may or may not reside within the buffer. It is the
+  //! user's responsibility to ensure that all objects within the buffer have been properly destroyed.
+  _CCCL_HIDE_FROM_ABI void destroy()
+  {
+    destroy(__stream_);
+  }
+
   //! @brief Destroys an \c uninitialized_async_buffer and deallocates the buffer in stream order on the stream
   //! that was used to create the buffer.
   //! @warning The destructor does not destroy any objects that may or may not reside within the buffer. It is the
@@ -331,6 +343,7 @@ public:
   }
 
   //! @brief Returns the stored stream
+  //! @note Stream used to allocate the buffer is initially stored in the buffer, but can be changed with `set_stream`
   [[nodiscard]] _CCCL_HIDE_FROM_ABI constexpr ::cuda::stream_ref stream() const noexcept
   {
     return __stream_;

cudax/test/containers/async_buffer/access.cu

@@ -95,12 +95,12 @@ C2H_CCCLRT_TEST("cudax::async_buffer access and stream", "[container][async_buff
 
     {
       cudax::stream other_stream{cuda::device_ref{0}};
-      buf.set_stream_unsynchronized(other_stream);
+      buf.set_stream(other_stream);
       CUDAX_CHECK(buf.stream() == other_stream);
-      // TODO swap to synchronized setter one cudax::stream_ref is moved to cuda namespace
-      buf.set_stream_unsynchronized(stream);
+      buf.set_stream(stream);
     }
 
     CUDAX_CHECK(buf.stream() == stream);
+    buf.destroy(stream);
   }
 }