From c63ac899ffd54d0eb39863b2283655f0f67934a5 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav <43375352+srinivasyadav18@users.noreply.github.com>
Date: Tue, 18 Jun 2024 19:39:30 -0700
Subject: [PATCH] make murmurhash3_x64_128 compatible with existing cuco data
 structures (#501)

---
 include/cuco/detail/probing_scheme_impl.inl |  4 +-
 include/cuco/detail/utils.cuh               | 47 +++++++++++++-
 tests/CMakeLists.txt                        |  1 +
 tests/static_map/hash_test.cu               | 72 +++++++++++++++++++++
 4 files changed, 119 insertions(+), 5 deletions(-)
 create mode 100644 tests/static_map/hash_test.cu

diff --git a/include/cuco/detail/probing_scheme_impl.inl b/include/cuco/detail/probing_scheme_impl.inl
index 50d7c4dcc..d0a67f87a 100644
--- a/include/cuco/detail/probing_scheme_impl.inl
+++ b/include/cuco/detail/probing_scheme_impl.inl
@@ -121,7 +121,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::sanitize_hash<size_type>(g, hash_(probe_key)) % upper_bound,
     cg_size,
     upper_bound};
 }
@@ -164,7 +164,7 @@ __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>(g, hash1_(probe_key)) % upper_bound,
     static_cast<size_type>(
       (cuco::detail::sanitize_hash<size_type>(hash2_(probe_key)) % (upper_bound / cg_size - 1) +
        1) *
diff --git a/include/cuco/detail/utils.cuh b/include/cuco/detail/utils.cuh
index 1cbe8fd26..21e4df759 100644
--- a/include/cuco/detail/utils.cuh
+++ b/include/cuco/detail/utils.cuh
@@ -17,11 +17,15 @@
 
 #include <cuco/detail/bitwise_compare.cuh>
 
+#include <cuda/std/array>
 #include <cuda/std/bit>
 #include <cuda/std/cmath>
+#include <cuda/std/limits>
 #include <cuda/std/type_traits>
 #include <thrust/tuple.h>
 
+#include <cstddef>
+
 namespace cuco {
 namespace detail {
 
@@ -81,6 +85,16 @@ struct slot_is_filled {
   }
 };
 
+template <typename SizeType, typename HashType>
+__host__ __device__ constexpr SizeType to_positive(HashType hash)
+{
+  if constexpr (cuda::std::is_signed_v<SizeType>) {
+    return cuda::std::abs(static_cast<SizeType>(hash));
+  } else {
+    return static_cast<SizeType>(hash);
+  }
+}
+
 /**
  * @brief Converts a given hash value into a valid (positive) size type.
  *
@@ -92,12 +106,39 @@ struct slot_is_filled {
 template <typename SizeType, typename HashType>
 __host__ __device__ constexpr SizeType sanitize_hash(HashType hash) noexcept
 {
-  if constexpr (cuda::std::is_signed_v<SizeType>) {
-    return cuda::std::abs(static_cast<SizeType>(hash));
+  if constexpr (cuda::std::is_same_v<HashType, cuda::std::array<std::uint64_t, 2>>) {
+#if !defined(CUCO_HAS_INT128)
+    static_assert(false,
+                  "CUCO_HAS_INT128 undefined. Need unsigned __int128 type when sanitizing "
+                  "cuda::std::array<std::uint64_t, 2>");
+#endif
+    unsigned __int128 ret{};
+    memcpy(&ret, &hash, sizeof(unsigned __int128));
+    return to_positive<SizeType>(static_cast<SizeType>(ret));
   } else {
-    return static_cast<SizeType>(hash);
+    return to_positive<SizeType>(hash);
   }
 }
 
+/**
+ * @brief Converts a given hash value and cg_rank, into a valid (positive) size type.
+ *
+ * @tparam SizeType The target type
+ * @tparam CG Cooperative group type
+ * @tparam HashType The input type
+ *
+ * @return Converted hash value
+ */
+template <typename SizeType, typename CG, typename HashType>
+__device__ constexpr SizeType sanitize_hash(CG const& group, HashType hash) noexcept
+{
+  auto const base_hash = sanitize_hash<SizeType>(hash);
+  auto const max_size  = cuda::std::numeric_limits<SizeType>::max();
+  auto const cg_rank   = static_cast<SizeType>(group.thread_rank());
+
+  if (base_hash > (max_size - cg_rank)) { return cg_rank - (max_size - base_hash); }
+  return base_hash + cg_rank;
+}
+
 }  // namespace detail
 }  // namespace cuco
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index a37f2d4e2..9d75d7a0e 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -76,6 +76,7 @@ ConfigureTest(STATIC_MAP_TEST
     static_map/custom_type_test.cu
     static_map/duplicate_keys_test.cu
     static_map/erase_test.cu
+    static_map/hash_test.cu
     static_map/heterogeneous_lookup_test.cu
     static_map/insert_and_find_test.cu
     static_map/insert_or_assign_test.cu
diff --git a/tests/static_map/hash_test.cu b/tests/static_map/hash_test.cu
new file mode 100644
index 000000000..c22eae998
--- /dev/null
+++ b/tests/static_map/hash_test.cu
@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <test_utils.hpp>
+
+#include <cuco/hash_functions.cuh>
+#include <cuco/static_map.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/functional.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+
+#include <catch2/catch_template_test_macros.hpp>
+
+using size_type = std::size_t;
+
+template <typename Key, typename Hash>
+void test_hash_function()
+{
+  using Value = int64_t;
+
+  constexpr size_type num_keys{400};
+
+  auto map = cuco::static_map<Key,
+                              Value,
+                              cuco::extent<size_type>,
+                              cuda::thread_scope_device,
+                              thrust::equal_to<Key>,
+                              cuco::linear_probing<1, Hash>,
+                              cuco::cuda_allocator<std::byte>,
+                              cuco::storage<2>>{
+    num_keys, cuco::empty_key<Key>{-1}, cuco::empty_value<Value>{-1}};
+
+  auto keys_begin = thrust::counting_iterator<Key>(1);
+
+  auto pairs_begin = thrust::make_transform_iterator(
+    keys_begin, cuda::proclaim_return_type<cuco::pair<Key, Value>>([] __device__(auto i) {
+      return cuco::pair<Key, Value>(i, i);
+    }));
+
+  thrust::device_vector<bool> d_keys_exist(num_keys);
+
+  map.insert(pairs_begin, pairs_begin + num_keys);
+
+  REQUIRE(map.size() == num_keys);
+
+  map.contains(keys_begin, keys_begin + num_keys, d_keys_exist.begin());
+
+  REQUIRE(cuco::test::all_of(d_keys_exist.begin(), d_keys_exist.end(), thrust::identity{}));
+}
+
+TEMPLATE_TEST_CASE_SIG("static_map hash tests", "", ((typename Key)), (int32_t), (int64_t))
+{
+  test_hash_function<Key, cuco::murmurhash3_32<Key>>();
+  test_hash_function<Key, cuco::murmurhash3_x64_128<Key>>();
+  test_hash_function<Key, cuco::xxhash_32<Key>>();
+  test_hash_function<Key, cuco::xxhash_64<Key>>();
+}
\ No newline at end of file