From 9b6a002aa3ecd62d99fb8274487b2af77435cd8a Mon Sep 17 00:00:00 2001
From: Daniel Juenger <2955913+sleeepyjack@users.noreply.github.com>
Date: Thu, 31 Aug 2023 16:30:41 +0000
Subject: [PATCH] Fix sanitize_hash helper

---
 include/cuco/detail/probing_scheme_impl.inl | 12 ++++++------
 include/cuco/detail/utils.cuh               | 16 ++++++++++++----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/include/cuco/detail/probing_scheme_impl.inl b/include/cuco/detail/probing_scheme_impl.inl
index 3090d026e..abaf9e6ac 100644
--- a/include/cuco/detail/probing_scheme_impl.inl
+++ b/include/cuco/detail/probing_scheme_impl.inl
@@ -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::hash_to_index<size_type>(hash_(probe_key)) % upper_bound,
     1,  // step size is 1
     upper_bound};
 }
@@ -114,7 +114,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) + g.thread_rank()) % upper_bound,
+    cuco::detail::hash_to_index<size_type>(hash_(probe_key), g.thread_rank()) % upper_bound,
     cg_size,
     upper_bound};
 }
@@ -133,9 +133,9 @@ __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::hash_to_index<size_type>(hash1_(probe_key)) % upper_bound,
     max(size_type{1},
-        cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) %
+        cuco::detail::hash_to_index<size_type>(hash2_(probe_key)) %
           upper_bound),  // step size in range [1, prime - 1]
     upper_bound};
 }
@@ -149,8 +149,8 @@ __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,
-    static_cast<size_type>((cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) %
+    cuco::detail::hash_to_index<size_type>(hash1_(probe_key), g.thread_rank()) % upper_bound,
+    static_cast<size_type>((cuco::detail::hash_to_index<size_type>(hash2_(probe_key)) %
                               (upper_bound.value() / cg_size - 1) +
                             1) *
                            cg_size),
diff --git a/include/cuco/detail/utils.cuh b/include/cuco/detail/utils.cuh
index 22675d496..ca7160b04 100644
--- a/include/cuco/detail/utils.cuh
+++ b/include/cuco/detail/utils.cuh
@@ -112,15 +112,23 @@ struct strong_type {
  * @tparam SizeType The target type
  * @tparam HashType The input type
  *
+ * @param[in] hash The hash value
+ * @param[in] thread_rank Local thread offset during CG probing
+ *
  * @return Converted hash value
  */
 template <typename SizeType, typename HashType>
-__host__ __device__ constexpr SizeType sanitize_hash(HashType hash) noexcept
+__host__ __device__ constexpr SizeType hash_to_index(HashType hash,
+                                                     uint32_t thread_rank = 0) noexcept
 {
-  if constexpr (cuda::std::is_signed_v<SizeType>) {
-    return cuda::std::abs(static_cast<SizeType>(hash));
+  if constexpr (std::is_unsigned_v<SizeType>) {
+    constexpr auto max_size = static_cast<HashType>(cuda::std::numeric_limits<SizeType>::max());
+    return static_cast<SizeType>((hash + thread_rank) & max_size);
   } else {
-    return static_cast<SizeType>(hash);
+    using larger_type =
+      cuda::std::conditional_t<sizeof(HashType) >= sizeof(SizeType), HashType, SizeType>;
+    constexpr auto max_size = static_cast<larger_type>(cuda::std::numeric_limits<SizeType>::max());
+    return static_cast<SizeType>((hash & max_size + thread_rank) & max_size);
   }
 }