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.

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

@@ -0,0 +1,204 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/span>
+#include <cuda/stream_ref>
+
+#if _CCCL_STD_VER >= 2014 && !defined(_CCCL_COMPILER_MSVC_2017) \
+  && defined(LIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE)
+
+//! @file The uninitialized_async_buffer class provides a typed buffer allocated 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 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, we can only forward stateless propertie. If a user wants to use a stateful one, then they
+//! need to implement :ref:`get_property(const uninitialized_async_buffer&, Property)
+//! <libcudacxx-extended-api-memory-resources-properties>`.
+//!
+//! .. note::
+//!
+//!    ``uninitialized_async_buffer`` stores a reference to the provided memory `memory resource
+//!    <libcudacxx-extended-api-memory-resources-resource>`. It is the users resposibility to ensure the lifetime of the
+//!    resource exceeds the lifetime of the buffer.
+//!
+//! .. note::
+//!
+//!    ``uninitialized_async_buffer`` utilizes `stream ordered allocations
+//!    <https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-1/>`__. It is the users
+//!    resposibility to ensure the lifetime of the provided stream resource exceeds 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_VMR::async_resource_ref<_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 _CCCL_HOST_DEVICE 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 _CCCL_HOST_DEVICE constexpr _Tp* __get_data() const noexcept
+  {
+    constexpr size_t __alignment = alignof(_Tp);
+    size_t __space               = __get_allocation_size(__count_);
+    void* __ptr                  = __buf_;
+    return reinterpret_cast<_Tp*>(_CUDA_VSTD::align(__alignment, __count_ * sizeof(_Tp), __ptr, __space));
+  }
+
+public:
+  using value_type = _Tp;
+  using reference  = _Tp&;
+  using pointer    = _Tp*;
+  using size_type  = size_t;
+
+  //! @brief Constructs a \c uninitialized_async_buffer, allocating sufficient storage for \p count elements through
+  //! \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_VMR::async_resource_ref<_Properties...> __mr, const ::cuda::stream_ref __stream, const size_t __count)
+      : __mr_(__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_(__other.__mr_)
+      , __stream_(__other.__stream_)
+      , __count_(__other.__count_)
+      , __buf_(__other.__buf_)
+  {
+    __other.__stream_ = {};
+    __other.__count_  = 0;
+    __other.__buf_    = nullptr;
+  }
+
+  //! @brief Move assignment
+  //! @param other Another \c uninitialized_async_buffer
+  uninitialized_async_buffer& operator=(uninitialized_async_buffer&& __other) noexcept
+  {
+    if (__buf_)
+    {
+      __mr_.deallocate_async(__buf_, __get_allocation_size(__count_), __stream_);
+    }
+    __mr_             = __other.__mr_;
+    __stream_         = __other.__stream_;
+    __count_          = __other.__count_;
+    __buf_            = __other.__buf_;
+    __other.__stream_ = {};
+    __other.__count_  = 0;
+    __other.__buf_    = nullptr;
+    return *this;
+  }
+
+  //! @brief Destroys an \c uninitialized_async_buffer deallocating the buffer
+  //! @warning The destructor does not destroy any objects that may or may not reside within the buffer. It is the users
+  //! 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 _CCCL_HOST_DEVICE constexpr pointer begin() const noexcept
+  {
+    return __get_data();
+  }
+
+  //! @brief Returns an aligned pointer to end of the buffer
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE constexpr pointer end() const noexcept
+  {
+    return __get_data() + __count_;
+  }
+
+  //! @brief Returns an aligned pointer to the buffer
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE constexpr pointer data() const noexcept
+  {
+    return __get_data();
+  }
+
+  //! @brief Returns the size of the buffer
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE constexpr size_t size() const noexcept
+  {
+    return __count_;
+  }
+
+  //! @brief Returns the stream used to allocate
+  _CCCL_NODISCARD _CCCL_HOST_DEVICE constexpr ::cuda::stream_ref stream() const noexcept
+  {
+    return __stream_;
+  }
+
+#  ifndef DOXYGEN_SHOULD_SKIP_THIS // friend functions are currently brocken
+  //! @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/buffer

@@ -20,6 +20,7 @@
 #  pragma system_header
 #endif // no system header
 
+#include <cuda/experimental/__container/uninitialized_async_buffer.h>
 #include <cuda/experimental/__container/uninitialized_buffer.h>
 
 #endif //_CUDA_BUFFER

cudax/test/CMakeLists.txt

@@ -58,6 +58,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
   )
 
   target_compile_options(${test_target} PRIVATE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--extended-lambda>)

cudax/test/containers/uninitialized_async_buffer.cu

