From ef4ccd1d909844374b15efd89a5b2c16d29f3409 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20J=C3=BCnger?= Date: Fri, 11 Sep 2026 15:18:38 -0700 Subject: [PATCH 1/3] Restore aligned bucket loads Derive storage alignment from the bucket stride and expose it at native bucket-load sites. Preserve general slot access and retain demand loads where widening regresses lookups or count. Add coverage for alignment, bounds, custom storage and probing, allocator ownership, shared memory, and wraparound. --- include/cuco/bucket_storage.cuh | 19 +- .../open_addressing_ref_impl.cuh | 23 +- include/cuco/detail/storage/load_bucket.cuh | 79 +++++ include/cuco/static_map_ref.cuh | 2 +- include/cuco/static_multimap_ref.cuh | 2 +- include/cuco/static_multiset_ref.cuh | 2 +- include/cuco/static_set_ref.cuh | 2 +- tests/CMakeLists.txt | 1 + tests/static_map/shared_memory_test.cu | 3 +- tests/static_set/shared_memory_test.cu | 3 +- tests/utility/aligned_storage_test.cu | 316 ++++++++++++++++++ 11 files changed, 436 insertions(+), 16 deletions(-) create mode 100644 include/cuco/detail/storage/load_bucket.cuh create mode 100644 tests/utility/aligned_storage_test.cu diff --git a/include/cuco/bucket_storage.cuh b/include/cuco/bucket_storage.cuh index 280e69583..7e82839c5 100644 --- a/include/cuco/bucket_storage.cuh +++ b/include/cuco/bucket_storage.cuh @@ -9,9 +9,10 @@ #include #include +#include #include -#include #include +#include #include #include @@ -29,15 +30,18 @@ namespace cuco { */ template > class bucket_storage_ref { + static_assert(BucketSize > 0, "Bucket size must be positive"); + public: static constexpr int32_t bucket_size = BucketSize; ///< Number of elements per bucket - static constexpr std::size_t max_vector_load_bytes = 16; ///< Maximum vector load width in bytes + static constexpr std::size_t max_vector_load_bytes = 32; ///< Maximum vector load width in bytes using bucket_type = cuda::std::array; ///< Slot bucket type static constexpr std::size_t alignment = - cuda::std::min(cuda::std::bit_ceil(sizeof(bucket_type)), - max_vector_load_bytes); ///< Required alignment in bytes + cuda::std::max(alignof(T), + cuda::std::gcd(sizeof(T) * BucketSize, + max_vector_load_bytes)); ///< Required alignment in bytes using extent_type = Extent; ///< Storage extent type using size_type = typename extent_type::value_type; ///< Storage size type @@ -46,6 +50,9 @@ class bucket_storage_ref { /** * @brief Constructor of slot storage ref. * + * @note `slots` must be aligned to `alignment` bytes. This alignment is + * preserved at every bucket boundary, including for non-power-of-two buckets. + * * @param size Number of slots * @param slots Pointer to the slots array */ @@ -92,6 +99,10 @@ class bucket_storage_ref { /** * @brief Returns an array of slots (or a bucket) for a given index. * + * @pre The complete range `[index, index + bucket_size)` is within the storage. + * + * @note `index` need not be a multiple of `bucket_size`. + * * @param index Index of the slot * @return An array of slots */ 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..d90aaf78d 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -93,6 +94,13 @@ class open_addressing_ref_impl /// Determines if the container is a key/value or key-only store static constexpr auto has_payload = not cuda::std::is_same_v; + // First-match probing can stop within a bucket. Keep its loads incremental; + // full-bucket scans below can use the storage's wider alignment. + static constexpr auto probe_load_alignment = cuda::std::max( + alignof(storage_value_type), + cuda::std::min(std::size_t{16}, + cuda::std::bit_floor(sizeof(typename StorageRef::bucket_type) / 2))); + /// Flag indicating whether duplicate keys are allowed or not static constexpr auto allows_duplicates = AllowsDuplicates; @@ -808,7 +816,8 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = + detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); auto const state = [&]() { auto res = detail::equal_result::UNEQUAL; @@ -892,7 +901,8 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = + detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -1298,7 +1308,8 @@ class open_addressing_ref_impl while (active_flushing_tile.any(running)) { if (running) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->storage_ref_[*probing_iter]; + auto const bucket_slots = + detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); cuda::static_for([&] __device__(auto i) { equals[i()] = false; @@ -1419,7 +1430,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->storage_ref_[*probing_iter]; + auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); bool should_return = false; cuda::static_for([&] __device__(auto i) { @@ -1476,7 +1487,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->storage_ref_[*probing_iter]; + auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( @@ -1542,7 +1553,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->storage_ref_[*probing_iter]; + auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( diff --git a/include/cuco/detail/storage/load_bucket.cuh b/include/cuco/detail/storage/load_bucket.cuh new file mode 100644 index 000000000..b8898ef71 --- /dev/null +++ b/include/cuco/detail/storage/load_bucket.cuh @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +#include +#include + +namespace cuco::detail { + +/** + * @brief Identifies the built-in flat bucket storage. + * @tparam Storage Storage reference type + */ +template +inline constexpr bool is_bucket_storage_ref_v = false; + +template +/// Native bucket storage uses a fixed byte stride. +inline constexpr bool is_bucket_storage_ref_v> = true; + +/** + * @brief Identifies probing schemes whose iterators preserve bucket boundaries. + * @tparam Probe Probing scheme type + */ +template +inline constexpr bool is_bucket_aligned_probing_v = false; + +template +/// Linear probing preserves bucket boundaries. +inline constexpr bool is_bucket_aligned_probing_v> = true; + +template +/// Double hashing preserves bucket boundaries. +inline constexpr bool is_bucket_aligned_probing_v> = true; + +/** + * @brief Loads a bucket using the alignment guaranteed by native probing. + * + * Custom storage and probing schemes retain their ordinary slot-indexed access. + * The native schemes initialize, advance, and wrap in multiples of the bucket + * size, so no runtime alignment branch is necessary. + * + * @tparam MaxLoadBytes Maximum alignment to expose for this access + * @tparam Storage Storage reference type + * @tparam Probe Probing scheme type + * @param storage Slot storage + * @param index Slot index produced by the probing iterator + * @return The bucket at `index` + */ +template +[[nodiscard]] __device__ constexpr typename Storage::bucket_type load_bucket( + Storage const& storage, typename Storage::size_type index, Probe const&) noexcept +{ + static_assert(cuda::std::has_single_bit(MaxLoadBytes), "Load alignment must be a power of two"); + if constexpr (is_bucket_storage_ref_v && is_bucket_aligned_probing_v) { + assert(index % Storage::bucket_size == 0); + assert(index <= storage.capacity() && Storage::bucket_size <= storage.capacity() - index); + constexpr auto alignment = cuda::std::min(Storage::alignment, MaxLoadBytes); + if constexpr (alignment <= alignof(typename Storage::value_type)) { + return storage[index]; + } else { + auto const* ptr = __builtin_assume_aligned(storage.data() + index, alignment); + return *static_cast(ptr); + } + } else { + return storage[index]; + } +} + +} // namespace cuco::detail diff --git a/include/cuco/static_map_ref.cuh b/include/cuco/static_map_ref.cuh index 15c7bf71c..612ca470e 100644 --- a/include/cuco/static_map_ref.cuh +++ b/include/cuco/static_map_ref.cuh @@ -269,7 +269,7 @@ class static_map_ref * * @param tile The cooperative thread group used to copy the data structure * @param memory_to_use Array large enough to support `capacity` elements. Object does not take - * the ownership of the memory + * the ownership of the memory. Must satisfy the storage reference's alignment requirements. * @param scope The thread scope of the newly created device ref * * @return Copy of the current device ref diff --git a/include/cuco/static_multimap_ref.cuh b/include/cuco/static_multimap_ref.cuh index f8e65ee3a..0a4968297 100644 --- a/include/cuco/static_multimap_ref.cuh +++ b/include/cuco/static_multimap_ref.cuh @@ -271,7 +271,7 @@ class static_multimap_ref * * @param tile The cooperative thread group used to copy the data structure * @param memory_to_use Array large enough to support `capacity` elements. Object does not take - * the ownership of the memory + * the ownership of the memory. Must satisfy the storage reference's alignment requirements. * @param scope The thread scope of the newly created device ref * * @return Copy of the current device ref diff --git a/include/cuco/static_multiset_ref.cuh b/include/cuco/static_multiset_ref.cuh index 9becb5324..557c09df0 100644 --- a/include/cuco/static_multiset_ref.cuh +++ b/include/cuco/static_multiset_ref.cuh @@ -255,7 +255,7 @@ class static_multiset_ref * * @param tile The cooperative thread group used to copy the data structure * @param memory_to_use Array large enough to support `capacity` elements. Object does not take - * the ownership of the memory + * the ownership of the memory. Must satisfy the storage reference's alignment requirements. * @param scope The thread scope of the newly created device ref * * @return Copy of the current device ref diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh index 84d15163f..77a40cc48 100644 --- a/include/cuco/static_set_ref.cuh +++ b/include/cuco/static_set_ref.cuh @@ -253,7 +253,7 @@ class static_set_ref * * @param tile The cooperative thread group used to copy the data structure * @param memory_to_use Array large enough to support `capacity` elements. Object does not take - * the ownership of the memory + * the ownership of the memory. Must satisfy the storage reference's alignment requirements. * @param scope The thread scope of the newly created device ref * * @return Copy of the current device ref diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index ae98b1216..29a2cc9b4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -40,6 +40,7 @@ endfunction(ConfigureTest) ################################################################################################### # - utility tests --------------------------------------------------------------------------------- ConfigureTest(UTILITY_TEST + utility/aligned_storage_test.cu utility/extent_test.cu utility/next_prime_test.cu utility/storage_test.cu diff --git a/tests/static_map/shared_memory_test.cu b/tests/static_map/shared_memory_test.cu index 28da53012..7a1574569 100644 --- a/tests/static_map/shared_memory_test.cu +++ b/tests/static_map/shared_memory_test.cu @@ -32,7 +32,8 @@ __global__ void shared_memory_test_kernel(Ref* maps, size_t const map_id = blockIdx.x; size_t const offset = map_id * number_of_elements; - __shared__ typename Ref::value_type sm_buffer[ValidSize]; + using storage_ref_type = typename Ref::storage_ref_type; + alignas(storage_ref_type::alignment) __shared__ typename Ref::value_type sm_buffer[ValidSize]; auto g = cuco::test::cg::this_thread_block(); auto insert_ref = maps[map_id].make_copy(g, sm_buffer, cuco::thread_scope_block); diff --git a/tests/static_set/shared_memory_test.cu b/tests/static_set/shared_memory_test.cu index 3e8066e8e..3e519e3b7 100644 --- a/tests/static_set/shared_memory_test.cu +++ b/tests/static_set/shared_memory_test.cu @@ -31,7 +31,8 @@ __global__ void shared_memory_test_kernel(Ref* sets, size_t const set_id = blockIdx.x; size_t const offset = set_id * number_of_elements; - __shared__ typename Ref::value_type sm_buffer[ValidSize]; + using storage_ref_type = typename Ref::storage_ref_type; + alignas(storage_ref_type::alignment) __shared__ typename Ref::value_type sm_buffer[ValidSize]; auto g = cuco::test::cg::this_thread_block(); auto insert_ref = sets[set_id].make_copy(g, sm_buffer, cuco::thread_scope_block); diff --git a/tests/utility/aligned_storage_test.cu b/tests/utility/aligned_storage_test.cu new file mode 100644 index 000000000..7e84d3344 --- /dev/null +++ b/tests/utility/aligned_storage_test.cu @@ -0,0 +1,316 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 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 +#include +#include + +namespace { + +struct slot12 { + std::uint32_t words[3]; +}; +struct slot24 { + std::uint64_t words[3]; +}; +struct alignas(64) slot64 { + std::uint32_t words[16]; +}; + +template +__device__ T slot_value(std::size_t index) +{ + cuda::std::array words{}; + for (std::size_t i = 0; i < words.size(); ++i) { + words[i] = static_cast(index * 17 + i + 1); + } + return cuda::std::bit_cast(words); +} + +template +__device__ bool same_value(T const& value, std::size_t index) +{ + using words = cuda::std::array; + return cuda::std::bit_cast(value) == cuda::std::bit_cast(slot_value(index)); +} + +struct custom_probe {}; + +struct absolute_equal { + __host__ __device__ bool operator()(std::int32_t a, std::int32_t b) const + { + return (a < 0 ? -a : a) == (b < 0 ? -b : b); + } +}; +struct absolute_hash { + __host__ __device__ std::uint32_t operator()(std::int32_t value) const + { + return static_cast(value < 0 ? -value : value); + } +}; + +template +struct shifted_storage : Ref { + __device__ explicit shifted_storage(Ref const& ref) : Ref{ref} {} + + __device__ typename Ref::bucket_type operator[](typename Ref::size_type index) const + { + return Ref::operator[](index + 1); + } +}; + +template +__device__ void check_reads(Ref ref, unsigned* errors) +{ + using value = typename Ref::value_type; + using probe = cuco::linear_probing<1, cuco::identity_hash>; + constexpr auto bucket = Ref::bucket_size; + auto const n = ref.capacity(); + for (std::size_t i = threadIdx.x; i < n; i += blockDim.x) { + ref.data()[i] = slot_value(i); + } + __syncthreads(); + + unsigned wrong{}; + for (std::size_t index = threadIdx.x * bucket; index + bucket <= n; + index += blockDim.x * bucket) { + auto const values = cuco::detail::load_bucket(ref, index, probe{}); + for (int i = 0; i < bucket; ++i) { + wrong += !same_value(values[i], index + i); + } + } + for (std::size_t index = threadIdx.x; index + bucket <= n; index += blockDim.x) { + auto const values = ref[index]; + auto const custom = cuco::detail::load_bucket(ref, index, custom_probe{}); + for (int i = 0; i < bucket; ++i) { + wrong += !same_value(values[i], index + i); + wrong += !same_value(custom[i], index + i); + } + } + auto const shifted = + cuco::detail::load_bucket(shifted_storage{ref}, std::size_t{0}, probe{}); + for (int i = 0; i < bucket; ++i) { + wrong += !same_value(shifted[i], i + 1); + } + if (wrong) { atomicAdd(errors, wrong); } +} + +template +__global__ void check_global(Ref ref, unsigned* errors) +{ + check_reads(ref, errors); +} + +template +__global__ void check_shared(unsigned* errors) +{ + using ref_type = cuco::bucket_storage_ref; + alignas(ref_type::alignment) __shared__ T values[N]; + check_reads(ref_type{cuco::extent{N}, values}, errors); +} + +struct allocation_record { + void* raw; + void* returned; + std::size_t count; + cudaStream_t stream; + bool freed{}; +}; +struct allocation_state { + std::vector records; + bool correct = true; +}; + +template +struct offset_allocator { + using value_type = T; + std::shared_ptr state; + + explicit offset_allocator(std::shared_ptr state_) : state{std::move(state_)} {} + template + offset_allocator(offset_allocator const& other) : state{other.state} + { + } + + T* allocate(std::size_t count, cuda::stream_ref stream) + { + void* raw{}; + CUCO_CUDA_TRY(cudaMallocAsync(&raw, count * sizeof(T) + alignof(T), stream.get())); + auto* result = reinterpret_cast(static_cast(raw) + alignof(T)); + state->records.push_back({raw, result, count, stream.get()}); + return result; + } + void deallocate(T* ptr, std::size_t count, cuda::stream_ref stream) + { + auto it = std::find_if(state->records.begin(), state->records.end(), [ptr](auto const& record) { + return record.returned == ptr; + }); + if (it == state->records.end()) { + state->correct = false; + it = std::find_if(state->records.begin(), state->records.end(), [](auto const& record) { + return !record.freed; + }); + } + if (it != state->records.end()) { + state->correct &= !it->freed && it->count == count && it->stream == stream.get(); + CUCO_CUDA_TRY(cudaFreeAsync(it->raw, stream.get())); + it->freed = true; + } + } +}; + +} // namespace + +TEMPLATE_TEST_CASE_SIG("aligned bucket loads and general slot access", + "", + ((typename T, int B, std::size_t A), T, B, A), + (std::int32_t, 1, 4), + (std::int32_t, 3, 4), + (std::int32_t, 4, 16), + (std::int32_t, 5, 4), + (std::int32_t, 8, 32), + (std::int64_t, 2, 16), + (std::int64_t, 3, 8), + (std::int64_t, 4, 32), + (std::int64_t, 5, 8), + (std::int64_t, 8, 32), + (cuco::pair, 3, 16), + (cuco::pair, 4, 32), + (slot12, 1, 4), + (slot24, 1, 8), + (slot64, 1, 64)) +{ + using ref_type = cuco::bucket_storage_ref; + STATIC_REQUIRE(ref_type::alignment == A); + constexpr std::size_t n = 17 * B + 5; + thrust::device_vector errors(1, 0); + + SECTION("Borrowed global storage has exact trailing bounds and only the required alignment.") + { + void* raw{}; + CUCO_CUDA_TRY(cudaMalloc(&raw, n * sizeof(T) + A)); + auto* slots = reinterpret_cast(static_cast(raw) + A); + REQUIRE(reinterpret_cast(slots) % A == 0); + CUCO_CUDA_TRY(cudaMemset(raw, 0xa5, A)); + check_global<<<1, 128>>>(ref_type{cuco::extent{n}, slots}, + thrust::raw_pointer_cast(errors.data())); + CUCO_CUDA_TRY(cudaDeviceSynchronize()); + std::vector prefix(A); + CUCO_CUDA_TRY(cudaMemcpy(prefix.data(), raw, A, cudaMemcpyDeviceToHost)); + CUCO_CUDA_TRY(cudaFree(raw)); + REQUIRE(std::all_of(prefix.begin(), prefix.end(), [](auto byte) { return byte == 0xa5; })); + REQUIRE(errors[0] == 0); + } + + SECTION("Shared storage follows the same load contract.") + { + check_shared<<<1, 128>>>(thrust::raw_pointer_cast(errors.data())); + CUCO_CUDA_TRY(cudaDeviceSynchronize()); + REQUIRE(errors[0] == 0); + } +} + +TEST_CASE("bucket storage realignment preserves allocator ownership and stream", "") +{ + auto state = std::make_shared(); + auto other_state = std::make_shared(); + cudaStream_t stream{}; + CUCO_CUDA_TRY(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + { + using storage = + cuco::bucket_storage, offset_allocator>; + storage first{cuco::extent{128}, offset_allocator{state}, {stream}}; + auto* original = first.data(); + REQUIRE(reinterpret_cast(original) % 32 == 0); + REQUIRE(static_cast(original) != state->records[0].returned); + first.initialize(17, {stream}); + storage moved{std::move(first)}; + REQUIRE(moved.data() == original); + storage assigned{cuco::extent{64}, offset_allocator{other_state}, {stream}}; + assigned = std::move(moved); + REQUIRE(assigned.data() == original); + REQUIRE(assigned.capacity() == 128); + } + CUCO_CUDA_TRY(cudaStreamSynchronize(stream)); + CUCO_CUDA_TRY(cudaStreamDestroy(stream)); + REQUIRE(state->correct); + REQUIRE(other_state->correct); + REQUIRE(state->records.size() == 1); + REQUIRE(other_state->records.size() == 1); + REQUIRE(std::all_of( + state->records.begin(), state->records.end(), [](auto const& record) { return record.freed; })); + REQUIRE(other_state->records[0].freed); +} + +TEMPLATE_TEST_CASE_SIG("odd buckets preserve aligned probing through a full table", + "", + ((typename Key, int CG, int B), Key, CG, B), + (std::int32_t, 2, 3), + (std::int64_t, 1, 5)) +{ + using probe = cuco::double_hashing>; + using set_type = cuco::static_set, + cuda::thread_scope_device, + cuda::std::equal_to, + probe, + cuco::cuda_allocator, + cuco::storage>; + constexpr std::size_t n = 7 * CG * B; + set_type set{n, cuco::empty_key{-1}}; + REQUIRE(set.capacity() == n); + thrust::device_vector keys(n + 7); + thrust::sequence(keys.begin(), keys.end()); + set.insert(keys.begin(), keys.begin() + n); + thrust::device_vector contains(n + 7); + thrust::device_vector found(n + 7); + set.contains(keys.begin(), keys.end(), contains.begin()); + set.find(keys.begin(), keys.end(), found.begin()); + for (std::size_t i = 0; i < n + 7; ++i) { + REQUIRE(contains[i] == (i < n)); + REQUIRE(found[i] == (i < n ? static_cast(i) : Key{-1})); + } +} + +TEMPLATE_TEST_CASE_SIG( + "aligned bucket comparisons preserve custom equality and stored keys", "", ((int CG), CG), 1, 2) +{ + using key = std::int32_t; + using set_type = cuco::static_set, + cuda::thread_scope_device, + absolute_equal, + cuco::linear_probing, + cuco::cuda_allocator, + cuco::storage<8>>; + set_type set{64, cuco::empty_key{-999}}; + thrust::device_vector stored(16); + thrust::sequence(stored.begin(), stored.end(), key{-2}, key{-1}); + set.insert(stored.begin(), stored.end()); + thrust::device_vector queries(17); + thrust::sequence(queries.begin(), queries.end(), key{2}); + thrust::device_vector contains(17); + thrust::device_vector found(17); + set.contains(queries.begin(), queries.end(), contains.begin()); + set.find(queries.begin(), queries.end(), found.begin()); + for (int i = 0; i < 17; ++i) { + REQUIRE(contains[i] == (i < 16)); + REQUIRE(found[i] == (i < 16 ? -i - 2 : -999)); + } +} From 7b3b69bb0c48c7b5b15dc88a55ce926d8858df67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20J=C3=BCnger?= Date: Fri, 11 Sep 2026 16:27:10 -0700 Subject: [PATCH 2/3] Move bucket load policies into storage Implement aligned bucket loads and compile-time width selection on bucket_storage_ref. Let open addressing select the access policy without passing byte counts, and remove the standalone load_bucket header. Require bucket-aligned indices from all probing schemes while preserving custom storage access. Cover both load policies and aligned custom probing without changing native kernel codegen. --- include/cuco/bucket_storage.cuh | 30 +++++++ .../open_addressing_ref_impl.cuh | 41 +++++++--- .../probing_scheme/probing_scheme_base.cuh | 3 + .../cuco/detail/storage/bucket_storage.inl | 20 +++++ include/cuco/detail/storage/load_bucket.cuh | 79 ------------------- tests/utility/aligned_storage_test.cu | 79 ++++++++++++++++--- 6 files changed, 147 insertions(+), 105 deletions(-) delete mode 100644 include/cuco/detail/storage/load_bucket.cuh diff --git a/include/cuco/bucket_storage.cuh b/include/cuco/bucket_storage.cuh index 7e82839c5..38e18c114 100644 --- a/include/cuco/bucket_storage.cuh +++ b/include/cuco/bucket_storage.cuh @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -21,6 +22,16 @@ namespace cuco { +/** + * @brief Selects the alignment exposed when loading a bucket. + * + * Both policies return the complete bucket; the compiler chooses the load instructions. + */ +enum class bucket_load_policy { + FULL, ///< Allow the full bucket alignment for wide loads. + FIRST_MATCH ///< Favor incremental loads when consumption can stop within the bucket. +}; + /** * @brief Non-owning array of slots storage reference type. * @@ -108,6 +119,21 @@ class bucket_storage_ref { */ [[nodiscard]] __device__ constexpr bucket_type operator[](size_type index) const noexcept; + /** + * @brief Loads a bucket starting at a bucket-aligned slot index. + * + * Unlike `operator[]`, this access exposes the guaranteed bucket alignment to the compiler. + * + * @pre `index` is a multiple of `bucket_size`. + * @pre The complete range `[index, index + bucket_size)` is within the storage. + * + * @tparam Policy Bucket load policy + * @param index Index of the first slot in the bucket + * @return An array containing the bucket's slots + */ + template + [[nodiscard]] __device__ constexpr bucket_type load_bucket(size_type index) const noexcept; + /** * @brief Gets the total number of slot buckets in the current storage. * @@ -130,6 +156,10 @@ class bucket_storage_ref { [[nodiscard]] __host__ __device__ constexpr extent_type extent() const noexcept; private: + // Keep first-match loads incremental; full scans can use the wider storage alignment. + static constexpr auto first_match_load_alignment = cuda::std::max( + alignof(T), cuda::std::min(std::size_t{16}, cuda::std::bit_floor(sizeof(bucket_type) / 2))); + extent_type extent_; ///< Storage extent value_type* slots_; ///< Pointer to the slots array }; 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 d90aaf78d..2ba08ae13 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -5,10 +5,10 @@ #pragma once +#include #include #include #include -#include #include #include #include @@ -94,12 +94,27 @@ class open_addressing_ref_impl /// Determines if the container is a key/value or key-only store static constexpr auto has_payload = not cuda::std::is_same_v; - // First-match probing can stop within a bucket. Keep its loads incremental; - // full-bucket scans below can use the storage's wider alignment. - static constexpr auto probe_load_alignment = cuda::std::max( - alignof(storage_value_type), - cuda::std::min(std::size_t{16}, - cuda::std::bit_floor(sizeof(typename StorageRef::bucket_type) / 2))); + /** + * @brief Selects aligned bucket access for native storage and preserves custom slot access. + * + * Probing schemes must produce bucket-aligned slot indices. + * @tparam Policy Bucket load policy + * @param index Slot index produced by the probing iterator + * @return The bucket at `index` + */ + template + __device__ typename StorageRef::bucket_type load_bucket( + typename StorageRef::size_type index) const noexcept + { + using native_storage_ref = bucket_storage_ref; + if constexpr (cuda::std::is_same_v) { + return storage_ref_.template load_bucket(index); + } else { + return storage_ref_[index]; + } + } /// Flag indicating whether duplicate keys are allowed or not static constexpr auto allows_duplicates = AllowsDuplicates; @@ -817,7 +832,7 @@ class open_addressing_ref_impl while (true) { auto const bucket_slots = - detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + this->template load_bucket(*probing_iter); auto const state = [&]() { auto res = detail::equal_result::UNEQUAL; @@ -902,7 +917,7 @@ class open_addressing_ref_impl while (true) { auto const bucket_slots = - detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + this->template load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -1309,7 +1324,7 @@ class open_addressing_ref_impl if (running) { // TODO atomic_ref::load if insert operator is present auto const bucket_slots = - detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + this->template load_bucket(*probing_iter); cuda::static_for([&] __device__(auto i) { equals[i()] = false; @@ -1430,7 +1445,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + auto const bucket_slots = this->template load_bucket(*probing_iter); bool should_return = false; cuda::static_for([&] __device__(auto i) { @@ -1487,7 +1502,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + auto const bucket_slots = this->template load_bucket(*probing_iter); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( @@ -1553,7 +1568,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = detail::load_bucket(storage_ref_, *probing_iter, probing_scheme_); + auto const bucket_slots = this->template load_bucket(*probing_iter); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( diff --git a/include/cuco/detail/probing_scheme/probing_scheme_base.cuh b/include/cuco/detail/probing_scheme/probing_scheme_base.cuh index 77bc2e8d3..a5ce7ee02 100644 --- a/include/cuco/detail/probing_scheme/probing_scheme_base.cuh +++ b/include/cuco/detail/probing_scheme/probing_scheme_base.cuh @@ -15,6 +15,9 @@ namespace detail { * * This class should not be used directly. * + * @note Derived probing schemes must produce bucket-aligned slot indices. For a bucket size + * `B`, every index must be a multiple of `B` and the complete bucket must fit within the storage. + * * @tparam CGSize Size of CUDA Cooperative Groups */ template diff --git a/include/cuco/detail/storage/bucket_storage.inl b/include/cuco/detail/storage/bucket_storage.inl index cfd549e2b..4814c97ce 100644 --- a/include/cuco/detail/storage/bucket_storage.inl +++ b/include/cuco/detail/storage/bucket_storage.inl @@ -65,6 +65,26 @@ bucket_storage_ref::operator[](size_type index) const noe return *reinterpret_cast(this->data() + index); } +template +template +__device__ constexpr bucket_storage_ref::bucket_type +bucket_storage_ref::load_bucket(size_type index) const noexcept +{ + static_assert(Policy == bucket_load_policy::FULL || Policy == bucket_load_policy::FIRST_MATCH, + "Unsupported bucket load policy"); + assert(index % bucket_size == 0); + assert(index <= capacity() && bucket_size <= capacity() - index); + constexpr auto load_alignment = cuda::std::min( + alignment, + Policy == bucket_load_policy::FIRST_MATCH ? first_match_load_alignment : max_vector_load_bytes); + if constexpr (load_alignment <= alignof(value_type)) { + return (*this)[index]; + } else { + auto const* ptr = __builtin_assume_aligned(this->data() + index, load_alignment); + return *static_cast(ptr); + } +} + template __host__ __device__ constexpr typename bucket_storage_ref::size_type bucket_storage_ref::num_buckets() const noexcept diff --git a/include/cuco/detail/storage/load_bucket.cuh b/include/cuco/detail/storage/load_bucket.cuh deleted file mode 100644 index b8898ef71..000000000 --- a/include/cuco/detail/storage/load_bucket.cuh +++ /dev/null @@ -1,79 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include - -#include -#include - -#include -#include - -namespace cuco::detail { - -/** - * @brief Identifies the built-in flat bucket storage. - * @tparam Storage Storage reference type - */ -template -inline constexpr bool is_bucket_storage_ref_v = false; - -template -/// Native bucket storage uses a fixed byte stride. -inline constexpr bool is_bucket_storage_ref_v> = true; - -/** - * @brief Identifies probing schemes whose iterators preserve bucket boundaries. - * @tparam Probe Probing scheme type - */ -template -inline constexpr bool is_bucket_aligned_probing_v = false; - -template -/// Linear probing preserves bucket boundaries. -inline constexpr bool is_bucket_aligned_probing_v> = true; - -template -/// Double hashing preserves bucket boundaries. -inline constexpr bool is_bucket_aligned_probing_v> = true; - -/** - * @brief Loads a bucket using the alignment guaranteed by native probing. - * - * Custom storage and probing schemes retain their ordinary slot-indexed access. - * The native schemes initialize, advance, and wrap in multiples of the bucket - * size, so no runtime alignment branch is necessary. - * - * @tparam MaxLoadBytes Maximum alignment to expose for this access - * @tparam Storage Storage reference type - * @tparam Probe Probing scheme type - * @param storage Slot storage - * @param index Slot index produced by the probing iterator - * @return The bucket at `index` - */ -template -[[nodiscard]] __device__ constexpr typename Storage::bucket_type load_bucket( - Storage const& storage, typename Storage::size_type index, Probe const&) noexcept -{ - static_assert(cuda::std::has_single_bit(MaxLoadBytes), "Load alignment must be a power of two"); - if constexpr (is_bucket_storage_ref_v && is_bucket_aligned_probing_v) { - assert(index % Storage::bucket_size == 0); - assert(index <= storage.capacity() && Storage::bucket_size <= storage.capacity() - index); - constexpr auto alignment = cuda::std::min(Storage::alignment, MaxLoadBytes); - if constexpr (alignment <= alignof(typename Storage::value_type)) { - return storage[index]; - } else { - auto const* ptr = __builtin_assume_aligned(storage.data() + index, alignment); - return *static_cast(ptr); - } - } else { - return storage[index]; - } -} - -} // namespace cuco::detail diff --git a/tests/utility/aligned_storage_test.cu b/tests/utility/aligned_storage_test.cu index 7e84d3344..fd96922be 100644 --- a/tests/utility/aligned_storage_test.cu +++ b/tests/utility/aligned_storage_test.cu @@ -4,7 +4,6 @@ */ #include -#include #include #include #include @@ -14,6 +13,8 @@ #include #include +#include + #include #include @@ -50,8 +51,6 @@ __device__ bool same_value(T const& value, std::size_t index) return cuda::std::bit_cast(value) == cuda::std::bit_cast(slot_value(index)); } -struct custom_probe {}; - struct absolute_equal { __host__ __device__ bool operator()(std::int32_t a, std::int32_t b) const { @@ -67,7 +66,7 @@ struct absolute_hash { template struct shifted_storage : Ref { - __device__ explicit shifted_storage(Ref const& ref) : Ref{ref} {} + __host__ __device__ explicit constexpr shifted_storage(Ref const& ref) : Ref{ref} {} __device__ typename Ref::bucket_type operator[](typename Ref::size_type index) const { @@ -75,11 +74,47 @@ struct shifted_storage : Ref { } }; +struct zero_hash { + __host__ __device__ std::uint32_t operator()(int) const noexcept { return 0; } +}; + +using native_probe = cuco::linear_probing<2, zero_hash>; + +struct custom_probe : native_probe { + template + __host__ __device__ auto make_iterator(cooperative_groups::thread_block_tile<2, ParentCG> group, + ProbeKey key, + Extent capacity) const noexcept + { + auto iter = native_probe::template make_iterator(group, key, capacity); + ++iter; + return iter; + } +}; + +template +__global__ void check_fallback_loads(Ref storage, unsigned* errors, unsigned* matches) +{ + using storage_type = cuda::std::conditional_t, Ref>; + using probe_type = cuda::std::conditional_t; + auto const group = + cooperative_groups::tiled_partition<2>(cooperative_groups::this_thread_block()); + auto ref = cuco::static_set_ref{cuco::empty_key{-1}, + cuda::std::equal_to{}, + probe_type{}, + cuco::thread_scope_device, + storage_type{storage}}; + if (!ref.rebind_operators(cuco::contains).contains(group, 0)) { atomicAdd(errors, 1u); } + ref.rebind_operators(cuco::for_each).for_each(group, 0, [=] __device__(int value) { + if (value != 0) { atomicAdd(errors, 1u); } + atomicAdd(matches, 1u); + }); +} + template __device__ void check_reads(Ref ref, unsigned* errors) { using value = typename Ref::value_type; - using probe = cuco::linear_probing<1, cuco::identity_hash>; constexpr auto bucket = Ref::bucket_size; auto const n = ref.capacity(); for (std::size_t i = threadIdx.x; i < n; i += blockDim.x) { @@ -90,24 +125,19 @@ __device__ void check_reads(Ref ref, unsigned* errors) unsigned wrong{}; for (std::size_t index = threadIdx.x * bucket; index + bucket <= n; index += blockDim.x * bucket) { - auto const values = cuco::detail::load_bucket(ref, index, probe{}); + auto const values = ref.load_bucket(index); + auto const first_match = ref.template load_bucket(index); for (int i = 0; i < bucket; ++i) { wrong += !same_value(values[i], index + i); + wrong += !same_value(first_match[i], index + i); } } for (std::size_t index = threadIdx.x; index + bucket <= n; index += blockDim.x) { auto const values = ref[index]; - auto const custom = cuco::detail::load_bucket(ref, index, custom_probe{}); for (int i = 0; i < bucket; ++i) { wrong += !same_value(values[i], index + i); - wrong += !same_value(custom[i], index + i); } } - auto const shifted = - cuco::detail::load_bucket(shifted_storage{ref}, std::size_t{0}, probe{}); - for (int i = 0; i < bucket; ++i) { - wrong += !same_value(shifted[i], i + 1); - } if (wrong) { atomicAdd(errors, wrong); } } @@ -226,6 +256,29 @@ TEMPLATE_TEST_CASE_SIG("aligned bucket loads and general slot access", } } +TEST_CASE("aligned bucket access preserves custom storage and probing", "") +{ + cuco::bucket_storage storage{cuco::extent{80}, cuco::cuda_allocator{}}; + storage.initialize(-1); + int const key = 0; + thrust::device_vector result(2, 0); + auto* errors = thrust::raw_pointer_cast(result.data()); + auto* matches = errors + 1; + SECTION("A derived storage ref must retain its overridden slot access.") + { + CUCO_CUDA_TRY(cudaMemcpy(storage.data() + 1, &key, sizeof(key), cudaMemcpyHostToDevice)); + check_fallback_loads<<<1, 2>>>(storage.ref(), errors, matches); + } + SECTION("A custom probe can use the aligned storage load policies.") + { + CUCO_CUDA_TRY(cudaMemcpy(storage.data() + 16, &key, sizeof(key), cudaMemcpyHostToDevice)); + check_fallback_loads<<<1, 2>>>(storage.ref(), errors, matches); + } + CUCO_CUDA_TRY(cudaDeviceSynchronize()); + REQUIRE(result[0] == 0); + REQUIRE(result[1] == 1); +} + TEST_CASE("bucket storage realignment preserves allocator ownership and stream", "") { auto state = std::make_shared(); From b30e9ec4112e3299dc8b54e082d9d11dde344d9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20J=C3=BCnger?= Date: Fri, 11 Sep 2026 17:25:00 -0700 Subject: [PATCH 3/3] Use aligned whole-bucket loads for all probing Remove FIRST_MATCH and the operation-specific load policies. Restore a single storage-owned alignment hint following the pre-slot-indexing implementation, and route scalar, cooperative, count, and mutation probing through the same whole-bucket accessor. --- include/cuco/bucket_storage.cuh | 17 ----- .../open_addressing_ref_impl.cuh | 75 +++++++++---------- .../cuco/detail/static_map/static_map_ref.inl | 10 +-- .../cuco/detail/storage/bucket_storage.inl | 13 +--- tests/utility/aligned_storage_test.cu | 6 +- 5 files changed, 42 insertions(+), 79 deletions(-) diff --git a/include/cuco/bucket_storage.cuh b/include/cuco/bucket_storage.cuh index 38e18c114..c186dbf50 100644 --- a/include/cuco/bucket_storage.cuh +++ b/include/cuco/bucket_storage.cuh @@ -11,7 +11,6 @@ #include #include -#include #include #include #include @@ -22,16 +21,6 @@ namespace cuco { -/** - * @brief Selects the alignment exposed when loading a bucket. - * - * Both policies return the complete bucket; the compiler chooses the load instructions. - */ -enum class bucket_load_policy { - FULL, ///< Allow the full bucket alignment for wide loads. - FIRST_MATCH ///< Favor incremental loads when consumption can stop within the bucket. -}; - /** * @brief Non-owning array of slots storage reference type. * @@ -127,11 +116,9 @@ class bucket_storage_ref { * @pre `index` is a multiple of `bucket_size`. * @pre The complete range `[index, index + bucket_size)` is within the storage. * - * @tparam Policy Bucket load policy * @param index Index of the first slot in the bucket * @return An array containing the bucket's slots */ - template [[nodiscard]] __device__ constexpr bucket_type load_bucket(size_type index) const noexcept; /** @@ -156,10 +143,6 @@ class bucket_storage_ref { [[nodiscard]] __host__ __device__ constexpr extent_type extent() const noexcept; private: - // Keep first-match loads incremental; full scans can use the wider storage alignment. - static constexpr auto first_match_load_alignment = cuda::std::max( - alignof(T), cuda::std::min(std::size_t{16}, cuda::std::bit_floor(sizeof(bucket_type) / 2))); - extent_type extent_; ///< Storage extent value_type* slots_; ///< Pointer to the slots array }; 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 2ba08ae13..285c23b93 100644 --- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh +++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh @@ -94,28 +94,6 @@ class open_addressing_ref_impl /// Determines if the container is a key/value or key-only store static constexpr auto has_payload = not cuda::std::is_same_v; - /** - * @brief Selects aligned bucket access for native storage and preserves custom slot access. - * - * Probing schemes must produce bucket-aligned slot indices. - * @tparam Policy Bucket load policy - * @param index Slot index produced by the probing iterator - * @return The bucket at `index` - */ - template - __device__ typename StorageRef::bucket_type load_bucket( - typename StorageRef::size_type index) const noexcept - { - using native_storage_ref = bucket_storage_ref; - if constexpr (cuda::std::is_same_v) { - return storage_ref_.template load_bucket(index); - } else { - return storage_ref_[index]; - } - } - /// Flag indicating whether duplicate keys are allowed or not static constexpr auto allows_duplicates = AllowsDuplicates; @@ -141,6 +119,24 @@ class open_addressing_ref_impl storage_ref_type::bucket_size; ///< Number of elements handled per bucket static constexpr auto thread_scope = Scope; ///< CUDA thread scope + /** + * @brief Loads the complete bucket at the probing iterator's slot index. + * + * Probing schemes must produce bucket-aligned slot indices. + * + * @param index Slot index produced by the probing iterator + * @return The bucket at `index` + */ + [[nodiscard]] __device__ bucket_type load_bucket(size_type index) const noexcept + { + using native_storage_ref = bucket_storage_ref; + if constexpr (cuda::std::is_same_v) { + return storage_ref_.load_bucket(index); + } else { + return storage_ref_[index]; + } + } + /** * @brief Constructs open_addressing_ref_impl. * @@ -398,7 +394,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); for (auto& slot_content : bucket_slots) { auto const eq_res = this->predicate_.template operator()( @@ -451,7 +447,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -547,7 +543,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); for (auto i = 0; i < bucket_size; ++i) { auto const eq_res = this->predicate_.template operator()( @@ -613,7 +609,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -687,7 +683,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); for (auto& slot_content : bucket_slots) { auto const eq_res = @@ -732,7 +728,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -793,7 +789,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); for (auto i = 0; i < bucket_size; ++i) { switch (this->predicate_.template operator()( @@ -831,8 +827,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = - this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); auto const state = [&]() { auto res = detail::equal_result::UNEQUAL; @@ -874,7 +869,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); for (auto i = 0; i < bucket_size; ++i) { switch (this->predicate_.template operator()( @@ -916,8 +911,7 @@ class open_addressing_ref_impl auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = - this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -970,7 +964,7 @@ class open_addressing_ref_impl size_type count = 0; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); cuda::std::int32_t equals[bucket_size] = {0}; bool empty_found = false; @@ -1012,7 +1006,7 @@ class open_addressing_ref_impl size_type count = 0; while (true) { - auto const bucket_slots = storage_ref_[*probing_iter]; + auto const bucket_slots = this->load_bucket(*probing_iter); cuda::std::int32_t equals[bucket_size] = {0}; bool empty_found = false; @@ -1323,8 +1317,7 @@ class open_addressing_ref_impl while (active_flushing_tile.any(running)) { if (running) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = - this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); cuda::static_for([&] __device__(auto i) { equals[i()] = false; @@ -1445,7 +1438,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); bool should_return = false; cuda::static_for([&] __device__(auto i) { @@ -1502,7 +1495,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( @@ -1568,7 +1561,7 @@ class open_addressing_ref_impl while (true) { // TODO atomic_ref::load if insert operator is present - auto const bucket_slots = this->template load_bucket(*probing_iter); + auto const bucket_slots = this->load_bucket(*probing_iter); for (cuda::std::int32_t i = 0; i < bucket_size and !empty; ++i) { switch (this->predicate_.template operator()( diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl index a4abf4cff..1be3f8636 100644 --- a/include/cuco/detail/static_map/static_map_ref.inl +++ b/include/cuco/detail/static_map/static_map_ref.inl @@ -513,7 +513,7 @@ class operator_impl< auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref[*probing_iter]; + auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter); for (auto& slot_content : bucket_slots) { auto const eq_res = @@ -563,7 +563,7 @@ class operator_impl< auto const init_idx = *probing_iter; while (true) { - auto const bucket_slots = storage_ref[*probing_iter]; + auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { detail::bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -888,7 +888,7 @@ class operator_impl< auto constexpr wait_for_payload = (not UseDirectApply) and (sizeof(value_type) > 8); while (true) { - auto const bucket_slots = storage_ref[*probing_iter]; + auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter); for (auto& slot_content : bucket_slots) { auto const eq_res = @@ -966,7 +966,7 @@ class operator_impl< auto constexpr wait_for_payload = (not UseDirectApply) and (sizeof(value_type) > 8); while (true) { - auto const bucket_slots = storage_ref[*probing_iter]; + auto const bucket_slots = ref_.impl_.load_bucket(*probing_iter); auto const [state, intra_bucket_index] = [&]() { detail::bucket_probing_results result{detail::equal_result::UNEQUAL, -1}; @@ -1592,4 +1592,4 @@ class operator_impl< }; } // namespace detail -} // namespace cuco \ No newline at end of file +} // namespace cuco diff --git a/include/cuco/detail/storage/bucket_storage.inl b/include/cuco/detail/storage/bucket_storage.inl index 4814c97ce..55d2e7532 100644 --- a/include/cuco/detail/storage/bucket_storage.inl +++ b/include/cuco/detail/storage/bucket_storage.inl @@ -66,23 +66,12 @@ bucket_storage_ref::operator[](size_type index) const noe } template -template __device__ constexpr bucket_storage_ref::bucket_type bucket_storage_ref::load_bucket(size_type index) const noexcept { - static_assert(Policy == bucket_load_policy::FULL || Policy == bucket_load_policy::FIRST_MATCH, - "Unsupported bucket load policy"); assert(index % bucket_size == 0); assert(index <= capacity() && bucket_size <= capacity() - index); - constexpr auto load_alignment = cuda::std::min( - alignment, - Policy == bucket_load_policy::FIRST_MATCH ? first_match_load_alignment : max_vector_load_bytes); - if constexpr (load_alignment <= alignof(value_type)) { - return (*this)[index]; - } else { - auto const* ptr = __builtin_assume_aligned(this->data() + index, load_alignment); - return *static_cast(ptr); - } + return *reinterpret_cast(__builtin_assume_aligned(this->data() + index, alignment)); } template diff --git a/tests/utility/aligned_storage_test.cu b/tests/utility/aligned_storage_test.cu index fd96922be..3f6d82e17 100644 --- a/tests/utility/aligned_storage_test.cu +++ b/tests/utility/aligned_storage_test.cu @@ -125,11 +125,9 @@ __device__ void check_reads(Ref ref, unsigned* errors) unsigned wrong{}; for (std::size_t index = threadIdx.x * bucket; index + bucket <= n; index += blockDim.x * bucket) { - auto const values = ref.load_bucket(index); - auto const first_match = ref.template load_bucket(index); + auto const values = ref.load_bucket(index); for (int i = 0; i < bucket; ++i) { wrong += !same_value(values[i], index + i); - wrong += !same_value(first_match[i], index + i); } } for (std::size_t index = threadIdx.x; index + bucket <= n; index += blockDim.x) { @@ -269,7 +267,7 @@ TEST_CASE("aligned bucket access preserves custom storage and probing", "") CUCO_CUDA_TRY(cudaMemcpy(storage.data() + 1, &key, sizeof(key), cudaMemcpyHostToDevice)); check_fallback_loads<<<1, 2>>>(storage.ref(), errors, matches); } - SECTION("A custom probe can use the aligned storage load policies.") + SECTION("A custom probe uses the aligned whole-bucket load.") { CUCO_CUDA_TRY(cudaMemcpy(storage.data() + 16, &key, sizeof(key), cudaMemcpyHostToDevice)); check_fallback_loads<<<1, 2>>>(storage.ref(), errors, matches);