Skip to content

Make stream and memory resource explicit in fixed_capacity_map APIs #9719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PointKernel merged 5 commits into from
Jul 8, 2026
Merged

Make stream and memory resource explicit in fixed_capacity_map APIs #9719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
355d423
Make stream and memory resource explicit in fixed_capacity_map APIs
PointKernel Jul 6, 2026
9f5813a
Use stream.sync() after stream-ordered kernel launch in test
PointKernel Jul 6, 2026
1a18099
Use stream-ordered buffer
PointKernel Jul 7, 2026
7126ca5
Merge remote-tracking branch 'upstream/main' into fix-9695-explicit-s…
PointKernel Jul 7, 2026
30ffeb3
Merge branch 'main' into fix-9695-explicit-stream-mr
PointKernel Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_impl.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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}
Expand All @@ -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}
Expand Down Expand Up @@ -234,7 +234,7 @@ public:

//! @brief Inserts keys in `[first, last)` and returns the number of successful insertions.
template <class _InputIt, class _Ref>
_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)
Expand Down Expand Up @@ -262,7 +262,7 @@ public:
//!
//! @throws cuda_error if the insert operation fails to launch
template <class _InputIt, class _Ref>
_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)
Expand Down Expand Up @@ -291,7 +291,7 @@ public:
//! @throws cuda_error if the query operation fails to launch
template <class _InputIt, class _OutputIt, class _Ref>
_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)
Expand Down
123 changes: 55 additions & 68 deletions cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,6 @@
# pragma system_header
#endif // no system header

#include <cuda/__device/device_ref.h>
#include <cuda/__driver/driver_api.h>
#include <cuda/__memory_pool/device_memory_pool.h>
#include <cuda/std/__concepts/concept_macros.h>
#include <cuda/std/__cstddef/types.h>
Expand Down Expand Up @@ -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)}
{}

Expand All @@ -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);
}
Expand All @@ -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);
}
Expand All @@ -301,29 +293,29 @@ 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 <class _InputIt>
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)`.
//!
//! @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 <class _InputIt>
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 =====
Expand All @@ -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 <class _InputIt, class _OutputIt>
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);
}

Expand All @@ -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 <class _InputIt, class _OutputIt>
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 =====
Expand Down
Loading
Loading