Skip to content

Commit dddc7a3

Browse files
PointKernelPointKernel
authored
Optimize extent calculation for linear probing by eliminating prime number computation (#692)
Closes #581 This PR enhances the extent calculation for linear probing by eliminating the need to compute the prime number for the final extent. It also resolves an embarrassing bug where two prime numbers were missing in the existing prime number array.
1 parent e445e99 commit dddc7a3

File tree

9 files changed

+20170
-20146
lines changed

9 files changed

+20170
-20146
lines changed

‎include/cuco/detail/extent/extent.inl‎

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
2+
* Copyright (c) 2023-2025, NVIDIA CORPORATION.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -120,24 +120,21 @@ template <int32_t CGSize, int32_t BucketSize, typename SizeType>
120120
template <typename ProbingScheme, typename Storage, typename SizeType, std::size_t N>
121121
[[nodiscard]] auto constexpr make_bucket_extent(extent<SizeType, N> ext)
122122
{
123-
return make_bucket_extent<ProbingScheme::cg_size, Storage::bucket_size, SizeType, N>(ext);
124-
/*
125-
// TODO fix linear probing with exact capacity
126-
if constexpr (cuco::is_double_hashing<ProbingScheme>::value) {
127-
return make_bucket_extent<ProbingScheme::cg_size, Storage::bucket_size, SizeType, N>(ext);
128-
} else {
129-
auto const size = cuco::detail::int_div_ceil(
130-
cuda::std::max(static_cast<SizeType>(ext), static_cast<SizeType>(1)),
131-
ProbingScheme::cg_size * Storage::bucket_size) +
132-
cuda::std::min(static_cast<SizeType>(ext), static_cast<SizeType>(1));
133-
if constexpr (N == dynamic_extent) {
134-
return bucket_extent<SizeType>{size * ProbingScheme::cg_size};
123+
if constexpr (cuco::is_double_hashing<ProbingScheme>::value) {
124+
return make_bucket_extent<ProbingScheme::cg_size, Storage::bucket_size, SizeType, N>(ext);
135125
} else {
136-
return bucket_extent<SizeType, size * ProbingScheme::cg_size>{};
126+
auto const size =
127+
cuco::detail::int_div_ceil(std::max(static_cast<SizeType>(ext), static_cast<SizeType>(1)),
128+
ProbingScheme::cg_size * Storage::bucket_size) +
129+
static_cast<SizeType>(ext == 0);
130+
131+
if constexpr (N == dynamic_extent) {
132+
return bucket_extent<SizeType>{size * ProbingScheme::cg_size};
133+
} else {
134+
return bucket_extent<SizeType, size * ProbingScheme::cg_size>{};
135+
}
137136
}
138137
}
139-
*/
140-
}
141138

142139
template <typename ProbingScheme, typename Storage, typename SizeType>
143140
[[nodiscard]] auto constexpr make_bucket_extent(SizeType size)

‎include/cuco/detail/prime.hpp‎

Lines changed: 20108 additions & 20108 deletions
Large diffs are not rendered by default.

‎tests/static_map/capacity_test.cu‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
2+
* Copyright (c) 2023-2025, NVIDIA CORPORATION.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -157,9 +157,9 @@ TEST_CASE("static_map capacity test", "")
157157
REQUIRE(ref_capacity == gold_capacity);
158158
}
159159

