Add ref with_hash_function and with_key_eq APIs for set and multi…

…set (#467) Closes #457 This PR adds `with_hash_function` and `with_key_eq` to allow making copies of an existing ref with new hasher and comparator. This is required to enable custom build and key equal for hash table queries.
NVIDIA · Apr 26, 2024 · a22d881 · a22d881
1 parent 1c8b920
commit a22d881
Show file tree

Hide file tree

Showing 6 changed files with 179 additions and 0 deletions.
diff --git a/include/cuco/detail/probing_scheme_impl.inl b/include/cuco/detail/probing_scheme_impl.inl
@@ -92,6 +92,14 @@ __host__ __device__ constexpr linear_probing<CGSize, Hash>::linear_probing(Hash
 {
 }
 
+template <int32_t CGSize, typename Hash>
+template <typename NewHash>
+__host__ __device__ constexpr auto linear_probing<CGSize, Hash>::with_hash_function(
+  NewHash const& hash) const noexcept
+{
+  return linear_probing<cg_size, NewHash>{hash};
+}
+
 template <int32_t CGSize, typename Hash>
 template <typename ProbeKey, typename Extent>
 __host__ __device__ constexpr auto linear_probing<CGSize, Hash>::operator()(
@@ -125,6 +133,14 @@ __host__ __device__ constexpr double_hashing<CGSize, Hash1, Hash2>::double_hashi
 {
 }
 
+template <int32_t CGSize, typename Hash1, typename Hash2>
+template <typename NewHash1, typename NewHash2>
+__host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::with_hash_function(
+  NewHash1 const& hash1, NewHash2 const& hash2) const noexcept
+{
+  return double_hashing<cg_size, NewHash1, NewHash2>{hash1, hash2};
+}
+
 template <int32_t CGSize, typename Hash1, typename Hash2>
 template <typename ProbeKey, typename Extent>
 __host__ __device__ constexpr auto double_hashing<CGSize, Hash1, Hash2>::operator()(

diff --git a/include/cuco/detail/static_multiset/static_multiset_ref.inl b/include/cuco/detail/static_multiset/static_multiset_ref.inl
@@ -211,6 +211,49 @@ auto static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operat
     std::move(*this)};
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+template <typename NewKeyEqual>
+__host__ __device__ constexpr auto
+static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::with_key_eq(
+  NewKeyEqual const& key_equal) const noexcept
+{
+  return static_multiset_ref<Key, Scope, NewKeyEqual, ProbingScheme, StorageRef, Operators...>{
+    cuco::empty_key<Key>{this->empty_key_sentinel()},
+    key_equal,
+    this->impl_.probing_scheme(),
+    {},
+    this->impl_.storage_ref()};
+}
+
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+template <typename NewHash>
+__host__ __device__ constexpr auto
+static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::
+  with_hash_function(NewHash const& hash) const noexcept
+{
+  auto const probing_scheme = this->impl_.probing_scheme().with_hash_function(hash);
+  return static_multiset_ref<Key,
+                             Scope,
+                             KeyEqual,
+                             decltype(probing_scheme),
+                             StorageRef,
+                             Operators...>{cuco::empty_key<Key>{this->empty_key_sentinel()},
+                                           this->impl_.key_eq(),
+                                           probing_scheme,
+                                           {},
+                                           this->impl_.storage_ref()};
+}
+
 namespace detail {
 
 template <typename Key,

diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
@@ -210,6 +210,45 @@ auto static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators..
     std::move(*this)};
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+template <typename NewKeyEqual>
+__host__ __device__ constexpr auto
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::with_key_eq(
+  NewKeyEqual const& key_equal) const noexcept
+{
+  return static_set_ref<Key, Scope, NewKeyEqual, ProbingScheme, StorageRef, Operators...>{
+    cuco::empty_key<Key>{this->empty_key_sentinel()},
+    key_equal,
+    this->impl_.probing_scheme(),
+    {},
+    this->impl_.storage_ref()};
+}
+
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+template <typename NewHash>
+__host__ __device__ constexpr auto
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::with_hash_function(
+  NewHash const& hash) const noexcept
+{
+  auto const probing_scheme = this->impl_.probing_scheme().with_hash_function(hash);
+  return static_set_ref<Key, Scope, KeyEqual, decltype(probing_scheme), StorageRef, Operators...>{
+    cuco::empty_key<Key>{this->empty_key_sentinel()},
+    this->impl_.key_eq(),
+    probing_scheme,
+    {},
+    this->impl_.storage_ref()};
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,

diff --git a/include/cuco/probing_scheme.cuh b/include/cuco/probing_scheme.cuh
@@ -46,6 +46,19 @@ class linear_probing : private detail::probing_scheme_base<CGSize> {
    */
   __host__ __device__ constexpr linear_probing(Hash const& hash = {});
 
+  /**
+   *@brief Makes a copy of the current probing method with the given hasher
+   *
+   * @tparam NewHash New hasher type
+   *
+   * @param hash Hasher
+   *
+   * @return Copy of the current probing method
+   */
+  template <typename NewHash>
+  [[nodiscard]] __host__ __device__ constexpr auto with_hash_function(
+    NewHash const& hash) const noexcept;
+
   /**
    * @brief Operator to return a probing iterator
    *
@@ -110,6 +123,22 @@ class double_hashing : private detail::probing_scheme_base<CGSize> {
    */
   __host__ __device__ constexpr double_hashing(Hash1 const& hash1 = {}, Hash2 const& hash2 = {1});
 
+  /**
+   *@brief Makes a copy of the current probing method with the given hasher
+   *
+   * @tparam NewHash1 First new hasher type
+   * @tparam NewHash2 Second new hasher type
+   *
+   * @param hash1 First hasher
+   * @param hash2 second hasher
+   *
+   * @return Copy of the current probing method
+   */
+  template <typename NewHash1, typename NewHash2 = NewHash1>
+  [[nodiscard]] __host__ __device__ constexpr auto with_hash_function(NewHash1 const& hash1,
+                                                                      NewHash2 const& hash2 = {
+                                                                        1}) const noexcept;
+
   /**
    * @brief Operator to return a probing iterator
    *

diff --git a/include/cuco/static_multiset_ref.cuh b/include/cuco/static_multiset_ref.cuh
@@ -200,6 +200,32 @@ class static_multiset_ref
   template <typename... NewOperators>
   [[nodiscard]] __host__ __device__ auto with(NewOperators... ops) && noexcept;
 
+  /**
+   * @brief Makes a copy of the current device reference with given key comparator
+   *
+   * @tparam NewKeyEqual The new key equal type
+   *
+   * @param key_equal New key comparator
+   *
+   * @return Copy of the current device ref
+   */
+  template <typename NewKeyEqual>
+  [[nodiscard]] __host__ __device__ constexpr auto with_key_eq(
+    NewKeyEqual const& key_equal) const noexcept;
+
+  /**
+   * @brief Makes a copy of the current device reference with given hasher
+   *
+   * @tparam NewHash The new hasher type
+   *
+   * @param hash New hasher
+   *
+   * @return Copy of the current device ref
+   */
+  template <typename NewHash>
+  [[nodiscard]] __host__ __device__ constexpr auto with_hash_function(
+    NewHash const& hash) const noexcept;
+
  private:
   impl_type impl_;
 

diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
@@ -198,6 +198,32 @@ class static_set_ref
   template <typename... NewOperators>
   [[nodiscard]] __host__ __device__ auto with(NewOperators... ops) && noexcept;
 
+  /**
+   * @brief Makes a copy of the current device reference with given key comparator
+   *
+   * @tparam NewKeyEqual The new key equal type
+   *
+   * @param key_equal New key comparator
+   *
+   * @return Copy of the current device ref
+   */
+  template <typename NewKeyEqual>
+  [[nodiscard]] __host__ __device__ constexpr auto with_key_eq(
+    NewKeyEqual const& key_equal) const noexcept;
+
+  /**
+   * @brief Makes a copy of the current device reference with given hasher
+   *
+   * @tparam NewHash The new hasher type
+   *
+   * @param hash New hasher
+   *
+   * @return Copy of the current device ref
+   */
+  template <typename NewHash>
+  [[nodiscard]] __host__ __device__ constexpr auto with_hash_function(
+    NewHash const& hash) const noexcept;
+
   /**
    * @brief Makes a copy of the current device reference using non-owned memory
    *