Update detail memcpy_async to return CUDA error

Author: PointKernel

include/cuco/detail/hyperloglog/hyperloglog_impl.cuh

@@ -19,7 +19,7 @@
 #include <cuco/detail/error.hpp>
 #include <cuco/detail/hyperloglog/finalizer.cuh>
 #include <cuco/detail/hyperloglog/kernels.cuh>
-#include <cuco/detail/utility/memcpy_async.cuh>
+#include <cuco/detail/utility/memcpy_async.hpp>
 #include <cuco/detail/utils.hpp>
 #include <cuco/hash_functions.cuh>
 #include <cuco/types.cuh>
@@ -421,11 +421,11 @@ class hyperloglog_impl {
     std::vector<register_type> host_sketch(num_regs);
 
     // TODO check if storage is host accessible
-    cuco::detail::memcpy_async(host_sketch.data(),
-                               this->sketch_.data(),
-                               sizeof(register_type) * num_regs,
-                               cudaMemcpyDefault,
-                               stream);
+    CUCO_CUDA_TRY(cuco::detail::memcpy_async(host_sketch.data(),
+                                             this->sketch_.data(),
+                                             sizeof(register_type) * num_regs,
+                                             cudaMemcpyDefault,
+                                             stream));
 #if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
     stream.sync();
 #else

include/cuco/detail/open_addressing/open_addressing_impl.cuh

@@ -21,7 +21,7 @@
 #include <cuco/detail/open_addressing/kernels.cuh>
 #include <cuco/detail/storage/counter_storage.cuh>
 #include <cuco/detail/utility/cuda.hpp>
-#include <cuco/detail/utility/memcpy_async.cuh>
+#include <cuco/detail/utility/memcpy_async.hpp>
 #include <cuco/detail/utils.hpp>
 #include <cuco/extent.cuh>
 #include <cuco/operator.hpp>
@@ -883,8 +883,8 @@ class open_addressing_impl {
                                           stream.get()));
 
       size_type temp_count;
-      cuco::detail::memcpy_async(
-        &temp_count, d_num_out, sizeof(size_type), cudaMemcpyDeviceToHost, stream);
+      CUCO_CUDA_TRY(cuco::detail::memcpy_async(
+        &temp_count, d_num_out, sizeof(size_type), cudaMemcpyDeviceToHost, stream));
 #if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
       stream.sync();
 #else

include/cuco/detail/static_map.inl

@@ -16,7 +16,7 @@
 
 #include <cuco/detail/bitwise_compare.cuh>
 #include <cuco/detail/error.hpp>
-#include <cuco/detail/utility/memcpy_async.cuh>
+#include <cuco/detail/utility/memcpy_async.hpp>
 #include <cuco/detail/utils.cuh>
 #include <cuco/detail/utils.hpp>
 
@@ -109,11 +109,11 @@ void static_map<Key, Value, Scope, Allocator>::insert(
 
   detail::insert<block_size, tile_size>
     <<<grid_size, block_size, 0, stream>>>(first, num_keys, num_successes_, view, hash, key_equal);
-  cuco::detail::memcpy_async(&h_num_successes,
-                             num_successes_,
-                             sizeof(atomic_ctr_type),
-                             cudaMemcpyDeviceToHost,
-                             cuda::stream_ref{stream});
+  CUCO_CUDA_TRY(cuco::detail::memcpy_async(&h_num_successes,
+                                           num_successes_,
+                                           sizeof(atomic_ctr_type),
+                                           cudaMemcpyDeviceToHost,
+                                           cuda::stream_ref{stream}));
 
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));  // stream sync to ensure h_num_successes is updated
 
