Add ref key_eq

NVIDIA · Dec 5, 2023 · 6ecc9b4 · 6ecc9b4
1 parent fe6e0e0
commit 6ecc9b4
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 1 deletion.
diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh
@@ -216,7 +216,7 @@ class open_addressing_ref_impl {
    *
    * @return The comparator used to compare keys
    */
-  [[nodiscard]] __device__ constexpr key_equal key_eq() const noexcept
+  [[nodiscard]] __host__ __device__ constexpr key_equal key_eq() const noexcept
   {
     return this->predicate().equal_;
   }

diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
@@ -101,6 +101,26 @@ __host__ __device__ constexpr static_map_ref<Key,
 {
 }
 
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_map_ref<Key,
+                                             T,
+                                             Scope,
+                                             KeyEqual,
+                                             ProbingScheme,
+                                             StorageRef,
+                                             Operators...>::key_equal
+static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::key_eq()
+  const noexcept
+{
+  return this->impl_.key_eq();
+}
+
 template <typename Key,
           typename T,
           cuda::thread_scope Scope,

diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
@@ -86,6 +86,24 @@ __host__ __device__ constexpr static_set_ref<Key,
 {
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_set_ref<Key,
+                                             Scope,
+                                             KeyEqual,
+                                             ProbingScheme,
+                                             StorageRef,
+                                             Operators...>::key_equal
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::key_eq()
+  const noexcept
+{
+  return this->impl_.key_eq();
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,

diff --git a/include/cuco/static_map_ref.cuh b/include/cuco/static_map_ref.cuh
@@ -166,6 +166,13 @@ class static_map_ref
    */
   [[nodiscard]] __host__ __device__ constexpr mapped_type empty_value_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the key comparator.
+   *
+   * @return The comparator used to compare keys
+   */
+  [[nodiscard]] __host__ __device__ constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *

diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
@@ -144,6 +144,13 @@ class static_set_ref
    */
   [[nodiscard]] __host__ __device__ constexpr key_type empty_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the key comparator.
+   *
+   * @return The comparator used to compare keys
+   */
+  [[nodiscard]] __host__ __device__ constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *