Use pinned host memory for improved async performance (#727)

The discussion in rapidsai/cudf#18967 reminded
me that we’ve never consistently prioritized using pinned memory in cuco
where applicable. In most of our cases, we only need a small host
buffer, typically 4 or 8 bytes, to receive data copied back from the
device. While the overhead of using pageable memory is negligible for
such small transfers, it can introduce unintended global synchronization
in multi-stream scenarios due to implicit driver fallback.

To address this, this PR updates all instances of host memory in cuco to
use pinned (page-locked) memory to enable true asynchronous data
transfers and avoid unnecessary synchronization overhead in multi-stream
contexts.

Author: PointKernel

include/cuco/detail/dynamic_map.inl

@@ -157,14 +157,17 @@ void dynamic_map<Key, Value, Scope, Allocator>::insert(
                                                hash,
                                                key_equal);
 
-      std::size_t h_num_successes;
-      CUCO_CUDA_TRY(cudaMemcpyAsync(&h_num_successes,
+      std::size_t* h_num_successes;
+      CUCO_CUDA_TRY(cudaMallocHost(&h_num_successes, sizeof(std::size_t)));
+      CUCO_CUDA_TRY(cudaMemcpyAsync(h_num_successes,
                                     submap_num_successes_[submap_idx],
                                     sizeof(atomic_ctr_type),
                                     cudaMemcpyDeviceToHost,
                                     stream));
-      submaps_[submap_idx]->size_ += h_num_successes;
-      size_ += h_num_successes;
+      CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
+      submaps_[submap_idx]->size_ += *h_num_successes;
+      size_ += *h_num_successes;
+      CUCO_CUDA_TRY(cudaFreeHost(h_num_successes));
       first += n;
       num_to_insert -= n;
     }
@@ -205,14 +208,17 @@ void dynamic_map<Key, Value, Scope, Allocator>::erase(
                                                            key_equal);
 
   for (uint32_t i = 0; i < submaps_.size(); ++i) {
-    std::size_t h_submap_num_successes;
-    CUCO_CUDA_TRY(cudaMemcpyAsync(&h_submap_num_successes,
+    std::size_t* h_submap_num_successes;
+    CUCO_CUDA_TRY(cudaMallocHost(&h_submap_num_successes, sizeof(std::size_t)));
+    CUCO_CUDA_TRY(cudaMemcpyAsync(h_submap_num_successes,
                                   submap_num_successes_[i],
                                   sizeof(atomic_ctr_type),
                                   cudaMemcpyDeviceToHost,
                                   stream));
-    submaps_[i]->size_ -= h_submap_num_successes;
-    size_ -= h_submap_num_successes;
+    CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
+    submaps_[i]->size_ -= *h_submap_num_successes;
+    size_ -= *h_submap_num_successes;
+    CUCO_CUDA_TRY(cudaFreeHost(h_submap_num_successes));
   }
 }
 

include/cuco/detail/hyperloglog/hyperloglog_impl.cuh

@@ -405,10 +405,10 @@ class hyperloglog_impl {
   [[nodiscard]] __host__ size_t estimate(cuda::stream_ref stream) const
   {
     auto const num_regs = 1ull << this->precision_;
-    std::vector<register_type> host_sketch(num_regs);
+    register_type* host_sketch;
+    CUCO_CUDA_TRY(cudaMallocHost(&host_sketch, sizeof(register_type) * num_regs));
 
-    // TODO check if storage is host accessible
-    CUCO_CUDA_TRY(cudaMemcpyAsync(host_sketch.data(),
+    CUCO_CUDA_TRY(cudaMemcpyAsync(host_sketch,
                                   this->sketch_.data(),
                                   sizeof(register_type) * num_regs,
                                   cudaMemcpyDefault,
@@ -419,11 +419,14 @@ class hyperloglog_impl {
     int zeroes  = 0;
 
     // geometric mean computation + count registers with 0s
-    for (auto const reg : host_sketch) {
+    for (size_t i = 0; i < num_regs; i++) {
+      auto const reg = host_sketch[i];
       sum += fp_type{1} / static_cast<fp_type>(1ull << reg);
       zeroes += reg == 0;
     }
 
+    CUCO_CUDA_TRY(cudaFreeHost(host_sketch));
+
     auto const finalize = cuco::hyperloglog_ns::detail::finalizer(this->precision_);
 
     // pass intermediate result to finalizer for bias correction, etc.

include/cuco/detail/open_addressing/open_addressing_impl.cuh

@@ -876,11 +876,13 @@ class open_addressing_impl {
                                           is_filled,
                                           stream.get()));
 
-      size_type temp_count;
+      size_type* temp_count;
+      CUCO_CUDA_TRY(cudaMallocHost(&temp_count, sizeof(size_type)));
       CUCO_CUDA_TRY(cudaMemcpyAsync(
-        &temp_count, d_num_out, sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
+        temp_count, d_num_out, sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
       stream.wait();
-      h_num_out += temp_count;
+      h_num_out += *temp_count;
+      CUCO_CUDA_TRY(cudaFreeHost(temp_count));
       temp_allocator.deallocate(d_temp_storage, temp_storage_bytes);
     }
 

include/cuco/detail/static_map.inl

@@ -104,16 +104,19 @@ void static_map<Key, Value, Scope, Allocator>::insert(
   // TODO: memset an atomic variable is unsafe
   static_assert(sizeof(std::size_t) == sizeof(atomic_ctr_type));
   CUCO_CUDA_TRY(cudaMemsetAsync(num_successes_, 0, sizeof(atomic_ctr_type), stream));
-  std::size_t h_num_successes;
+
+  std::size_t* h_num_successes;
+  CUCO_CUDA_TRY(cudaMallocHost(&h_num_successes, sizeof(std::size_t)));
 
   detail::insert<block_size, tile_size>
     <<<grid_size, block_size, 0, stream>>>(first, num_keys, num_successes_, view, hash, key_equal);
   CUCO_CUDA_TRY(cudaMemcpyAsync(
-    &h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
+    h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
 
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));  // stream sync to ensure h_num_successes is updated
 
-  size_ += h_num_successes;
+  size_ += *h_num_successes;
+  CUCO_CUDA_TRY(cudaFreeHost(h_num_successes));
 }
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
@@ -142,15 +145,18 @@ void static_map<Key, Value, Scope, Allocator>::insert_if(InputIt first,
   // TODO: memset an atomic variable is unsafe
   static_assert(sizeof(std::size_t) == sizeof(atomic_ctr_type));
   CUCO_CUDA_TRY(cudaMemsetAsync(num_successes_, 0, sizeof(atomic_ctr_type), stream));
-  std::size_t h_num_successes;
+
+  std::size_t* h_num_successes;
+  CUCO_CUDA_TRY(cudaMallocHost(&h_num_successes, sizeof(std::size_t)));
 
   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_CUDA_TRY(cudaMemcpyAsync(
-    &h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
+    h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
 
-  size_ += h_num_successes;
+  size_ += *h_num_successes;
+  CUCO_CUDA_TRY(cudaFreeHost(h_num_successes));
 }
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
@@ -174,16 +180,19 @@ void static_map<Key, Value, Scope, Allocator>::erase(
   // TODO: memset an atomic variable is unsafe
   static_assert(sizeof(std::size_t) == sizeof(atomic_ctr_type));
   CUCO_CUDA_TRY(cudaMemsetAsync(num_successes_, 0, sizeof(atomic_ctr_type), stream));
-  std::size_t h_num_successes;
+
+  std::size_t* h_num_successes;
+  CUCO_CUDA_TRY(cudaMallocHost(&h_num_successes, sizeof(std::size_t)));
 
   detail::erase<block_size, tile_size>
     <<<grid_size, block_size, 0, stream>>>(first, num_keys, num_successes_, view, hash, key_equal);
   CUCO_CUDA_TRY(cudaMemcpyAsync(
-    &h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
+    h_num_successes, num_successes_, sizeof(atomic_ctr_type), cudaMemcpyDeviceToHost, stream));
 
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));  // stream sync to ensure h_num_successes is updated
 
-  size_ -= h_num_successes;
+  size_ -= *h_num_successes;
+  CUCO_CUDA_TRY(cudaFreeHost(h_num_successes));
 }
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
@@ -249,16 +258,21 @@ std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_a
                         filled,
                         stream);
 
-  std::size_t h_num_out;
+  std::size_t* h_num_out;
+  CUCO_CUDA_TRY(cudaMallocHost(&h_num_out, sizeof(std::size_t)));
   CUCO_CUDA_TRY(
-    cudaMemcpyAsync(&h_num_out, d_num_out, sizeof(std::size_t), cudaMemcpyDeviceToHost, stream));
+    cudaMemcpyAsync(h_num_out, d_num_out, sizeof(std::size_t), cudaMemcpyDeviceToHost, stream));
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
+
+  auto result = std::make_pair(keys_out + *h_num_out, values_out + *h_num_out);
+
+  CUCO_CUDA_TRY(cudaFreeHost(h_num_out));
   std::allocator_traits<temp_allocator_type>::deallocate(
     temp_allocator, reinterpret_cast<char*>(d_num_out), sizeof(std::size_t));
   std::allocator_traits<temp_allocator_type>::deallocate(
     temp_allocator, d_temp_storage, temp_storage_bytes);
 
-  return std::make_pair(keys_out + h_num_out, values_out + h_num_out);
+  return result;
 }
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>

include/cuco/detail/storage/counter_storage.cuh

@@ -92,11 +92,14 @@ class counter_storage : public storage_base<cuco::extent<SizeType, 1>> {
    */
   [[nodiscard]] constexpr size_type load_to_host(cuda::stream_ref stream) const
   {
-    size_type h_count;
+    size_type* h_count;
+    CUCO_CUDA_TRY(cudaMallocHost(&h_count, sizeof(size_type)));
     CUCO_CUDA_TRY(cudaMemcpyAsync(
-      &h_count, this->data(), sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
+      h_count, this->data(), sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
     stream.wait();
-    return h_count;
+    size_type result = *h_count;
+    CUCO_CUDA_TRY(cudaFreeHost(h_count));
+    return result;
   }
 
  private:

include/cuco/detail/trie/dynamic_bitset/dynamic_bitset.inl

@@ -208,15 +208,17 @@ constexpr void dynamic_bitset<Allocator>::build_ranks_and_selects(
                                        num_blocks,
                                        stream.get()));
 
-  size_type num_selects{};
+  size_type* h_num_selects;
+  CUCO_CUDA_TRY(cudaMallocHost(&h_num_selects, sizeof(size_type)));
   CUCO_CUDA_TRY(
-    cudaMemcpyAsync(&num_selects, d_sum, sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
+    cudaMemcpyAsync(h_num_selects, d_sum, sizeof(size_type), cudaMemcpyDeviceToHost, stream.get()));
   stream.wait();
   std::allocator_traits<temp_allocator_type>::deallocate(
     temp_allocator, thrust::device_ptr<char>{reinterpret_cast<char*>(d_sum)}, sizeof(size_type));
   temp_allocator.deallocate(d_temp_storage, temp_storage_bytes);
 
-  selects.resize(num_selects);
+  selects.resize(*h_num_selects);
+  CUCO_CUDA_TRY(cudaFreeHost(h_num_selects));
 
   auto const select_begin = thrust::raw_pointer_cast(selects.data());