From 505697e86ab578a8f5c5ed2829c53b9e77455e6f Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Tue, 5 Nov 2024 10:19:17 -0800
Subject: [PATCH] Add rehash for multiset and multimap

---
 .../static_multimap/static_multimap.inl       | 58 +++++++++++++++++++
 .../static_multiset/static_multiset.inl       | 54 +++++++++++++++++
 include/cuco/static_multimap.cuh              | 55 ++++++++++++++++++
 include/cuco/static_multiset.cuh              | 55 ++++++++++++++++++
 4 files changed, 222 insertions(+)

diff --git a/include/cuco/detail/static_multimap/static_multimap.inl b/include/cuco/detail/static_multimap/static_multimap.inl
index 965e14f3d..2df33b50b 100644
--- a/include/cuco/detail/static_multimap/static_multimap.inl
+++ b/include/cuco/detail/static_multimap/static_multimap.inl
@@ -398,6 +398,64 @@ static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stora
   return impl_->count(first, last, ref(op::count), stream);
 }
 
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash(
+  cuda::stream_ref stream)
+{
+  this->impl_->rehash(*this, stream);
+}
+
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash(
+  size_type capacity, cuda::stream_ref stream)
+{
+  auto const extent = make_bucket_extent<static_multimap>(capacity);
+  this->impl_->rehash(extent, *this, stream);
+}
+
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  rehash_async(cuda::stream_ref stream)
+{
+  this->impl_->rehash_async(*this, stream);
+}
+
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  rehash_async(size_type capacity, cuda::stream_ref stream)
+{
+  auto const extent = make_bucket_extent<static_multimap>(capacity);
+  this->impl_->rehash_async(extent, *this, stream);
+}
+
 template <class Key,
           class T,
           class Extent,
diff --git a/include/cuco/detail/static_multiset/static_multiset.inl b/include/cuco/detail/static_multiset/static_multiset.inl
index 6fa640804..2eb23287e 100644
--- a/include/cuco/detail/static_multiset/static_multiset.inl
+++ b/include/cuco/detail/static_multiset/static_multiset.inl
@@ -474,6 +474,60 @@ static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
   return impl_->retrieve_outer(first, last, output_probe, output_match, probe_ref, stream);
 }
 
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash(
+  cuda::stream_ref stream)
+{
+  this->impl_->rehash(*this, stream);
+}
+
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash(
+  size_type capacity, cuda::stream_ref stream)
+{
+  auto const extent = make_bucket_extent<static_multiset>(capacity);
+  this->impl_->rehash(extent, *this, stream);
+}
+
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash_async(
+  cuda::stream_ref stream)
+{
+  this->impl_->rehash_async(*this, stream);
+}
+
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+void static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::rehash_async(
+  size_type capacity, cuda::stream_ref stream)
+{
+  auto const extent = make_bucket_extent<static_multiset>(capacity);
+  this->impl_->rehash_async(extent, *this, stream);
+}
+
 template <class Key,
           class Extent,
           cuda::thread_scope Scope,
diff --git a/include/cuco/static_multimap.cuh b/include/cuco/static_multimap.cuh
index e6e2a139e..04fb03211 100644
--- a/include/cuco/static_multimap.cuh
+++ b/include/cuco/static_multimap.cuh
@@ -602,6 +602,61 @@ class static_multimap {
   template <typename InputIt>
   size_type count(InputIt first, InputIt last, cuda::stream_ref stream = {}) const;
 
+  /**
+   * @brief Regenerates the container.
+   *
+   * @note This function synchronizes the given stream. For asynchronous execution use
+   * `rehash_async`.
+   *
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash(cuda::stream_ref stream = {});
+
+  /**
+   * @brief Reserves at least the specified number of slots and regenerates the container
+   *
+   * @note Changes the number of slots to a value that is not less than `capacity`, then
+   * rehashes the container, i.e. puts the elements into appropriate slots considering
+   * that the total number of slots has changed.
+   *
+   * @note This function synchronizes the given stream. For asynchronous execution use
+   * `rehash_async`.
+   *
+   * @note Behavior is undefined if the desired `capacity` is insufficient to store all of the
+   * contained elements.
+   *
+   * @note This function is not available if the conatiner's `extent_type` is static.
+   *
+   * @param capacity New capacity of the container
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash(size_type capacity, cuda::stream_ref stream = {});
+
+  /**
+   * @brief Asynchronously regenerates the container.
+   *
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash_async(cuda::stream_ref stream = {});
+
+  /**
+   * @brief Asynchronously reserves at least the specified number of slots and regenerates the
+   * container
+   *
+   * @note Changes the number of slots to a value that is not less than `capacity`, then
+   * rehashes the container, i.e. puts the elements into appropriate slots considering
+   * that the total number of slots has changed.
+   *
+   * @note Behavior is undefined if the desired `capacity` is insufficient to store all of the
+   * contained elements.
+   *
+   * @note This function is not available if the conatiner's `extent_type` is static.
+   *
+   * @param capacity New capacity of the container
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash_async(size_type capacity, cuda::stream_ref stream = {});
+
   /**
    * @brief Gets the maximum number of elements the hash map can hold.
    *
diff --git a/include/cuco/static_multiset.cuh b/include/cuco/static_multiset.cuh
index 6a3843329..9ecbde9b7 100644
--- a/include/cuco/static_multiset.cuh
+++ b/include/cuco/static_multiset.cuh
@@ -738,6 +738,61 @@ class static_multiset {
                                                          OutputMatchIt output_match,
                                                          cuda::stream_ref stream = {}) const;
 
+  /**
+   * @brief Regenerates the container.
+   *
+   * @note This function synchronizes the given stream. For asynchronous execution use
+   * `rehash_async`.
+   *
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash(cuda::stream_ref stream = {});
+
+  /**
+   * @brief Reserves at least the specified number of slots and regenerates the container
+   *
+   * @note Changes the number of slots to a value that is not less than `capacity`, then
+   * rehashes the container, i.e. puts the elements into appropriate slots considering
+   * that the total number of slots has changed.
+   *
+   * @note This function synchronizes the given stream. For asynchronous execution use
+   * `rehash_async`.
+   *
+   * @note Behavior is undefined if the desired `capacity` is insufficient to store all of the
+   * contained elements.
+   *
+   * @note This function is not available if the conatiner's `extent_type` is static.
+   *
+   * @param capacity New capacity of the container
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash(size_type capacity, cuda::stream_ref stream = {});
+
+  /**
+   * @brief Asynchronously regenerates the container.
+   *
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash_async(cuda::stream_ref stream = {});
+
+  /**
+   * @brief Asynchronously reserves at least the specified number of slots and regenerates the
+   * container
+   *
+   * @note Changes the number of slots to a value that is not less than `capacity`, then
+   * rehashes the container, i.e. puts the elements into appropriate slots considering
+   * that the total number of slots has changed.
+   *
+   * @note Behavior is undefined if the desired `capacity` is insufficient to store all of the
+   * contained elements.
+   *
+   * @note This function is not available if the conatiner's `extent_type` is static.
+   *
+   * @param capacity New capacity of the container
+   * @param stream CUDA stream used for this operation
+   */
+  void rehash_async(size_type capacity, cuda::stream_ref stream = {});
+
   /**
    * @brief Gets the number of elements in the container.
    *