diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 1b28746ea..bc7cc697f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -114,6 +114,7 @@ ConfigureTest(STATIC_MULTIMAP_TEST static_multimap/count_test.cu static_multimap/custom_pair_retrieve_test.cu static_multimap/custom_type_test.cu + static_multimap/find_test.cu static_multimap/heterogeneous_lookup_test.cu static_multimap/insert_contains_test.cu static_multimap/insert_if_test.cu diff --git a/tests/static_multimap/find_test.cu b/tests/static_multimap/find_test.cu new file mode 100644 index 000000000..51456b088 --- /dev/null +++ b/tests/static_multimap/find_test.cu @@ -0,0 +1,106 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include + +#include +#include +#include +#include +#include + +#include + +using size_type = int32_t; + +template +void test_multimap_find(Map& map, size_type num_keys) +{ + using Key = typename Map::key_type; + using Value = typename Map::mapped_type; + + auto zip_equal = cuda::proclaim_return_type( + [] __device__(auto val) { return thrust::get<0>(val) == thrust::get<1>(val); }); + + auto const keys_begin = thrust::counting_iterator{0}; + + SECTION("Non-inserted keys have no matches") + { + thrust::device_vector found_vals(num_keys); + + map.find(keys_begin, keys_begin + num_keys, found_vals.begin()); + auto zip = thrust::make_zip_iterator(thrust::make_tuple( + found_vals.begin(), thrust::constant_iterator{map.empty_value_sentinel()})); + + REQUIRE(cuco::test::all_of(zip, zip + num_keys, zip_equal)); + } + + auto const pairs_begin = thrust::make_transform_iterator( + thrust::make_counting_iterator(0), + cuda::proclaim_return_type>( + [] __device__(auto i) { return cuco::pair{i, i * 2}; })); + + map.insert(pairs_begin, pairs_begin + num_keys); + + SECTION("All inserted keys should be correctly recovered during find") + { + thrust::device_vector found_vals(num_keys); + + map.find(keys_begin, keys_begin + num_keys, found_vals.begin()); + + auto const gold_vals_begin = thrust::make_transform_iterator( + thrust::make_counting_iterator(0), + cuda::proclaim_return_type([] __device__(auto i) { return Value{i * 2}; })); + auto zip = thrust::make_zip_iterator(thrust::make_tuple(found_vals.begin(), gold_vals_begin)); + + REQUIRE(cuco::test::all_of(zip, zip + num_keys, zip_equal)); + } +} + +TEMPLATE_TEST_CASE_SIG( + "static_multimap find tests", + "", + ((typename T, cuco::test::probe_sequence Probe, int CGSize), T, 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)) +{ + constexpr size_type num_keys{1'000}; + + using probe = std::conditional_t< + Probe == cuco::test::probe_sequence::linear_probing, + cuco::linear_probing>, + cuco::double_hashing, cuco::default_hash_function>>; + + auto map = cuco::experimental::static_multimap, + cuda::thread_scope_device, + thrust::equal_to, + probe, + cuco::cuda_allocator, + cuco::storage<2>>{ + num_keys, cuco::empty_key{-1}, cuco::empty_value{-2}}; + + test_multimap_find(map, num_keys); +}