Merge branch 'dev' into fix-docs

Author: PointKernel

include/cuco/detail/dynamic_map.inl

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include <cstdint>
+
 namespace cuco {
 
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
@@ -77,42 +79,52 @@ dynamic_map<Key, Value, Scope, Allocator>::dynamic_map(std::size_t initial_capac
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
 void dynamic_map<Key, Value, Scope, Allocator>::reserve(std::size_t n, cudaStream_t stream)
 {
-  int64_t num_elements_remaining = n;
-  uint32_t submap_idx            = 0;
-  while (num_elements_remaining > 0) {
-    std::size_t submap_capacity;
-
-    // if the submap already exists
-    if (submap_idx < submaps_.size()) {
-      submap_capacity = submaps_[submap_idx]->get_capacity();
+  // Calculate current total available capacity across all submaps
+  std::size_t total_available_capacity = 0;
+  for (std::size_t i = 0; i < submaps_.size(); ++i) {
+    std::size_t submap_usable_capacity =
+      static_cast<std::size_t>(max_load_factor_ * submaps_[i]->get_capacity());
+    // Only count capacity above the minimum insert threshold
+    if (submap_usable_capacity >= min_insert_size_) {
+      total_available_capacity += submap_usable_capacity - min_insert_size_;
     }
-    // if the submap does not exist yet, create it
-    else {
-      submap_capacity = capacity_;
-      if (erased_key_sentinel_ != empty_key_sentinel_) {
-        submaps_.push_back(std::make_unique<cuco::legacy::static_map<Key, Value, Scope, Allocator>>(
-          submap_capacity,
-          empty_key<Key>{empty_key_sentinel_},
-          empty_value<Value>{empty_value_sentinel_},
-          erased_key<Key>{erased_key_sentinel_},
-          alloc_,
-          stream));
-      } else {
-        submaps_.push_back(std::make_unique<cuco::legacy::static_map<Key, Value, Scope, Allocator>>(
-          submap_capacity,
-          empty_key<Key>{empty_key_sentinel_},
-          empty_value<Value>{empty_value_sentinel_},
-          alloc_,
-          stream));
-      }
-      submap_num_successes_.push_back(submaps_[submap_idx]->num_successes_);
-      submap_views_.push_back(submaps_[submap_idx]->get_device_view());
-      submap_mutable_views_.push_back(submaps_[submap_idx]->get_device_mutable_view());
-      capacity_ *= 2;
+  }
+
+  // Create new submaps until we have enough capacity
+  while (total_available_capacity < n) {
+    std::size_t new_submap_capacity = capacity_;
+
+    if (erased_key_sentinel_ != empty_key_sentinel_) {
+      submaps_.push_back(std::make_unique<cuco::legacy::static_map<Key, Value, Scope, Allocator>>(
+        new_submap_capacity,
+        empty_key<Key>{empty_key_sentinel_},
+        empty_value<Value>{empty_value_sentinel_},
+        erased_key<Key>{erased_key_sentinel_},
+        alloc_,
+        stream));
+    } else {
+      submaps_.push_back(std::make_unique<cuco::legacy::static_map<Key, Value, Scope, Allocator>>(
+        new_submap_capacity,
+        empty_key<Key>{empty_key_sentinel_},
+        empty_value<Value>{empty_value_sentinel_},
+        alloc_,
+        stream));
     }
 
-    num_elements_remaining -= max_load_factor_ * submap_capacity - min_insert_size_;
-    submap_idx++;
+    std::size_t submap_idx = submaps_.size() - 1;
+    submap_num_successes_.push_back(submaps_[submap_idx]->num_successes_);
+    submap_views_.push_back(submaps_[submap_idx]->get_device_view());
+    submap_mutable_views_.push_back(submaps_[submap_idx]->get_device_mutable_view());
+
+    // Add the new submap's usable capacity
+    std::size_t new_usable_capacity =
+      static_cast<std::size_t>(max_load_factor_ * new_submap_capacity);
+    if (new_usable_capacity >= min_insert_size_) {
+      total_available_capacity += new_usable_capacity - min_insert_size_;
+    }
+
+    // Update capacity for next submap (double the size)
+    capacity_ *= 2;
   }
 }
 
@@ -133,7 +145,7 @@ void dynamic_map<Key, Value, Scope, Allocator>::insert(
 
   reserve(size_ + num_to_insert, stream);
 
-  uint32_t submap_idx = 0;
+  std::size_t submap_idx = 0;
   while (num_to_insert > 0) {
     std::size_t capacity_remaining =
       max_load_factor_ * submaps_[submap_idx]->get_capacity() - submaps_[submap_idx]->get_size();
@@ -189,7 +201,7 @@ void dynamic_map<Key, Value, Scope, Allocator>::erase(
   auto const grid_size = (tile_size * num_keys + stride * block_size - 1) / (stride * block_size);
 
   // zero out submap success counters
-  for (uint32_t i = 0; i < submaps_.size(); ++i) {
+  for (std::size_t i = 0; i < submaps_.size(); ++i) {
     CUCO_CUDA_TRY(cudaMemsetAsync(submap_num_successes_[i], 0, sizeof(atomic_ctr_type), stream));
   }
 
@@ -204,7 +216,7 @@ void dynamic_map<Key, Value, Scope, Allocator>::erase(
                                                            hash,
                                                            key_equal);
 
-  for (uint32_t i = 0; i < submaps_.size(); ++i) {
+  for (std::size_t i = 0; i < submaps_.size(); ++i) {
     std::size_t h_submap_num_successes;
     CUCO_CUDA_TRY(cudaMemcpyAsync(&h_submap_num_successes,
                                   submap_num_successes_[i],

include/cuco/detail/dynamic_map/dynamic_map.inl

@@ -82,7 +82,7 @@ void dynamic_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stor
   auto num_to_insert = cuco::detail::distance(first, last);
   this->reserve(size_ + num_to_insert, stream);
 
-  uint32_t submap_idx = 0;
+  std::size_t submap_idx = 0;
   while (num_to_insert > 0) {
     auto& cur = submaps_[submap_idx];
 
@@ -114,7 +114,7 @@ void dynamic_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stor
   size_type n, cuda::stream_ref stream)
 {
   size_type num_elements_remaining = n;
-  uint32_t submap_idx              = 0;
+  std::size_t submap_idx           = 0;
   while (num_elements_remaining > 0) {
     std::size_t submap_capacity;
 
@@ -157,9 +157,9 @@ template <typename InputIt, typename OutputIt>
 void dynamic_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::contains(
   InputIt first, InputIt last, OutputIt output_begin, cuda::stream_ref stream) const
 {
-  auto num_keys         = cuco::detail::distance(first, last);
-  std::size_t traversed = 0;
-  uint32_t submap_idx   = 0;
+  auto num_keys          = cuco::detail::distance(first, last);
+  std::size_t traversed  = 0;
+  std::size_t submap_idx = 0;
   while (num_keys > 0 && submap_idx < submaps_.size()) {
     const auto& cur       = submaps_[submap_idx];
     const size_t cur_size = cur->size();

include/cuco/detail/open_addressing/kernels.cuh

@@ -55,8 +55,8 @@ CUCO_SUPPRESS_KERNEL_WARNINGS
  * @param num_successes Number of successful inserted elements
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize,
-          int32_t BlockSize,
+template <int CGSize,
+          int BlockSize,
           typename InputIt,
           typename StencilIt,
           typename Predicate,
@@ -123,8 +123,8 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void insert_if_n(InputIt first,
  * @param pred Predicate to test on every element in the range `[stencil, stencil + n)`
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize,
-          int32_t BlockSize,
+template <int CGSize,
+          int BlockSize,
           typename InputIt,
           typename StencilIt,
           typename Predicate,
@@ -164,7 +164,7 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void insert_if_n(
  * @param n Number of input elements
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize, int32_t BlockSize, typename InputIt, typename Ref>
+template <int CGSize, int BlockSize, typename InputIt, typename Ref>
 CUCO_KERNEL __launch_bounds__(BlockSize) void erase(InputIt first,
                                                     cuco::detail::index_type n,
                                                     Ref ref)
@@ -204,7 +204,7 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void erase(InputIt first,
  * @param callback_op Function to call on every matched slot found in the container
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize, int32_t BlockSize, typename InputIt, typename CallbackOp, typename Ref>
+template <int CGSize, int BlockSize, typename InputIt, typename CallbackOp, typename Ref>
 CUCO_KERNEL __launch_bounds__(BlockSize) void for_each_n(InputIt first,
                                                          cuco::detail::index_type n,
                                                          CallbackOp callback_op,
@@ -252,8 +252,8 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void for_each_n(InputIt first,
  * @param output_begin Beginning of the sequence of booleans for the presence of each key
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize,
-          int32_t BlockSize,
+template <int CGSize,
+          int BlockSize,
           typename InputIt,
           typename StencilIt,
           typename Predicate,
@@ -348,8 +348,8 @@ struct find_buffer<Container, cuda::std::void_t<typename Container::mapped_type>
  * @param output_begin Beginning of the sequence of matched payloads retrieved for each key
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize,
-          int32_t BlockSize,
+template <int CGSize,
+          int BlockSize,
           typename InputIt,
           typename StencilIt,
           typename Predicate,
@@ -446,8 +446,8 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void find_if_n(InputIt first,
  * @param inserted_begin Beginning of the sequence of booleans for the presence of each key
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <int32_t CGSize,
-          int32_t BlockSize,
+template <int CGSize,
+          int BlockSize,
           typename InputIt,
           typename FoundIt,
           typename InsertedIt,
@@ -532,12 +532,7 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void insert_and_find(InputIt first,
  * @param count Number of matches
  * @param ref Non-owning container device ref used to access the slot storage
  */
-template <bool IsOuter,
-          int32_t CGSize,
-          int32_t BlockSize,
-          typename InputIt,
-          typename AtomicT,
-          typename Ref>
+template <bool IsOuter, int CGSize, int BlockSize, typename InputIt, typename AtomicT, typename Ref>
 CUCO_KERNEL __launch_bounds__(BlockSize) void count(InputIt first,
                                                     cuco::detail::index_type n,
                                                     AtomicT* count,
@@ -598,8 +593,8 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count(InputIt first,
  * @param ref Non-owning container device ref used to access the slot storage
  */
 template <bool IsOuter,
-          int32_t CGSize,
-          int32_t BlockSize,
+          int CGSize,
+          int BlockSize,
           typename InputIt,
           typename OutputIt,
           typename Ref>
@@ -671,7 +666,7 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count_each(InputIt first,
  * @param ref Non-owning container device ref used to access the slot storage
  */
 template <bool IsOuter,
-          int32_t BlockSize,
+          int BlockSize,
           class InputProbeIt,
           class OutputProbeIt,
           class OutputMatchIt,
@@ -724,7 +719,7 @@ CUCO_KERNEL void retrieve(InputProbeIt input_probe,
  * @param is_filled Predicate indicating if the given slot is filled
  * @param count Number of filled slots
  */
-template <int32_t BlockSize, typename StorageRef, typename Predicate, typename AtomicT>
+template <int BlockSize, typename StorageRef, typename Predicate, typename AtomicT>
 CUCO_KERNEL __launch_bounds__(BlockSize) void size(StorageRef storage,
                                                    Predicate is_filled,
                                                    AtomicT* count)
@@ -749,7 +744,7 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void size(StorageRef storage,
   if (threadIdx.x == 0) { count->fetch_add(block_count, cuda::std::memory_order_relaxed); }
 }
 
-template <int32_t BlockSize, typename ContainerRef, typename Predicate>
+template <int BlockSize, typename ContainerRef, typename Predicate>
 CUCO_KERNEL __launch_bounds__(BlockSize) void rehash(
   typename ContainerRef::storage_ref_type storage_ref,
   ContainerRef container_ref,
@@ -796,4 +791,4 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void rehash(
     idx += loop_stride;
   }
 }
-}  // namespace cuco::detail::open_addressing_ns
+}  // namespace cuco::detail::open_addressing_ns

include/cuco/detail/open_addressing/open_addressing_impl.cuh

@@ -37,6 +37,7 @@
 #include <thrust/iterator/transform_iterator.h>
 
 #include <cmath>
+#include <cstdint>
 
 namespace cuco {
 namespace detail {
@@ -834,7 +835,7 @@ class open_addressing_impl {
     using temp_allocator_type =
       typename std::allocator_traits<allocator_type>::template rebind_alloc<char>;
 
-    cuco::detail::index_type constexpr stride = std::numeric_limits<int32_t>::max();
+    cuco::detail::index_type constexpr stride = std::numeric_limits<std::int32_t>::max();
 
     cuco::detail::index_type h_num_out{0};
     auto temp_allocator = temp_allocator_type{this->allocator()};
@@ -860,7 +861,7 @@ class open_addressing_impl {
                                           begin,
                                           output_begin + h_num_out,
                                           d_num_out,
-                                          static_cast<int32_t>(num_items),
+                                          static_cast<std::int32_t>(num_items),
                                           is_filled,
                                           stream.get()));
 
@@ -872,7 +873,7 @@ class open_addressing_impl {
                                           begin,
                                           output_begin + h_num_out,
                                           d_num_out,
-                                          static_cast<int32_t>(num_items),
+                                          static_cast<std::int32_t>(num_items),
                                           is_filled,
                                           stream.get()));
 
@@ -1281,7 +1282,7 @@ class open_addressing_impl {
     auto counter       = counter_type{this->allocator()};
     counter.reset(stream.get());
 
-    int32_t constexpr block_size = cuco::detail::default_block_size();
+    auto constexpr block_size = cuco::detail::default_block_size();
 
     auto constexpr grid_stride = 1;
     auto const grid_size       = cuco::detail::grid_size(n, cg_size, grid_stride, block_size);