From a93867550cd98a768eff07b6f035b3c94fd916b4 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Tue, 5 Dec 2023 11:40:53 -0800
Subject: [PATCH] Fix a static extent bug in probing scheme

---
 include/cuco/detail/probing_scheme_impl.inl | 16 +++++++++-------
 include/cuco/utility/fast_int.cuh           |  2 +-
 tests/static_map/unique_sequence_test.cu    |  6 ++++--
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/include/cuco/detail/probing_scheme_impl.inl b/include/cuco/detail/probing_scheme_impl.inl
index 3090d026e..2458315d5 100644
--- a/include/cuco/detail/probing_scheme_impl.inl
+++ b/include/cuco/detail/probing_scheme_impl.inl
@@ -64,7 +64,7 @@ class probing_iterator {
   {
     // 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_;
+    curr_index_ = (curr_index_ + step_size_) % static_cast<size_type>(upper_bound_);
     return *this;
   }
 
@@ -100,7 +100,7 @@ __host__ __device__ constexpr auto linear_probing<CGSize, Hash>::operator()(
 {
   using size_type = typename Extent::value_type;
   return detail::probing_iterator<Extent>{
-    cuco::detail::sanitize_hash<size_type>(hash_(probe_key)) % upper_bound,
+    cuco::detail::sanitize_hash<size_type>(hash_(probe_key)) % static_cast<size_type>(upper_bound),
     1,  // step size is 1
     upper_bound};
 }
@@ -114,7 +114,8 @@ __host__ __device__ constexpr auto linear_probing<CGSize, Hash>::operator()(
 {
   using size_type = typename Extent::value_type;
   return detail::probing_iterator<Extent>{
-    cuco::detail::sanitize_hash<size_type>(hash_(probe_key) + g.thread_rank()) % upper_bound,
+    cuco::detail::sanitize_hash<size_type>(hash_(probe_key) + g.thread_rank()) %
+      static_cast<size_type>(upper_bound),
     cg_size,
     upper_bound};
 }
@@ -133,10 +134,10 @@ __host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::operato
 {
   using size_type = typename Extent::value_type;
   return detail::probing_iterator<Extent>{
-    cuco::detail::sanitize_hash<size_type>(hash1_(probe_key)) % upper_bound,
+    cuco::detail::sanitize_hash<size_type>(hash1_(probe_key)) % static_cast<size_type>(upper_bound),
     max(size_type{1},
         cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) %
-          upper_bound),  // step size in range [1, prime - 1]
+          static_cast<size_type>(upper_bound)),  // step size in range [1, prime - 1]
     upper_bound};
 }
 
@@ -149,9 +150,10 @@ __host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::operato
 {
   using size_type = typename Extent::value_type;
   return detail::probing_iterator<Extent>{
-    cuco::detail::sanitize_hash<size_type>(hash1_(probe_key) + g.thread_rank()) % upper_bound,
+    cuco::detail::sanitize_hash<size_type>(hash1_(probe_key) + g.thread_rank()) %
+      static_cast<size_type>(upper_bound),
     static_cast<size_type>((cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) %
-                              (upper_bound.value() / cg_size - 1) +
+                              (static_cast<size_type>(upper_bound) / cg_size - 1) +
                             1) *
                            cg_size),
     upper_bound};  // TODO use fast_int operator
diff --git a/include/cuco/utility/fast_int.cuh b/include/cuco/utility/fast_int.cuh
index 6616e2c5c..26442c577 100644
--- a/include/cuco/utility/fast_int.cuh
+++ b/include/cuco/utility/fast_int.cuh
@@ -156,4 +156,4 @@ struct fast_int {
     return lhs - (lhs / rhs) * rhs.value_;
   }
 };
-}  // namespace cuco::utility
\ No newline at end of file
+}  // namespace cuco::utility
diff --git a/tests/static_map/unique_sequence_test.cu b/tests/static_map/unique_sequence_test.cu
index 69fa69fb0..389453372 100644
--- a/tests/static_map/unique_sequence_test.cu
+++ b/tests/static_map/unique_sequence_test.cu
@@ -171,6 +171,8 @@ TEMPLATE_TEST_CASE_SIG(
   constexpr size_type gold_capacity = CGSize == 1 ? 422   // 211 x 1 x 2
                                                   : 412;  // 103 x 2 x 2
 
+  // XXX: test static extent on purpose, DO NOT CHANGE
+  using extent_type = cuco::experimental::extent<size_type, num_keys>;
   using probe =
     std::conditional_t<Probe == cuco::test::probe_sequence::linear_probing,
                        cuco::experimental::linear_probing<CGSize, cuco::murmurhash3_32<Key>>,
@@ -180,13 +182,13 @@ TEMPLATE_TEST_CASE_SIG(
 
   auto map = cuco::experimental::static_map<Key,
                                             Value,
-                                            cuco::experimental::extent<size_type>,
+                                            extent_type,
                                             cuda::thread_scope_device,
                                             thrust::equal_to<Key>,
                                             probe,
                                             cuco::cuda_allocator<std::byte>,
                                             cuco::experimental::storage<2>>{
-    num_keys, cuco::empty_key<Key>{-1}, cuco::empty_value<Value>{-1}};
+    extent_type{}, cuco::empty_key<Key>{-1}, cuco::empty_value<Value>{-1}};
 
   REQUIRE(map.capacity() == gold_capacity);