Add set ref make_copy

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

Hide file tree

Showing 3 changed files with 42 additions and 5 deletions.
diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
@@ -212,13 +212,12 @@ static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>
 {
   this->impl_.make_copy(g, memory_to_use);
   return static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>{
-    cuco : empty_key<Key>{this->empty_key_sentinel()},
+    cuco::empty_key<Key>{this->empty_key_sentinel()},
     cuco::empty_value<T>{this->empty_value_sentinel()},
     cuco::erased_key<Key>{this->erased_key_sentinel()},
     this->key_eq(),
     this->probing_scheme(),
-    storage_ref_type{this->window_extent(), memory_to_use}
-  };
+    storage_ref_type{this->window_extent(), memory_to_use}};
 }
 
 namespace detail {

diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
@@ -158,8 +158,28 @@ template <typename... NewOperators>
 auto static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::with(
   NewOperators...) && noexcept
 {
-  return static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, NewOperators...>(
-    std::move(*this));
+  return static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, NewOperators...>{
+    std::move(*this)};
+}
+
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+template <typename CG>
+__device__ constexpr auto
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::make_copy(
+  CG const& g, window_type* const memory_to_use) const noexcept
+{
+  this->impl_.make_copy(g, memory_to_use);
+  return static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>{
+    cuco::empty_key<Key>{this->empty_key_sentinel()},
+    cuco::erased_key<Key>{this->erased_key_sentinel()},
+    this->key_eq(),
+    this->probing_scheme(),
+    storage_ref_type{this->window_extent(), memory_to_use}};
 }
 
 namespace detail {

diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
@@ -168,6 +168,24 @@ class static_set_ref
   template <typename... NewOperators>
   [[nodiscard]] __host__ __device__ auto with(NewOperators... ops) && noexcept;
 
+  /**
+   * @brief Makes a copy of the current device ref using non-owned memory
+   *
+   * This function is intended to be used to create shared memory copies of small static sets,
+   * although global memory can be used as well.
+   *
+   * @tparam CG The type of the cooperative thread group
+   *
+   * @param g The ooperative thread group used to copy the data structure
+   * @param memory_to_use Array large enough to support `capacity` elements. Object does not take
+   * the ownership of the memory
+   *
+   * @return Copy of the current device ref
+   */
+  template <typename CG>
+  [[nodiscard]] __device__ constexpr auto make_copy(
+    CG const& g, window_type* const memory_to_use) const noexcept;
+
  private:
   impl_type impl_;