[PoC]: Implement cuda::experimental::uninitialized_async_buffer (NV…

…IDIA#1854)

* Implement `cuda::uninitialized_async_buffer`

This uninitialized buffer provides a stream ordered allocation of N elements of type T utilitzing a cuda::mr::async_resource to allocate the storage.

The buffer takes care of alignment and deallocation of the storage. The user is required to ensure that the lifetime of the memory resource exceeds the lifetime of the buffer.

Co-authored-by: Mark Harris <[email protected]>

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

@@ -0,0 +1,258 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the CUDA Toolkit, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef __CUDAX__CONTAINERS_UNINITIALIZED_ASYNC_BUFFER_H
+#define __CUDAX__CONTAINERS_UNINITIALIZED_ASYNC_BUFFER_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/__memory_resource/properties.h>
+#include <cuda/__memory_resource/resource_ref.h>
+#include <cuda/std/__concepts/_One_of.h>
+#include <cuda/std/__memory/align.h>
+#include <cuda/std/__new/launder.h>
+#include <cuda/std/__utility/exchange.h>
+#include <cuda/std/__utility/move.h>
+#include <cuda/std/__utility/swap.h>
+#include <cuda/std/span>
+#include <cuda/stream_ref>
+
+#include <cuda/experimental/__memory_resource/any_resource.cuh>
+
+#if _CCCL_STD_VER >= 2014 && !defined(_CCCL_COMPILER_MSVC_2017) \
+  && defined(LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE)
+
+//! @file
+//! The \c uninitialized_async_buffer class provides a typed buffer allocated in stream-order from a given memory
+//! resource.
+namespace cuda::experimental
+{
+
+//! @rst
+//! .. _cudax-containers-uninitialized-async-buffer:
+//!
+//! Uninitialized stream-ordered type-safe memory storage
+//! ------------------------------------------------------
+//!
+//! ``uninitialized_async_buffer`` provides a typed buffer allocated in stream order from a given :ref:`async memory
+//! resource <libcudacxx-extended-api-memory-resources-resource>`. It handles alignment and release of the allocation.
+//! The memory is uninitialized, so that a user needs to ensure elements are properly constructed.
+//!
+//! In addition to being type safe, ``uninitialized_async_buffer`` also takes a set of :ref:`properties
+//! <libcudacxx-extended-api-memory-resources-properties>` to ensure that e.g. execution space constraints are checked
+//! at compile time. However, only stateless properties can be forwarded. To use a stateful property,
+//! implement :ref:`get_property(const uninitialized_async_buffer&, Property)
+//! <libcudacxx-extended-api-memory-resources-properties>`.
+//!
+//! .. warning::
+//!
+//!    ``uninitialized_async_buffer`` uses `stream-ordered allocation
+//!    <https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/>`__. It is the user's
+//!    resposibility to ensure the lifetime of both the provided async resource and the stream exceed the lifetime of
+//!    the buffer.
+//!
+//! @endrst
+//! @tparam _T the type to be stored in the buffer
+//! @tparam _Properties... The properties the allocated memory satisfies
+template <class _Tp, class... _Properties>
+class uninitialized_async_buffer
+{
+private:
+  ::cuda::experimental::mr::async_any_resource<_Properties...> __mr_;
+  ::cuda::stream_ref __stream_ = {};
+  size_t __count_              = 0;
+  void* __buf_                 = nullptr;
+
+  //! @brief Determines the allocation size given the alignment and size of `T`
+  _CCCL_NODISCARD static constexpr size_t __get_allocation_size(const size_t __count) noexcept
+  {
+    constexpr size_t __alignment = alignof(_Tp);
+    return (__count * sizeof(_Tp) + (__alignment - 1)) & ~(__alignment - 1);
+  }
+
+  //! @brief Determines the properly aligned start of the buffer given the alignment and size of `T`
+  _CCCL_NODISCARD constexpr _Tp* __get_data() const noexcept
+  {
+    constexpr size_t __alignment = alignof(_Tp);
+    size_t __space               = __get_allocation_size(__count_);
+    void* __ptr                  = __buf_;
+    return _CUDA_VSTD::launder(
+      reinterpret_cast<_Tp*>(_CUDA_VSTD::align(__alignment, __count_ * sizeof(_Tp), __ptr, __space)));
+  }
+
+  //! @brief Causes the buffer to be treated as a span when passed to cudax::launch.
+  //! @pre The buffer must have the cuda::mr::device_accessible property.
+  _CCCL_NODISCARD_FRIEND _CUDA_VSTD::span<_Tp>
+  __cudax_launch_transform(::cuda::stream_ref, uninitialized_async_buffer& __self) noexcept
+  {
+    static_assert(_CUDA_VSTD::_One_of<_CUDA_VMR::device_accessible, _Properties...>,
+                  "The buffer must be device accessible to be passed to `launch`");
+    return {__self.__get_data(), __self.size()};
+  }
+
+  //! @brief Causes the buffer to be treated as a span when passed to cudax::launch
+  //! @pre The buffer must have the cuda::mr::device_accessible property.
+  _CCCL_NODISCARD_FRIEND _CUDA_VSTD::span<const _Tp>
+  __cudax_launch_transform(::cuda::stream_ref, const uninitialized_async_buffer& __self) noexcept
+  {
+    static_assert(_CUDA_VSTD::_One_of<_CUDA_VMR::device_accessible, _Properties...>,
+                  "The buffer must be device accessible to be passed to `launch`");
+    return {__self.__get_data(), __self.size()};
+  }
+
+public:
+  using value_type = _Tp;
+  using reference  = _Tp&;
+  using pointer    = _Tp*;
+  using size_type  = size_t;
+
+  //! @brief Constructs an \c uninitialized_async_buffer, allocating sufficient storage for \p __count elements using
+  //! \p __mr
+  //! @param __mr The async memory resource to allocate the buffer with.
+  //! @param __stream The CUDA stream used for stream-ordered allocation.
+  //! @param __count The desired size of the buffer.
+  //! @note Depending on the alignment requirements of `T` the size of the underlying allocation might be larger
+  //! than `count * sizeof(T)`. Only allocates memory when \p __count > 0
+  uninitialized_async_buffer(::cuda::experimental::mr::async_any_resource<_Properties...> __mr,
+                             const ::cuda::stream_ref __stream,
+                             const size_t __count)
+      : __mr_(_CUDA_VSTD::move(__mr))
+      , __stream_(__stream)
+      , __count_(__count)
+      , __buf_(__count_ == 0 ? nullptr : __mr_.allocate_async(__get_allocation_size(__count_), __stream_))
+  {}
+
+  uninitialized_async_buffer(const uninitialized_async_buffer&)            = delete;
+  uninitialized_async_buffer& operator=(const uninitialized_async_buffer&) = delete;
+
+  //! @brief Move construction
+  //! @param __other Another \c uninitialized_async_buffer
+  uninitialized_async_buffer(uninitialized_async_buffer&& __other) noexcept
+      : __mr_(_CUDA_VSTD::move(__other.__mr_))
+      , __stream_(_CUDA_VSTD::exchange(__other.__stream_, {}))
+      , __count_(_CUDA_VSTD::exchange(__other.__count_, 0))
+      , __buf_(_CUDA_VSTD::exchange(__other.__buf_, nullptr))
+  {}
+
+  //! @brief Move assignment
+  //! @param __other Another \c uninitialized_async_buffer
+  uninitialized_async_buffer& operator=(uninitialized_async_buffer&& __other) noexcept
+  {
+    if (this == _CUDA_VSTD::addressof(__other))
+    {
+      return *this;
+    }
+
+    if (__buf_)
+    {
+      __mr_.deallocate_async(__buf_, __get_allocation_size(__count_), __stream_);
+    }
+    __mr_     = __other.__mr_;
+    __stream_ = _CUDA_VSTD::exchange(__other.__stream_, {});
+    __count_  = _CUDA_VSTD::exchange(__other.__count_, 0);
+    __buf_    = _CUDA_VSTD::exchange(__other.__buf_, nullptr);
+    return *this;
+  }
+  //! @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
+  //! user's responsibility to ensure that all objects within the buffer have been properly destroyed.
+  ~uninitialized_async_buffer()
+  {
+    if (__buf_)
+    {
+      __mr_.deallocate_async(__buf_, __get_allocation_size(__count_), __stream_);
+    }
+  }
+
+  //! @brief Returns an aligned pointer to the buffer
+  _CCCL_NODISCARD constexpr pointer begin() const noexcept
+  {
+    return __get_data();
+  }
+
+  //! @brief Returns an aligned pointer to the element following the last element of the buffer.
+  //! This element acts as a placeholder; attempting to access it results in undefined behavior.
+  _CCCL_NODISCARD constexpr pointer end() const noexcept
+  {
+    return __get_data() + __count_;
+  }
+
+  //! @brief Returns an aligned pointer to the buffer
+  _CCCL_NODISCARD constexpr pointer data() const noexcept
+  {
+    return __get_data();
+  }
+
+  //! @brief Returns the size of the buffer
+  _CCCL_NODISCARD constexpr size_t size() const noexcept
+  {
+    return __count_;
+  }
+
+  //! @rst
+  //! Returns an :ref:`asnyc_resource_ref <libcudacxx-extended-api-memory-resources-resource-ref>` to the resource used
+  //! to allocate the buffer
+  //! @endrst
+  _CCCL_NODISCARD _CUDA_VMR::async_resource_ref<_Properties...> get_resource() const noexcept
+  {
+    return _CUDA_VMR::async_resource_ref<_Properties...>{const_cast<uninitialized_async_buffer*>(this)->__mr_};
+  }
+
+  //! @brief Returns the stored stream
+  _CCCL_NODISCARD constexpr ::cuda::stream_ref get_stream() const noexcept
+  {
+    return __stream_;
+  }
+
+  //! @brief Replaces the stored stream
+  //! @param __new_stream the new stream
+  //! @note Always synchronizes with the old stream
+  constexpr void change_stream(::cuda::stream_ref __new_stream)
+  {
+    if (__new_stream != __stream_)
+    {
+      __stream_.wait();
+    }
+    __stream_ = __new_stream;
+  }
+
+  //! @brief Swaps the contents with those of another \c uninitialized_async_buffer
+  //! @param __other The other \c uninitialized_async_buffer.
+  constexpr void swap(uninitialized_async_buffer& __other) noexcept
+  {
+    _CUDA_VSTD::swap(__mr_, __other.__mr_);
+    _CUDA_VSTD::swap(__count_, __other.__count_);
+    _CUDA_VSTD::swap(__buf_, __other.__buf_);
+  }
+
+#  ifndef DOXYGEN_SHOULD_SKIP_THIS // friend functions are currently broken
+  //! @brief Forwards the passed properties
+  _LIBCUDACXX_TEMPLATE(class _Property)
+  _LIBCUDACXX_REQUIRES((!property_with_value<_Property>) _LIBCUDACXX_AND _CUDA_VSTD::_One_of<_Property, _Properties...>)
+  friend constexpr void get_property(const uninitialized_async_buffer&, _Property) noexcept {}
+#  endif // DOXYGEN_SHOULD_SKIP_THIS
+};
+
+template <class _Tp>
+using uninitialized_async_device_buffer = uninitialized_async_buffer<_Tp, _CUDA_VMR::device_accessible>;
+
+} // namespace cuda::experimental
+
+#endif // _CCCL_STD_VER >= 2014 && !_CCCL_COMPILER_MSVC_2017 && LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE
+
+#endif //__CUDAX__CONTAINERS_UNINITIALIZED_ASYNC_BUFFER_H

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

@@ -25,6 +25,7 @@
 #include <cuda/std/__concepts/_One_of.h>
 #include <cuda/std/__memory/align.h>
 #include <cuda/std/__new/launder.h>
+#include <cuda/std/__utility/exchange.h>
 #include <cuda/std/__utility/move.h>
 #include <cuda/std/__utility/swap.h>
 #include <cuda/std/span>
@@ -128,12 +129,9 @@ public:
   //! @param __other Another \c uninitialized_buffer
   uninitialized_buffer(uninitialized_buffer&& __other) noexcept
       : __mr_(_CUDA_VSTD::move(__other.__mr_))
-      , __count_(__other.__count_)
-      , __buf_(__other.__buf_)
-  {
-    __other.__count_ = 0;
-    __other.__buf_   = nullptr;
-  }
+      , __count_(_CUDA_VSTD::exchange(__other.__count_, 0))
+      , __buf_(_CUDA_VSTD::exchange(__other.__buf_, nullptr))
+  {}
 
   //! @brief Move assignment
   //! @param __other Another \c uninitialized_buffer
@@ -148,11 +146,9 @@ public:
     {
       __mr_.deallocate(__buf_, __get_allocation_size(__count_));
     }
-    __mr_            = _CUDA_VSTD::move(__other.__mr_);
-    __count_         = __other.__count_;
-    __buf_           = __other.__buf_;
-    __other.__count_ = 0;
-    __other.__buf_   = nullptr;
+    __mr_    = _CUDA_VSTD::move(__other.__mr_);
+    __count_ = _CUDA_VSTD::exchange(__other.__count_, 0);
+    __buf_   = _CUDA_VSTD::exchange(__other.__buf_, nullptr);
     return *this;
   }
 
