make uninitialized[_async]_buffer's range accessors const-correct

NVIDIA · Jan 30, 2025 · e1515cf · e1515cf
1 parent cea61a3
commit e1515cf
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 16 deletions.
diff --git a/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh b/cudax/include/cuda/experimental/__container/uninitialized_async_buffer.cuh
@@ -108,7 +108,7 @@ private:
     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)));
+      static_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.
@@ -136,10 +136,12 @@ private:
   }
 
 public:
-  using value_type = _Tp;
-  using reference  = _Tp&;
-  using pointer    = _Tp*;
-  using size_type  = size_t;
+  using value_type      = _Tp;
+  using reference       = _Tp&;
+  using const_reference = const _Tp&;
+  using pointer         = _Tp*;
+  using const_pointer   = const _Tp*;
+  using size_type       = size_t;
 
   //! @brief Constructs an \c uninitialized_async_buffer, allocating sufficient storage for \p __count elements through
   //! \p __mr
@@ -215,20 +217,38 @@ public:
   }
 
   //! @brief Returns an aligned pointer to the first element in the buffer
-  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr pointer begin() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr pointer begin() noexcept
+  {
+    return __get_data();
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr const_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 _CCCL_HIDE_FROM_ABI constexpr pointer end() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr pointer end() noexcept
+  {
+    return __get_data() + __count_;
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr const_pointer end() const noexcept
   {
     return __get_data() + __count_;
   }
 
   //! @brief Returns an aligned pointer to the first element in the buffer
-  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr pointer data() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr pointer data() noexcept
+  {
+    return __get_data();
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI constexpr const_pointer data() const noexcept
   {
     return __get_data();
   }

diff --git a/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh b/cudax/include/cuda/experimental/__container/uninitialized_buffer.cuh
@@ -98,7 +98,7 @@ private:
     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)));
+      static_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.
@@ -124,10 +124,12 @@ private:
   }
 
 public:
-  using value_type = _Tp;
-  using reference  = _Tp&;
-  using pointer    = _Tp*;
-  using size_type  = size_t;
+  using value_type      = _Tp;
+  using reference       = _Tp&;
+  using const_reference = const _Tp&;
+  using pointer         = _Tp*;
+  using const_pointer   = const _Tp*;
+  using size_type       = size_t;
 
   //! @brief Constructs an \c uninitialized_buffer and allocates sufficient storage for \p __count elements through
   //! \p __mr
@@ -198,20 +200,38 @@ public:
   }
 
   //! @brief Returns an aligned pointer to the first element in the buffer
-  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI pointer begin() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI pointer begin() noexcept
+  {
+    return __get_data();
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI const_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 _CCCL_HIDE_FROM_ABI pointer end() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI pointer end() noexcept
+  {
+    return __get_data() + __count_;
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI const_pointer end() const noexcept
   {
     return __get_data() + __count_;
   }
 
   //! @brief Returns an aligned pointer to the first element in the buffer
-  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI pointer data() const noexcept
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI pointer data() noexcept
+  {
+    return __get_data();
+  }
+
+  //! @overload
+  _CCCL_NODISCARD _CCCL_HIDE_FROM_ABI const_pointer data() const noexcept
   {
     return __get_data();
   }

diff --git a/cudax/test/containers/uninitialized_async_buffer.cu b/cudax/test/containers/uninitialized_async_buffer.cu
@@ -58,6 +58,18 @@ TEMPLATE_TEST_CASE(
   cuda::experimental::device_memory_resource resource{};
   cuda::experimental::stream stream{};
 
+  if (false)
+  {
+    uninitialized_async_buffer buf{resource, stream, 42};
+    uninitialized_async_buffer const& cbuf = buf;
+    static_assert(cuda::std::is_same<decltype(buf.begin()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.begin()), TestType const*>::value, "");
+    static_assert(cuda::std::is_same<decltype(buf.end()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.end()), TestType const*>::value, "");
+    static_assert(cuda::std::is_same<decltype(buf.data()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.data()), TestType const*>::value, "");
+  }
+
   SECTION("construction")
   {
     {

diff --git a/cudax/test/containers/uninitialized_buffer.cu b/cudax/test/containers/uninitialized_buffer.cu
@@ -83,6 +83,18 @@ TEMPLATE_TEST_CASE(
 
   cudax::device_memory_resource resource{};
 
+  if (false)
+  {
+    uninitialized_buffer buf{resource, 42};
+    uninitialized_buffer const& cbuf = buf;
+    static_assert(cuda::std::is_same<decltype(buf.begin()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.begin()), TestType const*>::value, "");
+    static_assert(cuda::std::is_same<decltype(buf.end()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.end()), TestType const*>::value, "");
+    static_assert(cuda::std::is_same<decltype(buf.data()), TestType*>::value, "");
+    static_assert(cuda::std::is_same<decltype(cbuf.data()), TestType const*>::value, "");
+  }
+
   SECTION("construction")
   {
     static_assert(!cuda::std::is_copy_constructible<uninitialized_buffer>::value, "");