Add static_multimap::pair_contains

include/cuco/detail/static_multimap/device_view_impl.inl

@@ -559,29 +559,36 @@ class static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view_
   }
 
   /**
-   * @brief Indicates whether the key `k` exists in the map using vector loads.
+   * @brief Indicates whether `element` exists in the map using vector loads.
    *
-   * If the key `k` was inserted into the map, `contains` returns
-   * true. Otherwise, it returns false. Uses the CUDA Cooperative Groups API to
-   * to leverage multiple threads to perform a single `contains` operation. This provides a
-   * significant boost in throughput compared to the non Cooperative Group based
-   * `contains` at moderate to high load factors.
+   * If `element` was inserted into the map, `contains` returns true. Otherwise, it returns false.
+   * Uses the CUDA Cooperative Groups API to leverage multiple threads to perform a single
+   * `contains` operation. This provides a significant boost in throughput compared to the non
+   * Cooperative Group based `contains` at moderate to high load factors.
    *
+   * @tparam is_pair_contains `true` if it's a `pair_contains` implementation
    * @tparam uses_vector_load Boolean flag indicating whether vector loads are used
-   * @tparam CG Cooperative Group type
-   * @tparam KeyEqual Binary callable type
+   * @tparam Element The element type to search for
+   * @tparam Equal Binary callable type
+   *
    * @param g The Cooperative Group used to perform the contains operation
-   * @param k The key to search for
-   * @param key_equal The binary callable used to compare two keys
-   * for equality
-   * @return A boolean indicating whether the key/value pair
-   * containing `k` was inserted
+   * @param element The element to search for
+   * @param equal The binary function to compare input element and slot content for equality
+   * @return A boolean indicating whether the key/value pair represented by `element` was inserted
    */
