Implement OA retrieve(_outer) and its multiset API

benchmarks/CMakeLists.txt

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2018-2023, NVIDIA CORPORATION.
+# Copyright (c) 2018-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -69,6 +69,7 @@ ConfigureBench(STATIC_MAP_BENCH
 # - static_multiset benchmarks --------------------------------------------------------------------
 ConfigureBench(STATIC_MULTISET_BENCH
  hash_table/static_multiset/contains_bench.cu
+ hash_table/static_multiset/retrieve_bench.cu
  hash_table/static_multiset/count_bench.cu
  hash_table/static_multiset/find_bench.cu
  hash_table/static_multiset/insert_bench.cu)

benchmarks/hash_table/static_multiset/retrieve_bench.cu

@@ -0,0 +1,87 @@
+/*
+ * Copyright (c) 2024, 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 <benchmark_defaults.hpp>
+#include <benchmark_utils.hpp>
+
+#include <cuco/static_multiset.cuh>
+#include <cuco/utility/key_generator.cuh>
+
+#include <nvbench/nvbench.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/transform.h>
+
+using namespace cuco::benchmark;
+using namespace cuco::utility;
+
+/**
+ * @brief A benchmark evaluating `cuco::static_multiset::retrieve` performance
+ */
+template <typename Key, typename Dist>
+void static_multiset_retrieve(nvbench::state& state, nvbench::type_list<Key, Dist>)
+{
+ auto const num_keys = state.get_int64_or_default("NumInputs", defaults::N);
+ auto const occupancy = state.get_float64_or_default("Occupancy", defaults::OCCUPANCY);
+ auto const matching_rate = state.get_float64_or_default("MatchingRate", defaults::MATCHING_RATE);
+
+ std::size_t const size = num_keys / occupancy;
+
+ thrust::device_vector<Key> keys(num_keys);
+
+ key_generator gen;
+ gen.generate(dist_from_state<Dist>(state), keys.begin(), keys.end());
+
+ gen.dropout(keys.begin(), keys.end(), matching_rate);
+
+ state.add_element_count(num_keys);
+
+ cuco::static_multiset<Key> set{size, cuco::empty_key<Key>{-1}};
+ set.insert(keys.begin(), keys.end());
+
+ auto const output_size = set.count(keys.begin(), keys.end());
+ thrust::device_vector<Key> output_match(output_size);
+ auto output_probe_begin = thrust::discard_iterator{};
+
+ state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
+ set.retrieve(
+ keys.begin(), keys.end(), output_probe_begin, output_match.begin(), {launch.get_stream()});
+ });
+}
+
+NVBENCH_BENCH_TYPES(static_multiset_retrieve,
+ NVBENCH_TYPE_AXES(defaults::KEY_TYPE_RANGE,
+ nvbench::type_list<distribution::uniform>))
+ .set_name("static_multiset_retrieve_uniform_occupancy")
+ .set_type_axes_names({"Key", "Distribution"})
+ .set_max_noise(defaults::MAX_NOISE)
+ .add_float64_axis("Occupancy", defaults::OCCUPANCY_RANGE);
+
+NVBENCH_BENCH_TYPES(static_multiset_retrieve,
+ NVBENCH_TYPE_AXES(defaults::KEY_TYPE_RANGE,
+ nvbench::type_list<distribution::uniform>))
+ .set_name("static_multiset_retrieve_uniform_matching_rate")
+ .set_type_axes_names({"Key", "Distribution"})
+ .set_max_noise(defaults::MAX_NOISE)
+ .add_float64_axis("MatchingRate", defaults::MATCHING_RATE_RANGE);
+
+NVBENCH_BENCH_TYPES(static_multiset_retrieve,
+ NVBENCH_TYPE_AXES(defaults::KEY_TYPE_RANGE,
+ nvbench::type_list<distribution::uniform>))
+ .set_name("static_multiset_retrieve_uniform_multiplicity")
+ .set_type_axes_names({"Key", "Distribution"})
+ .set_max_noise(defaults::MAX_NOISE)
+ .add_int64_axis("Multiplicity", defaults::MULTIPLICITY_RANGE);