@@ -0,0 +1,129 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDA Experimental in CUDA C++ Core Libraries,
+// 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.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cuda/experimental/buffer>
+#include <cuda/experimental/memory_resource>
+#include <cuda/std/cassert>
+#include <cuda/std/cstdint>
+#include <cuda/std/type_traits>
+#include <cuda/std/utility>
+#include <cuda/stream_ref>
+
+#include <catch2/catch.hpp>
+
+struct do_not_construct
+{
+  do_not_construct()
+  {
+    CHECK(false);
+  }
+};
+
+struct my_property
+{
+  using value_type = int;
+};
+constexpr int get_property(const cuda::experimental::uninitialized_async_buffer<int, my_property>&, my_property)
+{
+  return 42;
+}
+
+TEMPLATE_TEST_CASE(
+  "uninitialized_async_buffer", "[memory_resource]", char, short, int, long, long long, float, double, do_not_construct)
+{
+  using uninitialized_async_buffer = cuda::experimental::uninitialized_async_buffer<TestType>;
+  static_assert(!cuda::std::is_default_constructible<uninitialized_async_buffer>::value, "");
+  static_assert(!cuda::std::is_copy_constructible<uninitialized_async_buffer>::value, "");
+  static_assert(!cuda::std::is_copy_assignable<uninitialized_async_buffer>::value, "");
+
+  cuda::experimental::mr::cuda_async_memory_resource resource{};
+
+  cudaStream_t raw_stream;
+  cudaStreamCreate(&raw_stream);
+  cuda::stream_ref stream{raw_stream};
+
+  SECTION("construction")
+  {
+    {
+      uninitialized_async_buffer from_stream_count{resource, stream, 42};
+      CHECK(from_stream_count.data() != nullptr);
+      CHECK(from_stream_count.size() == 42);
+    }
+    {
+      uninitialized_async_buffer input{resource, stream, 42};
+      const TestType* ptr = input.data();
+
+      uninitialized_async_buffer from_rvalue{cuda::std::move(input)};
+      CHECK(from_rvalue.data() == ptr);
+      CHECK(from_rvalue.size() == 42);
+      CHECK(from_rvalue.stream() == stream);
+
+      // Ensure that we properly reset the input buffer
+      CHECK(input.data() == nullptr);
+      CHECK(input.size() == 0);
+      CHECK(input.stream() == cuda::stream_ref{});
+    }
+
+    cudaStream_t other_raw_stream;
+    cudaStreamCreate(&other_raw_stream);
+    cuda::stream_ref other_stream{other_raw_stream};
+    {
+      uninitialized_async_buffer input{resource, other_stream, 42};
+      const TestType* ptr = input.data();
+
+      uninitialized_async_buffer assign_rvalue{resource, stream, 1337};
+      assign_rvalue = cuda::std::move(input);
+      CHECK(assign_rvalue.data() == ptr);
+      CHECK(assign_rvalue.size() == 42);
+      CHECK(assign_rvalue.stream() == other_stream);
+
+      // Ensure that we properly reset the input buffer
+      CHECK(input.data() == nullptr);
+      CHECK(input.size() == 0);
+      CHECK(input.stream() == cuda::stream_ref{});
+    }
+    cudaStreamDestroy(other_raw_stream);
+  }
+
+  SECTION("access")
+  {
+    uninitialized_async_buffer buf{resource, stream, 42};
+    CHECK(buf.data() != nullptr);
+    CHECK(buf.size() == 42);
+    CHECK(buf.begin() == buf.data());
+    CHECK(buf.end() == buf.begin() + buf.size());
+    CHECK(buf.stream() == stream);
+
+    CHECK(cuda::std::as_const(buf).data() != nullptr);
+    CHECK(cuda::std::as_const(buf).size() == 42);
+    CHECK(cuda::std::as_const(buf).begin() == buf.data());
+    CHECK(cuda::std::as_const(buf).end() == buf.begin() + buf.size());
+    CHECK(cuda::std::as_const(buf).stream() == stream);
+  }
+
+  SECTION("properties")
+  {
+    static_assert(cuda::has_property<cuda::experimental::uninitialized_async_buffer<int, cuda::mr::device_accessible>,
+                                     cuda::mr::device_accessible>,
+                  "");
+    static_assert(cuda::has_property<cuda::experimental::uninitialized_async_buffer<int, my_property>, my_property>,
+                  "");
+  }
+
+  SECTION("convertion to span")
+  {
+    uninitialized_async_buffer buf{resource, stream, 42};
+    const cuda::std::span<TestType> as_span{buf};
+    CHECK(as_span.data() == buf.data());
+    CHECK(as_span.size() == 42);
+  }
+
+  cudaStreamDestroy(raw_stream);
+}

docs/cudax/container.rst

@@ -13,8 +13,9 @@ annotations are checked by the type system.
 Uninitialized buffers
 ---------------------
 
-The ``<cuda/experimental/buffer>`` header contains facilities, that provide *heterogeneous* allocations to store objects
-in uninitialized memory. This is a common request in HPC due to the high cost of initialization of large arrays.
+The ``<cuda/experimental/buffer>`` header contains facilities, that provide *heterogeneous* and potentially
+*stream ordered* allocations to store objects in uninitialized memory. This is a common request in HPC due to the
+high cost of initialization of large arrays.
 
 .. warning::
 
@@ -25,3 +26,4 @@ in uninitialized memory. This is a common request in HPC due to the high cost of
    :maxdepth: 3
 
    container/uninitialized_buffer
+   container/uninitialized_async_buffer

docs/cudax/container/uninitialized_async_buffer.rst

@@ -0,0 +1,5 @@
+uninitialized_async_buffer
+==========================
+
+.. doxygenclass:: cuda::experimental::uninitialized_async_buffer
+   :members: