Fix open-addressing index generation and wraparound

include/cuco/detail/extent/extent.inl

@@ -14,9 +14,48 @@
 
 #include <cuda/std/type_traits>
 
+#include <cmath>
 #include <cstdint>
 
 namespace cuco {
+namespace detail {
+
+constexpr std::uint64_t extent_div_ceil(std::uint64_t dividend, std::uint64_t divisor)
+{
+  return dividend / divisor + static_cast<std::uint64_t>(dividend % divisor != 0);
+}
+
+template <typename SizeType>
+constexpr std::uint64_t max_extent_value()
+{
+  static_assert(cuda::std::is_integral_v<SizeType>);
+  static_assert(sizeof(SizeType) <= sizeof(std::uint64_t));
+  return static_cast<std::uint64_t>(cuda::std::numeric_limits<SizeType>::max());
+}
+
+template <typename SizeType, std::size_t N>
+constexpr bool is_static_extent_representable()
+{
+  if constexpr (sizeof(SizeType) > sizeof(std::size_t) ||
+                (sizeof(SizeType) == sizeof(std::size_t) && cuda::std::is_unsigned_v<SizeType>)) {
+    return true;
+  } else {
+    return N <= static_cast<std::size_t>(cuda::std::numeric_limits<SizeType>::max());
+  }
+}
+
+template <typename SizeType>
+constexpr std::uint64_t normalize_extent(SizeType size)
+{
+  if constexpr (cuda::std::is_signed_v<SizeType>) {
+    return size > 0 ? static_cast<std::uint64_t>(size) : 1ull;
+  } else {
+    return size == 0 ? 1ull : static_cast<std::uint64_t>(size);
+  }
+}
+
+}  // namespace detail
+
 template <typename SizeType, std::size_t N>
 struct valid_extent {
   using value_type = SizeType;  ///< Extent value type
@@ -102,17 +141,26 @@ struct valid_extent<SizeType, dynamic_extent> : cuco::utility::fast_int<SizeType
 template <int32_t CGSize, int32_t BucketSize, typename SizeType, std::size_t N>
 [[nodiscard]] auto constexpr make_valid_extent(extent<SizeType, N> ext)
 {
-  auto constexpr stride = CGSize * BucketSize;
-  auto const size       = cuco::detail::int_div_ceil(
-    cuda::std::max(static_cast<SizeType>(ext), static_cast<SizeType>(1)), stride);
+  static_assert(CGSize > 0);
+  static_assert(BucketSize > 0);
+
+  constexpr auto stride     = static_cast<std::uint64_t>(CGSize) * BucketSize;
+  constexpr auto max_groups = detail::max_extent_value<SizeType>() / stride;
 
   if constexpr (N == dynamic_extent) {
-    return valid_extent<SizeType, dynamic_extent>{
-      static_cast<SizeType>(cuco::detail::next_prime(static_cast<std::uint64_t>(size)) * stride)};
+    auto const requested = detail::normalize_extent(static_cast<SizeType>(ext));
+    auto const groups    = detail::extent_div_ceil(requested, stride);
+    auto const prime     = cuco::detail::next_prime(groups, max_groups);
+    if (prime == 0) { CUCO_FAIL("Requested extent exceeds the representable capacity"); }
+    return valid_extent<SizeType, dynamic_extent>{static_cast<SizeType>(prime * stride)};
   } else {
-    return valid_extent<SizeType,
-                        static_cast<std::size_t>(
-                          cuco::detail::next_prime(static_cast<std::uint64_t>(size)) * stride)>{};
+    static_assert(detail::is_static_extent_representable<SizeType, N>(),
+                  "Static extent must be representable by its size type");
+    constexpr auto requested = N == 0 ? 1 : N;
+    constexpr auto groups    = detail::extent_div_ceil(requested, stride);
+    constexpr auto prime     = cuco::detail::next_prime(groups, max_groups);
+    static_assert(prime != 0, "Requested extent exceeds the representable capacity");
+    return valid_extent<SizeType, static_cast<std::size_t>(prime * stride)>{};
   }
 }
 
@@ -130,16 +178,27 @@ template <typename ProbingScheme, typename Storage, typename SizeType, std::size
   if constexpr (cuco::is_double_hashing<ProbingScheme>::value) {
     return make_valid_extent<ProbingScheme::cg_size, Storage::bucket_size, SizeType, N>(ext);
   } else {
-    auto constexpr stride = ProbingScheme::cg_size * Storage::bucket_size;
-    auto const size =
-      cuco::detail::int_div_ceil(
-        cuda::std::max(static_cast<SizeType>(ext), static_cast<SizeType>(1)), stride) +
-      static_cast<SizeType>(ext == 0);
+    static_assert(ProbingScheme::cg_size > 0);
+    static_assert(Storage::bucket_size > 0);
+
+    constexpr auto stride =
+      static_cast<std::uint64_t>(ProbingScheme::cg_size) * Storage::bucket_size;
+    constexpr auto max_groups = detail::max_extent_value<SizeType>() / stride;
 
     if constexpr (N == dynamic_extent) {
-      return valid_extent<SizeType, dynamic_extent>{size * stride};
+      auto const value = static_cast<SizeType>(ext);
+      auto groups      = detail::extent_div_ceil(detail::normalize_extent(value), stride);
+      if (value == 0) { ++groups; }
+      if (groups > max_groups) { CUCO_FAIL("Requested extent exceeds the representable capacity"); }
+      return valid_extent<SizeType, dynamic_extent>{static_cast<SizeType>(groups * stride)};
     } else {
-      return valid_extent<SizeType, size * stride>{};
+      static_assert(detail::is_static_extent_representable<SizeType, N>(),
+                    "Static extent must be representable by its size type");
+      constexpr auto requested = N == 0 ? 1 : N;
+      constexpr auto groups =
+        detail::extent_div_ceil(requested, stride) + static_cast<std::uint64_t>(N == 0);
+      static_assert(groups <= max_groups, "Requested extent exceeds the representable capacity");
+      return valid_extent<SizeType, static_cast<std::size_t>(groups * stride)>{};
     }
   }
 }
@@ -151,8 +210,16 @@ template <typename ProbingScheme, typename Storage, typename SizeType>
   CUCO_EXPECTS(desired_load_factor > 0., "Desired occupancy must be larger than zero");
   CUCO_EXPECTS(desired_load_factor <= 1., "Desired occupancy must be no larger than one");
 
-  auto const temp = cuda::std::ceil(static_cast<double>(SizeType{ext}) / desired_load_factor);
-  if (temp > static_cast<double>(cuda::std::numeric_limits<SizeType>::max())) {
+  auto const value = static_cast<SizeType>(ext);
+  if constexpr (cuda::std::is_signed_v<SizeType>) {
+    if (value <= 0) { return make_valid_extent<ProbingScheme, Storage>(ext); }
+  } else {
+    if (value == 0) { return make_valid_extent<ProbingScheme, Storage>(ext); }
+  }
+
+  auto const temp =
+    std::ceil(static_cast<long double>(value) / static_cast<long double>(desired_load_factor));
+  if (temp > static_cast<long double>(cuda::std::numeric_limits<SizeType>::max())) {
     CUCO_FAIL(
       "Invalid load factor: requested extent divided by load factor exceeds maximum representable "
       "value");

include/cuco/detail/prime.hpp

@@ -8,6 +8,7 @@
 #include <cuco/detail/__config>
 
 #include <cstdint>
+#include <limits>
 
 namespace cuco {
 namespace detail {
@@ -102,22 +103,39 @@ constexpr bool is_prime(std::uint64_t n)
 }
 
 /**
- * @brief Returns the smallest prime >= n.
+ * @brief Returns the smallest prime in `[n, upper_bound]`.
  *
- * For n <= 2, returns 2. Otherwise searches odd numbers starting
- * from n (or n+1 if n is even).
+ * @param n Lower bound of the search range
+ * @param upper_bound Upper bound of the search range
+ *
+ * @return The smallest prime in `[n, upper_bound]`, or zero if none exists
  */
-constexpr std::uint64_t next_prime(std::uint64_t n)
+constexpr std::uint64_t next_prime(std::uint64_t n, std::uint64_t upper_bound)
 {
+  if (upper_bound < 2ull || n > upper_bound) { return 0ull; }
   if (n <= 2ull) { return 2ull; }
 
-  n |= 1;  // make odd
+  if ((n & 1ull) == 0) { ++n; }
 
-  while (!is_prime(n)) {
+  while (n <= upper_bound) {
+    if (is_prime(n)) { return n; }
+    if (upper_bound - n < 2ull) { break; }
     n += 2ull;
   }
 
-  return n;
+  return 0ull;
+}
+
+/**
+ * @brief Returns the smallest representable prime greater than or equal to `n`.
+ *
+ * @param n Lower bound of the search range
+ *
+ * @return The smallest representable prime greater than or equal to `n`, or zero if none exists
+ */
+constexpr std::uint64_t next_prime(std::uint64_t n)
+{
+  return next_prime(n, std::numeric_limits<std::uint64_t>::max());
 }
 
 }  // namespace detail

include/cuco/detail/probing_scheme/probing_scheme_impl.inl

@@ -54,9 +54,10 @@ class probing_iterator {
    */
   __host__ __device__ constexpr auto operator++() noexcept
   {
-    // TODO: step_size_ can be a build time constant (e.g. linear probing)
-    //  Worth passing another extent type?
-    curr_index_ = (curr_index_ + step_size_) % upper_bound_;
+    // Probe construction guarantees curr_index_ < upper_bound_ and step_size_ <= upper_bound_,
+    // therefore advancing can wrap at most once.
+    auto const remaining = upper_bound_ - curr_index_;
+    curr_index_ = step_size_ >= remaining ? step_size_ - remaining : curr_index_ + step_size_;
     return *this;
   }
 
@@ -98,9 +99,10 @@ template <int32_t BucketSize, typename ProbeKey, typename Extent>
 __host__ __device__ constexpr auto linear_probing<CGSize, Hash>::make_iterator(
   ProbeKey probe_key, Extent upper_bound) const noexcept
 {
-  using size_type      = typename Extent::value_type;
-  size_type const init = cuco::detail::sanitize_hash<size_type>(hash_(probe_key)) %
-                         (upper_bound / BucketSize) * BucketSize;
+  using size_type             = typename Extent::value_type;
+  size_type const num_buckets = upper_bound / BucketSize;
+  size_type const init =
+    cuco::detail::sanitize_hash<size_type>(hash_(probe_key), num_buckets) * BucketSize;
   return detail::probing_iterator<Extent>{init, static_cast<size_type>(BucketSize), upper_bound};
 }
 
@@ -113,8 +115,9 @@ __host__ __device__ constexpr auto linear_probing<CGSize, Hash>::make_iterator(
 {
   using size_type            = typename Extent::value_type;
   size_type constexpr stride = cg_size * BucketSize;
+  size_type const num_groups = upper_bound / stride;
   size_type const init =
-    cuco::detail::sanitize_hash<size_type>(hash_(probe_key)) % (upper_bound / stride) * stride +
+    cuco::detail::sanitize_hash<size_type>(hash_(probe_key), num_groups) * stride +
     g.thread_rank() * BucketSize;
   return detail::probing_iterator<Extent>{init, stride, upper_bound};
 }
@@ -159,14 +162,15 @@ template <int32_t BucketSize, typename ProbeKey, typename Extent>
 __host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::make_iterator(
   ProbeKey probe_key, Extent upper_bound) const noexcept
 {
-  using size_type = typename Extent::value_type;
+  using size_type             = typename Extent::value_type;
+  size_type const num_buckets = upper_bound / BucketSize;
   return detail::probing_iterator<Extent>{
-    static_cast<size_type>(cuco::detail::sanitize_hash<size_type>(hash1_(probe_key)) %
-                           (upper_bound / BucketSize) * BucketSize),
-    static_cast<size_type>(
-      (cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) % (upper_bound / BucketSize - 1) +
-       1) *
-      BucketSize),  // step size in range [1, prime - 1]
+    static_cast<size_type>(cuco::detail::sanitize_hash<size_type>(hash1_(probe_key), num_buckets) *
+                           BucketSize),
+    static_cast<size_type>((cuco::detail::sanitize_hash<size_type>(
+                              hash2_(probe_key), static_cast<size_type>(num_buckets - 1)) +
+                            1) *
+                           BucketSize),  // step size in range [1, prime - 1]
     upper_bound};
 }
 
@@ -177,16 +181,18 @@ __host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::make_it
   ProbeKey probe_key,
   Extent upper_bound) const noexcept
 {
-  int32_t const stride = cg_size * BucketSize;
-  using size_type      = typename Extent::value_type;
+  int32_t const stride       = cg_size * BucketSize;
+  using size_type            = typename Extent::value_type;
+  size_type const num_groups = upper_bound / stride;
 
   return detail::probing_iterator<Extent>{
-    static_cast<size_type>(cuco::detail::sanitize_hash<size_type>(hash1_(probe_key)) %
-                             (upper_bound / stride) * stride +
+    static_cast<size_type>(cuco::detail::sanitize_hash<size_type>(hash1_(probe_key), num_groups) *
+                             stride +
                            g.thread_rank() * BucketSize),
-    static_cast<size_type>(
-      (cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) % (upper_bound / stride - 1) + 1) *
-      stride),
+    static_cast<size_type>((cuco::detail::sanitize_hash<size_type>(
+                              hash2_(probe_key), static_cast<size_type>(num_groups - 1)) +
+                            1) *
+                           stride),
     upper_bound};  // TODO use fast_int operator
 }
 

include/cuco/detail/utils.cuh

@@ -10,6 +10,7 @@
 #include <cuda/std/array>
 #include <cuda/std/cmath>
 #include <cuda/std/cstdint>
+#include <cuda/std/limits>
 #include <cuda/std/type_traits>
 
 namespace cuco {
@@ -26,26 +27,40 @@ __device__ __forceinline__ cuda::std::int32_t count_least_significant_bits(cuda:
 }
 
 template <typename SizeType, typename HashType>
-__host__ __device__ constexpr SizeType to_positive(HashType hash)
+__host__ __device__ constexpr cuda::std::make_unsigned_t<SizeType> to_positive(HashType hash)
 {
+  using unsigned_size_type = cuda::std::make_unsigned_t<SizeType>;
+  auto const value         = static_cast<unsigned_size_type>(hash);
+
   if constexpr (cuda::std::is_signed_v<SizeType>) {
-    return cuda::std::abs(static_cast<SizeType>(hash));
+    auto constexpr max =
+      static_cast<unsigned_size_type>(cuda::std::numeric_limits<SizeType>::max());
+    return value > max ? unsigned_size_type{0} - value : value;
   } else {
-    return static_cast<SizeType>(hash);
+    return value;
   }
 }
 
 /**
- * @brief Converts a given hash value into a valid (positive) size type.
+ * @brief Converts a hash value into a valid index for the given modulus.
+ *
+ * @note Hash values wider than `SizeType` are narrowed before reduction, preserving the existing
+ * low-bit mapping policy.
  *
  * @tparam SizeType The target type
  * @tparam HashType The input type
  *
- * @return Converted hash value
+ * @param hash The hash value
+ * @param modulus Exclusive upper bound for the returned index
+ *
+ * @return An index in `[0, modulus)`
  */
 template <typename SizeType, typename HashType>
-__host__ __device__ constexpr SizeType sanitize_hash(HashType hash) noexcept
+__host__ __device__ constexpr SizeType sanitize_hash(HashType hash, SizeType modulus) noexcept
 {
+  using unsigned_size_type = cuda::std::make_unsigned_t<SizeType>;
+
+  unsigned_size_type magnitude;
   if constexpr (cuda::std::is_same_v<HashType, cuda::std::array<std::uint64_t, 2>>) {
 #if !defined(CUCO_HAS_INT128)
     static_assert(false,
@@ -54,10 +69,12 @@ __host__ __device__ constexpr SizeType sanitize_hash(HashType hash) noexcept
 #endif
     unsigned __int128 ret{};
     memcpy(&ret, &hash, sizeof(unsigned __int128));
-    return to_positive<SizeType>(static_cast<SizeType>(ret));
+    magnitude = to_positive<SizeType>(ret);
   } else {
-    return to_positive<SizeType>(hash);
+    magnitude = to_positive<SizeType>(hash);
   }
+
+  return static_cast<SizeType>(magnitude % static_cast<unsigned_size_type>(modulus));
 }
 
 }  // namespace detail