160-
SECTION("Dynamic extent is evaluated at run time.")
160+
SECTION("Dynamic extent of linear probing is evaluated at run time.")
161161
{
162-
auto constexpr gold_capacity = 412; // 103 x 2 x 2
162+
auto constexpr gold_capacity = 400;
163163

164164
using probe = cuco::linear_probing<2, cuco::default_hash_function<Key>>;
165165
auto map =

‎tests/static_map/erase_test.cu‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,10 @@ TEMPLATE_TEST_CASE_SIG(
118118
thrust::equal_to<Key>,
119119
probe,
120120
cuco::cuda_allocator<cuda::std::byte>,
121-
cuco::storage<2>>{
122-
num_keys, cuco::empty_key<Key>{-1}, cuco::empty_value<Value>{-1}, cuco::erased_key<Key>{-2}};
121+
cuco::storage<2>>{num_keys * 2,
122+
cuco::empty_key<Key>{-1},
123+
cuco::empty_value<Value>{-1},
124+
cuco::erased_key<Key>{-2}};
123125

124126
test_erase(map, num_keys);
125127
}

‎tests/static_map/find_test.cu‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -155,9 +155,7 @@ TEMPLATE_TEST_CASE_SIG(
155155
(int64_t, int32_t, cuco::test::probe_sequence::linear_probing, 2),
156156
(int64_t, int64_t, cuco::test::probe_sequence::linear_probing, 2))
157157
{
158-
constexpr size_type num_keys{400};
159-
constexpr size_type gold_capacity = CGSize == 1 ? 422 // 211 x 1 x 2
160-
: 412; // 103 x 2 x 2
158+
constexpr size_type num_keys{301};
161159

162160
// XXX: testing static extent is intended, DO NOT CHANGE
163161
using extent_type = cuco::extent<size_type, num_keys>;
@@ -166,6 +164,15 @@ TEMPLATE_TEST_CASE_SIG(
166164
cuco::linear_probing<CGSize, cuco::murmurhash3_32<Key>>,
167165
cuco::double_hashing<CGSize, cuco::murmurhash3_32<Key>, cuco::murmurhash3_32<Key>>>;
168166

167+
constexpr size_type gold_capacity = [&]() {
168+
if constexpr (cuco::is_double_hashing<probe>::value) {
169+
return (CGSize == 1) ? 302 // 151 x 1 x 2
170+
: 316; // 79 x 2 x 2
171+
} else {
172+
return (CGSize == 1) ? 302 : 304;
173+
}
174+
}();
175+
169176
auto map = cuco::static_map<Key,
170177
Value,
171178
extent_type,

‎tests/static_multiset/insert_test.cu‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2024, NVIDIA CORPORATION.
2+
* Copyright (c) 2024-2025, NVIDIA CORPORATION.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -89,13 +89,20 @@ TEMPLATE_TEST_CASE_SIG(
8989
(int64_t, cuco::test::probe_sequence::linear_probing, 2))
9090
{
9191
constexpr size_type num_keys{400};
92-
constexpr size_type gold_capacity = CGSize == 1 ? 422 // 211 x 1 x 2
93-
: 412; // 103 x 2 x 2
9492

9593
using probe = std::conditional_t<Probe == cuco::test::probe_sequence::linear_probing,
9694
cuco::linear_probing<CGSize, cuco::default_hash_function<Key>>,
9795
cuco::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
9896

97+
constexpr size_type gold_capacity = [&]() {
98+
if constexpr (cuco::is_double_hashing<probe>::value) {
99+
return (CGSize == 1) ? 422 // 211 x 1 x 2
100+
: 412; // 103 x 2 x 2
101+
} else {
102+
return 400;
103+
}
104+
}();
105+
99106
auto set =
100107
cuco::static_multiset{num_keys, cuco::empty_key<Key>{-1}, {}, probe{}, {}, cuco::storage<2>{}};
101108

‎tests/static_set/capacity_test.cu‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
2+
* Copyright (c) 2023-2025, NVIDIA CORPORATION.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -120,9 +120,9 @@ TEST_CASE("static_set capacity test", "")
120120
REQUIRE(ref_capacity == gold_capacity);
121121
}
122122

123-
SECTION("Dynamic extent is evaluated at run time.")
123+
SECTION("Dynamic extent of linear probing is evaluated at run time.")
124124
{
125-
auto constexpr gold_capacity = 412; // 103 x 2 x 2
125+
auto constexpr gold_capacity = 400;
126126

127127
using probe = cuco::linear_probing<2, cuco::default_hash_function<Key>>;
128128
auto set = cuco::static_set<Key,

‎tests/static_set/retrieve_all_test.cu‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
2+
* Copyright (c) 2023-2025, NVIDIA CORPORATION.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -72,14 +72,20 @@ TEMPLATE_TEST_CASE_SIG(
7272
{
7373
constexpr std::size_t num_keys{400};
7474
constexpr double desired_load_factor = 1.;
75-
auto constexpr gold_capacity = CGSize == 1 ? 409 // 409 x 1 x 1
76-
: 422 // 211 x 2 x 1
77-
;
7875

7976
using probe = std::conditional_t<Probe == cuco::test::probe_sequence::linear_probing,
8077
cuco::linear_probing<CGSize, cuco::default_hash_function<Key>>,
8178
cuco::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
8279

80+
constexpr std::size_t gold_capacity = [&]() {
81+
if constexpr (cuco::is_double_hashing<probe>::value) {
82+
return (CGSize == 1) ? 409 // 409 x 1 x 2
83+
: 422; // 211 x 2 x 2
84+
} else {
85+
return 400;
86+
}
87+
}();
88+
8389
auto set = cuco::static_set{num_keys, desired_load_factor, cuco::empty_key<Key>{-1}, {}, probe{}};
8490

8591
REQUIRE(set.capacity() == gold_capacity);

‎tests/static_set/unique_sequence_test.cu‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -148,14 +148,19 @@ TEMPLATE_TEST_CASE_SIG(
148148
(int64_t, cuco::test::probe_sequence::linear_probing, 2))
149149
{
150150
constexpr size_type num_keys{400};
151-
constexpr size_type gold_capacity = CGSize == 1 ? 422 // 211 x 1 x 2
152-
: 412 // 103 x 2 x 2
153-
;
154-
155151
using probe = std::conditional_t<Probe == cuco::test::probe_sequence::linear_probing,
156152
cuco::linear_probing<CGSize, cuco::default_hash_function<Key>>,
157153
cuco::double_hashing<CGSize, cuco::default_hash_function<Key>>>;
158154

155+
constexpr size_type gold_capacity = [&]() {
156+
if constexpr (cuco::is_double_hashing<probe>::value) {
157+
return (CGSize == 1) ? 422 // 211 x 1 x 2
158+
: 412; // 103 x 2 x 2
159+
} else {
160+
return 400;
161+
}
162+
}();
163+
159164
auto set =
160165
cuco::static_set{num_keys, cuco::empty_key<Key>{SENTINEL}, {}, probe{}, {}, cuco::storage<2>{}};
161166

0 commit comments

Comments
 (0)