From d39d59ae8b6be942e125b4103ee2500114bb8534 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Mon, 11 Nov 2024 17:31:42 -0800
Subject: [PATCH] Add `retrieve_all` for multiset and multimap (#635)

This PR adds `retrieve_all` for multiset and multimap
---
 .../static_multimap/static_multimap.inl       | 20 +++++++++++++++
 .../static_multiset/static_multiset.inl       | 15 +++++++++++
 include/cuco/static_map.cuh                   |  2 +-
 include/cuco/static_multimap.cuh              | 25 +++++++++++++++++++
 include/cuco/static_multiset.cuh              | 20 +++++++++++++++
 5 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/include/cuco/detail/static_multimap/static_multimap.inl b/include/cuco/detail/static_multimap/static_multimap.inl
index 7236d3175..649cbc749 100644
--- a/include/cuco/detail/static_multimap/static_multimap.inl
+++ b/include/cuco/detail/static_multimap/static_multimap.inl
@@ -398,6 +398,26 @@ 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>
+template <typename KeyOut, typename ValueOut>
+std::pair<KeyOut, ValueOut>
+static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::retrieve_all(
+  KeyOut keys_out, ValueOut values_out, cuda::stream_ref stream) const
+{
+  auto const zipped_out_begin = thrust::make_zip_iterator(thrust::make_tuple(keys_out, values_out));
+  auto const zipped_out_end   = impl_->retrieve_all(zipped_out_begin, stream);
+  auto const num              = std::distance(zipped_out_begin, zipped_out_end);
+
+  return std::make_pair(keys_out + num, values_out + num);
+}
+
 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 23702695a..844a1e4f8 100644
--- a/include/cuco/detail/static_multiset/static_multiset.inl
+++ b/include/cuco/detail/static_multiset/static_multiset.inl
@@ -474,6 +474,21 @@ 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>
+template <typename OutputIt>
+OutputIt
+static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::retrieve_all(
+  OutputIt output_begin, cuda::stream_ref stream) const
+{
+  return impl_->retrieve_all(output_begin, stream);
+}
+
 template <class Key,
           class Extent,
           cuda::thread_scope Scope,
diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
index 8b6b92e3a..c219538bc 100644
--- a/include/cuco/static_map.cuh
+++ b/include/cuco/static_map.cuh
@@ -852,7 +852,7 @@ class static_map {
   size_type count(InputIt first, InputIt last, cuda::stream_ref stream = {}) const;
 
   /**
-   * @brief Retrieves all of the keys and their associated values.
+   * @brief Retrieves all of the keys and their associated values contained in the map
    *
    * @note This API synchronizes the given stream.
    * @note The order in which keys are returned is implementation defined and not guaranteed to be
diff --git a/include/cuco/static_multimap.cuh b/include/cuco/static_multimap.cuh
index bf3e5998f..8d81b90ac 100644
--- a/include/cuco/static_multimap.cuh
+++ b/include/cuco/static_multimap.cuh
@@ -602,6 +602,31 @@ class static_multimap {
   template <typename InputIt>
   size_type count(InputIt first, InputIt last, cuda::stream_ref stream = {}) const;
 
+  /**
+   * @brief Retrieves all of the keys and their associated values contained in the multimap
+   *
+   * @note This API synchronizes the given stream.
+   * @note The order in which keys are returned is implementation defined and not guaranteed to be
+   * consistent between subsequent calls to `retrieve_all`.
+   * @note Behavior is undefined if the range beginning at `keys_out` or `values_out` is smaller
+   * than the return value of `size()`.
+   *
+   * @tparam KeyOut Device accessible random access output iterator whose `value_type` is
+   * convertible from `key_type`.
+   * @tparam ValueOut Device accesible random access output iterator whose `value_type` is
+   * convertible from `mapped_type`.
+   *
+   * @param keys_out Beginning output iterator for keys
+   * @param values_out Beginning output iterator for associated values
+   * @param stream CUDA stream used for this operation
+   *
+   * @return Pair of iterators indicating the last elements in the output
+   */
+  template <typename KeyOut, typename ValueOut>
+  std::pair<KeyOut, ValueOut> retrieve_all(KeyOut keys_out,
+                                           ValueOut values_out,
+                                           cuda::stream_ref stream = {}) const;
+
   /**
    * @brief Regenerates the container.
    *
diff --git a/include/cuco/static_multiset.cuh b/include/cuco/static_multiset.cuh
index 9ecbde9b7..044f5d8be 100644
--- a/include/cuco/static_multiset.cuh
+++ b/include/cuco/static_multiset.cuh
@@ -738,6 +738,26 @@ class static_multiset {
                                                          OutputMatchIt output_match,
                                                          cuda::stream_ref stream = {}) const;
 
+  /**
+   * @brief Retrieves all keys contained in the multiset
+   *
+   * @note This API synchronizes the given stream.
+   * @note The order in which keys are returned is implementation defined and not guaranteed to be
+   * consistent between subsequent calls to `retrieve_all`.
+   * @note Behavior is undefined if the range beginning at `output_begin` is smaller than the return
+   * value of `size()`.
+   *
+   * @tparam OutputIt Device accessible random access output iterator whose `value_type` is
+   * convertible from the container's `key_type`.
+   *
+   * @param output_begin Beginning output iterator for keys
+   * @param stream CUDA stream used for this operation
+   *
+   * @return Iterator indicating the end of the output
+   */
+  template <typename OutputIt>
+  OutputIt retrieve_all(OutputIt output_begin, cuda::stream_ref stream = {}) const;
+
   /**
    * @brief Regenerates the container.
    *