Skip to content

Commit 0898ffa

Browse files
committed
Add stream tests
1 parent c99311d commit 0898ffa

File tree

3 files changed

+269
-1
lines changed

3 files changed

+269
-1
lines changed

tests/CMakeLists.txt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ ConfigureTest(STATIC_SET_TEST
7575
static_set/rehash_test.cu
7676
static_set/size_test.cu
7777
static_set/shared_memory_test.cu
78+
static_set/stream_test.cu
7879
static_set/unique_sequence_test.cu)
7980

8081
###################################################################################################
@@ -120,7 +121,8 @@ ConfigureTest(STATIC_MULTISET_TEST
120121
static_multiset/retrieve_test.cu
121122
static_multiset/retrieve_if_test.cu
122123
static_multiset/large_input_test.cu
123-
static_multiset/load_factor_test.cu)
124+
static_multiset/load_factor_test.cu
125+
static_multiset/stream_test.cu)
124126

125127
###################################################################################################
126128
# - static_multimap tests -------------------------------------------------------------------------
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
/*
2+
* Copyright (c) 2025, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include <test_utils.hpp>
18+
19+
#include <cuco/static_multiset.cuh>
20+
21+
#include <cuda/functional>
22+
#include <cuda/std/iterator>
23+
#include <cuda/std/tuple>
24+
#include <thrust/device_vector.h>
25+
#include <thrust/execution_policy.h>
26+
#include <thrust/iterator/counting_iterator.h>
27+
#include <thrust/iterator/zip_iterator.h>
28+
#include <thrust/sequence.h>
29+
30+
#include <catch2/catch_template_test_macros.hpp>
31+
32+
TEMPLATE_TEST_CASE_SIG("static_multiset: operations on different stream than constructor",
33+
"",
34+
((typename Key), Key),
35+
(int32_t),
36+
(int64_t))
37+
{
38+
cudaStream_t constructor_stream;
39+
cudaStream_t operation_stream;
40+
CUCO_CUDA_TRY(cudaStreamCreate(&constructor_stream));
41+
CUCO_CUDA_TRY(cudaStreamCreate(&operation_stream));
42+
43+
{ // Scope ensures multiset is destroyed before streams
44+
constexpr std::size_t num_keys{500'000};
45+
auto multiset =
46+
cuco::static_multiset{num_keys * 2,
47+
cuco::empty_key<Key>{-1},
48+
{},
49+
cuco::linear_probing<1, cuco::default_hash_function<Key>>{},
50+
{},
51+
{},
52+
{},
53+
constructor_stream};
54+
55+
thrust::device_vector<Key> d_keys(num_keys);
56+
thrust::sequence(thrust::device, d_keys.begin(), d_keys.end());
57+
58+
SECTION("Insert and count_each on different stream than constructor")
59+
{
60+
multiset.insert(d_keys.begin(), d_keys.end(), operation_stream);
61+
62+
thrust::device_vector<typename cuco::static_multiset<Key>::size_type> d_counts(num_keys);
63+
multiset.count_each(d_keys.begin(),
64+
d_keys.end(),
65+
cuda::std::equal_to<Key>{},
66+
cuco::default_hash_function<Key>{},
67+
d_counts.begin(),
68+
operation_stream);
69+
70+
REQUIRE(cuco::test::all_of(
71+
d_counts.begin(),
72+
d_counts.end(),
73+
cuda::proclaim_return_type<bool>([] __device__(auto const& count) { return count == 1; }),
74+
operation_stream));
75+
}
76+
77+
SECTION("Insert and contains on different stream than constructor")
78+
{
79+
multiset.insert(d_keys.begin(), d_keys.end(), operation_stream);
80+
81+
thrust::device_vector<bool> d_contained(num_keys);
82+
multiset.contains(d_keys.begin(), d_keys.end(), d_contained.begin(), operation_stream);
83+
84+
REQUIRE(cuco::test::all_of(
85+
d_contained.begin(), d_contained.end(), cuda::std::identity{}, operation_stream));
86+
}
87+
88+
SECTION("Insert and find on different stream than constructor")
89+
{
90+
multiset.insert(d_keys.begin(), d_keys.end(), operation_stream);
91+
92+
thrust::device_vector<Key> d_results(num_keys);
93+
multiset.find(d_keys.begin(), d_keys.end(), d_results.begin(), operation_stream);
94+
95+
auto zip =
96+
thrust::make_zip_iterator(cuda::std::make_tuple(d_results.begin(), d_keys.begin()));
97+
REQUIRE(cuco::test::all_of(zip,
98+
zip + num_keys,
99+
cuda::proclaim_return_type<bool>([] __device__(auto const& p) {
100+
return cuda::std::get<0>(p) == cuda::std::get<1>(p);
101+
}),
102+
operation_stream));
103+
}
104+
105+
SECTION("Insert and retrieve on different stream than constructor")
106+
{
107+
multiset.insert(d_keys.begin(), d_keys.end(), operation_stream);
108+
109+
thrust::device_vector<Key> d_probe_results(num_keys);
110+
thrust::device_vector<Key> d_match_results(num_keys);
111+
auto [d_probe_end, d_match_end] = multiset.retrieve(d_keys.begin(),
112+
d_keys.end(),
113+
d_probe_results.begin(),
114+
d_match_results.begin(),
115+
operation_stream);
116+
117+
auto const num_retrieved = cuda::std::distance(d_probe_results.begin(), d_probe_end);
118+
REQUIRE(num_retrieved == num_keys);
119+
}
120+
121+
SECTION("Insert and size on different stream than constructor")
122+
{
123+
multiset.insert(d_keys.begin(), d_keys.end(), operation_stream);
124+
125+
auto const size = multiset.size(operation_stream);
126+
REQUIRE(size == num_keys);
127+
}
128+
129+
SECTION("Insert on constructor stream and query on different stream")
130+
{
131+
multiset.insert(d_keys.begin(), d_keys.end(), constructor_stream);
132+
133+
thrust::device_vector<bool> d_contained(num_keys);
134+
multiset.contains(d_keys.begin(), d_keys.end(), d_contained.begin(), operation_stream);
135+
136+
REQUIRE(cuco::test::all_of(
137+
d_contained.begin(), d_contained.end(), cuda::std::identity{}, operation_stream));
138+
}
139+
} // multiset is destroyed here
140+
141+
CUCO_CUDA_TRY(cudaStreamDestroy(operation_stream));
142+
CUCO_CUDA_TRY(cudaStreamDestroy(constructor_stream));
143+
}

tests/static_set/stream_test.cu

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/*
2+
* Copyright (c) 2025, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include <test_utils.hpp>
18+
19+
#include <cuco/static_set.cuh>
20+
21+
#include <cuda/functional>
22+
#include <cuda/std/iterator>
23+
#include <cuda/std/tuple>
24+
#include <thrust/device_vector.h>
25+
#include <thrust/execution_policy.h>
26+
#include <thrust/iterator/counting_iterator.h>
27+
#include <thrust/iterator/zip_iterator.h>
28+
#include <thrust/sequence.h>
29+
30+
#include <catch2/catch_template_test_macros.hpp>
31+
32+
TEMPLATE_TEST_CASE_SIG("static_set: operations on different stream than constructor",
33+
"",
34+
((typename Key), Key),
35+
(int32_t),
36+
(int64_t))
37+
{
38+
cudaStream_t constructor_stream;
39+
cudaStream_t operation_stream;
40+
CUCO_CUDA_TRY(cudaStreamCreate(&constructor_stream));
41+
CUCO_CUDA_TRY(cudaStreamCreate(&operation_stream));
42+
43+
{ // Scope ensures set is destroyed before streams
44+
constexpr std::size_t num_keys{500'000};
45+
auto set = cuco::static_set{num_keys * 2,
46+
cuco::empty_key<Key>{-1},
47+
{},
48+
cuco::linear_probing<1, cuco::default_hash_function<Key>>{},
49+
{},
50+
{},
51+
{},
52+
constructor_stream};
53+
54+
thrust::device_vector<Key> d_keys(num_keys);
55+
thrust::sequence(thrust::device, d_keys.begin(), d_keys.end());
56+
57+
SECTION("Insert and contains on different stream than constructor")
58+
{
59+
set.insert(d_keys.begin(), d_keys.end(), operation_stream);
60+
61+
thrust::device_vector<bool> d_contained(num_keys);
62+
set.contains(d_keys.begin(), d_keys.end(), d_contained.begin(), operation_stream);
63+
64+
REQUIRE(cuco::test::all_of(
65+
d_contained.begin(), d_contained.end(), cuda::std::identity{}, operation_stream));
66+
}
67+
68+
SECTION("Insert and find on different stream than constructor")
69+
{
70+
set.insert(d_keys.begin(), d_keys.end(), operation_stream);
71+
72+
thrust::device_vector<Key> d_results(num_keys);
73+
set.find(d_keys.begin(), d_keys.end(), d_results.begin(), operation_stream);
74+
75+
auto zip =
76+
thrust::make_zip_iterator(cuda::std::make_tuple(d_results.begin(), d_keys.begin()));
77+
REQUIRE(cuco::test::all_of(zip,
78+
zip + num_keys,
79+
cuda::proclaim_return_type<bool>([] __device__(auto const& p) {
80+
return cuda::std::get<0>(p) == cuda::std::get<1>(p);
81+
}),
82+
operation_stream));
83+
}
84+
85+
SECTION("Insert and retrieve on different stream than constructor")
86+
{
87+
set.insert(d_keys.begin(), d_keys.end(), operation_stream);
88+
89+
thrust::device_vector<Key> d_probe_results(num_keys);
90+
thrust::device_vector<Key> d_match_results(num_keys);
91+
auto [d_probe_end, d_match_end] = set.retrieve(d_keys.begin(),
92+
d_keys.end(),
93+
d_probe_results.begin(),
94+
d_match_results.begin(),
95+
operation_stream);
96+
97+
auto const num_retrieved = cuda::std::distance(d_probe_results.begin(), d_probe_end);
98+
REQUIRE(num_retrieved == num_keys);
99+
}
100+
101+
SECTION("Insert and size on different stream than constructor")
102+
{
103+
set.insert(d_keys.begin(), d_keys.end(), operation_stream);
104+
105+
auto const size = set.size(operation_stream);
106+
REQUIRE(size == num_keys);
107+
}
108+
109+
SECTION("Insert on constructor stream and query on different stream")
110+
{
111+
set.insert(d_keys.begin(), d_keys.end(), constructor_stream);
112+
113+
thrust::device_vector<bool> d_contained(num_keys);
114+
set.contains(d_keys.begin(), d_keys.end(), d_contained.begin(), operation_stream);
115+
116+
REQUIRE(cuco::test::all_of(
117+
d_contained.begin(), d_contained.end(), cuda::std::identity{}, operation_stream));
118+
}
119+
} // set is destroyed here
120+
121+
CUCO_CUDA_TRY(cudaStreamDestroy(operation_stream));
122+
CUCO_CUDA_TRY(cudaStreamDestroy(constructor_stream));
123+
}

0 commit comments

Comments
 (0)