examples/CMakeLists.txt

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2018-2023, NVIDIA CORPORATION.
+# Copyright (c) 2018-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -38,6 +38,7 @@ ConfigureExample(STATIC_SET_DEVICE_REF_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/stat
 ConfigureExample(STATIC_SET_DEVICE_SUBSETS_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_set/device_subsets_example.cu")
 ConfigureExample(STATIC_SET_SHARED_MEMORY_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_set/shared_memory_example.cu")
 ConfigureExample(STATIC_SET_MAPPING_TABLE_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_set/mapping_table_example.cu")
+ConfigureExample(STATIC_MULTISET_HOST_BULK_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_multiset/host_bulk_example.cu")
 ConfigureExample(STATIC_MAP_HOST_BULK_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_map/host_bulk_example.cu")
 ConfigureExample(STATIC_MAP_DEVICE_SIDE_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_map/device_ref_example.cu")
 ConfigureExample(STATIC_MAP_CUSTOM_TYPE_EXAMPLE "${CMAKE_CURRENT_SOURCE_DIR}/static_map/custom_type_example.cu")

examples/static_multiset/host_bulk_example.cu

@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2024, 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 <cuco/static_multiset.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/functional.h>
+#include <thrust/logical.h>
+#include <thrust/sequence.h>
+
+#include <iostream>
+#include <limits>
+
+/**
+ * @file host_bulk_example.cu
+ * @brief Demonstrates usage of the static_multiset "bulk" host APIs.
+ *
+ * The bulk APIs are only invocable from the host and are used for doing operations like `insert` or
+ * `retrieve` on a multiset of keys.
+ *
+ */
+int main(void)
+{
+ using key_type = int;
+
+ // Empty slots are represented by reserved "sentinel" values. These values should be selected such
+ // that they never occur in your input data.
+ key_type constexpr empty_key_sentinel = -1;
+
+ // Number of keys to be inserted
+ std::size_t constexpr num_keys = 50'000;
+
+ // Compute capacity based on a 50% load factor
+ auto constexpr load_factor = 0.5;
+ std::size_t const capacity = std::ceil(num_keys / load_factor);
+
+ // Constructs a set with at least `capacity` slots using -1 as the empty keys sentinel.
+ cuco::static_multiset<key_type> multiset{capacity, cuco::empty_key{empty_key_sentinel}};
+
+ // Create a sequence of keys {0, 1, 2, .., i}
+ // We're going to insert each key twice so we only need 'num_keys / 2' distinct keys.
+ thrust::device_vector<key_type> keys(num_keys / 2);
+ thrust::sequence(keys.begin(), keys.end(), 0);
+
+ // Inserts all keys into the hash set
+ multiset.insert(keys.begin(), keys.end());
+ // Insert the same set of keys again, so each distinct key should occur twice in the multiset
+ multiset.insert(keys.begin(), keys.end());
+
+ // Counts the occurrences of matching keys contained in the multiset.
+ std::size_t const counted_output_size = multiset.count(keys.begin(), keys.end());
+
+ // Storage for result
+ thrust::device_vector<key_type> output_probes(counted_output_size);
+ thrust::device_vector<key_type> output_matches(counted_output_size);
+
+ // Retrieve all matching keys
+ auto const [output_probes_end, _] =
+ multiset.retrieve(keys.begin(), keys.end(), output_probes.begin(), output_matches.begin());
+ std::size_t const retrieved_output_size = output_probes_end - output_probes.begin();
+
+ if ((retrieved_output_size == counted_output_size) and (retrieved_output_size == num_keys)) {
+ std::cout << "Success! Found all keys.\n";
+ } else {
+ std::cout << "Fail! Something went wrong.\n";
+ }
+
+ return 0;
+}

include/cuco/detail/open_addressing/kernels.cuh

@@ -359,6 +359,50 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void find(InputIt first,
  }
 }
 
+// TODO docs
+template <bool IsOuter,
+ int32_t BlockSize,
+ class InputProbeIt,
+ class OutputProbeIt,
+ class OutputMatchIt,
+ class AtomicCounter,
+ class Ref>
+CUCO_KERNEL __launch_bounds__(BlockSize) void retrieve(InputProbeIt input_probe,
+ cuco::detail::index_type n,
+ OutputProbeIt output_probe,
+ OutputMatchIt output_match,
+ AtomicCounter* atomic_counter,
+ Ref ref)
+{
+ namespace cg = cooperative_groups;
+
+ auto const block = cg::this_thread_block();
+ auto constexpr tiles_in_block = BlockSize / Ref::cg_size;
+ // make sure all but the last block are always occupied
+ auto const items_per_block = detail::int_div_ceil(n, tiles_in_block * gridDim.x) * tiles_in_block;
+
+ auto const block_begin_offset = block.group_index().x * items_per_block;
+ auto const block_end_offset = min(n, block_begin_offset + items_per_block);
+
+ if (block_begin_offset < block_end_offset) {
+ if constexpr (IsOuter) {
+ ref.retrieve_outer<BlockSize>(block,
+ input_probe + block_begin_offset,
+ input_probe + block_end_offset,
+ output_probe,
+ output_match,
+ *atomic_counter);
+ } else {
+ ref.retrieve<BlockSize>(block,
+ input_probe + block_begin_offset,
+ input_probe + block_end_offset,
+ output_probe,
+ output_match,
+ *atomic_counter);
+ }
+ }
+}
+
 /**
  * @brief Inserts all elements in the range `[first, last)`.
  *