Add retrieve_all for multiset and multimap (#635)

This PR adds `retrieve_all` for multiset and multimap

Author: PointKernel

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,

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,

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

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.
    *

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.
    *