Merge pull request #150 from PointKernel/static-map-retrieve-all

Add `static_map::retrieve_all`

Author: PointKernel

include/cuco/detail/static_map.inl

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,11 @@
  */
 
 #include <cuco/detail/bitwise_compare.cuh>
+#include <cuco/detail/utils.cuh>
+
+#include <thrust/copy.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/zip_iterator.h>
 
 namespace cuco {
 
@@ -136,6 +141,25 @@ void static_map<Key, Value, Scope, Allocator>::find(InputIt first,
     <<<grid_size, block_size, 0, stream>>>(first, last, output_begin, view, hash, key_equal);
 }
 
+template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
+template <typename KeyOut, typename ValueOut>
+std::pair<KeyOut, ValueOut> static_map<Key, Value, Scope, Allocator>::retrieve_all(
+  KeyOut keys_out, ValueOut values_out, cudaStream_t stream)
+{
+  static_assert(sizeof(pair_atomic_type) == sizeof(value_type));
+  auto slots_begin = reinterpret_cast<value_type*>(slots_);
+
+  auto begin  = thrust::make_transform_iterator(slots_begin, detail::slot_to_tuple<Key, Value>{});
+  auto end    = begin + get_capacity();
+  auto filled = detail::slot_is_filled<Key>{get_empty_key_sentinel()};
+  auto zipped_out_begin = thrust::make_zip_iterator(thrust::make_tuple(keys_out, values_out));
+
+  auto const zipped_out_end =
+    thrust::copy_if(thrust::cuda::par.on(stream), begin, end, zipped_out_begin, filled);
+  auto const num = std::distance(zipped_out_begin, zipped_out_end);
+  return std::make_pair(keys_out + num, values_out + num);
+}
+
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
 template <typename InputIt, typename OutputIt, typename Hash, typename KeyEqual>
 void static_map<Key, Value, Scope, Allocator>::contains(InputIt first,

include/cuco/detail/utils.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,5 +27,30 @@ __device__ __forceinline__ int32_t count_least_significant_bits(uint32_t x, int3
   return __popc(x & (1 << n) - 1);
 }
 
+/**
+ * @brief Converts `cuco::pair` to `thrust::tuple` to allow assigning to a zip iterator.
+ */
+template <typename Key, typename Value>
+struct slot_to_tuple {
+  template <typename S>
+  __device__ thrust::tuple<Key, Value> operator()(S const& s)
+  {
+    return thrust::tuple<Key, Value>(s.first, s.second);
+  }
+};
+
+/**
+ * @brief Device functor returning whether the input slot `s` is filled.
+ */
+template <typename Key>
+struct slot_is_filled {
+  Key empty_key_sentinel;
+  template <typename S>
+  __device__ bool operator()(S const& s)
+  {
+    return thrust::get<0>(s) != empty_key_sentinel;
+  }
+};
+
 }  // namespace detail
 }  // namespace cuco

include/cuco/static_map.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -296,6 +296,29 @@ class static_map {
             KeyEqual key_equal  = KeyEqual{},
             cudaStream_t stream = 0);
 
+  /**
+   * @brief Retrieves all of the keys and their associated values.
+   *
+   * The order in which keys are returned is implementation defined and not guaranteed to be
+   * consistent between subsequent calls to `retrieve_all`.
+   *
+   * Behavior is undefined if the range beginning at `keys_out` or `values_out` is less than
+   * `get_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 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,
+                                           cudaStream_t stream = 0);
+
   /**
    * @brief Indicates whether the keys in the range
    * `[first, last)` are contained in the map.

tests/CMakeLists.txt

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2018-2021, NVIDIA CORPORATION.
+# Copyright (c) 2018-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -55,6 +55,7 @@ endfunction(ConfigureTest)
 # - static_map tests ------------------------------------------------------------------------------
 ConfigureTest(STATIC_MAP_TEST
     static_map/custom_type_test.cu
+    static_map/duplicate_keys_test.cu
     static_map/key_sentinel_test.cu
     static_map/shared_memory_test.cu
     static_map/stream_test.cu

tests/static_map/duplicate_keys_test.cu

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <utils.hpp>
+
+#include <cuco/static_map.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/functional.h>
+#include <thrust/iterator/discard_iterator.h>
+#include <thrust/sort.h>
+
+#include <catch2/catch.hpp>
+
+TEMPLATE_TEST_CASE_SIG("Duplicate keys",
+                       "",
+                       ((typename Key, typename Value), Key, Value),
+                       (int32_t, int32_t),
+                       (int32_t, int64_t),
+                       (int64_t, int32_t),
+                       (int64_t, int64_t))
+{
+  constexpr std::size_t num_keys{500'000};
+  cuco::static_map<Key, Value> map{num_keys * 2, -1, -1};
+
+  auto m_view = map.get_device_mutable_view();
+  auto view   = map.get_device_view();
+
+  thrust::device_vector<Key> d_keys(num_keys);
+  thrust::device_vector<Value> d_values(num_keys);
+
+  thrust::sequence(thrust::device, d_keys.begin(), d_keys.end());
+  thrust::sequence(thrust::device, d_values.begin(), d_values.end());
+
+  auto pairs_begin = thrust::make_transform_iterator(
+    thrust::make_counting_iterator<int>(0),
+    [] __device__(auto i) { return cuco::pair_type<Key, Value>(i / 2, i / 2); });
+
+  thrust::device_vector<Value> d_results(num_keys);
+  thrust::device_vector<bool> d_contained(num_keys);
+
+  SECTION("Retrieve all entries")
+  {
+    auto constexpr gold = num_keys / 2;
+    thrust::device_vector<Key> unique_keys(gold);
+    thrust::device_vector<Key> unique_values(gold);
+
+    // Retrieve all from an empty map
+    auto [empty_key_end, empty_value_end] =
+      map.retrieve_all(unique_keys.begin(), unique_values.begin());
+    REQUIRE(std::distance(unique_keys.begin(), empty_key_end) == 0);
+    REQUIRE(std::distance(unique_values.begin(), empty_value_end) == 0);
+
+    map.insert(pairs_begin, pairs_begin + num_keys);
+
+    auto const num_entries = map.get_size();
+    REQUIRE(num_entries == gold);
+
+    auto [key_out_end, value_out_end] =
+      map.retrieve_all(unique_keys.begin(), unique_values.begin());
+    REQUIRE(std::distance(unique_keys.begin(), key_out_end) == gold);
+    REQUIRE(std::distance(unique_values.begin(), value_out_end) == gold);
+
+    thrust::sort(thrust::device, unique_keys.begin(), unique_keys.end());
+    REQUIRE(cuco::test::equal(unique_keys.begin(),
+                              unique_keys.end(),
+                              thrust::make_counting_iterator<Key>(0),
+                              thrust::equal_to<Key>{}));
+  }
+
+  SECTION("Tests of contains")
+  {
+    map.insert(pairs_begin, pairs_begin + num_keys);
+    map.contains(d_keys.begin(), d_keys.end(), d_contained.begin());
+
+    REQUIRE(cuco::test::all_of(d_contained.begin(),
+                               d_contained.begin() + num_keys / 2,
+                               [] __device__(bool const& b) { return b; }));
+
+    REQUIRE(cuco::test::none_of(d_contained.begin() + num_keys / 2,
+                                d_contained.end(),
+                                [] __device__(bool const& b) { return b; }));
+  }
+}

tests/utils.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,9 +16,11 @@
 
 #pragma once
 
+#include <utils.cuh>
+
 #include <thrust/functional.h>
 
-#include <utils.cuh>
+#include <cooperative_groups.h>
 
 namespace cuco {
 namespace test {