@@ -196,7 +192,7 @@ public:
   //! allocate the buffer
   //! @endrst
   _CCCL_EXEC_CHECK_DISABLE
-  _CCCL_NODISCARD _CCCL_HOST_DEVICE _CUDA_VMR::resource_ref<_Properties...> resource() const noexcept
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE _CUDA_VMR::resource_ref<_Properties...> get_resource() const noexcept
   {
     return _CUDA_VMR::resource_ref<_Properties...>{const_cast<uninitialized_buffer*>(this)->__mr_};
   }

cudax/include/cuda/experimental/__memory_resource/async_memory_resource.cuh

@@ -44,6 +44,9 @@
 
 #  if _CCCL_STD_VER >= 2014
 
+//! @file
+//! The \c async_memory_pool class provides an asynchronous memory resource that allocates device memory in stream
+//! order.
 namespace cuda::experimental::mr
 {
 

cudax/include/cuda/experimental/buffer.cuh

@@ -7,8 +7,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef __CUDAX_BUFFER
-#define __CUDAX_BUFFER
+#ifndef __CUDAX_BUFFER__
+#define __CUDAX_BUFFER__
 
 #include <cuda/std/detail/__config>
 
@@ -20,6 +20,7 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cuda/experimental/__container/uninitialized_async_buffer.cuh>
 #include <cuda/experimental/__container/uninitialized_buffer.cuh>
 
 #endif // __CUDAX_BUFFER

cudax/test/CMakeLists.txt

@@ -88,6 +88,7 @@ foreach(cn_target IN LISTS cudax_TARGETS)
 
   cudax_add_catch2_test(test_target containers ${cn_target}
     containers/uninitialized_buffer.cu
+    containers/uninitialized_async_buffer.cu
   )
 
   cudax_add_catch2_test(test_target memory_resource ${cn_target}