From d21aaed3171579eeec945e7d2799798266682294 Mon Sep 17 00:00:00 2001 From: William Fan Date: Fri, 18 Sep 2026 11:56:38 -0400 Subject: [PATCH 01/11] Add count retrieve inner and outer if functions --- .../cuco/detail/open_addressing/kernels.cuh | 124 +++ .../open_addressing/open_addressing_impl.cuh | 159 +++- .../open_addressing_ref_impl.cuh | 23 + .../static_multiset/static_multiset.inl | 186 +++++ .../static_multiset/static_multiset_ref.inl | 27 + include/cuco/static_multiset.cuh | 90 +++ tests/CMakeLists.txt | 1 + tests/static_multiset/count_test.cu | 180 +++++ tests/static_multiset/custom_count_test.cu | 368 +++++++++ .../static_multiset/retrieve_if_impl_test.cu | 752 ++++++++++++++++++ 10 files changed, 1877 insertions(+), 33 deletions(-) create mode 100644 tests/static_multiset/retrieve_if_impl_test.cu diff --git a/include/cuco/detail/open_addressing/kernels.cuh b/include/cuco/detail/open_addressing/kernels.cuh index 7c41e89d6..141aaa53b 100644 --- a/include/cuco/detail/open_addressing/kernels.cuh +++ b/include/cuco/detail/open_addressing/kernels.cuh @@ -565,6 +565,78 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count(InputIt first, if (threadIdx.x == 0) { count->fetch_add(block_count, cuda::std::memory_order_relaxed); } } +template +CUCO_KERNEL __launch_bounds__(BlockSize) void count_if(InputIt first, + cuco::detail::index_type n, + StencilIt stencil, + Predicate pred, + AtomicT* count, + Ref ref) +{ + using size_type = typename Ref::size_type; + + size_type constexpr outer_min_count = 1; + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + size_type thread_count = 0; + + auto const loop_stride = cuco::detail::grid_stride() / CGSize; + auto idx = cuco::detail::global_thread_id() / CGSize; + + while (idx < n) { + if constexpr (CGSize == 1) { + if (pred(*(stencil + idx))) { + typename cuda::std::iterator_traits::value_type const key = *(first + idx); + + if constexpr (IsOuter) { + thread_count += max(ref.count(key), outer_min_count); + } else { + thread_count += ref.count(key); + } + } else if constexpr (IsOuter) { + thread_count += outer_min_count; + } + } else { + auto const tile = + cooperative_groups::tiled_partition( + cooperative_groups::this_thread_block()); + + // bool const selected = pred(*(stencil + idx)); + + if (pred(*(stencil + idx))) { + typename cuda::std::iterator_traits::value_type const key = *(first + idx); + + if constexpr (IsOuter) { + auto temp_count = ref.count(tile, key); + + if (tile.all(temp_count == 0) && tile.thread_rank() == 0) { ++temp_count; } + + thread_count += temp_count; + } else { + thread_count += ref.count(tile, key); + } + } else if constexpr (IsOuter) { + if (tile.thread_rank() == 0) { thread_count += outer_min_count; } + } + } + + idx += loop_stride; + } + + auto const block_count = BlockReduce(temp_storage).Sum(thread_count); + + if (threadIdx.x == 0) { count->fetch_add(block_count, cuda::std::memory_order_relaxed); } +} + /** * @brief Counts the occurrences of each key in `[first, last)` contained in the container * and stores the counts in the output array. @@ -699,6 +771,58 @@ CUCO_KERNEL void retrieve(InputProbeIt input_probe, } } +template +CUCO_KERNEL void retrieve_if(InputProbeIt input_probe, + cuco::detail::index_type n, + StencilIt stencil, + Predicate pred, + 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; + auto constexpr tiles_per_block = TileStride * tiles_in_block; + + auto const block_begin_offset = block.group_index().x * tiles_per_block; + auto const block_end_offset = + min(n, static_cast(block_begin_offset + tiles_per_block)); + + if (block_begin_offset < block_end_offset) { + if constexpr (IsOuter) { + ref.template retrieve_outer_if(block, + input_probe + block_begin_offset, + input_probe + block_end_offset, + stencil + block_begin_offset, + pred, + output_probe, + output_match, + *atomic_counter); + } else { + ref.template retrieve_if(block, + input_probe + block_begin_offset, + input_probe + block_end_offset, + stencil + block_begin_offset, + pred, + output_probe, + output_match, + *atomic_counter); + } + } +} + template CUCO_KERNEL __launch_bounds__(BlockSize) void rehash( typename ContainerRef::storage_ref_type storage_ref, diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh index cfae2d38f..cd71f29df 100644 --- a/include/cuco/detail/open_addressing/open_addressing_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_impl.cuh @@ -686,6 +686,46 @@ class open_addressing_impl : private open_addressing_compatible + std::pair retrieve_if(InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + Ref container_ref, + cuda::stream_ref stream) const + { + auto constexpr is_outer = false; + return this->retrieve_if_impl( + first, last, stencil, pred, output_probe, output_match, container_ref, stream); + } + + template + std::pair retrieve_outer_if(InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + Ref container_ref, + cuda::stream_ref stream) const + { + auto constexpr is_outer = true; + return this->retrieve_if_impl( + first, last, stencil, pred, output_probe, output_match, container_ref, stream); + } + /** * @brief Counts the occurrences of keys in `[first, last)` contained in the container * @@ -708,6 +748,18 @@ class open_addressing_impl : private open_addressing_compatiblecount(first, last, container_ref, stream); } + template + [[nodiscard]] size_type count_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate pred, + Ref container_ref, + cuda::stream_ref stream) const + { + auto constexpr is_outer = false; + return this->count_if(first, last, stencil, pred, container_ref, stream); + } + /** * @brief Counts the occurrences of keys in `[first, last)` contained in the container * @@ -732,6 +784,44 @@ class open_addressing_impl : private open_addressing_compatiblecount(first, last, container_ref, stream); } + template + [[nodiscard]] size_type count_outer_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate pred, + Ref container_ref, + cuda::stream_ref stream) const + { + auto constexpr is_outer = true; + + return this->count_if(first, last, stencil, pred, container_ref, stream); + } + + template + [[nodiscard]] size_type count_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate pred, + Ref container_ref, + cuda::stream_ref stream) const + { + auto const num_keys = cuco::detail::distance(first, last); + if (num_keys == 0) { return 0; } + + auto counter = + detail::counter_storage{this->allocator(), stream}; + + counter.reset(stream); + + auto const grid_size = cuco::detail::grid_size(num_keys, cg_size); + + detail::open_addressing_ns::count_if + <<>>( + first, num_keys, stencil, pred, counter.data(), container_ref); + + return counter.load_to_host(stream); + } + /** * @brief Counts the number of occurrences of each query key in the container * @@ -1264,39 +1354,6 @@ class open_addressing_impl : private open_addressing_compatible std::pair retrieve_impl(InputProbeIt first, InputProbeIt last, @@ -1326,6 +1383,42 @@ class open_addressing_impl : private open_addressing_compatible + std::pair retrieve_if_impl(InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + Ref container_ref, + cuda::stream_ref stream) const + { + auto const n = detail::distance(first, last); + if (n == 0) { return {output_probe, output_match}; } + + using counter_type = detail::counter_storage; + auto counter = counter_type{this->allocator(), stream}; + counter.reset(stream.get()); + + auto constexpr block_size = cuco::detail::default_block_size(); + auto constexpr grid_stride = 4; + auto const grid_size = cuco::detail::grid_size(n, cg_size, grid_stride, block_size); + + detail::open_addressing_ns::retrieve_if + <<>>( + first, n, stencil, pred, output_probe, output_match, counter.data(), container_ref); + + auto const num_retrieved = counter.load_to_host(stream.get()); + + return {output_probe + num_retrieved, output_match + num_retrieved}; + } + /** * @brief Extracts the key from a given slot. * diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh index b84d07624..62082cb36 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -1175,6 +1175,29 @@ class open_addressing_ref_impl block, input_probe_begin, n, stencil, pred, output_probe, output_match, atomic_counter); } + template + __device__ void retrieve_outer_if(cooperative_groups::thread_block const& block, + InputProbeIt input_probe_begin, + InputProbeIt input_probe_end, + StencilIt stencil, + Predicate pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + AtomicCounter& atomic_counter) const + { + auto constexpr is_outer = true; + auto const n = cuco::detail::distance(input_probe_begin, input_probe_end); + + this->retrieve_impl( + block, input_probe_begin, n, stencil, pred, output_probe, output_match, atomic_counter); + } + /** * @brief Retrieves all the slots corresponding to all keys in the range `[input_probe_begin, * input_probe_end)`. diff --git a/include/cuco/detail/static_multiset/static_multiset.inl b/include/cuco/detail/static_multiset/static_multiset.inl index 811c07781..e39eb830f 100644 --- a/include/cuco/detail/static_multiset/static_multiset.inl +++ b/include/cuco/detail/static_multiset/static_multiset.inl @@ -457,6 +457,100 @@ static_multiset stream); } +template +template +static_multiset::size_type +static_multiset::count_if( + InputIt first, InputIt last, StencilIt stencil, Predicate const& pred, cuda::stream_ref stream) + const +{ + return this->count_if(first, last, stencil, pred, key_eq(), hash_function(), stream); +} + +template +template +static_multiset::size_type +static_multiset::count_if( + InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + ProbeKeyEqual const& probe_key_equal, + ProbeHash const& probe_hash, + cuda::stream_ref stream) const +{ + return impl_->count_if( + first, + last, + stencil, + pred, + ref(op::count).rebind_key_eq(probe_key_equal).rebind_hash_function(probe_hash), + stream); +} + +template +template +static_multiset::size_type +static_multiset::count_outer_if( + InputIt first, InputIt last, StencilIt stencil, Predicate const& pred, cuda::stream_ref stream) + const +{ + return this->count_outer_if(first, last, stencil, pred, key_eq(), hash_function(), stream); +} + +template +template +static_multiset::size_type +static_multiset::count_outer_if( + InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + ProbeKeyEqual const& probe_key_equal, + ProbeHash const& probe_hash, + cuda::stream_ref stream) const +{ + return impl_->count_outer_if( + first, + last, + stencil, + pred, + ref(op::count).rebind_key_eq(probe_key_equal).rebind_hash_function(probe_hash), + stream); +} + template return impl_->retrieve_outer(first, last, output_probe, output_match, probe_ref, stream); } +template +template +std::pair +static_multiset::retrieve_if( + InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream) const +{ + return impl_->retrieve_if( + first, last, stencil, pred, output_probe, output_match, this->ref(op::retrieve), stream); +} + +template +template +std::pair +static_multiset::retrieve_if( + InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + ProbeEqual const& probe_equal, + ProbeHash const& probe_hash, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream) const +{ + auto const probe_ref = + this->ref(op::retrieve).rebind_key_eq(probe_equal).rebind_hash_function(probe_hash); + + return impl_->retrieve_if( + first, last, stencil, pred, output_probe, output_match, probe_ref, stream); +} + +template +template +std::pair +static_multiset::retrieve_outer_if( + InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + ProbeEqual const& probe_equal, + ProbeHash const& probe_hash, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream) const +{ + auto const probe_ref = + this->ref(op::retrieve).rebind_key_eq(probe_equal).rebind_hash_function(probe_hash); + + return impl_->retrieve_outer_if( + first, last, stencil, pred, output_probe, output_match, probe_ref, stream); +} + template + __device__ void retrieve_outer_if(cooperative_groups::thread_block const& block, + InputProbeIt input_probe_begin, + InputProbeIt input_probe_end, + StencilIt stencil, + Predicate pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + AtomicCounter& atomic_counter) const + { + auto const& ref_ = static_cast(*this); + ref_.impl_.template retrieve_outer_if(block, + input_probe_begin, + input_probe_end, + stencil, + pred, + output_probe, + output_match, + atomic_counter); + } }; template + size_type count_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + + template + size_type count_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + ProbeKeyEqual const& probe_key_equal, + ProbeHash const& probe_hash, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + /** * @brief Counts the occurrences of keys in `[first, last)` contained in the multiset * @@ -680,6 +700,26 @@ class static_multiset { ProbeHash const& probe_hash, cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + template + size_type count_outer_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + + template + size_type count_outer_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate const& pred, + ProbeKeyEqual const& probe_key_equal, + ProbeHash const& probe_hash, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + /** * @brief Counts the number of occurrences of each query key in the multiset * @@ -873,6 +913,56 @@ class static_multiset { cuda::stream_ref stream = cuda::stream_ref{ cudaStream_t{nullptr}}) const; + template + std::pair retrieve_if(InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream = cuda::stream_ref{ + cudaStream_t{nullptr}}) const; + + template + std::pair retrieve_if(InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + ProbeEqual const& probe_equal, + ProbeHash const& probe_hash, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream = cuda::stream_ref{ + cudaStream_t{nullptr}}) const; + + template + std::pair retrieve_outer_if( + InputProbeIt first, + InputProbeIt last, + StencilIt stencil, + Predicate const& pred, + ProbeEqual const& probe_equal, + ProbeHash const& probe_hash, + OutputProbeIt output_probe, + OutputMatchIt output_match, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + /** * @brief Retrieves all keys contained in the multiset * diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index dd146341e..0e073f859 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -106,6 +106,7 @@ ConfigureTest(STATIC_MULTISET_TEST static_multiset/insert_test.cu static_multiset/for_each_test.cu static_multiset/retrieve_test.cu + static_multiset/retrieve_if_impl_test.cu static_multiset/retrieve_if_test.cu static_multiset/large_input_test.cu static_multiset/load_factor_test.cu diff --git a/tests/static_multiset/count_test.cu b/tests/static_multiset/count_test.cu index e9e4e664d..b3b393c86 100644 --- a/tests/static_multiset/count_test.cu +++ b/tests/static_multiset/count_test.cu @@ -168,6 +168,182 @@ void test_count_each_outer(Set& set, size_type num_keys) } } +template +void test_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + + set.insert(keys_begin, keys_begin + num_keys); + + SECTION("Count_if with all elements selected should match count.") + { + auto const count = + set.count_if(keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == num_keys); + } + + SECTION("Count_if with no elements selected should return zero.") + { + auto const count = + set.count_if(keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == 0); + } + + SECTION("Count_if with alternating predicate should count selected keys.") + { + auto const count = set.count_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + REQUIRE(count == (num_keys + 1) / 2); + } +} + +template +void test_count_outer_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + + set.insert(keys_begin, keys_begin + num_keys); + SECTION("Count_outer_if with all elements selected should match count_outer.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == num_keys); + } + + SECTION("Count_outer_if with no elements selected should return one per input.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == num_keys); + } + + SECTION( + "Count_outer_if with alternating predicate should count selected matches and unselected rows.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + REQUIRE(count == num_keys); + } +} + +template +void test_count_if_stencil(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return i + 1000; })); + + set.clear(); + set.insert(keys_begin, keys_begin + num_keys); + + SECTION("Count_if should apply the predicate to the stencil.") + { + auto const count = set.count_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); + + REQUIRE(count == 100); + } + + SECTION("Count_outer_if should apply the predicate to the stencil.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); + + REQUIRE(count == num_keys); + } +} + +template +void test_count_if_overloads(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + set.insert(keys_begin, keys_begin + num_keys); + + auto const pred = + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }); + + SECTION("Count_if explicit default key equality/hash matches overload.") + { + auto const count_default = set.count_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); + + auto const count_explicit = set.count_if( + keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); + + REQUIRE(count_explicit == count_default); + } + + SECTION("Count_outer_if explicit default key equality/hash matches overload.") + { + auto const count_default = + set.count_outer_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); + + auto const count_explicit = set.count_outer_if( + keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); + + REQUIRE(count_explicit == count_default); + } +} + TEMPLATE_TEST_CASE_SIG( "static_multiset count tests", "", @@ -201,4 +377,8 @@ TEMPLATE_TEST_CASE_SIG( test_unique_sequence(set, num_keys); test_count_each(set, num_keys); test_count_each_outer(set, num_keys); + test_count_if(set, num_keys); + test_count_outer_if(set, num_keys); + test_count_if_stencil(set, num_keys); + test_count_if_overloads(set, num_keys); } diff --git a/tests/static_multiset/custom_count_test.cu b/tests/static_multiset/custom_count_test.cu index 6449f7022..1d6743800 100644 --- a/tests/static_multiset/custom_count_test.cu +++ b/tests/static_multiset/custom_count_test.cu @@ -108,6 +108,369 @@ void test_custom_count(Set& set, size_type num_keys) } } +template +void test_count_if_duplicates(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto constexpr multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if with duplicates and all keys selected returns total multiplicity.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_if with duplicates and no keys selected returns zero.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == 0); + } + + SECTION("Count_if with duplicates counts only selected keys.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + auto const expected = ((query_size + 1) / 2) * multiplicity; + + REQUIRE(count == expected); + } + + SECTION("Count_if with duplicates counts a single selected key by its multiplicity.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); + + REQUIRE(count == multiplicity); + } +} + +template +void test_count_outer_if_duplicates(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto constexpr multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + // Query each unique key once. + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_outer_if with duplicates and all keys selected returns total multiplicity.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_outer_if with duplicates and no keys selected returns one per query.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == query_size); + } + + SECTION( + "Count_outer_if with duplicates counts selected matches and one for each unselected query.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + auto const expected = selected_count * multiplicity + unselected_count; + + REQUIRE(count == expected); + } + + SECTION("Count_outer_if with a selected key counts its multiplicity.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); + + auto const expected = multiplicity + (query_size - 1); + + REQUIRE(count == expected); + } +} + +template +void test_custom_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + constexpr auto multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if custom key equality/hash overload counts selected duplicates.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + REQUIRE(count == selected_count * multiplicity); + } + + SECTION("Count_outer_if custom key equality/hash overload counts selected duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + REQUIRE(count == selected_count * multiplicity + unselected_count); + } +} + +template +void test_custom_count_if_overloads(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + constexpr auto multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto query_size = num_keys / multiplicity; + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if explicit key equality/hash overload selects all duplicates.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_if explicit key equality/hash overload selects no duplicates.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; }), + custom_key_eq{}, + hash); + + REQUIRE(count == 0); + } + + SECTION("Count_if explicit key equality/hash overload selects alternating duplicates.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + REQUIRE(count == selected_count * multiplicity); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects all duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects no duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects alternating duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + REQUIRE(count == selected_count * multiplicity + unselected_count); + } +} + +template +void test_custom_hash_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + auto const iter = cuda::counting_iterator{0}; + set.clear(); + set.insert(iter, iter + num_keys); + + auto query_begin = cuda::make_transform_iterator( + cuda::make_counting_iterator(0), + cuda::proclaim_return_type([] __device__(auto i) { return static_cast(i * XXX); })); + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if uses custom key equality and hash.") + { + auto const count = set.count_if( + query_begin, + query_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + REQUIRE(count == (num_keys + 1) / 2); + } + + SECTION("Count_outer_if uses custom key equality and hash.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + REQUIRE(count == num_keys); + } +} + TEMPLATE_TEST_CASE_SIG( "static_multiset custom count tests", "", @@ -139,4 +502,9 @@ TEMPLATE_TEST_CASE_SIG( cuco::static_multiset{num_keys, cuco::empty_key{-1}, {}, probe{}, {}, cuco::storage<2>{}}; test_custom_count(set, num_keys); + test_count_if_duplicates(set, num_keys); + test_count_outer_if_duplicates(set, num_keys); + test_custom_count_if(set, num_keys); + test_custom_count_if_overloads(set, num_keys); + test_custom_hash_count_if(set, num_keys); } diff --git a/tests/static_multiset/retrieve_if_impl_test.cu b/tests/static_multiset/retrieve_if_impl_test.cu new file mode 100644 index 000000000..9eb9bedea --- /dev/null +++ b/tests/static_multiset/retrieve_if_impl_test.cu @@ -0,0 +1,752 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +template +void test_multiplicity(Container& container, std::size_t num_keys, std::size_t multiplicity) +{ + using key_type = typename Container::key_type; + + container.clear(); + + auto const num_unique_keys = num_keys / multiplicity; + REQUIRE(num_unique_keys > 0); + + auto const num_actual_keys = num_unique_keys * multiplicity; + REQUIRE(num_actual_keys <= num_keys); + + thrust::device_vector probed_keys(num_actual_keys); + thrust::device_vector matched_keys(num_actual_keys); + + auto const keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator(0), + cuda::proclaim_return_type([multiplicity] __device__(auto const& i) { + return static_cast(i / multiplicity); + })); + + container.insert(keys_begin, keys_begin + num_actual_keys); + REQUIRE(container.size() == num_actual_keys); + + SECTION("All inserted keys should be contained.") + { + auto const [probed_end, matched_end] = container.retrieve( + keys_begin, keys_begin + num_actual_keys, probed_keys.begin(), matched_keys.begin()); + + thrust::sort(probed_keys.begin(), probed_end); + thrust::sort(matched_keys.begin(), matched_end); + + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_keys.end(), keys_begin, cuda::std::equal_to{})); + + REQUIRE(cuco::test::equal( + matched_keys.begin(), matched_keys.end(), keys_begin, cuda::std::equal_to{})); + } +} + +template +void test_outer(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + auto const empty_key_sentinel = container.empty_key_sentinel(); + + container.clear(); + + auto const keys_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys * 2ull; + + thrust::device_vector probed_keys(query_size); + thrust::device_vector matched_keys(query_size); + + SECTION("Non-inserted keys should output sentinels.") + { + auto const [probed_end, matched_end] = container.retrieve_outer(keys_begin, + keys_begin + query_size, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == + query_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + query_size); + + REQUIRE(cuco::test::all_of( + matched_keys.begin(), + matched_keys.end(), + cuda::proclaim_return_type([empty_key_sentinel] __device__(auto const& k) { + return static_cast(k == static_cast(empty_key_sentinel)); + }))); + } + + container.insert(keys_begin, keys_begin + num_keys); + + SECTION("All inserted keys should be contained.") + { + auto const [probed_end, matched_end] = container.retrieve_outer(keys_begin, + keys_begin + query_size, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == + query_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + query_size); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_keys.end(), keys_begin, cuda::std::equal_to{})); + + REQUIRE(cuco::test::equal( + matched_keys.begin(), + matched_keys.begin() + num_keys, + keys_begin, + cuda::std::equal_to{})); + + REQUIRE(cuco::test::all_of( + matched_keys.begin() + num_keys, + matched_keys.end(), + cuda::proclaim_return_type([empty_key_sentinel] __device__(auto const& k) { + return static_cast(k == static_cast(empty_key_sentinel)); + }))); + } +} + +template +void test_retrieve_if(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + + container.clear(); + + auto const keys_begin = cuda::counting_iterator{0}; + + container.insert(keys_begin, keys_begin + num_keys); + + thrust::device_vector probed_keys(num_keys); + thrust::device_vector matched_keys(num_keys); + thrust::device_vector stencil(num_keys); + + SECTION("retrieve_if should predicate on the stencil, not the probe.") + { + // Make the stencil intentionally different from the probe sequence. + // Probes are 0..N-1, while stencil values are 1..N. + thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); + + auto const pred = [] __device__(key_type key) { + return key % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + auto const expected_size = num_keys / 2; + + REQUIRE(num_results == expected_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + expected_size); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + // Even stencil values correspond to odd probe values. + for (std::size_t i = 0; i < expected_size; ++i) + { + auto const expected = static_cast(i * 2 + 1); + + REQUIRE(probed_keys[i] == expected); + REQUIRE(matched_keys[i] == expected); + } + } + + SECTION("retrieve_if should retrieve only elements satisfying the predicate.") + { + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + auto const pred = [] __device__(key_type key) { + return key % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + REQUIRE(num_results == (num_keys + 1) / 2); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + num_results); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + for (std::size_t i = 0; i < num_results; ++i) + { + auto const expected = static_cast(i * 2); + + REQUIRE(probed_keys[i] == expected); + REQUIRE(matched_keys[i] == expected); + } + } + + SECTION("retrieve_if should return nothing when the predicate is always false.") + { + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + auto const pred = [] __device__(key_type) { + return false; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); + } + + SECTION("retrieve_if should retrieve everything when the predicate is always true.") + { + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + auto const pred = [] __device__(key_type) { + return true; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == num_keys); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == num_keys); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_end, keys_begin, cuda::std::equal_to{})); + + REQUIRE(cuco::test::equal( + matched_keys.begin(), matched_end, keys_begin, cuda::std::equal_to{})); + } +} + +template +void test_retrieve_if_with_probe(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + + container.clear(); + + auto const keys_begin = cuda::counting_iterator{0}; + + container.insert(keys_begin, keys_begin + num_keys); + + thrust::device_vector probed_keys(num_keys); + thrust::device_vector matched_keys(num_keys); + thrust::device_vector stencil(num_keys); + + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + SECTION("retrieve_if should accept explicit equality and hash functions.") + { + auto const pred = [] __device__(key_type key) { + return key % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + REQUIRE(num_results == (num_keys + 1) / 2); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + num_results); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + for (std::size_t i = 0; i < num_results; ++i) + { + auto const expected = static_cast(i * 2); + + REQUIRE(probed_keys[i] == expected); + REQUIRE(matched_keys[i] == expected); + } + } + + SECTION("retrieve_if with explicit equality and hash should return nothing for false predicate.") + { + auto const pred = [] __device__(key_type) { + return false; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); + } +} + +template +void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + + constexpr std::size_t multiplicity = 2; + + container.clear(); + + auto const num_unique_keys = num_keys / multiplicity; + auto const num_actual_keys = num_unique_keys * multiplicity; + + auto const keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator(0), + cuda::proclaim_return_type([multiplicity] __device__(auto const& i) { + return static_cast(i / multiplicity); + })); + + container.insert(keys_begin, keys_begin + num_actual_keys); + REQUIRE(container.size() == num_actual_keys); + + thrust::device_vector stencil(num_actual_keys); + + // A probe can match `multiplicity` container slots, so the output capacity + // must account for duplicate matches. + thrust::device_vector probed_keys(num_actual_keys * multiplicity); + thrust::device_vector matched_keys(num_actual_keys * multiplicity); + + thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); + + SECTION("retrieve_if should filter duplicate matches using the stencil predicate.") + { + auto const pred = [] __device__(key_type value) { + return value % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + // There are num_actual_keys / 2 selected input probes. Each selected + // probe has multiplicity matching slots in the multiset. + auto const expected_results = (num_actual_keys / 2) * multiplicity; + + REQUIRE(num_results == expected_results); + REQUIRE(static_cast( + std::distance(matched_keys.begin(), matched_end)) == expected_results); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + for (std::size_t i = 0; i < expected_results; ++i) + { + auto const input_index = i / multiplicity; + auto const expected_key = + static_cast((input_index * 2) / multiplicity); + + REQUIRE(probed_keys[i] == expected_key); + REQUIRE(matched_keys[i] == expected_key); + } + } + + SECTION("retrieve_if should return nothing when the predicate is always false.") + { + auto const pred = [] __device__(key_type) { return false; }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast( + std::distance(probed_keys.begin(), probed_end)) == 0); + REQUIRE(static_cast( + std::distance(matched_keys.begin(), matched_end)) == 0); + } + + SECTION("retrieve_if should return all matches when the predicate is always true.") + { + auto const pred = [] __device__(key_type) { return true; }; + + auto const [probed_end, matched_end] = container.retrieve_if( + keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + // There are num_actual_keys input probes and every probe has + // `multiplicity` matching slots. + auto const expected_results = num_actual_keys * multiplicity; + + REQUIRE(static_cast( + std::distance(probed_keys.begin(), probed_end)) == expected_results); + REQUIRE(static_cast( + std::distance(matched_keys.begin(), matched_end)) == expected_results); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + // Each unique key occurs `multiplicity` times in the input and + // `multiplicity` times in the container, producing multiplicity^2 + // output pairs for each unique key. + for (std::size_t key = 0; key < num_unique_keys; ++key) + { + auto const expected_key = static_cast(key); + + auto const expected_count = multiplicity * multiplicity; + + for (std::size_t j = 0; j < expected_count; ++j) + { + auto const output_index = key * expected_count + j; + + REQUIRE(probed_keys[output_index] == expected_key); + REQUIRE(matched_keys[output_index] == expected_key); + } + } + } +} + +template +void test_retrieve_outer_if(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + auto const empty_key_sentinel = container.empty_key_sentinel(); + + container.clear(); + + auto const keys_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys * 2ull; + + container.insert(keys_begin, keys_begin + num_keys); + + thrust::device_vector probes(query_size); + thrust::device_vector stencil(query_size); + thrust::device_vector probed_keys(query_size); + thrust::device_vector matched_keys(query_size); + + thrust::sequence(probes.begin(), probes.end(), key_type{0}); + + SECTION("retrieve_outer_if should return matches and sentinels for misses.") + { + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + auto const pred = [] __device__(key_type) { + return true; + }; + + auto const [probed_end, matched_end] = container.retrieve_outer_if( + probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == + query_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + query_size); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_keys.end(), probes.begin(), cuda::std::equal_to{})); + + REQUIRE(cuco::test::equal( + matched_keys.begin(), + matched_keys.begin() + num_keys, + keys_begin, + cuda::std::equal_to{})); + + REQUIRE(cuco::test::all_of( + matched_keys.begin() + num_keys, + matched_keys.end(), + cuda::proclaim_return_type([empty_key_sentinel] __device__(auto const& k) { + return static_cast(k == static_cast(empty_key_sentinel)); + }))); + } + + SECTION("retrieve_outer_if should predicate on the stencil, not the probe.") + { + // probes = 0..N-1 + // stencil = 1..N + // Even stencil values select odd probes. + thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); + + auto const pred = [] __device__(key_type key) { + return key % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_outer_if( + probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + auto const expected_size = query_size / 2; + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == + expected_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + expected_size); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + for (std::size_t i = 0; i < expected_size; ++i) + { + auto const expected_probe = static_cast(i * 2 + 1); + + REQUIRE(probed_keys[i] == expected_probe); + + if (expected_probe < static_cast(num_keys)) + { + REQUIRE(matched_keys[i] == expected_probe); + } + else + { + REQUIRE(matched_keys[i] == static_cast(empty_key_sentinel)); + } + } + } + + SECTION("retrieve_outer_if should return nothing for an always-false predicate.") + { + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); + + auto const pred = [] __device__(key_type) { + return false; + }; + + auto const [probed_end, matched_end] = container.retrieve_outer_if( + probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); + } +} + +template +void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + auto const empty_key_sentinel = container.empty_key_sentinel(); + + constexpr std::size_t multiplicity = 2; + + container.clear(); + + auto const num_unique_keys = num_keys / multiplicity; + auto const num_actual_keys = num_unique_keys * multiplicity; + + auto const keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator(0), + cuda::proclaim_return_type([multiplicity] __device__(auto const& i) { + return static_cast(i / multiplicity); + })); + + container.insert(keys_begin, keys_begin + num_actual_keys); + + auto const query_size = num_unique_keys * 2ull; + + thrust::device_vector probes(query_size); + thrust::device_vector stencil(query_size); + thrust::device_vector probed_keys(query_size * multiplicity); + thrust::device_vector matched_keys(query_size * multiplicity); + + thrust::sequence(probes.begin(), probes.end(), key_type{0}); + + // Select odd probes by using stencil values 1..N. + thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); + + auto const pred = [] __device__(key_type key) { + return key % 2 == 0; + }; + + auto const [probed_end, matched_end] = container.retrieve_outer_if( + probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + auto const num_matching_probes = query_size / 2; + auto const num_matching_unique_keys = num_unique_keys / 2; + auto const num_missing_probes = num_matching_probes - num_matching_unique_keys; + + // Each selected probe that exists in the multiset has `multiplicity` + // matches. Each selected probe that does not exist produces one sentinel. + auto const expected_results = + num_matching_unique_keys * multiplicity + num_missing_probes; + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + REQUIRE(num_results == expected_results); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + expected_results); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + std::size_t output_index = 0; + + for (std::size_t probe = 1; probe < query_size; probe += 2) + { + auto const expected_probe = static_cast(probe); + + if (probe < num_unique_keys) + { + for (std::size_t j = 0; j < multiplicity; ++j) + { + REQUIRE(probed_keys[output_index] == expected_probe); + REQUIRE(matched_keys[output_index] == expected_probe); + ++output_index; + } + } + else + { + REQUIRE(probed_keys[output_index] == expected_probe); + REQUIRE(matched_keys[output_index] == + static_cast(empty_key_sentinel)); + ++output_index; + } + } + + REQUIRE(output_index == expected_results); +} + +TEMPLATE_TEST_CASE_SIG( + "static_multiset retrieve if tests", + "", + ((typename Key, cuco::test::probe_sequence Probe, int CGSize), Key, Probe, CGSize), + (int32_t, cuco::test::probe_sequence::double_hashing, 1), + (int32_t, cuco::test::probe_sequence::double_hashing, 2), + (int64_t, cuco::test::probe_sequence::double_hashing, 1), + (int64_t, cuco::test::probe_sequence::double_hashing, 2), + (int32_t, cuco::test::probe_sequence::linear_probing, 1), + (int32_t, cuco::test::probe_sequence::linear_probing, 2), + (int64_t, cuco::test::probe_sequence::linear_probing, 1), + (int64_t, cuco::test::probe_sequence::linear_probing, 2) +#if defined(CUCO_HAS_128BIT_ATOMICS) + , + (__int128_t, cuco::test::probe_sequence::double_hashing, 1), + (__int128_t, cuco::test::probe_sequence::double_hashing, 2), + (__int128_t, cuco::test::probe_sequence::linear_probing, 1), + (__int128_t, cuco::test::probe_sequence::linear_probing, 2) +#endif +) +{ + constexpr std::size_t num_keys{400}; + constexpr double desired_load_factor = 0.5; + constexpr auto empty_key_sentinel = std::numeric_limits::max(); + + using probe = std::conditional_t>, + cuco::double_hashing>>; + + auto set = cuco::static_multiset{ + num_keys, desired_load_factor, cuco::empty_key{empty_key_sentinel}, {}, probe{}}; + + test_multiplicity(set, num_keys, 1); + test_multiplicity(set, num_keys, 2); + test_multiplicity(set, num_keys, 11); + + test_outer(set, num_keys); + + test_retrieve_if(set, num_keys); + test_retrieve_if_with_probe(set, num_keys); + test_retrieve_if_multiplicity(set, num_keys); + + test_retrieve_outer_if(set, num_keys); + test_retrieve_outer_if_multiplicity(set, num_keys); +} From e8b40812ad8b3f540a49143087b7a24bb7a3577a Mon Sep 17 00:00:00 2001 From: William Fan Date: Fri, 18 Sep 2026 11:58:25 -0400 Subject: [PATCH 02/11] Reformatted --- tests/static_multiset/retrieve_if_impl_test.cu | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/tests/static_multiset/retrieve_if_impl_test.cu b/tests/static_multiset/retrieve_if_impl_test.cu index 9eb9bedea..7c17c14df 100644 --- a/tests/static_multiset/retrieve_if_impl_test.cu +++ b/tests/static_multiset/retrieve_if_impl_test.cu @@ -151,8 +151,6 @@ void test_retrieve_if(Container& container, std::size_t num_keys) SECTION("retrieve_if should predicate on the stencil, not the probe.") { - // Make the stencil intentionally different from the probe sequence. - // Probes are 0..N-1, while stencil values are 1..N. thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); auto const pred = [] __device__(key_type key) { @@ -179,7 +177,6 @@ void test_retrieve_if(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - // Even stencil values correspond to odd probe values. for (std::size_t i = 0; i < expected_size; ++i) { auto const expected = static_cast(i * 2 + 1); @@ -370,8 +367,6 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) thrust::device_vector stencil(num_actual_keys); - // A probe can match `multiplicity` container slots, so the output capacity - // must account for duplicate matches. thrust::device_vector probed_keys(num_actual_keys * multiplicity); thrust::device_vector matched_keys(num_actual_keys * multiplicity); @@ -394,8 +389,6 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); - // There are num_actual_keys / 2 selected input probes. Each selected - // probe has multiplicity matching slots in the multiset. auto const expected_results = (num_actual_keys / 2) * multiplicity; REQUIRE(num_results == expected_results); @@ -446,8 +439,6 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) probed_keys.begin(), matched_keys.begin()); - // There are num_actual_keys input probes and every probe has - // `multiplicity` matching slots. auto const expected_results = num_actual_keys * multiplicity; REQUIRE(static_cast( @@ -458,9 +449,6 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - // Each unique key occurs `multiplicity` times in the input and - // `multiplicity` times in the container, producing multiplicity^2 - // output pairs for each unique key. for (std::size_t key = 0; key < num_unique_keys; ++key) { auto const expected_key = static_cast(key); @@ -543,9 +531,6 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) SECTION("retrieve_outer_if should predicate on the stencil, not the probe.") { - // probes = 0..N-1 - // stencil = 1..N - // Even stencil values select odd probes. thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); auto const pred = [] __device__(key_type key) { @@ -642,7 +627,6 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k thrust::sequence(probes.begin(), probes.end(), key_type{0}); - // Select odd probes by using stencil values 1..N. thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); auto const pred = [] __device__(key_type key) { @@ -663,8 +647,6 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k auto const num_matching_unique_keys = num_unique_keys / 2; auto const num_missing_probes = num_matching_probes - num_matching_unique_keys; - // Each selected probe that exists in the multiset has `multiplicity` - // matches. Each selected probe that does not exist produces one sentinel. auto const expected_results = num_matching_unique_keys * multiplicity + num_missing_probes; From ea3d38264d5f5fdb184cac622978d288897b7e2e Mon Sep 17 00:00:00 2001 From: William Fan Date: Fri, 18 Sep 2026 12:10:51 -0400 Subject: [PATCH 03/11] Fixed formatting --- .../open_addressing/open_addressing_impl.cuh | 33 ++ .../static_multiset/retrieve_if_impl_test.cu | 353 ++++++++---------- 2 files changed, 179 insertions(+), 207 deletions(-) diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh index cd71f29df..509fe6b9f 100644 --- a/include/cuco/detail/open_addressing/open_addressing_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_impl.cuh @@ -1354,6 +1354,39 @@ class open_addressing_impl : private open_addressing_compatible std::pair retrieve_impl(InputProbeIt first, InputProbeIt last, diff --git a/tests/static_multiset/retrieve_if_impl_test.cu b/tests/static_multiset/retrieve_if_impl_test.cu index 7c17c14df..32782115e 100644 --- a/tests/static_multiset/retrieve_if_impl_test.cu +++ b/tests/static_multiset/retrieve_if_impl_test.cu @@ -84,8 +84,7 @@ void test_outer(Container& container, std::size_t num_keys) probed_keys.begin(), matched_keys.begin()); - REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == - query_size); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == query_size); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == query_size); @@ -108,8 +107,7 @@ void test_outer(Container& container, std::size_t num_keys) probed_keys.begin(), matched_keys.begin()); - REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == - query_size); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == query_size); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == query_size); @@ -119,11 +117,10 @@ void test_outer(Container& container, std::size_t num_keys) REQUIRE(cuco::test::equal( probed_keys.begin(), probed_keys.end(), keys_begin, cuda::std::equal_to{})); - REQUIRE(cuco::test::equal( - matched_keys.begin(), - matched_keys.begin() + num_keys, - keys_begin, - cuda::std::equal_to{})); + REQUIRE(cuco::test::equal(matched_keys.begin(), + matched_keys.begin() + num_keys, + keys_begin, + cuda::std::equal_to{})); REQUIRE(cuco::test::all_of( matched_keys.begin() + num_keys, @@ -153,17 +150,14 @@ void test_retrieve_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); - auto const pred = [] __device__(key_type key) { - return key % 2 == 0; - }; + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); @@ -177,8 +171,7 @@ void test_retrieve_if(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < expected_size; ++i) - { + for (std::size_t i = 0; i < expected_size; ++i) { auto const expected = static_cast(i * 2 + 1); REQUIRE(probed_keys[i] == expected); @@ -190,17 +183,14 @@ void test_retrieve_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); - auto const pred = [] __device__(key_type key) { - return key % 2 == 0; - }; + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); @@ -212,8 +202,7 @@ void test_retrieve_if(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < num_results; ++i) - { + for (std::size_t i = 0; i < num_results; ++i) { auto const expected = static_cast(i * 2); REQUIRE(probed_keys[i] == expected); @@ -225,17 +214,14 @@ void test_retrieve_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); - auto const pred = [] __device__(key_type) { - return false; - }; + auto const pred = [] __device__(key_type) { return false; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); @@ -245,17 +231,14 @@ void test_retrieve_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); - auto const pred = [] __device__(key_type) { - return true; - }; + auto const pred = [] __device__(key_type) { return true; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == num_keys); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == num_keys); @@ -290,19 +273,16 @@ void test_retrieve_if_with_probe(Container& container, std::size_t num_keys) SECTION("retrieve_if should accept explicit equality and hash functions.") { - auto const pred = [] __device__(key_type key) { - return key % 2 == 0; - }; - - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; + + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); @@ -314,8 +294,7 @@ void test_retrieve_if_with_probe(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < num_results; ++i) - { + for (std::size_t i = 0; i < num_results; ++i) { auto const expected = static_cast(i * 2); REQUIRE(probed_keys[i] == expected); @@ -325,19 +304,16 @@ void test_retrieve_if_with_probe(Container& container, std::size_t num_keys) SECTION("retrieve_if with explicit equality and hash should return nothing for false predicate.") { - auto const pred = [] __device__(key_type) { - return false; - }; - - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_keys, - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); + auto const pred = [] __device__(key_type) { return false; }; + + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_keys, + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); @@ -374,17 +350,14 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) SECTION("retrieve_if should filter duplicate matches using the stencil predicate.") { - auto const pred = [] __device__(key_type value) { - return value % 2 == 0; - }; - - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_actual_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const pred = [] __device__(key_type value) { return value % 2 == 0; }; + + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); @@ -392,17 +365,15 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) auto const expected_results = (num_actual_keys / 2) * multiplicity; REQUIRE(num_results == expected_results); - REQUIRE(static_cast( - std::distance(matched_keys.begin(), matched_end)) == expected_results); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + expected_results); thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < expected_results; ++i) - { - auto const input_index = i / multiplicity; - auto const expected_key = - static_cast((input_index * 2) / multiplicity); + for (std::size_t i = 0; i < expected_results; ++i) { + auto const input_index = i / multiplicity; + auto const expected_key = static_cast((input_index * 2) / multiplicity); REQUIRE(probed_keys[i] == expected_key); REQUIRE(matched_keys[i] == expected_key); @@ -413,50 +384,44 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) { auto const pred = [] __device__(key_type) { return false; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_actual_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); - - REQUIRE(static_cast( - std::distance(probed_keys.begin(), probed_end)) == 0); - REQUIRE(static_cast( - std::distance(matched_keys.begin(), matched_end)) == 0); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); } SECTION("retrieve_if should return all matches when the predicate is always true.") { auto const pred = [] __device__(key_type) { return true; }; - auto const [probed_end, matched_end] = container.retrieve_if( - keys_begin, - keys_begin + num_actual_keys, - stencil.begin(), - pred, - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_if(keys_begin, + keys_begin + num_actual_keys, + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); auto const expected_results = num_actual_keys * multiplicity; - REQUIRE(static_cast( - std::distance(probed_keys.begin(), probed_end)) == expected_results); - REQUIRE(static_cast( - std::distance(matched_keys.begin(), matched_end)) == expected_results); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == + expected_results); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + expected_results); thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t key = 0; key < num_unique_keys; ++key) - { + for (std::size_t key = 0; key < num_unique_keys; ++key) { auto const expected_key = static_cast(key); auto const expected_count = multiplicity * multiplicity; - for (std::size_t j = 0; j < expected_count; ++j) - { + for (std::size_t j = 0; j < expected_count; ++j) { auto const output_index = key * expected_count + j; REQUIRE(probed_keys[output_index] == expected_key); @@ -490,22 +455,18 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); - auto const pred = [] __device__(key_type) { - return true; - }; + auto const pred = [] __device__(key_type) { return true; }; - auto const [probed_end, matched_end] = container.retrieve_outer_if( - probes.begin(), - probes.end(), - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_outer_if(probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); - REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == - query_size); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == query_size); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == query_size); @@ -515,11 +476,10 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) REQUIRE(cuco::test::equal( probed_keys.begin(), probed_keys.end(), probes.begin(), cuda::std::equal_to{})); - REQUIRE(cuco::test::equal( - matched_keys.begin(), - matched_keys.begin() + num_keys, - keys_begin, - cuda::std::equal_to{})); + REQUIRE(cuco::test::equal(matched_keys.begin(), + matched_keys.begin() + num_keys, + keys_begin, + cuda::std::equal_to{})); REQUIRE(cuco::test::all_of( matched_keys.begin() + num_keys, @@ -533,19 +493,16 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); - auto const pred = [] __device__(key_type key) { - return key % 2 == 0; - }; + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; - auto const [probed_end, matched_end] = container.retrieve_outer_if( - probes.begin(), - probes.end(), - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_outer_if(probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); auto const expected_size = query_size / 2; @@ -557,18 +514,14 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < expected_size; ++i) - { + for (std::size_t i = 0; i < expected_size; ++i) { auto const expected_probe = static_cast(i * 2 + 1); REQUIRE(probed_keys[i] == expected_probe); - if (expected_probe < static_cast(num_keys)) - { + if (expected_probe < static_cast(num_keys)) { REQUIRE(matched_keys[i] == expected_probe); - } - else - { + } else { REQUIRE(matched_keys[i] == static_cast(empty_key_sentinel)); } } @@ -578,19 +531,16 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); - auto const pred = [] __device__(key_type) { - return false; - }; + auto const pred = [] __device__(key_type) { return false; }; - auto const [probed_end, matched_end] = container.retrieve_outer_if( - probes.begin(), - probes.end(), - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); + auto const [probed_end, matched_end] = container.retrieve_outer_if(probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); @@ -629,29 +579,24 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); - auto const pred = [] __device__(key_type key) { - return key % 2 == 0; - }; - - auto const [probed_end, matched_end] = container.retrieve_outer_if( - probes.begin(), - probes.end(), - stencil.begin(), - pred, - container.key_eq(), - container.hash_function(), - probed_keys.begin(), - matched_keys.begin()); - - auto const num_matching_probes = query_size / 2; + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; + + auto const [probed_end, matched_end] = container.retrieve_outer_if(probes.begin(), + probes.end(), + stencil.begin(), + pred, + container.key_eq(), + container.hash_function(), + probed_keys.begin(), + matched_keys.begin()); + + auto const num_matching_probes = query_size / 2; auto const num_matching_unique_keys = num_unique_keys / 2; - auto const num_missing_probes = num_matching_probes - num_matching_unique_keys; + auto const num_missing_probes = num_matching_probes - num_matching_unique_keys; - auto const expected_results = - num_matching_unique_keys * multiplicity + num_missing_probes; + auto const expected_results = num_matching_unique_keys * multiplicity + num_missing_probes; - auto const num_results = - static_cast(std::distance(probed_keys.begin(), probed_end)); + auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); REQUIRE(num_results == expected_results); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == @@ -662,24 +607,18 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k std::size_t output_index = 0; - for (std::size_t probe = 1; probe < query_size; probe += 2) - { + for (std::size_t probe = 1; probe < query_size; probe += 2) { auto const expected_probe = static_cast(probe); - if (probe < num_unique_keys) - { - for (std::size_t j = 0; j < multiplicity; ++j) - { + if (probe < num_unique_keys) { + for (std::size_t j = 0; j < multiplicity; ++j) { REQUIRE(probed_keys[output_index] == expected_probe); REQUIRE(matched_keys[output_index] == expected_probe); ++output_index; } - } - else - { + } else { REQUIRE(probed_keys[output_index] == expected_probe); - REQUIRE(matched_keys[output_index] == - static_cast(empty_key_sentinel)); + REQUIRE(matched_keys[output_index] == static_cast(empty_key_sentinel)); ++output_index; } } From 1c6881526d0a74d77de7aaefe77d1e4095cb31d9 Mon Sep 17 00:00:00 2001 From: William Fan Date: Sat, 19 Sep 2026 15:22:30 -0400 Subject: [PATCH 04/11] Add documentation --- .../cuco/detail/open_addressing/kernels.cuh | 63 +++++- .../open_addressing/open_addressing_impl.cuh | 202 ++++++++++++++--- .../static_multiset/static_multiset_ref.inl | 39 ++++ include/cuco/static_multiset.cuh | 214 ++++++++++++++++++ 4 files changed, 489 insertions(+), 29 deletions(-) diff --git a/include/cuco/detail/open_addressing/kernels.cuh b/include/cuco/detail/open_addressing/kernels.cuh index 141aaa53b..a5ef9bee5 100644 --- a/include/cuco/detail/open_addressing/kernels.cuh +++ b/include/cuco/detail/open_addressing/kernels.cuh @@ -565,6 +565,28 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count(InputIt first, if (threadIdx.x == 0) { count->fetch_add(block_count, cuda::std::memory_order_relaxed); } } +/** + * @brief Counts the occurrences of keys in `[first, last)` contained in the container + * if `pred` of the corresponding stencil returns true. + * + * @tparam IsOuter Flag indicating whether it's an outer count or not + * @tparam CGSize Number of threads in each CG + * @tparam BlockSize Number of threads in each block + * @tparam InputIt Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` + * and argument type is convertible from `std::iterator_traits::value_type` + * @tparam AtomicT Atomic counter type + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the sequence of input elements + * @param n Number of input elements + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + n)` + * @param count Number of matches + * @param ref Non-owning container device ref used to access the slot storage + */ template ::value_type` + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam AtomicCounter Integral atomic type that follows the same semantics as + * `cuda::(std::)atomic(_ref)` + * @tparam Ref Type of non-owning device ref allowing access to storage + * + * @param input_probe Beginning of the sequence of input keys + * @param n Number of the keys to query + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + n)` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param atomic_counter Pointer to an atomic object of integral type that is used to count the + * number of output elements + * @param ref Non-owning container device ref used to access the slot storage + */ template std::iterator_traits::value_type + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template std::iterator_traits::value_type + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template count(first, last, container_ref, stream); } + /** + * @brief Counts the occurrences of keys in `[first, last)` contained in the container + * if `pred` of the corresponding stencil returns true. + * + * @tparam Input Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` and + * argument type is convertible from std::iterator_traits::value_type + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` + */ template [[nodiscard]] size_type count_if(InputIt first, InputIt last, @@ -770,6 +869,7 @@ class open_addressing_impl : private open_addressing_compatiblecount(first, last, container_ref, stream); } + /** + * @brief Counts the occurrences of keys in `[first, last)` contained in the container + * if `pred` of the corresponding stencil returns true. + * + * @note If a given key has no matches, or `pred` of its corresponding stencil is false, + * its occurrence is 1. + * + * @tparam Input Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` and + * argument type is convertible from std::iterator_traits::value_type + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` + */ template [[nodiscard]] size_type count_outer_if(InputIt first, InputIt last, @@ -797,31 +921,6 @@ class open_addressing_impl : private open_addressing_compatiblecount_if(first, last, stencil, pred, container_ref, stream); } - template - [[nodiscard]] size_type count_if(InputIt first, - InputIt last, - StencilIt stencil, - Predicate pred, - Ref container_ref, - cuda::stream_ref stream) const - { - auto const num_keys = cuco::detail::distance(first, last); - if (num_keys == 0) { return 0; } - - auto counter = - detail::counter_storage{this->allocator(), stream}; - - counter.reset(stream); - - auto const grid_size = cuco::detail::grid_size(num_keys, cg_size); - - detail::open_addressing_ns::count_if - <<>>( - first, num_keys, stencil, pred, counter.data(), container_ref); - - return counter.load_to_host(stream); - } - /** * @brief Counts the number of occurrences of each query key in the container * @@ -1281,11 +1380,12 @@ class open_addressing_impl : private open_addressing_compatiblestd::iterator_traits::value_type + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` + */ + template + [[nodiscard]] size_type count_if(InputIt first, + InputIt last, + StencilIt stencil, + Predicate pred, + Ref container_ref, + cuda::stream_ref stream) const + { + auto const num_keys = cuco::detail::distance(first, last); + if (num_keys == 0) { return 0; } + + auto counter = + detail::counter_storage{this->allocator(), stream}; + + counter.reset(stream); + + auto const grid_size = cuco::detail::grid_size(num_keys, cg_size); + + detail::open_addressing_ns::count_if + <<>>( + first, num_keys, stencil, pred, counter.data(), container_ref); + + return counter.load_to_host(stream); + } + /** * @brief Counts the number of occurrences of each query key in the container * diff --git a/include/cuco/detail/static_multiset/static_multiset_ref.inl b/include/cuco/detail/static_multiset/static_multiset_ref.inl index 483810cc1..a760374e8 100644 --- a/include/cuco/detail/static_multiset/static_multiset_ref.inl +++ b/include/cuco/detail/static_multiset/static_multiset_ref.inl @@ -708,6 +708,45 @@ class operator_impl< atomic_counter); } + /** + * @brief Retrieves all the slots corresponding to all keys in the range `[input_probe_begin, + * input_probe_end)` if `pred` of the corresponding stencil returns true. + * + * If key `k = *(first + i)` exists in the container and `pred` of its corresponding stencil + * is true, copies `k` to `output_probe` and associated slot contents to `output_match`, + * respectively. The output order is unspecified. + * + * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. + * Use `count_outer()` to determine the size of the output range. + * + * If a key `k` has no matches in the container, or `pred` of its corresponding stencil is + * false, then `{key, empty_slot_sentinel}` will be added to the output sequence. + * + * @tparam BlockSize Size of the thread block this operation is executed in + * @tparam InputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the container's `key_type` + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` + * and argument type is convertible from `std::iterator_traits::value_type` + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the container's `key_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam AtomicCounter Atomic counter type that follows the same semantics as + * `cuda::atomic(_ref)` + * + * @param block Thread block this operation is executed in + * @param input_probe_begin Beginning of the input sequence of keys + * @param input_probe_end End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + n)` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param atomic_counter Counter that is used to determine the next free position in the output + * sequences + */ template ::value_type` + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` + */ template size_type count_if(InputIt first, InputIt last, @@ -661,6 +682,31 @@ class static_multiset { Predicate const& pred, cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + /** + * @brief Counts the occurrences of keys in `[first, last)` contained in the multiset + * if `pred` of the corresponding stencil returns true. + * + * @note This function synchronizes the given stream. + * + * @tparam Input Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` and + * argument type is convertible from `std::iterator_traits::value_type` + * @tparam ProbeKeyEqual Binary callable + * @tparam ProbeHash Unary hash callable + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param probe_key_equal Binary callable to compare two keys for equality + * @param probe_hash Unary callable to hash a given key + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` + */ template ::value_type` + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` where keys have no matches + * are considered to have a single occurrence. + */ template size_type count_outer_if(InputIt first, InputIt last, @@ -707,6 +777,34 @@ class static_multiset { Predicate const& pred, cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}) const; + /** + * @brief Counts the occurrences of keys in `[first, last)` contained in the multiset + * if `pred` of the corresponding stencil returns true. + * + * @note This function synchronizes the given stream. + * @note If a given key has no matches or `pred` of its corresponding stencil is false, + * its occurrence is 1. + * + * @tparam Input Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` and + * argument type is convertible from `std::iterator_traits::value_type` + * @tparam ProbeKeyEqual Binary callable + * @tparam ProbeHash Unary hash callable + * + * @param first Beginning of the sequence of keys to count + * @param last End of the sequence of keys to count + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param probe_key_equal Binary callable to compare two keys for equality + * @param probe_hash Unary callable to hash a given key + * @param stream CUDA stream used for count + * + * @return The sum of total occurrences of all keys in `[first, last)` where keys have no matches + * are considered to have a single occurrence. + */ template ::value_type` + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template ::value_type` + * @tparam ProbeEqual Binary callable equal type + * @tparam ProbeHash Unary callable hasher type that can be constructed from + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param probe_equal The binary function to compare set keys and probe keys for equality + * @param probe_hash The unary function to hash probe keys + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template ::value_type` + * @tparam ProbeEqual Binary callable equal type + * @tparam ProbeHash Unary callable hasher type that can be constructed from + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param probe_equal The binary function to compare set keys and probe keys for equality + * @param probe_hash The unary function to hash probe keys + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template Date: Sat, 19 Sep 2026 15:29:18 -0400 Subject: [PATCH 05/11] Fix documentation --- .../cuco/detail/open_addressing/kernels.cuh | 12 ++--- .../open_addressing/open_addressing_impl.cuh | 45 ++++++++++++++++++- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/include/cuco/detail/open_addressing/kernels.cuh b/include/cuco/detail/open_addressing/kernels.cuh index a5ef9bee5..9d008dec0 100644 --- a/include/cuco/detail/open_addressing/kernels.cuh +++ b/include/cuco/detail/open_addressing/kernels.cuh @@ -724,9 +724,9 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count_each(InputIt first, * @brief Retrieves the equivalent container elements of all keys in the range `[input_probe, * input_probe + n)`. * - * If key `k = *(input_probe + i)` has one or more matches in the container and `pred` of - * its corresponding stencil is true, copies `k` to `output_probe` and associated slot - * contents to `output_match`, respectively. The output order is unspecified. + * If key `k = *(input_probe + i)` has one or more matches in the container, copies `k` to + * `output_probe` and associated slot contents to `output_match`, respectively. The output order is + * unspecified. * * @tparam IsOuter Flag indicating whether it's an outer count or not * @tparam BlockSize The size of the thread block @@ -797,9 +797,9 @@ CUCO_KERNEL void retrieve(InputProbeIt input_probe, * @brief Retrieves the equivalent container elements of all keys in the range `[input_probe, * input_probe + n)` if `pred` of the corresponding stencil returns true. * - * If key `k = *(input_probe + i)` has one or more matches in the container, copies `k` to - * `output_probe` and associated slot contents to `output_match`, respectively. The output order is - * unspecified. + * If key `k = *(input_probe + i)` has one or more matches in the container and `pred` of + * its corresponding stencil is true, copies `k` to `output_probe` and associated slot + * contents to `output_match`, respectively. The output order is unspecified. * * @tparam IsOuter Flag indicating whether it's an outer count or not * @tparam BlockSize The size of the thread block diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh index b253f96e0..a4d852a7c 100644 --- a/include/cuco/detail/open_addressing/open_addressing_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_impl.cuh @@ -1425,7 +1425,7 @@ class open_addressing_impl : private open_addressing_compatiblestd::iterator_traits::value_type + * argument type is convertible from `std::iterator_traits::value_type` * @tparam Ref Type of non-owning device container ref allowing access to storage * * @param first Beginning of the sequence of keys to count @@ -1566,6 +1566,49 @@ class open_addressing_impl : private open_addressing_compatible::value_type` + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the container's `key_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam Ref Type of non-owning device container ref allowing access to storage + * + * @param first Beginning of the input sequence of keys + * @param last End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + + * std::distance(first, last))` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param container_ref Non-owning device reference to the container + * @param stream CUDA stream this operation is executed in + * + * @return Iterator pair indicating the the end of the output sequences + */ template Date: Sat, 19 Sep 2026 15:33:04 -0400 Subject: [PATCH 06/11] Add documentation --- .../open_addressing_ref_impl.cuh | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh index 62082cb36..a9a795538 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -1175,6 +1175,44 @@ class open_addressing_ref_impl block, input_probe_begin, n, stencil, pred, output_probe, output_match, atomic_counter); } + /** + * @brief Retrieves all the slots corresponding to all keys in the range `[input_probe_begin, + * input_probe_end)` if `pred` of the corresponding stencil returns true. + * + * If key `k = *(first + i)` exists in the container and `pred( *(stencil + i) )` returns true, + * copies `k` to `output_probe` and associated slot contents to `output_match`, + * respectively. The output order is unspecified. + * + * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. + * Use `count()` to determine the size of the output range. + * + * If a key `k` has no matches in the container, or `pred` of the corresponding stencil is + * false, then `{key, empty_slot_sentinel}` will be added to the output sequence. + * + * @tparam BlockSize Size of the thread block this operation is executed in + * @tparam InputProbeIt Device accessible input iterator + * @tparam StencilIt Device accessible random access iterator whose value_type is + * convertible to Predicate's argument type + * @tparam Predicate Unary predicate callable whose return type must be convertible to `bool` + * and argument type is convertible from `std::iterator_traits::value_type` + * @tparam OutputProbeIt Device accessible input iterator whose `value_type` is + * convertible to the `InputProbeIt`'s `value_type` + * @tparam OutputMatchIt Device accessible input iterator whose `value_type` is + * convertible to the container's `value_type` + * @tparam AtomicCounter Integral atomic counter type that follows the same semantics as + * `cuda::(std::)atomic(_ref)` + * + * @param block Thread block this operation is executed in + * @param input_probe_begin Beginning of the input sequence of keys + * @param input_probe_end End of the input sequence of keys + * @param stencil Beginning of the stencil sequence + * @param pred Predicate to test on every element in the range `[stencil, stencil + n)` + * @param output_probe Beginning of the sequence of keys corresponding to matching elements in + * `output_match` + * @param output_match Beginning of the sequence of matching elements + * @param atomic_counter Atomic object of integral type that is used to count the + * number of output elements + */ template Date: Sat, 19 Sep 2026 16:16:59 -0400 Subject: [PATCH 07/11] Reorganized unit tests --- tests/CMakeLists.txt | 2 + tests/static_multiset/count_if_test.cu | 230 +++++++++ tests/static_multiset/count_test.cu | 180 ------- tests/static_multiset/custom_count_if_test.cu | 443 ++++++++++++++++++ tests/static_multiset/custom_count_test.cu | 368 --------------- 5 files changed, 675 insertions(+), 548 deletions(-) create mode 100644 tests/static_multiset/count_if_test.cu create mode 100644 tests/static_multiset/custom_count_if_test.cu diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0e073f859..3a639fba4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -100,7 +100,9 @@ ConfigureTest(DYNAMIC_MAP_TEST # - static_multiset tests ------------------------------------------------------------------------- ConfigureTest(STATIC_MULTISET_TEST static_multiset/contains_test.cu + static_multiset/count_if_test.cu static_multiset/count_test.cu + static_multiset/custom_count_if_test.cu static_multiset/custom_count_test.cu static_multiset/find_test.cu static_multiset/insert_test.cu diff --git a/tests/static_multiset/count_if_test.cu b/tests/static_multiset/count_if_test.cu new file mode 100644 index 000000000..9826e36d4 --- /dev/null +++ b/tests/static_multiset/count_if_test.cu @@ -0,0 +1,230 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include + +#include +#include +#include + +#include + +using size_type = int32_t; + +template +void test_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + + set.insert(keys_begin, keys_begin + num_keys); + + SECTION("Count_if with all elements selected should match count.") + { + auto const count = + set.count_if(keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == num_keys); + } + + SECTION("Count_if with no elements selected should return zero.") + { + auto const count = + set.count_if(keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == 0); + } + + SECTION("Count_if with alternating predicate should count selected keys.") + { + auto const count = set.count_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + REQUIRE(count == (num_keys + 1) / 2); + } +} + +template +void test_count_outer_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + + set.insert(keys_begin, keys_begin + num_keys); + + SECTION("Count_outer_if with all elements selected should match count_outer.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == num_keys); + } + + SECTION("Count_outer_if with no elements selected should return one per input.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == num_keys); + } + + SECTION( + "Count_outer_if with alternating predicate should count selected matches and unselected rows.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + REQUIRE(count == num_keys); + } +} + +template +void test_count_if_stencil(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return i + 1000; })); + + set.clear(); + set.insert(keys_begin, keys_begin + num_keys); + + SECTION("Count_if should apply the predicate to the stencil.") + { + auto const count = set.count_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); + + REQUIRE(count == 100); + } + + SECTION("Count_outer_if should apply the predicate to the stencil.") + { + auto const count = set.count_outer_if( + keys_begin, + keys_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); + + REQUIRE(count == num_keys); + } +} + +template +void test_count_if_overloads(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto keys_begin = cuda::make_transform_iterator( + cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); + + auto stencil_begin = cuda::counting_iterator{0}; + + set.clear(); + set.insert(keys_begin, keys_begin + num_keys); + + auto const pred = + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }); + + SECTION("Count_if explicit default key equality/hash matches overload.") + { + auto const count_default = set.count_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); + + auto const count_explicit = set.count_if( + keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); + + REQUIRE(count_explicit == count_default); + } + + SECTION("Count_outer_if explicit default key equality/hash matches overload.") + { + auto const count_default = + set.count_outer_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); + + auto const count_explicit = set.count_outer_if( + keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); + + REQUIRE(count_explicit == count_default); + } +} + +TEMPLATE_TEST_CASE_SIG( + "static_multiset count_if tests", + "", + ((typename Key, cuco::test::probe_sequence Probe, int CGSize), Key, Probe, CGSize), + (int32_t, cuco::test::probe_sequence::double_hashing, 1), + (int32_t, cuco::test::probe_sequence::double_hashing, 2), + (int64_t, cuco::test::probe_sequence::double_hashing, 1), + (int64_t, cuco::test::probe_sequence::double_hashing, 2), + (int32_t, cuco::test::probe_sequence::linear_probing, 1), + (int32_t, cuco::test::probe_sequence::linear_probing, 2), + (int64_t, cuco::test::probe_sequence::linear_probing, 1), + (int64_t, cuco::test::probe_sequence::linear_probing, 2) +#if defined(CUCO_HAS_128BIT_ATOMICS) + , + (__int128_t, cuco::test::probe_sequence::double_hashing, 1), + (__int128_t, cuco::test::probe_sequence::double_hashing, 2), + (__int128_t, cuco::test::probe_sequence::linear_probing, 1), + (__int128_t, cuco::test::probe_sequence::linear_probing, 2) +#endif +) +{ + constexpr size_type num_keys{666}; + + using probe = std::conditional_t>, + cuco::double_hashing>>; + + auto set = + cuco::static_multiset{num_keys, cuco::empty_key{-1}, {}, probe{}, {}, cuco::storage<2>{}}; + + test_count_if(set, num_keys); + test_count_outer_if(set, num_keys); + test_count_if_stencil(set, num_keys); + test_count_if_overloads(set, num_keys); +} diff --git a/tests/static_multiset/count_test.cu b/tests/static_multiset/count_test.cu index b3b393c86..e9e4e664d 100644 --- a/tests/static_multiset/count_test.cu +++ b/tests/static_multiset/count_test.cu @@ -168,182 +168,6 @@ void test_count_each_outer(Set& set, size_type num_keys) } } -template -void test_count_if(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto keys_begin = cuda::make_transform_iterator( - cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); - - auto stencil_begin = cuda::counting_iterator{0}; - - set.clear(); - - set.insert(keys_begin, keys_begin + num_keys); - - SECTION("Count_if with all elements selected should match count.") - { - auto const count = - set.count_if(keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; })); - - REQUIRE(count == num_keys); - } - - SECTION("Count_if with no elements selected should return zero.") - { - auto const count = - set.count_if(keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; })); - - REQUIRE(count == 0); - } - - SECTION("Count_if with alternating predicate should count selected keys.") - { - auto const count = set.count_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); - - REQUIRE(count == (num_keys + 1) / 2); - } -} - -template -void test_count_outer_if(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto keys_begin = cuda::make_transform_iterator( - cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); - - auto stencil_begin = cuda::counting_iterator{0}; - - set.clear(); - - set.insert(keys_begin, keys_begin + num_keys); - SECTION("Count_outer_if with all elements selected should match count_outer.") - { - auto const count = set.count_outer_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; })); - - REQUIRE(count == num_keys); - } - - SECTION("Count_outer_if with no elements selected should return one per input.") - { - auto const count = set.count_outer_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; })); - - REQUIRE(count == num_keys); - } - - SECTION( - "Count_outer_if with alternating predicate should count selected matches and unselected rows.") - { - auto const count = set.count_outer_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); - - REQUIRE(count == num_keys); - } -} - -template -void test_count_if_stencil(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto keys_begin = cuda::make_transform_iterator( - cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); - - auto stencil_begin = cuda::make_transform_iterator( - cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(auto i) { return i + 1000; })); - - set.clear(); - set.insert(keys_begin, keys_begin + num_keys); - - SECTION("Count_if should apply the predicate to the stencil.") - { - auto const count = set.count_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); - - REQUIRE(count == 100); - } - - SECTION("Count_outer_if should apply the predicate to the stencil.") - { - auto const count = set.count_outer_if( - keys_begin, - keys_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type value) { return value < 1100; })); - - REQUIRE(count == num_keys); - } -} - -template -void test_count_if_overloads(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto keys_begin = cuda::make_transform_iterator( - cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(auto i) { return Key{i}; })); - - auto stencil_begin = cuda::counting_iterator{0}; - - set.clear(); - set.insert(keys_begin, keys_begin + num_keys); - - auto const pred = - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }); - - SECTION("Count_if explicit default key equality/hash matches overload.") - { - auto const count_default = set.count_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); - - auto const count_explicit = set.count_if( - keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); - - REQUIRE(count_explicit == count_default); - } - - SECTION("Count_outer_if explicit default key equality/hash matches overload.") - { - auto const count_default = - set.count_outer_if(keys_begin, keys_begin + num_keys, stencil_begin, pred); - - auto const count_explicit = set.count_outer_if( - keys_begin, keys_begin + num_keys, stencil_begin, pred, set.key_eq(), set.hash_function()); - - REQUIRE(count_explicit == count_default); - } -} - TEMPLATE_TEST_CASE_SIG( "static_multiset count tests", "", @@ -377,8 +201,4 @@ TEMPLATE_TEST_CASE_SIG( test_unique_sequence(set, num_keys); test_count_each(set, num_keys); test_count_each_outer(set, num_keys); - test_count_if(set, num_keys); - test_count_outer_if(set, num_keys); - test_count_if_stencil(set, num_keys); - test_count_if_overloads(set, num_keys); } diff --git a/tests/static_multiset/custom_count_if_test.cu b/tests/static_multiset/custom_count_if_test.cu new file mode 100644 index 000000000..d5fb6b82b --- /dev/null +++ b/tests/static_multiset/custom_count_if_test.cu @@ -0,0 +1,443 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include + +#include +#include + +#include + +using size_type = std::size_t; + +static auto constexpr XXX = 111; + +template +struct identity_hash { + __host__ __device__ identity_hash() {}; + __host__ __device__ identity_hash([[maybe_unused]] int i) {} + __device__ T operator()(T k) const { return k; } +}; + +struct custom_hash { + __host__ __device__ custom_hash() {} + __host__ __device__ custom_hash([[maybe_unused]] int i) {} + template + __device__ custom_type operator()(custom_type k) const + { + return k / XXX; + }; +}; + +struct custom_key_eq { + template + __device__ bool operator()(lhs_type lhs, rhs_type rhs) const + { + return lhs / XXX == rhs; + } +}; + +template +void test_count_if_duplicates(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto constexpr multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if with duplicates and all keys selected returns total multiplicity.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_if with duplicates and no keys selected returns zero.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == 0); + } + + SECTION("Count_if with duplicates counts only selected keys.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + auto const expected = ((query_size + 1) / 2) * multiplicity; + + REQUIRE(count == expected); + } + + SECTION("Count_if with duplicates counts a single selected key by its multiplicity.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); + + REQUIRE(count == multiplicity); + } +} + +template +void test_count_outer_if_duplicates(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto constexpr multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + // Query each unique key once. + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_outer_if with duplicates and all keys selected returns total multiplicity.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; })); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_outer_if with duplicates and no keys selected returns one per query.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; })); + + REQUIRE(count == query_size); + } + + SECTION( + "Count_outer_if with duplicates counts selected matches and one for each unselected query.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + auto const expected = selected_count * multiplicity + unselected_count; + + REQUIRE(count == expected); + } + + SECTION("Count_outer_if with a selected key counts its multiplicity.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); + + auto const expected = multiplicity + (query_size - 1); + + REQUIRE(count == expected); + } +} + +template +void test_custom_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + constexpr auto multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto const query_size = num_keys / multiplicity; + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if custom key equality/hash overload counts selected duplicates.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + REQUIRE(count == selected_count * multiplicity); + } + + SECTION("Count_outer_if custom key equality/hash overload counts selected duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + REQUIRE(count == selected_count * multiplicity + unselected_count); + } +} + +template +void test_custom_count_if_overloads(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + constexpr auto multiplicity = 3; + + auto duplicate_keys_begin = + cuda::make_transform_iterator(cuda::counting_iterator{0}, + cuda::proclaim_return_type([] __device__(size_type i) { + return static_cast(i / multiplicity); + })); + + set.clear(); + set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); + + auto query_begin = cuda::counting_iterator{0}; + auto query_size = num_keys / multiplicity; + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if explicit key equality/hash overload selects all duplicates.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_if explicit key equality/hash overload selects no duplicates.") + { + auto const count = + set.count_if(query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; }), + custom_key_eq{}, + hash); + + REQUIRE(count == 0); + } + + SECTION("Count_if explicit key equality/hash overload selects alternating duplicates.") + { + auto const count = set.count_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + REQUIRE(count == selected_count * multiplicity); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects all duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return true; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size * multiplicity); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects no duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type) { return false; }), + custom_key_eq{}, + hash); + + REQUIRE(count == query_size); + } + + SECTION("Count_outer_if explicit key equality/hash overload selects alternating duplicates.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + query_size, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + auto const selected_count = (query_size + 1) / 2; + auto const unselected_count = query_size / 2; + + REQUIRE(count == selected_count * multiplicity + unselected_count); + } +} + +template +void test_custom_hash_count_if(Set& set, size_type num_keys) +{ + using Key = typename Set::key_type; + + auto const hash = []() { + if constexpr (cuco::is_double_hashing::value) { + return cuda::std::tuple{custom_hash{}, custom_hash{}}; + } else { + return custom_hash{}; + } + }(); + + auto const iter = cuda::counting_iterator{0}; + set.clear(); + set.insert(iter, iter + num_keys); + + auto query_begin = cuda::make_transform_iterator( + cuda::make_counting_iterator(0), + cuda::proclaim_return_type([] __device__(auto i) { return static_cast(i * XXX); })); + + auto stencil_begin = cuda::counting_iterator{0}; + + SECTION("Count_if uses custom key equality and hash.") + { + auto const count = set.count_if( + query_begin, + query_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + REQUIRE(count == (num_keys + 1) / 2); + } + + SECTION("Count_outer_if uses custom key equality and hash.") + { + auto const count = set.count_outer_if( + query_begin, + query_begin + num_keys, + stencil_begin, + cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), + custom_key_eq{}, + hash); + + REQUIRE(count == num_keys); + } +} + +TEMPLATE_TEST_CASE_SIG( + "static_multiset custom count_if tests", + "", + ((typename Key, cuco::test::probe_sequence Probe, int CGSize), Key, Probe, CGSize), + (int32_t, cuco::test::probe_sequence::double_hashing, 1), + (int32_t, cuco::test::probe_sequence::double_hashing, 2), + (int64_t, cuco::test::probe_sequence::double_hashing, 1), + (int64_t, cuco::test::probe_sequence::double_hashing, 2), + (int32_t, cuco::test::probe_sequence::linear_probing, 1), + (int32_t, cuco::test::probe_sequence::linear_probing, 2), + (int64_t, cuco::test::probe_sequence::linear_probing, 1), + (int64_t, cuco::test::probe_sequence::linear_probing, 2) +#if defined(CUCO_HAS_128BIT_ATOMICS) + , + (__int128_t, cuco::test::probe_sequence::double_hashing, 1), + (__int128_t, cuco::test::probe_sequence::double_hashing, 2), + (__int128_t, cuco::test::probe_sequence::linear_probing, 1), + (__int128_t, cuco::test::probe_sequence::linear_probing, 2) +#endif +) +{ + constexpr size_type num_keys{555}; + + using probe = std::conditional_t>, + cuco::double_hashing>>; + + auto set = + cuco::static_multiset{num_keys, cuco::empty_key{-1}, {}, probe{}, {}, cuco::storage<2>{}}; + + test_count_if_duplicates(set, num_keys); + test_count_outer_if_duplicates(set, num_keys); + test_custom_count_if(set, num_keys); + test_custom_count_if_overloads(set, num_keys); + test_custom_hash_count_if(set, num_keys); +} diff --git a/tests/static_multiset/custom_count_test.cu b/tests/static_multiset/custom_count_test.cu index 1d6743800..6449f7022 100644 --- a/tests/static_multiset/custom_count_test.cu +++ b/tests/static_multiset/custom_count_test.cu @@ -108,369 +108,6 @@ void test_custom_count(Set& set, size_type num_keys) } } -template -void test_count_if_duplicates(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto constexpr multiplicity = 3; - - auto duplicate_keys_begin = - cuda::make_transform_iterator(cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(size_type i) { - return static_cast(i / multiplicity); - })); - - set.clear(); - set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); - - auto query_begin = cuda::counting_iterator{0}; - auto const query_size = num_keys / multiplicity; - - auto stencil_begin = cuda::counting_iterator{0}; - - SECTION("Count_if with duplicates and all keys selected returns total multiplicity.") - { - auto const count = - set.count_if(query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; })); - - REQUIRE(count == query_size * multiplicity); - } - - SECTION("Count_if with duplicates and no keys selected returns zero.") - { - auto const count = - set.count_if(query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; })); - - REQUIRE(count == 0); - } - - SECTION("Count_if with duplicates counts only selected keys.") - { - auto const count = set.count_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); - - auto const expected = ((query_size + 1) / 2) * multiplicity; - - REQUIRE(count == expected); - } - - SECTION("Count_if with duplicates counts a single selected key by its multiplicity.") - { - auto const count = - set.count_if(query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); - - REQUIRE(count == multiplicity); - } -} - -template -void test_count_outer_if_duplicates(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto constexpr multiplicity = 3; - - auto duplicate_keys_begin = - cuda::make_transform_iterator(cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(size_type i) { - return static_cast(i / multiplicity); - })); - - set.clear(); - set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); - - // Query each unique key once. - auto query_begin = cuda::counting_iterator{0}; - auto const query_size = num_keys / multiplicity; - - auto stencil_begin = cuda::counting_iterator{0}; - - SECTION("Count_outer_if with duplicates and all keys selected returns total multiplicity.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; })); - - REQUIRE(count == query_size * multiplicity); - } - - SECTION("Count_outer_if with duplicates and no keys selected returns one per query.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; })); - - REQUIRE(count == query_size); - } - - SECTION( - "Count_outer_if with duplicates counts selected matches and one for each unselected query.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; })); - - auto const selected_count = (query_size + 1) / 2; - auto const unselected_count = query_size / 2; - - auto const expected = selected_count * multiplicity + unselected_count; - - REQUIRE(count == expected); - } - - SECTION("Count_outer_if with a selected key counts its multiplicity.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return i == 0; })); - - auto const expected = multiplicity + (query_size - 1); - - REQUIRE(count == expected); - } -} - -template -void test_custom_count_if(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto const hash = []() { - if constexpr (cuco::is_double_hashing::value) { - return cuda::std::tuple{custom_hash{}, custom_hash{}}; - } else { - return custom_hash{}; - } - }(); - - constexpr auto multiplicity = 3; - - auto duplicate_keys_begin = - cuda::make_transform_iterator(cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(size_type i) { - return static_cast(i / multiplicity); - })); - - set.clear(); - set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); - - auto query_begin = cuda::counting_iterator{0}; - auto const query_size = num_keys / multiplicity; - auto stencil_begin = cuda::counting_iterator{0}; - - SECTION("Count_if custom key equality/hash overload counts selected duplicates.") - { - auto const count = set.count_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - auto const selected_count = (query_size + 1) / 2; - REQUIRE(count == selected_count * multiplicity); - } - - SECTION("Count_outer_if custom key equality/hash overload counts selected duplicates.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - auto const selected_count = (query_size + 1) / 2; - auto const unselected_count = query_size / 2; - - REQUIRE(count == selected_count * multiplicity + unselected_count); - } -} - -template -void test_custom_count_if_overloads(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto const hash = []() { - if constexpr (cuco::is_double_hashing::value) { - return cuda::std::tuple{custom_hash{}, custom_hash{}}; - } else { - return custom_hash{}; - } - }(); - - constexpr auto multiplicity = 3; - - auto duplicate_keys_begin = - cuda::make_transform_iterator(cuda::counting_iterator{0}, - cuda::proclaim_return_type([] __device__(size_type i) { - return static_cast(i / multiplicity); - })); - - set.clear(); - set.insert(duplicate_keys_begin, duplicate_keys_begin + num_keys); - - auto query_begin = cuda::counting_iterator{0}; - auto query_size = num_keys / multiplicity; - auto stencil_begin = cuda::counting_iterator{0}; - - SECTION("Count_if explicit key equality/hash overload selects all duplicates.") - { - auto const count = - set.count_if(query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; }), - custom_key_eq{}, - hash); - - REQUIRE(count == query_size * multiplicity); - } - - SECTION("Count_if explicit key equality/hash overload selects no duplicates.") - { - auto const count = - set.count_if(query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; }), - custom_key_eq{}, - hash); - - REQUIRE(count == 0); - } - - SECTION("Count_if explicit key equality/hash overload selects alternating duplicates.") - { - auto const count = set.count_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - auto const selected_count = (query_size + 1) / 2; - REQUIRE(count == selected_count * multiplicity); - } - - SECTION("Count_outer_if explicit key equality/hash overload selects all duplicates.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return true; }), - custom_key_eq{}, - hash); - - REQUIRE(count == query_size * multiplicity); - } - - SECTION("Count_outer_if explicit key equality/hash overload selects no duplicates.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type) { return false; }), - custom_key_eq{}, - hash); - - REQUIRE(count == query_size); - } - - SECTION("Count_outer_if explicit key equality/hash overload selects alternating duplicates.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + query_size, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - auto const selected_count = (query_size + 1) / 2; - auto const unselected_count = query_size / 2; - - REQUIRE(count == selected_count * multiplicity + unselected_count); - } -} - -template -void test_custom_hash_count_if(Set& set, size_type num_keys) -{ - using Key = typename Set::key_type; - - auto const hash = []() { - if constexpr (cuco::is_double_hashing::value) { - return cuda::std::tuple{custom_hash{}, custom_hash{}}; - } else { - return custom_hash{}; - } - }(); - - auto const iter = cuda::counting_iterator{0}; - set.clear(); - set.insert(iter, iter + num_keys); - - auto query_begin = cuda::make_transform_iterator( - cuda::make_counting_iterator(0), - cuda::proclaim_return_type([] __device__(auto i) { return static_cast(i * XXX); })); - - auto stencil_begin = cuda::counting_iterator{0}; - - SECTION("Count_if uses custom key equality and hash.") - { - auto const count = set.count_if( - query_begin, - query_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - REQUIRE(count == (num_keys + 1) / 2); - } - - SECTION("Count_outer_if uses custom key equality and hash.") - { - auto const count = set.count_outer_if( - query_begin, - query_begin + num_keys, - stencil_begin, - cuda::proclaim_return_type([] __device__(size_type i) { return (i % 2) == 0; }), - custom_key_eq{}, - hash); - - REQUIRE(count == num_keys); - } -} - TEMPLATE_TEST_CASE_SIG( "static_multiset custom count tests", "", @@ -502,9 +139,4 @@ TEMPLATE_TEST_CASE_SIG( cuco::static_multiset{num_keys, cuco::empty_key{-1}, {}, probe{}, {}, cuco::storage<2>{}}; test_custom_count(set, num_keys); - test_count_if_duplicates(set, num_keys); - test_count_outer_if_duplicates(set, num_keys); - test_custom_count_if(set, num_keys); - test_custom_count_if_overloads(set, num_keys); - test_custom_hash_count_if(set, num_keys); } From dffdac8b53b730798ede15f1475c90315eab3fa7 Mon Sep 17 00:00:00 2001 From: William Fan Date: Sat, 19 Sep 2026 17:06:21 -0400 Subject: [PATCH 08/11] Removed comment --- include/cuco/detail/open_addressing/kernels.cuh | 2 -- 1 file changed, 2 deletions(-) diff --git a/include/cuco/detail/open_addressing/kernels.cuh b/include/cuco/detail/open_addressing/kernels.cuh index 9d008dec0..1505041fe 100644 --- a/include/cuco/detail/open_addressing/kernels.cuh +++ b/include/cuco/detail/open_addressing/kernels.cuh @@ -632,8 +632,6 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void count_if(InputIt first, cooperative_groups::tiled_partition( cooperative_groups::this_thread_block()); - // bool const selected = pred(*(stencil + idx)); - if (pred(*(stencil + idx))) { typename cuda::std::iterator_traits::value_type const key = *(first + idx); From 2921184a1901676f319edc5ab6ff47b28d8c66fa Mon Sep 17 00:00:00 2001 From: William Fan Date: Sun, 20 Sep 2026 15:32:00 -0400 Subject: [PATCH 09/11] Fix retrieve_outer_if to include sentinel --- .../open_addressing_ref_impl.cuh | 27 ++++- .../static_multiset/retrieve_if_impl_test.cu | 109 ++++++++++++------ 2 files changed, 97 insertions(+), 39 deletions(-) diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh index a9a795538..b60950366 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -1446,8 +1446,31 @@ class open_addressing_ref_impl // onto the next probing bucket ++probing_iter; if (*probing_iter == init_idx) { running = false; } - } // while running - } // if active_flag + } + } else if constexpr (IsOuter) { + // Predicate rejected this key. It is already known to be a miss, + // so do not probe the hash table. Emit the outer sentinel directly. + if (idx < n and probing_tile.thread_rank() == 0) { + auto ref = cuda::atomic_ref{ + counters[flushing_tile_id]}; + auto const output_idx = ref.fetch_add(1, cuda::memory_order_relaxed); + probe_type const probe_key = *(input_probe + idx); + // printf("sentinel = %lld\n", + // static_cast(this->empty_slot_sentinel())); + + buffers[flushing_tile_id][output_idx] = {probe_key, this->empty_slot_sentinel()}; + } + active_flushing_tile.sync(); + // if the buffer has not enough empty slots for the next iteration + if (counters[flushing_tile_id] > (buffer_size - max_matches_per_step)) { + flush_buffers(active_flushing_tile); + active_flushing_tile.sync(); + + // reset buffer counter + if (active_flushing_tile.thread_rank() == 0) { counters[flushing_tile_id] = 0; } + active_flushing_tile.sync(); + } + } // onto the next key idx += stride; diff --git a/tests/static_multiset/retrieve_if_impl_test.cu b/tests/static_multiset/retrieve_if_impl_test.cu index 32782115e..5c1cc4b28 100644 --- a/tests/static_multiset/retrieve_if_impl_test.cu +++ b/tests/static_multiset/retrieve_if_impl_test.cu @@ -75,7 +75,7 @@ void test_outer(Container& container, std::size_t num_keys) thrust::device_vector probed_keys(query_size); thrust::device_vector matched_keys(query_size); - SECTION("Non-inserted keys should output sentinels.") + SECTION("Non-inserted keys should return the empty-key sentinel.") { auto const [probed_end, matched_end] = container.retrieve_outer(keys_begin, keys_begin + query_size, @@ -88,9 +88,12 @@ void test_outer(Container& container, std::size_t num_keys) REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == query_size); + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_end, keys_begin, cuda::std::equal_to{})); + REQUIRE(cuco::test::all_of( matched_keys.begin(), - matched_keys.end(), + matched_end, cuda::proclaim_return_type([empty_key_sentinel] __device__(auto const& k) { return static_cast(k == static_cast(empty_key_sentinel)); }))); @@ -98,7 +101,7 @@ void test_outer(Container& container, std::size_t num_keys) container.insert(keys_begin, keys_begin + num_keys); - SECTION("All inserted keys should be contained.") + SECTION("All inserted keys should be contained and missing keys should return the sentinel.") { auto const [probed_end, matched_end] = container.retrieve_outer(keys_begin, keys_begin + query_size, @@ -346,7 +349,7 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) thrust::device_vector probed_keys(num_actual_keys * multiplicity); thrust::device_vector matched_keys(num_actual_keys * multiplicity); - thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); + thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); SECTION("retrieve_if should filter duplicate matches using the stencil predicate.") { @@ -371,12 +374,14 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < expected_results; ++i) { - auto const input_index = i / multiplicity; - auto const expected_key = static_cast((input_index * 2) / multiplicity); + for (std::size_t key = 0; key < num_unique_keys; ++key) { + auto const expected_key = static_cast(key); + auto const output_offset = key * multiplicity; - REQUIRE(probed_keys[i] == expected_key); - REQUIRE(matched_keys[i] == expected_key); + for (std::size_t j = 0; j < multiplicity; ++j) { + REQUIRE(probed_keys[output_offset + j] == expected_key); + REQUIRE(matched_keys[output_offset + j] == expected_key); + } } } @@ -417,8 +422,7 @@ void test_retrieve_if_multiplicity(Container& container, std::size_t num_keys) probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); for (std::size_t key = 0; key < num_unique_keys; ++key) { - auto const expected_key = static_cast(key); - + auto const expected_key = static_cast(key); auto const expected_count = multiplicity * multiplicity; for (std::size_t j = 0; j < expected_count; ++j) { @@ -504,30 +508,34 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) probed_keys.begin(), matched_keys.begin()); - auto const expected_size = query_size / 2; - - REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == - expected_size); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == query_size); REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == - expected_size); + query_size); thrust::sort_by_key( probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); - for (std::size_t i = 0; i < expected_size; ++i) { - auto const expected_probe = static_cast(i * 2 + 1); + for (std::size_t i = 0; i < query_size; ++i) { + auto const expected_probe = static_cast(i); REQUIRE(probed_keys[i] == expected_probe); - if (expected_probe < static_cast(num_keys)) { - REQUIRE(matched_keys[i] == expected_probe); - } else { + if (i % 2 == 0) { + // stencil[i] = i + 1 is odd -> predicate false REQUIRE(matched_keys[i] == static_cast(empty_key_sentinel)); + } else { + // stencil[i] = i + 1 is even -> predicate true + if (i < num_keys) { + REQUIRE(matched_keys[i] == expected_probe); + } else { + REQUIRE(matched_keys[i] == static_cast(empty_key_sentinel)); + } } } } - SECTION("retrieve_outer_if should return nothing for an always-false predicate.") + SECTION( + "retrieve_outer_if should return every probe with the sentinel for an always-false predicate.") { thrust::sequence(stencil.begin(), stencil.end(), key_type{0}); @@ -542,8 +550,22 @@ void test_retrieve_outer_if(Container& container, std::size_t num_keys) probed_keys.begin(), matched_keys.begin()); - REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == 0); - REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 0); + REQUIRE(static_cast(std::distance(probed_keys.begin(), probed_end)) == query_size); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == + query_size); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less{}); + + REQUIRE(cuco::test::equal( + probed_keys.begin(), probed_end, probes.begin(), cuda::std::equal_to{})); + + REQUIRE(cuco::test::all_of( + matched_keys.begin(), + matched_end, + cuda::proclaim_return_type([empty_key_sentinel] __device__(auto const& k) { + return k == static_cast(empty_key_sentinel); + }))); } } @@ -572,11 +594,11 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k thrust::device_vector probes(query_size); thrust::device_vector stencil(query_size); + thrust::device_vector probed_keys(query_size * multiplicity); thrust::device_vector matched_keys(query_size * multiplicity); thrust::sequence(probes.begin(), probes.end(), key_type{0}); - thrust::sequence(stencil.begin(), stencil.end(), key_type{1}); auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; @@ -590,11 +612,16 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k probed_keys.begin(), matched_keys.begin()); - auto const num_matching_probes = query_size / 2; - auto const num_matching_unique_keys = num_unique_keys / 2; - auto const num_missing_probes = num_matching_probes - num_matching_unique_keys; + auto const num_predicate_false_probes = query_size / 2; + auto const num_predicate_true_probes = query_size / 2; - auto const expected_results = num_matching_unique_keys * multiplicity + num_missing_probes; + auto const num_matching_unique_keys = (num_unique_keys > 1) ? (num_unique_keys / 2) : 0; + + auto const num_missing_selected_probes = num_predicate_true_probes - num_matching_unique_keys; + + auto const expected_results = num_predicate_false_probes + + num_matching_unique_keys * multiplicity + + num_missing_selected_probes; auto const num_results = static_cast(std::distance(probed_keys.begin(), probed_end)); @@ -607,19 +634,27 @@ void test_retrieve_outer_if_multiplicity(Container& container, std::size_t num_k std::size_t output_index = 0; - for (std::size_t probe = 1; probe < query_size; probe += 2) { + for (std::size_t probe = 0; probe < query_size; ++probe) { auto const expected_probe = static_cast(probe); - if (probe < num_unique_keys) { - for (std::size_t j = 0; j < multiplicity; ++j) { - REQUIRE(probed_keys[output_index] == expected_probe); - REQUIRE(matched_keys[output_index] == expected_probe); - ++output_index; - } - } else { + if (probe % 2 == 0) { + // Predicate is false for even stencil values. REQUIRE(probed_keys[output_index] == expected_probe); REQUIRE(matched_keys[output_index] == static_cast(empty_key_sentinel)); ++output_index; + } else { + // Predicate is true for odd probe positions. + if (probe < num_unique_keys) { + for (std::size_t j = 0; j < multiplicity; ++j) { + REQUIRE(probed_keys[output_index] == expected_probe); + REQUIRE(matched_keys[output_index] == expected_probe); + ++output_index; + } + } else { + REQUIRE(probed_keys[output_index] == expected_probe); + REQUIRE(matched_keys[output_index] == static_cast(empty_key_sentinel)); + ++output_index; + } } } From d2d12deff14e08628294a029130d2724dd50614b Mon Sep 17 00:00:00 2001 From: William Fan Date: Sun, 20 Sep 2026 15:35:29 -0400 Subject: [PATCH 10/11] Fixed comments --- .../detail/open_addressing/open_addressing_ref_impl.cuh | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh index b60950366..cd7b4c720 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -1446,8 +1446,9 @@ class open_addressing_ref_impl // onto the next probing bucket ++probing_iter; if (*probing_iter == init_idx) { running = false; } - } - } else if constexpr (IsOuter) { + } // while running + } // if active_flag + else if constexpr (IsOuter) { // Predicate rejected this key. It is already known to be a miss, // so do not probe the hash table. Emit the outer sentinel directly. if (idx < n and probing_tile.thread_rank() == 0) { @@ -1455,8 +1456,6 @@ class open_addressing_ref_impl counters[flushing_tile_id]}; auto const output_idx = ref.fetch_add(1, cuda::memory_order_relaxed); probe_type const probe_key = *(input_probe + idx); - // printf("sentinel = %lld\n", - // static_cast(this->empty_slot_sentinel())); buffers[flushing_tile_id][output_idx] = {probe_key, this->empty_slot_sentinel()}; } From ffa9c31cee9d86833362ad26d0610e124895d5b4 Mon Sep 17 00:00:00 2001 From: William Fan Date: Sun, 20 Sep 2026 16:10:29 -0400 Subject: [PATCH 11/11] Fix documentation and unit tests --- .../open_addressing_ref_impl.cuh | 2 +- .../static_multiset/static_multiset_ref.inl | 4 +- include/cuco/static_multiset.cuh | 6 +-- .../static_multiset/retrieve_if_impl_test.cu | 49 +++++++++++++++++++ 4 files changed, 55 insertions(+), 6 deletions(-) diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh index cd7b4c720..03466f888 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -1184,7 +1184,7 @@ class open_addressing_ref_impl * respectively. The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count()` to determine the size of the output range. + * Use `count_outer_if()` to determine the size of the output range. * * If a key `k` has no matches in the container, or `pred` of the corresponding stencil is * false, then `{key, empty_slot_sentinel}` will be added to the output sequence. diff --git a/include/cuco/detail/static_multiset/static_multiset_ref.inl b/include/cuco/detail/static_multiset/static_multiset_ref.inl index a760374e8..f1117d117 100644 --- a/include/cuco/detail/static_multiset/static_multiset_ref.inl +++ b/include/cuco/detail/static_multiset/static_multiset_ref.inl @@ -654,7 +654,7 @@ class operator_impl< * The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count()` to determine the size of the output range. + * Use `count_if()` to determine the size of the output range. * * @tparam BlockSize Size of the thread block this operation is executed in * @tparam InputProbeIt Device accessible input iterator whose `value_type` is @@ -717,7 +717,7 @@ class operator_impl< * respectively. The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count_outer()` to determine the size of the output range. + * Use `count_outer_if()` to determine the size of the output range. * * If a key `k` has no matches in the container, or `pred` of its corresponding stencil is * false, then `{key, empty_slot_sentinel}` will be added to the output sequence. diff --git a/include/cuco/static_multiset.cuh b/include/cuco/static_multiset.cuh index b9f834dbb..6f6f26eb1 100644 --- a/include/cuco/static_multiset.cuh +++ b/include/cuco/static_multiset.cuh @@ -1020,7 +1020,7 @@ class static_multiset { * respectively. The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count()` to determine the size of the output range. + * Use `count_if()` to determine the size of the output range. * * This function synchronizes the given CUDA stream. * @@ -1069,7 +1069,7 @@ class static_multiset { * respectively. The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count()` to determine the size of the output range. + * Use `count_if()` to determine the size of the output range. * * This function synchronizes the given CUDA stream. * @@ -1126,7 +1126,7 @@ class static_multiset { * respectively. The output order is unspecified. * * Behavior is undefined if the size of the output range exceeds the number of retrieved slots. - * Use `count_outer()` to determine the size of the output range. + * Use `count_outer_if()` to determine the size of the output range. * * If a key `k` has no matches in the container, or `pred` of its corresponding stencil is * false, then `{key, empty_slot_sentinel}` will be added to the output sequence. diff --git a/tests/static_multiset/retrieve_if_impl_test.cu b/tests/static_multiset/retrieve_if_impl_test.cu index 5c1cc4b28..fcb799e0a 100644 --- a/tests/static_multiset/retrieve_if_impl_test.cu +++ b/tests/static_multiset/retrieve_if_impl_test.cu @@ -257,6 +257,54 @@ void test_retrieve_if(Container& container, std::size_t num_keys) } } +template +void test_retrieve_if_stencil(Container& container, std::size_t num_keys) +{ + using key_type = typename Container::key_type; + + container.clear(); + + auto const keys_begin = cuda::counting_iterator{0}; + + container.insert(keys_begin, keys_begin + num_keys); + + thrust::device_vector probes{2, 1, 4, 3}; + thrust::device_vector stencil{1, 2, 3, 4}; + thrust::device_vector probed_keys(4); + thrust::device_vector matched_keys(4); + + SECTION("retrieve_if should predicate on the stencil, not the probe.") + { + auto const pred = [] __device__(key_type key) { return key % 2 == 0; }; + + auto const [probed_end, matched_end] = container.retrieve_if(probes.begin(), + probes.end(), + stencil.begin(), + pred, + probed_keys.begin(), + matched_keys.begin()); + + auto const num_results = + static_cast(std::distance(probed_keys.begin(), probed_end)); + + // stencil: [1, 2, 3, 4] + // probe: [2, 1, 4, 3] + // + // Only stencil values 2 and 4 satisfy the predicate, corresponding + // to probes 1 and 3. + REQUIRE(num_results == 2); + REQUIRE(static_cast(std::distance(matched_keys.begin(), matched_end)) == 2); + + thrust::sort_by_key( + probed_keys.begin(), probed_end, matched_keys.begin(), cuda::std::less()); + + REQUIRE(probed_keys[0] == key_type{1}); + REQUIRE(probed_keys[1] == key_type{3}); + REQUIRE(matched_keys[0] == key_type{1}); + REQUIRE(matched_keys[1] == key_type{3}); + } +} + template void test_retrieve_if_with_probe(Container& container, std::size_t num_keys) { @@ -701,6 +749,7 @@ TEMPLATE_TEST_CASE_SIG( test_retrieve_if(set, num_keys); test_retrieve_if_with_probe(set, num_keys); + test_retrieve_if_stencil(set, num_keys); test_retrieve_if_multiplicity(set, num_keys); test_retrieve_outer_if(set, num_keys);