@@ -150,11 +150,11 @@ void static_map<Key, Value, Scope, Allocator>::insert_if(InputIt first,
 
   detail::insert_if_n<block_size, tile_size><<<grid_size, block_size, 0, stream>>>(
     first, num_keys, num_successes_, view, stencil, pred, hash, key_equal);
-  cuco::detail::memcpy_async(&h_num_successes,
-                             num_successes_,
-                             sizeof(atomic_ctr_type),
-                             cudaMemcpyDeviceToHost,
-                             cuda::stream_ref{stream});
+  CUCO_CUDA_TRY(cuco::detail::memcpy_async(&h_num_successes,
+                                           num_successes_,
+                                           sizeof(atomic_ctr_type),
+                                           cudaMemcpyDeviceToHost,
+                                           cuda::stream_ref{stream}));
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
 
   size_ += h_num_successes;
@@ -185,11 +185,11 @@ void static_map<Key, Value, Scope, Allocator>::erase(
 
   detail::erase<block_size, tile_size>
     <<<grid_size, block_size, 0, stream>>>(first, num_keys, num_successes_, view, hash, key_equal);
-  cuco::detail::memcpy_async(&h_num_successes,
-                             num_successes_,
-                             sizeof(atomic_ctr_type),
-                             cudaMemcpyDeviceToHost,
-                             cuda::stream_ref{stream});
+  CUCO_CUDA_TRY(cuco::detail::memcpy_async(&h_num_successes,
+                                           num_successes_,
+                                           sizeof(atomic_ctr_type),
+                                           cudaMemcpyDeviceToHost,
+                                           cuda::stream_ref{stream}));
 
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));  // stream sync to ensure h_num_successes is updated
 
@@ -259,8 +259,8 @@ std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_a
                         stream);
 
   std::size_t h_num_out;
-  cuco::detail::memcpy_async(
-    &h_num_out, d_num_out, sizeof(std::size_t), cudaMemcpyDeviceToHost, cuda::stream_ref{stream});
+  CUCO_CUDA_TRY(cuco::detail::memcpy_async(
+    &h_num_out, d_num_out, sizeof(std::size_t), cudaMemcpyDeviceToHost, cuda::stream_ref{stream}));
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
   temp_allocator.deallocate(
     reinterpret_cast<char*>(d_num_out), sizeof(std::size_t), cuda::stream_ref{stream});

include/cuco/detail/utility/memcpy_async.cuh

@@ -33,17 +33,15 @@ namespace cuco::detail {
  * @param count Number of bytes to copy
  * @param kind Memory copy direction
  * @param stream CUDA stream for the operation
+ * @return cudaError_t Error code from the memory copy operation
  */
-inline void memcpy_async(
+[[nodiscard]] inline cudaError_t memcpy_async(
   void* dst, void const* src, size_t count, cudaMemcpyKind kind, cuda::stream_ref stream)
 {
-  if (dst == nullptr || src == nullptr || count == 0) { return; }
+  if (dst == nullptr || src == nullptr || count == 0) { return cudaSuccess; }
 
 #if CUDART_VERSION >= 12080
-  if (stream.get() == 0) {
-    CUCO_CUDA_TRY(cudaMemcpyAsync(dst, src, count, kind, stream.get()));
-    return;
-  }
+  if (stream.get() == 0) { return cudaMemcpyAsync(dst, src, count, kind, stream.get()); }
 
   void* dsts[1]             = {dst};
   void* srcs[1]             = {const_cast<void*>(src)};
@@ -55,15 +53,14 @@ inline void memcpy_async(
   attrs[0].flags                = cudaMemcpyFlagPreferOverlapWithCompute;
 
 #if CUDART_VERSION >= 13000
-  CUCO_CUDA_TRY(cudaMemcpyBatchAsync(dsts, srcs, sizes, 1, attrs, attrs_idxs, 1, stream.get()));
+  return cudaMemcpyBatchAsync(dsts, srcs, sizes, 1, attrs, attrs_idxs, 1, stream.get());
 #else
   std::size_t fail_idx;
-  CUCO_CUDA_TRY(
-    cudaMemcpyBatchAsync(dsts, srcs, sizes, 1, attrs, attrs_idxs, 1, &fail_idx, stream.get()));
+  return cudaMemcpyBatchAsync(dsts, srcs, sizes, 1, attrs, attrs_idxs, 1, &fail_idx, stream.get());
 #endif  // CUDART_VERSION >= 13000
 #else
   // CUDA < 12.8 - use regular cudaMemcpyAsync
-  CUCO_CUDA_TRY(cudaMemcpyAsync(dst, src, count, kind, stream.get()));
+  return cudaMemcpyAsync(dst, src, count, kind, stream.get());
 #endif  // CUDART_VERSION >= 12080
 }
 }  // namespace cuco::detail