-  template <bool uses_vector_load, typename CG, typename KeyEqual>
+  template <bool is_pair_contains, bool uses_vector_load, typename Element, typename Equal>
   __device__ __forceinline__ std::enable_if_t<uses_vector_load, bool> contains(
-    CG g, Key const& k, KeyEqual key_equal) noexcept
+    cooperative_groups::thread_block_tile<ProbeSequence::cg_size> const& g,
+    Element const& element,
+    Equal equal) noexcept
   {
-    auto current_slot = initial_slot(g, k);
+    auto current_slot = [&]() {
+      if constexpr (is_pair_contains) {
+        return initial_slot(g, element.first);
+      } else {
+        return initial_slot(g, element);
+      }
+    }();
 
     while (true) {
       value_type arr[2];
@@ -591,8 +598,20 @@ class static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view_
         detail::bitwise_compare(arr[0].first, this->get_empty_key_sentinel());
       auto const second_slot_is_empty =
         detail::bitwise_compare(arr[1].first, this->get_empty_key_sentinel());
-      auto const first_equals  = (not first_slot_is_empty and key_equal(arr[0].first, k));
-      auto const second_equals = (not second_slot_is_empty and key_equal(arr[1].first, k));
+      auto const first_equals = [&]() {
+        if constexpr (is_pair_contains) {
+          return not first_slot_is_empty and equal(arr[0], element);
+        } else {
+          return not first_slot_is_empty and equal(arr[0].first, element);
+        }
+      }();
+      auto const second_equals = [&]() {
+        if constexpr (is_pair_contains) {
+          return not second_slot_is_empty and equal(arr[1], element);
+        } else {
+          return not second_slot_is_empty and equal(arr[1].first, element);
+        }
+      }();
 
       // the key we were searching for was found by one of the threads, so we return true
       if (g.any(first_equals or second_equals)) { return true; }
@@ -607,29 +626,36 @@ class static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view_
   }
 
   /**
-   * @brief Indicates whether the key `k` exists in the map using scalar loads.
+   * @brief Indicates whether `element` exists in the map using scalar loads.
    *
-   * If the key `k` was inserted into the map, `contains` returns
-   * true. Otherwise, it returns false. Uses the CUDA Cooperative Groups API to
-   * to leverage multiple threads to perform a single `contains` operation. This provides a
-   * significant boost in throughput compared to the non Cooperative Group
-   * `contains` at moderate to high load factors.
+   * If `element` was inserted into the map, `contains` returns true. Otherwise, it returns false.
+   * Uses the CUDA Cooperative Groups API to leverage multiple threads to perform a single
+   * `contains` operation. This provides a significant boost in throughput compared to the non
+   * Cooperative Group `contains` at moderate to high load factors.
    *
+   * @tparam is_pair_contains `true` if it's a `pair_contains` implementation
    * @tparam uses_vector_load Boolean flag indicating whether vector loads are used
-   * @tparam CG Cooperative Group type
-   * @tparam KeyEqual Binary callable type
+   * @tparam Element The element type to search for
+   * @tparam Equal Binary callable type
+   *
    * @param g The Cooperative Group used to perform the contains operation
-   * @param k The key to search for
-   * @param key_equal The binary callable used to compare two keys
-   * for equality
-   * @return A boolean indicating whether the key/value pair
-   * containing `k` was inserted
+   * @param element The element to search for
+   * @param equal The binary function to compare input element and slot content for equality
+   * @return A boolean indicating whether the key/value pair represented by `element` was inserted
    */
-  template <bool uses_vector_load, typename CG, typename KeyEqual>
+  template <bool is_pair_contains, bool uses_vector_load, typename Element, typename Equal>
   __device__ __forceinline__ std::enable_if_t<not uses_vector_load, bool> contains(
-    CG g, Key const& k, KeyEqual key_equal) noexcept
+    cooperative_groups::thread_block_tile<ProbeSequence::cg_size> const& g,
+    Element const& element,
+    Equal equal) noexcept
   {
-    auto current_slot = initial_slot(g, k);
+    auto current_slot = [&]() {
+      if constexpr (is_pair_contains) {
+        return initial_slot(g, element.first);
+      } else {
+        return initial_slot(g, element);
+      }
+    }();
 
     while (true) {
       value_type slot_contents = *reinterpret_cast<value_type const*>(current_slot);
@@ -640,7 +666,13 @@ class static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view_
       auto const slot_is_empty =
         detail::bitwise_compare(existing_key, this->get_empty_key_sentinel());
 
-      auto const equals = (not slot_is_empty and key_equal(existing_key, k));
+      auto const equals = [&]() {
+        if constexpr (is_pair_contains) {
+          return not slot_is_empty and equal(slot_contents, element);
+        } else {
+          return not slot_is_empty and equal(existing_key, element);
+        }
+      }();
 
       // the key we were searching for was found by one of the threads, so we return true
       if (g.any(equals)) { return true; }

include/cuco/detail/static_multimap/kernels.cuh

@@ -144,45 +144,55 @@ __global__ void insert_if_n(InputIt first, StencilIt s, std::size_t n, viewT vie
 }
 
 /**
- * @brief Indicates whether the keys in the range `[first, last)` are contained in the map.
+ * @brief Indicates whether the elements in the range `[first, last)` are contained in the map.
  *
- * Stores `true` or `false` to `(output + i)` indicating if the key `*(first + i)` exists in the
+ * Stores `true` or `false` to `(output + i)` indicating if the element `*(first + i)` exists in the
  * map.
  *
  * Uses the CUDA Cooperative Groups API to leverage groups of multiple threads to perform the
- * contains operation for each key. This provides a significant boost in throughput compared
+ * contains operation for each element. This provides a significant boost in throughput compared
  * to the non Cooperative Group `contains` at moderate to high load factors.
  *
+ * @tparam is_pair_contains `true` if it's a `pair_contains` implementation
  * @tparam block_size The size of the thread block
  * @tparam tile_size The number of threads in the Cooperative Groups
- * @tparam InputIt Device accessible input iterator whose `value_type` is
- * convertible to the map's `key_type`
+ * @tparam InputIt Device accessible input iterator whose `value_type` is convertible to:
+ *         - the map's `key_type` if `is_pair_contains` is `false`
+ *         - the map's `value_type` if `is_pair_contains` is `true`
  * @tparam OutputIt Device accessible output iterator whose `value_type` is convertible from `bool`
  * @tparam viewT Type of device view allowing access of hash map storage
- * @tparam KeyEqual Binary callable type
- * @param first Beginning of the sequence of keys
- * @param last End of the sequence of keys
- * @param output_begin Beginning of the sequence of booleans for the presence of each key
+ * @tparam Equal Binary callable type
+ *
+ * @param first Beginning of the sequence of elements
+ * @param last End of the sequence of elements
+ * @param output_begin Beginning of the sequence of booleans for the presence of each element
  * @param view Device view used to access the hash map's slot storage
- * @param key_equal The binary function to compare two keys for equality
+ * @param equal The binary function to compare input element and slot content for equality
  */
-template <uint32_t block_size,
+template <bool is_pair_contains,
+          uint32_t block_size,
           uint32_t tile_size,
           typename InputIt,
           typename OutputIt,
           typename viewT,
-          typename KeyEqual>
+          typename Equal>
 __global__ void contains(
-  InputIt first, InputIt last, OutputIt output_begin, viewT view, KeyEqual key_equal)
+  InputIt first, InputIt last, OutputIt output_begin, viewT view, Equal equal)
 {
-  auto tile    = cg::tiled_partition<tile_size>(cg::this_thread_block());
-  auto tid     = block_size * blockIdx.x + threadIdx.x;
-  auto key_idx = tid / tile_size;
+  auto tile = cg::tiled_partition<tile_size>(cg::this_thread_block());
+  auto tid  = block_size * blockIdx.x + threadIdx.x;
+  auto idx  = tid / tile_size;
   __shared__ bool writeBuffer[block_size];
 
-  while (first + key_idx < last) {
-    auto key   = *(first + key_idx);
-    auto found = view.contains(tile, key, key_equal);
+  while (first + idx < last) {
+    auto element = *(first + idx);
+    auto found   = [&]() {
+      if constexpr (is_pair_contains) {
+        return view.pair_contains(tile, element, equal);
+      } else {
+        return view.contains(tile, element, equal);
+      }
+    }();
 
     /*
      * The ld.relaxed.gpu instruction used in view.find causes L1 to
@@ -193,10 +203,8 @@ __global__ void contains(
      */
     if (tile.thread_rank() == 0) { writeBuffer[threadIdx.x / tile_size] = found; }
     __syncthreads();
-    if (tile.thread_rank() == 0) {
-      *(output_begin + key_idx) = writeBuffer[threadIdx.x / tile_size];
-    }
-    key_idx += (gridDim.x * block_size) / tile_size;
+    if (tile.thread_rank() == 0) { *(output_begin + idx) = writeBuffer[threadIdx.x / tile_size]; }
+    idx += (gridDim.x * block_size) / tile_size;
   }
 }
 

include/cuco/detail/static_multimap/static_multimap.inl

@@ -108,21 +108,46 @@ template <typename Key,
           class ProbeSequence>
 template <typename InputIt, typename OutputIt, typename KeyEqual>
 void static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::contains(
-  InputIt first, InputIt last, OutputIt output_begin, cudaStream_t stream, KeyEqual key_equal) const
+  InputIt first, InputIt last, OutputIt output_begin, KeyEqual key_equal, cudaStream_t stream) const
 {
   auto const num_keys = std::distance(first, last);
   if (num_keys == 0) { return; }
 
-  auto constexpr block_size = 128;
-  auto constexpr stride     = 1;
+  auto constexpr is_pair_contains = false;
+  auto constexpr block_size       = 128;
+  auto constexpr stride           = 1;
   auto const grid_size = (cg_size() * num_keys + stride * block_size - 1) / (stride * block_size);
   auto view            = get_device_view();
 
-  detail::contains<block_size, cg_size()>
+  detail::contains<is_pair_contains, block_size, cg_size()>
     <<<grid_size, block_size, 0, stream>>>(first, last, output_begin, view, key_equal);
   CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
 }
 
+template <typename Key,
+          typename Value,
+          cuda::thread_scope Scope,
+          typename Allocator,
+          class ProbeSequence>
+template <typename InputIt, typename OutputIt, typename PairEqual>
+void static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::pair_contains(
+  InputIt first, InputIt last, OutputIt output_begin, PairEqual pair_equal, cudaStream_t stream)
+  const
+{
+  auto const num_pairs = std::distance(first, last);
+  if (num_pairs == 0) { return; }
+
+  auto constexpr is_pair_contains = true;
+  auto constexpr block_size       = 128;
+  auto constexpr stride           = 1;
+  auto const grid_size = (cg_size() * num_pairs + stride * block_size - 1) / (stride * block_size);
+  auto view            = get_device_view();
+
+  detail::contains<is_pair_contains, block_size, cg_size()>
+    <<<grid_size, block_size, 0, stream>>>(first, last, output_begin, view, pair_equal);
+  CUCO_CUDA_TRY(cudaStreamSynchronize(stream));
+}
+
 template <typename Key,
           typename Value,
           cuda::thread_scope Scope,
@@ -543,7 +568,24 @@ static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view::conta
   Key const& k,
   KeyEqual key_equal) noexcept
 {
-  return impl_.contains<uses_vector_load()>(g, k, key_equal);
+  constexpr bool is_pair_contains = false;
+  return impl_.contains<is_pair_contains, uses_vector_load()>(g, k, key_equal);
+}
+
+template <typename Key,
+          typename Value,
+          cuda::thread_scope Scope,
+          typename Allocator,
+          class ProbeSequence>
+template <typename PairEqual>
+__device__ __forceinline__ bool
+static_multimap<Key, Value, Scope, Allocator, ProbeSequence>::device_view::pair_contains(
+  cooperative_groups::thread_block_tile<ProbeSequence::cg_size> const& g,
+  value_type const& p,
+  PairEqual pair_equal) noexcept
+{
+  constexpr bool is_pair_contains = true;
+  return impl_.contains<is_pair_contains, uses_vector_load()>(g, p, pair_equal);
 }
 
 template <typename Key,

