From 355d423a9afa1d4490e11fa181cda84a6c40a7b1 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Mon, 6 Jul 2026 18:27:13 +0000 Subject: [PATCH 1/3] Make stream and memory resource explicit in fixed_capacity_map APIs --- .../open_addressing/open_addressing_impl.cuh | 24 ++-- .../__cuco/fixed_capacity_map.cuh | 123 ++++++++---------- .../cuco/fixed_capacity_map/test_capacity.cu | 17 ++- .../test_insert_and_contains.cu | 18 ++- .../fixed_capacity_map/test_key_sentinel.cu | 16 ++- .../fixed_capacity_map/test_shared_memory.cu | 11 +- 6 files changed, 115 insertions(+), 94 deletions(-) diff --git a/cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_impl.cuh b/cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_impl.cuh index 9830203b812..66d4fc1f14d 100644 --- a/cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_impl.cuh +++ b/cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_impl.cuh @@ -147,12 +147,12 @@ private: public: //! @brief Constructs an open addressing implementation with the given capacity. _CCCL_HOST_API __open_addressing_impl( + ::cuda::stream_ref __stream, + _MemoryResource __mr, __size_type __capacity, __value_type __empty_slot_sentinel, const _KeyEqual& __pred, - const _ProbingScheme& __probing_scheme, - _MemoryResource __mr, - ::cuda::stream_ref __stream) + const _ProbingScheme& __probing_scheme) : __empty_slot_sentinel{__empty_slot_sentinel} , __erased_key_sentinel{__extract_key(__empty_slot_sentinel)} , __predicate{__pred} @@ -166,13 +166,13 @@ public: //! @brief Constructs an open addressing implementation with capacity derived from desired load //! factor. _CCCL_HOST_API __open_addressing_impl( + ::cuda::stream_ref __stream, + _MemoryResource __mr, __size_type __n, double __desired_load_factor, __value_type __empty_slot_sentinel, const _KeyEqual& __pred, - const _ProbingScheme& __probing_scheme, - _MemoryResource __mr, - ::cuda::stream_ref __stream) + const _ProbingScheme& __probing_scheme) : __empty_slot_sentinel{__empty_slot_sentinel} , __erased_key_sentinel{__extract_key(__empty_slot_sentinel)} , __predicate{__pred} @@ -185,13 +185,13 @@ public: //! @brief Constructs an open addressing implementation with erasure support. _CCCL_HOST_API __open_addressing_impl( + ::cuda::stream_ref __stream, + _MemoryResource __mr, __size_type __capacity, __value_type __empty_slot_sentinel, __key_type __erased_key_sentinel, const _KeyEqual& __pred, - const _ProbingScheme& __probing_scheme, - _MemoryResource __mr, - ::cuda::stream_ref __stream) + const _ProbingScheme& __probing_scheme) : __empty_slot_sentinel{__empty_slot_sentinel} , __erased_key_sentinel{__erased_key_sentinel} , __predicate{__pred} @@ -234,7 +234,7 @@ public: //! @brief Inserts keys in `[first, last)` and returns the number of successful insertions. template - _CCCL_HOST_API __size_type insert(_InputIt __first, _InputIt __last, _Ref __container_ref, ::cuda::stream_ref __stream) + _CCCL_HOST_API __size_type insert(::cuda::stream_ref __stream, _InputIt __first, _InputIt __last, _Ref __container_ref) { const auto __num_keys = detail::__distance(__first, __last); if (__num_keys == 0) @@ -262,7 +262,7 @@ public: //! //! @throws cuda_error if the insert operation fails to launch template - _CCCL_HOST_API void insert_async(_InputIt __first, _InputIt __last, _Ref __container_ref, ::cuda::stream_ref __stream) + _CCCL_HOST_API void insert_async(::cuda::stream_ref __stream, _InputIt __first, _InputIt __last, _Ref __container_ref) { const auto __num_keys = detail::__distance(__first, __last); if (__num_keys == 0) @@ -291,7 +291,7 @@ public: //! @throws cuda_error if the query operation fails to launch template _CCCL_HOST_API void contains_async( - _InputIt __first, _InputIt __last, _OutputIt __output_begin, _Ref __container_ref, ::cuda::stream_ref __stream) const + ::cuda::stream_ref __stream, _InputIt __first, _InputIt __last, _OutputIt __output_begin, _Ref __container_ref) const { const auto __num_keys = detail::__distance(__first, __last); if (__num_keys == 0) diff --git a/cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh b/cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh index 3b1dd8e91a4..230f07ea1b9 100644 --- a/cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh +++ b/cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh @@ -21,8 +21,6 @@ # pragma system_header #endif // no system header -#include -#include #include #include #include @@ -114,160 +112,154 @@ private: __stream.sync(); } - //! @brief Returns the default memory pool of the current device. - [[nodiscard]] _CCCL_HOST_API static ::cuda::device_memory_pool_ref __default_memory_resource() - { - return ::cuda::device_default_memory_pool(::cuda::device_ref{::cuda::__driver::__ctxGetDevice()}); - } - public: //! @brief Constructs a map with static capacity (encoded in `_Capacity`) and no erasure. //! + //! @param __stream Stream used for allocation and initialization + //! @param __mr Memory resource for device storage //! @param __empty_key_sentinel Sentinel indicating an empty key slot //! @param __empty_value_sentinel Sentinel indicating an empty payload //! @param __pred Key equality binary callable //! @param __probing_scheme Probing scheme - //! @param __mr Memory resource for device storage - //! @param __stream Stream used for allocation and initialization _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity) _CCCL_REQUIRES((_C != ::cuda::std::dynamic_extent)) _CCCL_HOST_API fixed_capacity_map( + ::cuda::stream_ref __stream, + _MemoryResource __mr, empty_key<_Key> __empty_key_sentinel, empty_value<_Tp> __empty_value_sentinel, const _KeyEqual& __pred = {}, - const _ProbingScheme& __probing_scheme = {}, - _MemoryResource __mr = __default_memory_resource(), - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + const _ProbingScheme& __probing_scheme = {}) : __impl{::cuda::std::make_unique<__impl_type>( + __stream, + __mr, _Capacity, value_type{key_type(__empty_key_sentinel), mapped_type(__empty_value_sentinel)}, __pred, - __probing_scheme, - __mr, - __stream)} + __probing_scheme)} , __empty_value_sentinel{mapped_type(__empty_value_sentinel)} {} //! @brief Constructs a map with dynamic capacity and no erasure. //! + //! @param __stream Stream used for allocation and initialization + //! @param __mr Memory resource for device storage //! @param __capacity Requested slot count (prime/stride-adjusted internally) //! @param __empty_key_sentinel Sentinel indicating an empty key slot //! @param __empty_value_sentinel Sentinel indicating an empty payload //! @param __pred Key equality binary callable //! @param __probing_scheme Probing scheme - //! @param __mr Memory resource for device storage - //! @param __stream Stream used for allocation and initialization _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity) _CCCL_REQUIRES((_C == ::cuda::std::dynamic_extent)) _CCCL_HOST_API fixed_capacity_map( + ::cuda::stream_ref __stream, + _MemoryResource __mr, size_type __capacity, empty_key<_Key> __empty_key_sentinel, empty_value<_Tp> __empty_value_sentinel, const _KeyEqual& __pred = {}, - const _ProbingScheme& __probing_scheme = {}, - _MemoryResource __mr = __default_memory_resource(), - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + const _ProbingScheme& __probing_scheme = {}) : __impl{::cuda::std::make_unique<__impl_type>( + __stream, + __mr, __capacity, value_type{key_type(__empty_key_sentinel), mapped_type(__empty_value_sentinel)}, __pred, - __probing_scheme, - __mr, - __stream)} + __probing_scheme)} , __empty_value_sentinel{mapped_type(__empty_value_sentinel)} {} //! @brief Constructs a map sized by a target load factor (dynamic capacity only). //! + //! @param __stream Stream used for allocation and initialization + //! @param __mr Memory resource for device storage //! @param __n Expected number of keys //! @param __desired_load_factor Target load factor in (0, 1] //! @param __empty_key_sentinel Sentinel indicating an empty key slot //! @param __empty_value_sentinel Sentinel indicating an empty payload //! @param __pred Key equality binary callable //! @param __probing_scheme Probing scheme - //! @param __mr Memory resource for device storage - //! @param __stream Stream used for allocation and initialization _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity) _CCCL_REQUIRES((_C == ::cuda::std::dynamic_extent)) _CCCL_HOST_API fixed_capacity_map( + ::cuda::stream_ref __stream, + _MemoryResource __mr, size_type __n, double __desired_load_factor, empty_key<_Key> __empty_key_sentinel, empty_value<_Tp> __empty_value_sentinel, const _KeyEqual& __pred = {}, - const _ProbingScheme& __probing_scheme = {}, - _MemoryResource __mr = __default_memory_resource(), - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + const _ProbingScheme& __probing_scheme = {}) : __impl{::cuda::std::make_unique<__impl_type>( + __stream, + __mr, __n, __desired_load_factor, value_type{key_type(__empty_key_sentinel), mapped_type(__empty_value_sentinel)}, __pred, - __probing_scheme, - __mr, - __stream)} + __probing_scheme)} , __empty_value_sentinel{mapped_type(__empty_value_sentinel)} {} //! @brief Constructs a map with static capacity and erasure support. //! + //! @param __stream Stream used for allocation and initialization + //! @param __mr Memory resource for device storage //! @param __empty_key_sentinel Sentinel indicating an empty key slot //! @param __empty_value_sentinel Sentinel indicating an empty payload //! @param __erased_key_sentinel Sentinel indicating an erased key slot //! @param __pred Key equality binary callable //! @param __probing_scheme Probing scheme - //! @param __mr Memory resource for device storage - //! @param __stream Stream used for allocation and initialization _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity) _CCCL_REQUIRES((_C != ::cuda::std::dynamic_extent)) _CCCL_HOST_API fixed_capacity_map( + ::cuda::stream_ref __stream, + _MemoryResource __mr, empty_key<_Key> __empty_key_sentinel, empty_value<_Tp> __empty_value_sentinel, erased_key<_Key> __erased_key_sentinel, const _KeyEqual& __pred = {}, - const _ProbingScheme& __probing_scheme = {}, - _MemoryResource __mr = __default_memory_resource(), - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + const _ProbingScheme& __probing_scheme = {}) : __impl{::cuda::std::make_unique<__impl_type>( + __stream, + __mr, _Capacity, value_type{key_type(__empty_key_sentinel), mapped_type(__empty_value_sentinel)}, key_type(__erased_key_sentinel), __pred, - __probing_scheme, - __mr, - __stream)} + __probing_scheme)} , __empty_value_sentinel{mapped_type(__empty_value_sentinel)} {} //! @brief Constructs a map with dynamic capacity and erasure support. //! + //! @param __stream Stream used for allocation and initialization + //! @param __mr Memory resource for device storage //! @param __capacity Requested slot count (prime/stride-adjusted internally) //! @param __empty_key_sentinel Sentinel indicating an empty key slot //! @param __empty_value_sentinel Sentinel indicating an empty payload //! @param __erased_key_sentinel Sentinel indicating an erased key slot //! @param __pred Key equality binary callable //! @param __probing_scheme Probing scheme - //! @param __mr Memory resource for device storage - //! @param __stream Stream used for allocation and initialization _CCCL_TEMPLATE(::cuda::std::size_t _C = _Capacity) _CCCL_REQUIRES((_C == ::cuda::std::dynamic_extent)) _CCCL_HOST_API fixed_capacity_map( + ::cuda::stream_ref __stream, + _MemoryResource __mr, size_type __capacity, empty_key<_Key> __empty_key_sentinel, empty_value<_Tp> __empty_value_sentinel, erased_key<_Key> __erased_key_sentinel, const _KeyEqual& __pred = {}, - const _ProbingScheme& __probing_scheme = {}, - _MemoryResource __mr = __default_memory_resource(), - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + const _ProbingScheme& __probing_scheme = {}) : __impl{::cuda::std::make_unique<__impl_type>( + __stream, + __mr, __capacity, value_type{key_type(__empty_key_sentinel), mapped_type(__empty_value_sentinel)}, key_type(__erased_key_sentinel), __pred, - __probing_scheme, - __mr, - __stream)} + __probing_scheme)} , __empty_value_sentinel{mapped_type(__empty_value_sentinel)} {} @@ -276,7 +268,7 @@ public: //! @brief Erases all elements from the container. After this call, `size()` returns zero. //! //! @param __stream CUDA stream this operation is executed in - void clear(::cuda::stream_ref __stream = cudaStream_t{nullptr}) + void clear(::cuda::stream_ref __stream) { __impl->clear(__stream); } @@ -285,7 +277,7 @@ public: //! returns zero. //! //! @param __stream CUDA stream this operation is executed in - void clear_async(::cuda::stream_ref __stream = cudaStream_t{nullptr}) noexcept + void clear_async(::cuda::stream_ref __stream) noexcept { __impl->clear_async(__stream); } @@ -301,15 +293,15 @@ public: //! @tparam _InputIt Device accessible random access input iterator whose `value_type` is //! convertible to the map's `value_type` //! + //! @param __stream CUDA stream used for insert //! @param __first Beginning of the sequence of keys //! @param __last End of the sequence of keys - //! @param __stream CUDA stream used for insert //! //! @return Number of successful insertions template - size_type insert(_InputIt __first, _InputIt __last, ::cuda::stream_ref __stream = cudaStream_t{nullptr}) + size_type insert(::cuda::stream_ref __stream, _InputIt __first, _InputIt __last) { - return __impl->insert(__first, __last, ref(), __stream); + return __impl->insert(__stream, __first, __last, ref()); } //! @brief Asynchronously inserts all keys in the range `[__first, __last)`. @@ -317,13 +309,13 @@ public: //! @tparam _InputIt Device accessible random access input iterator whose `value_type` is //! convertible to the map's `value_type` //! + //! @param __stream CUDA stream used for insert //! @param __first Beginning of the sequence of keys //! @param __last End of the sequence of keys - //! @param __stream CUDA stream used for insert template - void insert_async(_InputIt __first, _InputIt __last, ::cuda::stream_ref __stream = cudaStream_t{nullptr}) noexcept + void insert_async(::cuda::stream_ref __stream, _InputIt __first, _InputIt __last) noexcept { - __impl->insert_async(__first, __last, ref(), __stream); + __impl->insert_async(__stream, __first, __last, ref()); } // ===== Contains ===== @@ -336,17 +328,14 @@ public: //! @tparam _InputIt Device accessible input iterator //! @tparam _OutputIt Device accessible output iterator assignable from `bool` //! + //! @param __stream CUDA stream used for executing the kernels //! @param __first Beginning of the sequence of keys //! @param __last End of the sequence of keys //! @param __output_begin Beginning of the output sequence of booleans - //! @param __stream CUDA stream used for executing the kernels template - void contains(_InputIt __first, - _InputIt __last, - _OutputIt __output_begin, - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) const + void contains(::cuda::stream_ref __stream, _InputIt __first, _InputIt __last, _OutputIt __output_begin) const { - contains_async(__first, __last, __output_begin, __stream); + contains_async(__stream, __first, __last, __output_begin); __sync(__stream); } @@ -355,17 +344,15 @@ public: //! @tparam _InputIt Device accessible input iterator //! @tparam _OutputIt Device accessible output iterator assignable from `bool` //! + //! @param __stream CUDA stream used for executing the kernels //! @param __first Beginning of the sequence of keys //! @param __last End of the sequence of keys //! @param __output_begin Beginning of the output sequence of booleans - //! @param __stream CUDA stream used for executing the kernels template - void contains_async(_InputIt __first, - _InputIt __last, - _OutputIt __output_begin, - ::cuda::stream_ref __stream = cudaStream_t{nullptr}) const noexcept + void contains_async( + ::cuda::stream_ref __stream, _InputIt __first, _InputIt __last, _OutputIt __output_begin) const noexcept { - __impl->contains_async(__first, __last, __output_begin, ref(), __stream); + __impl->contains_async(__stream, __first, __last, __output_begin, ref()); } // ===== Accessors ===== diff --git a/cudax/test/cuco/fixed_capacity_map/test_capacity.cu b/cudax/test/cuco/fixed_capacity_map/test_capacity.cu index 9f4b5be51a4..5a7f02c62b7 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_capacity.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_capacity.cu @@ -13,7 +13,9 @@ # pragma nv_diag_suppress 20011 #endif +#include #include +#include #include #include @@ -38,7 +40,10 @@ C2H_TEST("fixed_capacity_map dynamic capacity — capacity() reflects the valid const auto valid = cudax::cuco::make_valid_capacity(requested); - dyn_map_t map{requested, cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + + dyn_map_t map{stream, mr, requested, cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; REQUIRE(map.capacity() == valid); REQUIRE(map.capacity() >= requested); } @@ -59,7 +64,10 @@ C2H_TEST("fixed_capacity_map static capacity — valid capacity and capacity_v", static_assert(smap_t::capacity_v == valid, "the map type carries the valid capacity, not the request"); static_assert(smap_t::ref_type::capacity_v == valid, "the ref carries the same valid capacity"); - smap_t map{cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + + smap_t map{stream, mr, cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; REQUIRE(map.capacity() == valid); } @@ -68,7 +76,12 @@ C2H_TEST("fixed_capacity_map dynamic extent — load factor constructor", "[capa constexpr int num_elements = 500; constexpr double load_factor = 0.5; + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + cudax::cuco::fixed_capacity_map map{ + stream, + mr, static_cast<::cuda::std::size_t>(num_elements), load_factor, cudax::cuco::empty_key{empty_key}, diff --git a/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu b/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu index f8582176552..252399d4b06 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu @@ -19,9 +19,11 @@ #include #include +#include #include #include #include +#include #include #include @@ -71,16 +73,22 @@ C2H_TEST("fixed_capacity_map insert and contains", "[container]", key_types, cg_ constexpr int num_keys = 400; - map_type map{static_cast<::cuda::std::size_t>(num_keys * 2), + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + + map_type map{stream, + mr, + static_cast<::cuda::std::size_t>(num_keys * 2), cudax::cuco::empty_key{key_type{-1}}, cudax::cuco::empty_value{key_type{-1}}}; auto __pairs = cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair{}); - map.insert(__pairs, __pairs + num_keys); + map.insert(stream, __pairs, __pairs + num_keys); // Query present keys [0, num_keys) and absent keys [num_keys, 2 * num_keys) ::thrust::device_vector found(2 * num_keys, 0); - map.contains(cuda::counting_iterator{0}, cuda::counting_iterator{2 * num_keys}, found.begin()); + map.contains( + stream, cuda::counting_iterator{0}, cuda::counting_iterator{2 * num_keys}, found.begin()); REQUIRE(cudaDeviceSynchronize() == cudaSuccess); ::thrust::host_vector h_found(found); @@ -93,9 +101,9 @@ C2H_TEST("fixed_capacity_map insert and contains", "[container]", key_types, cg_ REQUIRE(mismatches == 0); // After clear the map is empty, so none of the previously inserted keys are found - map.clear(); + map.clear(stream); ::thrust::fill(found.begin(), found.end(), 1); - map.contains(cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); + map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); REQUIRE(cudaDeviceSynchronize() == cudaSuccess); REQUIRE(::thrust::none_of(found.begin(), found.begin() + num_keys, [] __device__(int v) { return v != 0; diff --git a/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu b/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu index 0f17849df35..5623d3be9d2 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu @@ -17,7 +17,9 @@ #include #include +#include #include +#include #include #include @@ -51,14 +53,20 @@ C2H_TEST("fixed_capacity_map — empty and erased key sentinels", "[sentinel]") cudax::cuco::make_valid_capacity(static_cast<::cuda::std::size_t>(num_keys * 2)); using map_type = cudax::cuco::fixed_capacity_map; - map_type map{ - cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}, cudax::cuco::erased_key{erased_sentinel}}; + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + + map_type map{stream, + mr, + cudax::cuco::empty_key{empty_key}, + cudax::cuco::empty_value{empty_value}, + cudax::cuco::erased_key{erased_sentinel}}; auto __pairs = cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair<::cuda::std::pair>{}); - map.insert(__pairs, __pairs + num_keys); + map.insert(stream, __pairs, __pairs + num_keys); ::thrust::device_vector found(num_keys, 0); - map.contains(cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); + map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); REQUIRE(cudaDeviceSynchronize() == cudaSuccess); REQUIRE(::thrust::all_of(found.begin(), found.end(), [] __device__(int v) { return v != 0; diff --git a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu index d9675dd2cd8..7dbc57d990c 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu @@ -18,7 +18,9 @@ #include #include +#include #include +#include #include #include @@ -71,12 +73,15 @@ C2H_TEST("fixed_capacity_map static extent — shared memory sizing via capacity { constexpr int num_keys = 64; - fixed_capacity_map_512_type map{cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; + ::cuda::stream stream{::cuda::device_ref{0}}; + auto mr = ::cuda::device_default_memory_pool(::cuda::device_ref{0}); + + fixed_capacity_map_512_type map{stream, mr, cudax::cuco::empty_key{empty_key}, cudax::cuco::empty_value{empty_value}}; const int block_size = 128; const int grid_size = (num_keys + block_size - 1) / block_size; - insert_shmem_kernel<<>>( + insert_shmem_kernel<<>>( map.ref(), cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair{}), num_keys); @@ -84,7 +89,7 @@ C2H_TEST("fixed_capacity_map static extent — shared memory sizing via capacity // Verify the insertions actually landed in the global map ::thrust::device_vector found(num_keys, 0); - map.contains(cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); + map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); REQUIRE(cudaDeviceSynchronize() == cudaSuccess); REQUIRE(::thrust::all_of(found.begin(), found.end(), [] __device__(int v) { return v != 0; From 9f5813af003c3159015196b7ae9ceec6cc84bed5 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Mon, 6 Jul 2026 18:47:12 +0000 Subject: [PATCH 2/3] Use stream.sync() after stream-ordered kernel launch in test --- cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu index 7dbc57d990c..7b0b52b20bb 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu @@ -85,7 +85,7 @@ C2H_TEST("fixed_capacity_map static extent — shared memory sizing via capacity map.ref(), cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair{}), num_keys); - REQUIRE(cudaDeviceSynchronize() == cudaSuccess); + stream.sync(); // Verify the insertions actually landed in the global map ::thrust::device_vector found(num_keys, 0); From 1a18099b6785a35885870daed4d5c44ea98b8c99 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Tue, 7 Jul 2026 18:14:22 +0000 Subject: [PATCH 3/3] Use stream-ordered buffer --- .../test_insert_and_contains.cu | 52 ++++++++++++------- .../fixed_capacity_map/test_key_sentinel.cu | 18 ++++--- .../fixed_capacity_map/test_shared_memory.cu | 20 ++++--- 3 files changed, 58 insertions(+), 32 deletions(-) diff --git a/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu b/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu index 252399d4b06..49506daf10f 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu @@ -13,11 +13,10 @@ # pragma nv_diag_suppress 20011 #endif -#include -#include -#include +#include #include +#include #include #include #include @@ -49,6 +48,26 @@ struct iota_pair } }; +// Present keys [0, num_keys) are found, absent keys [num_keys, ...) are not +struct match_expected +{ + const int* found; + int num_keys; + + __device__ bool operator()(int i) const noexcept + { + return static_cast(found[i]) == (i < num_keys); + } +}; + +struct is_nonzero +{ + __device__ bool operator()(int v) const noexcept + { + return v != 0; + } +}; + C2H_TEST("fixed_capacity_map insert and contains", "[container]", key_types, cg_sizes, bucket_sizes, probing_kinds) { using key_type = c2h::get<0, TestType>; @@ -86,26 +105,21 @@ C2H_TEST("fixed_capacity_map insert and contains", "[container]", key_types, cg_ map.insert(stream, __pairs, __pairs + num_keys); // Query present keys [0, num_keys) and absent keys [num_keys, 2 * num_keys) - ::thrust::device_vector found(2 * num_keys, 0); + auto found = ::cuda::make_buffer(stream, mr, 2 * num_keys, 0); map.contains( stream, cuda::counting_iterator{0}, cuda::counting_iterator{2 * num_keys}, found.begin()); - REQUIRE(cudaDeviceSynchronize() == cudaSuccess); - ::thrust::host_vector h_found(found); - int mismatches = 0; - for (int i = 0; i < 2 * num_keys; ++i) - { - const bool expected = i < num_keys; // present keys found, absent keys not - mismatches += (static_cast(h_found[i]) != expected); - } - REQUIRE(mismatches == 0); + REQUIRE(::thrust::all_of( + ::thrust::cuda::par.on(stream.get()), + cuda::counting_iterator{0}, + cuda::counting_iterator{2 * num_keys}, + match_expected{found.data(), num_keys})); // After clear the map is empty, so none of the previously inserted keys are found map.clear(stream); - ::thrust::fill(found.begin(), found.end(), 1); - map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); - REQUIRE(cudaDeviceSynchronize() == cudaSuccess); - REQUIRE(::thrust::none_of(found.begin(), found.begin() + num_keys, [] __device__(int v) { - return v != 0; - })); + auto cleared = ::cuda::make_buffer(stream, mr, num_keys, 1); + map.contains( + stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, cleared.begin()); + REQUIRE( + ::thrust::none_of(::thrust::cuda::par.on(stream.get()), cleared.data(), cleared.data() + num_keys, is_nonzero{})); } diff --git a/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu b/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu index 5623d3be9d2..adbd59eaee3 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_key_sentinel.cu @@ -13,9 +13,10 @@ # pragma nv_diag_suppress 20011 #endif -#include +#include #include +#include #include #include #include @@ -42,6 +43,14 @@ struct iota_pair } }; +struct is_nonzero +{ + __device__ bool operator()(int v) const noexcept + { + return v != 0; + } +}; + C2H_TEST("fixed_capacity_map — empty and erased key sentinels", "[sentinel]") { constexpr int erased_sentinel = -2; @@ -65,10 +74,7 @@ C2H_TEST("fixed_capacity_map — empty and erased key sentinels", "[sentinel]") auto __pairs = cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair<::cuda::std::pair>{}); map.insert(stream, __pairs, __pairs + num_keys); - ::thrust::device_vector found(num_keys, 0); + auto found = ::cuda::make_buffer(stream, mr, num_keys, 0); map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); - REQUIRE(cudaDeviceSynchronize() == cudaSuccess); - REQUIRE(::thrust::all_of(found.begin(), found.end(), [] __device__(int v) { - return v != 0; - })); + REQUIRE(::thrust::all_of(::thrust::cuda::par.on(stream.get()), found.data(), found.data() + num_keys, is_nonzero{})); } diff --git a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu index 7b0b52b20bb..0d1aa919b00 100644 --- a/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu +++ b/cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu @@ -13,9 +13,10 @@ # pragma nv_diag_suppress 20011 #endif -#include +#include #include +#include #include #include #include @@ -50,6 +51,14 @@ struct iota_pair } }; +struct is_nonzero +{ + __device__ bool operator()(int v) const noexcept + { + return v != 0; + } +}; + // Demonstrates compile-time __shared__ sizing via ref_type::capacity_v. template __global__ void insert_shmem_kernel(fixed_capacity_map_512_type::ref_type global_ref, _PairIt pairs, int n) @@ -85,13 +94,10 @@ C2H_TEST("fixed_capacity_map static extent — shared memory sizing via capacity map.ref(), cuda::transform_iterator(cuda::counting_iterator{0}, iota_pair{}), num_keys); - stream.sync(); + REQUIRE(cudaGetLastError() == cudaSuccess); // Verify the insertions actually landed in the global map - ::thrust::device_vector found(num_keys, 0); + auto found = ::cuda::make_buffer(stream, mr, num_keys, 0); map.contains(stream, cuda::counting_iterator{0}, cuda::counting_iterator{num_keys}, found.begin()); - REQUIRE(cudaDeviceSynchronize() == cudaSuccess); - REQUIRE(::thrust::all_of(found.begin(), found.end(), [] __device__(int v) { - return v != 0; - })); + REQUIRE(::thrust::all_of(::thrust::cuda::par.on(stream.get()), found.data(), found.data() + num_keys, is_nonzero{})); }