From 09000e93594bf5703fb9a8d1c243932634443fd9 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Fri, 13 Nov 2026 10:37:26 -0800 Subject: [PATCH] Migrate rebind APIs for fixed_capacity_set_ref --- .../__cuco/fixed_capacity_set_ref.cuh | 36 ++++++ .../test/cuco/fixed_capacity_set/test_ref.cu | 105 ++++++++++++++++++ 2 files changed, 141 insertions(+) diff --git a/cudax/include/cuda/experimental/__cuco/fixed_capacity_set_ref.cuh b/cudax/include/cuda/experimental/__cuco/fixed_capacity_set_ref.cuh index 5c96d997c53a..eac26d2a544e 100644 --- a/cudax/include/cuda/experimental/__cuco/fixed_capacity_set_ref.cuh +++ b/cudax/include/cuda/experimental/__cuco/fixed_capacity_set_ref.cuh @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -178,6 +179,41 @@ public: return __impl_.probing_scheme(); } + //! @brief Returns a copy of this ref that compares keys with `__predicate`. + //! + //! The copy references the same slot storage. + //! @tparam _NewKeyEqual Key equality predicate type + //! @param[in] __predicate Key equality predicate + //! @return A ref over the same storage using `__predicate` + template + [[nodiscard]] _CCCL_HOST_DEVICE_API auto rebind_key_eq(const _NewKeyEqual& __predicate) const + { + using __rebound_ref = fixed_capacity_set_ref<_Key, _Scope, _NewKeyEqual, _ProbingScheme, _BucketSize, _Capacity>; + return __rebound_ref{ + empty_key<_Key>{empty_key_sentinel()}, __predicate, probing_scheme(), storage_span_type{data(), capacity()}}; + } + + //! @brief Returns a copy of this ref that hashes keys with `__hash`. + //! + //! The copy references the same slot storage and keeps the probing scheme kind and CG size. + //! @tparam _NewHash Hash function type, or a tuple-like of two hash functions for double hashing + //! @param[in] __hash Hash function + //! @return A ref over the same storage using `__hash` + template + [[nodiscard]] _CCCL_HOST_DEVICE_API auto rebind_hash_function(const _NewHash& __hash) const + { + const auto __probing_scheme = probing_scheme().rebind_hash_function(__hash); + using __rebound_ref = + fixed_capacity_set_ref<_Key, + _Scope, + _KeyEqual, + ::cuda::std::decay_t, + _BucketSize, + _Capacity>; + return __rebound_ref{ + empty_key<_Key>{empty_key_sentinel()}, key_eq(), __probing_scheme, storage_span_type{data(), capacity()}}; + } + //! @brief Returns a pointer to the slot storage, including empty slots. //! @return The slot storage pointer; stored keys must not be modified directly [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr value_type* data() const noexcept diff --git a/cudax/test/cuco/fixed_capacity_set/test_ref.cu b/cudax/test/cuco/fixed_capacity_set/test_ref.cu index 6e75e05c922a..9e8fac32250b 100644 --- a/cudax/test/cuco/fixed_capacity_set/test_ref.cu +++ b/cudax/test/cuco/fixed_capacity_set/test_ref.cu @@ -7,15 +7,18 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include +#include #include #include @@ -228,3 +231,105 @@ C2H_TEST("fixed_capacity_set CUDA-trivially-copyable key", "[container][ref]", c results.data()); REQUIRE(context.all_equal(results.data(), 1, 1)); } + +struct offset_hash +{ + int offset = 0; + int seed = 0; + + [[nodiscard]] _CCCL_HOST_DEVICE_API cuda::std::uint64_t operator()(int key) const noexcept + { + // Index reduction uses the high hash bits, so spread them with a Fibonacci multiplier. + return (static_cast(key - offset) * 0x9E3779B97F4A7C15ULL) + + static_cast(seed); + } +}; + +struct offset_equal +{ + int offset; + + [[nodiscard]] _CCCL_HOST_DEVICE_API bool operator()(int probe_key, int stored_key) const noexcept + { + return probe_key - offset == stored_key; + } +}; + +// Queries `i + offset` for each `i`, which only the rebound predicate and hash map back to stored key `i`. +struct contains_offset_keys +{ + template + _CCCL_DEVICE_API void operator()(Ref ref, Hash hash, int offset, int* found) const + { + const auto group = cooperative_groups::tiled_partition(cooperative_groups::this_thread_block()); + const int index = static_cast(threadIdx.x) / Ref::cg_size; + const auto rebound = ref.rebind_key_eq(offset_equal{offset}).rebind_hash_function(hash); + const bool contains = rebound.contains(group, index + offset); + if (group.thread_rank() == 0) + { + found[index] = static_cast(contains); + } + } +}; + +using probing_kinds = c2h::type_list, int_c<1>>; + +C2H_TEST("fixed_capacity_set_ref rebind", "[container][ref][rebind]", key_types, cg_sizes, probing_kinds) +{ + using key_type = c2h::get<0, TestType>; + [[maybe_unused]] constexpr int cg_size = c2h::get<1, TestType>::value; + [[maybe_unused]] constexpr int probing_kind = c2h::get<2, TestType>::value; + using probing_type = + cuda::std::conditional_t, + cudax::cuco::double_hashing>; + using set_type = + cudax::cuco::fixed_capacity_set, + probing_type>; + + constexpr int num_keys = 32; + constexpr int offset = 64; + test_context context; + const auto make_hash = [](int hash_offset) { + if constexpr (probing_kind == 0) + { + return offset_hash{hash_offset, 0}; + } + else + { + return cuda::std::tuple{offset_hash{hash_offset, 0}, offset_hash{hash_offset, 1}}; + } + }; + set_type set{ + context.stream, + context.mr, + cuda::std::size_t{2 * num_keys}, + cudax::cuco::empty_key{static_cast(-1)}, + {}, + probing_type{make_hash(0)}}; + const auto first = cuda::counting_iterator{0}; + REQUIRE(set.insert(context.stream, first, first + num_keys) == num_keys); + + const auto hash = make_hash(offset); + const auto ref = set.ref(); + const auto rebound = ref.rebind_key_eq(offset_equal{offset}).rebind_hash_function(hash); + static_assert(cuda::std::is_same_v); + static_assert(cuda::std::is_same_v); + REQUIRE(rebound.data() == ref.data()); + REQUIRE(rebound.capacity() == ref.capacity()); + REQUIRE(rebound.empty_key_sentinel() == ref.empty_key_sentinel()); + + auto found = cuda::make_buffer(context.stream, context.mr, 2 * num_keys, -1); + cuda::launch( + context.stream, + cuda::make_config(cuda::grid_dims<1>(), cuda::block_dims<2 * num_keys * cg_size>()), + contains_offset_keys{}, + ref, + hash, + offset, + found.data()); + REQUIRE(context.matches(found.data(), 2 * num_keys, num_keys)); +}