Merge remote-tracking branch 'upstream/dev' into header-test

Author: PointKernel

README.md

@@ -159,7 +159,8 @@ int main()
   auto const total_num_slots = offsets.back();
 
   // Create a single bulk storage used by all subsets
-  auto set_storage = storage_type{total_num_slots};
+  auto set_storage = storage_type{
+    total_num_slots, cuco::cuda_allocator<key_type>{}, cuda::stream_ref{cudaStream_t{nullptr}}};
   // Initializes the storage with the given sentinel
   set_storage.initialize(empty_key_sentinel);
 

examples/static_set/device_subsets_example.cu

@@ -164,8 +164,12 @@ class bucket_storage {
    *
    * @param size Number of slots to (de)allocate
    * @param allocator Allocator used for (de)allocating device storage
+   * @param stream Stream to use for (de)allocating device storage
    */
-  explicit constexpr bucket_storage(Extent size, Allocator const& allocator = {});
+  explicit constexpr bucket_storage(Extent size,
+                                    Allocator const& allocator,
+                                    cuda::stream_ref stream = cuda::stream_ref{
+                                      cudaStream_t{nullptr}});
 
   bucket_storage(bucket_storage&&) = default;  ///< Move constructor
   /**

include/cuco/bucket_storage.cuh

@@ -34,8 +34,8 @@ __host__ constexpr bloom_filter<Key, Extent, Scope, Policy, Allocator>::bloom_fi
   Allocator const& alloc,
   cuda::stream_ref stream)
   : allocator_{alloc},
-    data_{allocator_.allocate(num_blocks),
-          detail::custom_deleter<std::size_t, allocator_type>{num_blocks, allocator_}},
+    data_{allocator_.allocate(num_blocks, stream),
+          detail::custom_deleter<std::size_t, allocator_type>{num_blocks, allocator_, stream}},
     ref_{data_.get(), num_blocks, {}, policy}
 {
   this->clear_async(stream);

include/cuco/detail/bloom_filter/bloom_filter.inl

@@ -271,7 +271,8 @@ std::pair<KeyOut, ValueOut> dynamic_map<Key, Value, Scope, Allocator>::retrieve_
     size_t{0});
   thrust::device_vector<size_t> submap_cap_prefix_d(submap_cap_prefix);
 
-  auto counter = detail::counter_storage<size_t, Scope, Allocator>{this->alloc_};
+  auto counter =
+    detail::counter_storage<size_t, Scope, Allocator>{this->alloc_, cuda::stream_ref{stream}};
   counter.reset({stream});
 
   detail::retrieve_all<block_size>

include/cuco/detail/dynamic_map.inl

@@ -22,9 +22,9 @@ constexpr hyperloglog<T, Scope, Hash, Allocator>::hyperloglog(cuco::sketch_size_
                                                               Allocator const& alloc,
                                                               cuda::stream_ref stream)
   : allocator_{alloc},
-    sketch_{
-      allocator_.allocate(sketch_bytes(sketch_size_kb) / sizeof(register_type)),
-      detail::custom_deleter{sketch_bytes(sketch_size_kb) / sizeof(register_type), allocator_}},
+    sketch_{allocator_.allocate(sketch_bytes(sketch_size_kb) / sizeof(register_type), stream),
+            detail::custom_deleter{
+              sketch_bytes(sketch_size_kb) / sizeof(register_type), allocator_, stream}},
     ref_{cuda::std::span{reinterpret_cast<cuda::std::byte*>(sketch_.get()),
                          sketch_bytes(sketch_size_kb)},
          hash}
@@ -39,9 +39,9 @@ constexpr hyperloglog<T, Scope, Hash, Allocator>::hyperloglog(
   Allocator const& alloc,
   cuda::stream_ref stream)
   : allocator_{alloc},
-    sketch_{
-      allocator_.allocate(sketch_bytes(standard_deviation) / sizeof(register_type)),
-      detail::custom_deleter{sketch_bytes(standard_deviation) / sizeof(register_type), allocator_}},
+    sketch_{allocator_.allocate(sketch_bytes(standard_deviation) / sizeof(register_type), stream),
+            detail::custom_deleter{
+              sketch_bytes(standard_deviation) / sizeof(register_type), allocator_, stream}},
     ref_{cuda::std::span{reinterpret_cast<cuda::std::byte*>(sketch_.get()),
                          sketch_bytes(standard_deviation)},
          hash}

include/cuco/detail/hyperloglog/hyperloglog.inl

@@ -136,7 +136,7 @@ class open_addressing_impl {
       erased_key_sentinel_{this->extract_key(empty_slot_sentinel)},
       predicate_{pred},
       probing_scheme_{probing_scheme},
-      storage_{make_valid_extent<probing_scheme_type, Storage>(capacity), alloc}
+      storage_{make_valid_extent<probing_scheme_type, Storage>(capacity), alloc, stream}
   {
     this->clear_async(stream);
   }
@@ -182,7 +182,8 @@ class open_addressing_impl {
       erased_key_sentinel_{this->extract_key(empty_slot_sentinel)},
       predicate_{pred},
       probing_scheme_{probing_scheme},
-      storage_{make_valid_extent<probing_scheme_type, Storage>(n, desired_load_factor), alloc}
+      storage_{
+        make_valid_extent<probing_scheme_type, Storage>(n, desired_load_factor), alloc, stream}
   {
     this->clear_async(stream);
   }
@@ -219,7 +220,7 @@ class open_addressing_impl {
       erased_key_sentinel_{erased_key_sentinel},
       predicate_{pred},
       probing_scheme_{probing_scheme},
-      storage_{make_valid_extent<probing_scheme_type, Storage>(capacity), alloc}
+      storage_{make_valid_extent<probing_scheme_type, Storage>(capacity), alloc, stream}
   {
     CUCO_EXPECTS(this->empty_key_sentinel() != this->erased_key_sentinel(),
                  "The empty key sentinel and erased key sentinel cannot be the same value.",
@@ -334,7 +335,7 @@ class open_addressing_impl {
     if (num_keys == 0) { return 0; }
 
     auto counter =
-      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator()};
+      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator(), stream};
     counter.reset(stream);
 
     auto const grid_size = cuco::detail::grid_size(num_keys, cg_size);
@@ -842,8 +843,8 @@ class open_addressing_impl {
 
     cuco::detail::index_type h_num_out{0};
     auto temp_allocator = temp_allocator_type{this->allocator()};
-    auto d_num_out      = reinterpret_cast<size_type*>(
-      std::allocator_traits<temp_allocator_type>::allocate(temp_allocator, sizeof(size_type)));
+    auto d_num_out =
+      reinterpret_cast<size_type*>(temp_allocator.allocate(sizeof(size_type), stream));
 
     // TODO: PR #580 to be reverted once https://github.com/NVIDIA/cccl/issues/1422 is resolved
     for (cuco::detail::index_type offset = 0;
@@ -869,7 +870,7 @@ class open_addressing_impl {
                                           stream.get()));
 
       // Allocate temporary storage
-      auto d_temp_storage = temp_allocator.allocate(temp_storage_bytes);
+      auto d_temp_storage = temp_allocator.allocate(temp_storage_bytes, stream);
 
       CUCO_CUDA_TRY(cub::DeviceSelect::If(d_temp_storage,
                                           temp_storage_bytes,
@@ -889,11 +890,10 @@ class open_addressing_impl {
       stream.wait();
 #endif
       h_num_out += temp_count;
-      temp_allocator.deallocate(d_temp_storage, temp_storage_bytes);
+      temp_allocator.deallocate(d_temp_storage, temp_storage_bytes, stream);
     }
 
-    std::allocator_traits<temp_allocator_type>::deallocate(
-      temp_allocator, reinterpret_cast<char*>(d_num_out), sizeof(size_type));
+    temp_allocator.deallocate(reinterpret_cast<char*>(d_num_out), sizeof(size_type), stream);
 
     return output_begin + h_num_out;
   }
@@ -969,7 +969,7 @@ class open_addressing_impl {
   [[nodiscard]] size_type size(cuda::stream_ref stream) const
   {
     auto counter =
-      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator()};
+      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator(), stream};
     counter.reset(stream);
 
     auto const grid_size = cuco::detail::grid_size(this->capacity());
@@ -1079,7 +1079,7 @@ class open_addressing_impl {
   void rehash_async(extent_type extent, Container const& container, cuda::stream_ref stream)
   {
     auto const old_storage = std::move(this->storage_);
-    new (&storage_) storage_type{extent, this->allocator()};
+    new (&storage_) storage_type{extent, this->allocator(), stream};
     this->clear_async(stream);
 
     auto const num_buckets = old_storage.num_buckets();
@@ -1196,7 +1196,7 @@ class open_addressing_impl {
     if (num_keys == 0) { return 0; }
 
     auto counter =
-      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator()};
+      detail::counter_storage<size_type, thread_scope, allocator_type>{this->allocator(), stream};
     counter.reset(stream);
 
     auto const grid_size = cuco::detail::grid_size(num_keys, cg_size);
@@ -1294,7 +1294,7 @@ class open_addressing_impl {
     if (n == 0) { return {output_probe, output_match}; }
 
     using counter_type = detail::counter_storage<size_type, thread_scope, allocator_type>;
-    auto counter       = counter_type{this->allocator()};
+    auto counter       = counter_type{this->allocator(), stream};
     counter.reset(stream.get());
 
     auto constexpr block_size = cuco::detail::default_block_size();

include/cuco/detail/open_addressing/open_addressing_impl.cuh

@@ -266,9 +266,9 @@ class roaring_bitmap_storage<cuda::std::uint32_t, Allocator> {
                          cuda::stream_ref stream)
     : allocator_{alloc},
       metadata_{bitmap},
-      data_{allocator_.allocate(metadata_.size_bytes),
-            cuco::detail::custom_deleter<cuda::std::size_t, allocator_type>{metadata_.size_bytes,
-                                                                            allocator_}},
+      data_{allocator_.allocate(metadata_.size_bytes, stream),
+            cuco::detail::custom_deleter<cuda::std::size_t, allocator_type>{
+              metadata_.size_bytes, allocator_, stream}},
       ref_{data_.get(), metadata_}
   {
     CUCO_CUDA_TRY(cudaMemcpyAsync(
@@ -362,12 +362,12 @@ class roaring_bitmap_storage<cuda::std::uint64_t, Allocator> {
         [bitmap](std::vector<typename ref_type::metadata_type::bucket_metadata>& bucket_metadata) {
           return typename ref_type::metadata_type{bitmap, bucket_metadata};
         }(bucket_metadata_)},
-      data_{allocator_.allocate(metadata_.size_bytes),
-            cuco::detail::custom_deleter<cuda::std::size_t, allocator_type>{metadata_.size_bytes,
-                                                                            allocator_}},
-      buckets_{bucket_allocator_.allocate(metadata_.num_buckets),
+      data_{allocator_.allocate(metadata_.size_bytes, stream),
+            cuco::detail::custom_deleter<cuda::std::size_t, allocator_type>{
+              metadata_.size_bytes, allocator_, stream}},
+      buckets_{bucket_allocator_.allocate(metadata_.num_buckets, stream),
                cuco::detail::custom_deleter<cuda::std::size_t, bucket_allocator_type>{
-                 metadata_.num_buckets, bucket_allocator_}},
+                 metadata_.num_buckets, bucket_allocator_, stream}},
       ref_{data_.get(), metadata_, buckets_.get()}
   {
     assert(metadata_.valid);

include/cuco/detail/roaring_bitmap/roaring_bitmap_storage.cuh

@@ -40,8 +40,8 @@ static_map<Key, Value, Scope, Allocator>::static_map(std::size_t capacity,
     slot_allocator_{alloc},
     counter_allocator_{alloc}
 {
-  slots_         = std::allocator_traits<slot_allocator_type>::allocate(slot_allocator_, capacity_);
-  num_successes_ = std::allocator_traits<counter_allocator_type>::allocate(counter_allocator_, 1);
+  slots_         = slot_allocator_.allocate(capacity_, cuda::stream_ref{stream});
+  num_successes_ = counter_allocator_.allocate(1, cuda::stream_ref{stream});
 
   auto constexpr block_size = 256;
   auto constexpr stride     = 4;
@@ -69,8 +69,8 @@ static_map<Key, Value, Scope, Allocator>::static_map(std::size_t capacity,
                "The empty key sentinel and erased key sentinel cannot be the same value.",
                std::runtime_error);
 
-  slots_         = std::allocator_traits<slot_allocator_type>::allocate(slot_allocator_, capacity_);
-  num_successes_ = std::allocator_traits<counter_allocator_type>::allocate(counter_allocator_, 1);
+  slots_         = slot_allocator_.allocate(capacity_, cuda::stream_ref{stream});
+  num_successes_ = counter_allocator_.allocate(1, cuda::stream_ref{stream});
 
   auto constexpr block_size = 256;
   auto constexpr stride     = 4;
@@ -83,8 +83,8 @@ static_map<Key, Value, Scope, Allocator>::static_map(std::size_t capacity,
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
 static_map<Key, Value, Scope, Allocator>::~static_map()
 {
-  std::allocator_traits<slot_allocator_type>::deallocate(slot_allocator_, slots_, capacity_);
-  std::allocator_traits<counter_allocator_type>::deallocate(counter_allocator_, num_successes_, 1);
+  slot_allocator_.deallocate(slots_, capacity_, cuda::stream_ref{cudaStream_t{nullptr}});
+  counter_allocator_.deallocate(num_successes_, 1, cuda::stream_ref{cudaStream_t{nullptr}});
 }
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
@@ -226,7 +226,7 @@ std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_a
     typename std::allocator_traits<Allocator>::template rebind_alloc<char>;
   auto temp_allocator = temp_allocator_type{slot_allocator_};
   auto d_num_out      = reinterpret_cast<std::size_t*>(
-    std::allocator_traits<temp_allocator_type>::allocate(temp_allocator, sizeof(std::size_t)));
+    temp_allocator.allocate(sizeof(std::size_t), cuda::stream_ref{stream}));
   cub::DeviceSelect::If(nullptr,
                         temp_storage_bytes,
                         begin,
@@ -237,8 +237,7 @@ std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_a
                         stream);
 
   // Allocate temporary storage
-  auto d_temp_storage =
-    std::allocator_traits<temp_allocator_type>::allocate(temp_allocator, temp_storage_bytes);
+  auto d_temp_storage = temp_allocator.allocate(temp_storage_bytes, cuda::stream_ref{stream});
 
   cub::DeviceSelect::If(d_temp_storage,
                         temp_storage_bytes,
@@ -253,10 +252,9 @@ std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_a
   CUCO_CUDA_TRY(
     cudaMemcpyAsync(&h_num_out, d_num_out, sizeof(std::size_t), cudaMemcpyDeviceToHost, stream));
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
-  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);
+  temp_allocator.deallocate(
+    reinterpret_cast<char*>(d_num_out), sizeof(std::size_t), cuda::stream_ref{stream});
+  temp_allocator.deallocate(d_temp_storage, temp_storage_bytes, cuda::stream_ref{stream});
 
   return std::make_pair(keys_out + h_num_out, values_out + h_num_out);
 }

include/cuco/detail/static_map.inl

@@ -787,8 +787,8 @@ static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::static_multimap(
     empty_key_sentinel_{empty_key_sentinel.value},
     empty_value_sentinel_{empty_value_sentinel.value},
     allocator_{alloc},
-    delete_slots_{allocator_, capacity_},
-    slots_{allocator_.allocate(capacity_), delete_slots_}
+    delete_slots_{allocator_, capacity_, cuda::stream_ref{stream}},
+    slots_{allocator_.allocate(capacity_, cuda::stream_ref{stream}), delete_slots_}
 {
   auto constexpr block_size = 128;
   auto constexpr stride     = 4;
@@ -909,7 +909,7 @@ std::size_t static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::count(
   auto view            = get_device_view();
   auto const grid_size = (cg_size() * num_keys + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::count<block_size, cg_size(), is_outer>
@@ -937,7 +937,7 @@ std::size_t static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::count_
   auto view            = get_device_view();
   auto const grid_size = (cg_size() * num_keys + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::count<block_size, cg_size(), is_outer>
@@ -965,7 +965,7 @@ std::size_t static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::pair_c
   auto view            = get_device_view();
   auto const grid_size = (cg_size() * num_pairs + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::pair_count<block_size, cg_size(), is_outer>
@@ -993,7 +993,7 @@ std::size_t static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::pair_c
   auto view            = get_device_view();
   auto const grid_size = (cg_size() * num_pairs + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::pair_count<block_size, cg_size(), is_outer>
@@ -1026,7 +1026,7 @@ OutputIt static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::retrieve(
 
   auto const grid_size = detail::grid_size(num_keys, cg_size());
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::retrieve<detail::default_block_size(), flushing_cg_size, cg_size(), buffer_size, is_outer>
@@ -1060,7 +1060,7 @@ OutputIt static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::retrieve_
 
   auto const grid_size = detail::grid_size(num_keys, cg_size());
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::retrieve<detail::default_block_size(), flushing_cg_size, cg_size(), buffer_size, is_outer>
@@ -1101,7 +1101,7 @@ static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::pair_retrieve(
   }();
   auto const grid_size = (cg_size() * num_pairs + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::pair_retrieve<block_size, flushing_cg_size, cg_size(), buffer_size, is_outer>
@@ -1148,7 +1148,7 @@ static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::pair_retrieve_oute
   }();
   auto const grid_size = (cg_size() * num_pairs + stride * block_size - 1) / (stride * block_size);
 
-  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_};
+  auto counter = detail::counter_storage<size_type, Scope, allocator_type>{allocator_, stream};
   counter.reset(stream);
 
   detail::pair_retrieve<block_size, flushing_cg_size, cg_size(), buffer_size, is_outer>