include/cuco/static_multimap.cuh

@@ -280,15 +280,41 @@ class static_multimap {
    * @param first Beginning of the sequence of keys
    * @param last End of the sequence of keys
    * @param output_begin Beginning of the output sequence indicating whether each key is present
-   * @param stream CUDA stream used for contains
    * @param key_equal The binary function to compare two keys for equality
+   * @param stream CUDA stream used for contains
    */
   template <typename InputIt, typename OutputIt, typename KeyEqual = thrust::equal_to<key_type>>
   void contains(InputIt first,
                 InputIt last,
                 OutputIt output_begin,
-                cudaStream_t stream = 0,
-                KeyEqual key_equal  = KeyEqual{}) const;
+                KeyEqual key_equal  = KeyEqual{},
+                cudaStream_t stream = 0) const;
+
+  /**
+   * @brief Indicates whether the pairs in the range `[first, last)` are contained in the map.
+   *
+   * Stores `true` or `false` to `(output + i)` indicating if the pair `*(first + i)` exists in
+   * the map.
+   *
+   * @tparam InputIt Device accessible random access input iterator where
+   * `std::is_convertible<std::iterator_traits<InputIt>::value_type,
+   * static_multimap<K, V>::value_type>` is `true`
+   * @tparam OutputIt Device accessible output iterator whose `value_type` is convertible from
+   * `bool`
+   * @tparam PairEqual Binary callable type used to compare input pair and slot content for equality
+   *
+   * @param first Beginning of the sequence of pairs
+   * @param last End of the sequence of pairs
+   * @param output_begin Beginning of the output sequence indicating whether each pair is present
+   * @param pair_equal The binary function to compare input pair and slot content for equality
+   * @param stream CUDA stream used for contains
+   */
+  template <typename InputIt, typename OutputIt, typename PairEqual = thrust::equal_to<value_type>>
+  void pair_contains(InputIt first,
+                     InputIt last,
+                     OutputIt output_begin,
+                     PairEqual pair_equal = PairEqual{},
+                     cudaStream_t stream  = 0) const;
 
   /**
    * @brief Counts the occurrences of keys in `[first, last)` contained in the multimap.
@@ -834,6 +860,28 @@ class static_multimap {
       Key const& k,
       KeyEqual key_equal = KeyEqual{}) noexcept;
 
+    /**
+     * @brief Indicates whether the pair `p` exists in the map.
+     *
+     * If the pair `p` was inserted into the map, `contains` returns
+     * true. Otherwise, it returns false. Uses the CUDA Cooperative Groups API to
+     * to leverage multiple threads to perform a single `contains` operation. This provides a
+     * significant boost in throughput compared to the non Cooperative Group
+     * `contains` at moderate to high load factors.
+     *
+     * @tparam PairEqual Binary callable type
+     * @param g The Cooperative Group used to perform the contains operation
+     * @param k The pair to search for
+     * @param pair_equal The binary callable used to compare input pair and slot content
+     * for equality
+     * @return A boolean indicating whether the input pair was inserted in the map
+     */
+    template <typename PairEqual = thrust::equal_to<value_type>>
+    __device__ __forceinline__ bool pair_contains(
+      cooperative_groups::thread_block_tile<ProbeSequence::cg_size> const& g,
+      value_type const& p,
+      PairEqual pair_equal = PairEqual{}) noexcept;
+
     /**
      * @brief Counts the occurrence of a given key contained in multimap.
      *