Skip to content

Replace cuda::atomic_ref with plain CUDA atomics #10001

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

Draft
PointKernel wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Replace cuda::atomic_ref with plain CUDA atomics #10001

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 2 additions & 2 deletions cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@

#include <cuda/experimental/__cuco/detail/hyperloglog/finalizer.cuh>
#include <cuda/experimental/__cuco/detail/hyperloglog/kernels.cuh>
#include <cuda/experimental/__cuco/detail/utility/atomic.cuh>
#include <cuda/experimental/__cuco/detail/utility/strong_type.cuh>
#include <cuda/experimental/__cuco/hash_functions.cuh>

Expand Down Expand Up @@ -565,8 +566,7 @@ private:
//! @param __value New value
_CCCL_DEVICE_API constexpr void __update_max(int __i, __register_type __value) noexcept
{
::cuda::atomic_ref<__register_type, _Scope> __register_ref(__sketch[__i]);
__register_ref.fetch_max(__value, ::cuda::memory_order_relaxed);
(void) ::cuda::experimental::cuco::detail::__atomic_fetch_max<_Scope>(&__sketch[__i], __value);
}

//! @brief Try expanding the shmem partition for a given kernel beyond 48KB if necessary.
Expand Down
5 changes: 2 additions & 3 deletions cudax/include/cuda/experimental/__cuco/detail/open_addressing/kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@

#include <cub/block/block_reduce.cuh>

#include <cuda/__atomic/atomic.h>
#include <cuda/std/__iterator/iterator_traits.h>

#include <cuda/experimental/__cuco/detail/utility/atomic.cuh>
#include <cuda/experimental/__cuco/detail/utility/cuda.cuh>

#include <cooperative_groups.h>
Expand Down Expand Up @@ -128,8 +128,7 @@ _CCCL_KERNEL_ATTRIBUTES _CCCL_LAUNCH_BOUNDS(_BlockSize) void __insert_if_n(
const auto __block_num_successes = __block_reduce(__temp_storage).Sum(__thread_num_successes);
if (threadIdx.x == 0)
{
::cuda::atomic_ref<typename _Ref::size_type, _Ref::thread_scope>{*__num_successes}.fetch_add(
__block_num_successes, ::cuda::std::memory_order_relaxed);
::cuda::experimental::cuco::detail::__atomic_fetch_add<_Ref::thread_scope>(__num_successes, __block_num_successes);
}
}

Expand Down
37 changes: 15 additions & 22 deletions cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@

#include <thrust/device_reference.h>

#include <cuda/__atomic/atomic.h>
#include <cuda/__type_traits/is_bitwise_comparable.h>
#include <cuda/std/__concepts/concept_macros.h>
#include <cuda/std/__functional/operations.h>
Expand All @@ -36,6 +35,7 @@

#include <cuda/experimental/__cuco/detail/bitwise_compare.cuh>
#include <cuda/experimental/__cuco/detail/equal_wrapper.cuh>
#include <cuda/experimental/__cuco/detail/utility/atomic.cuh>
#include <cuda/experimental/__cuco/detail/utility/cuda.cuh>
#include <cuda/experimental/__cuco/detail/utility/traits.cuh>
#include <cuda/experimental/__cuco/probing_scheme.cuh>
Expand Down Expand Up @@ -706,10 +706,8 @@ public:
auto* __expected_ptr = reinterpret_cast<packed_type*>(&__expected);
auto* __desired_ptr = reinterpret_cast<packed_type*>(&__desired);

auto __slot_ref = ::cuda::atomic_ref<packed_type, _Scope>{*__slot_ptr};

const auto success =
__slot_ref.compare_exchange_strong(*__expected_ptr, *__desired_ptr, ::cuda::memory_order_relaxed);
const auto success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>(
__slot_ptr, *__expected_ptr, *__desired_ptr);

if (success)
{
Expand Down Expand Up @@ -745,28 +743,26 @@ public:
auto __expected_key = __expected.first;
auto __expected_payload = empty_value_sentinel();

::cuda::atomic_ref<__key_type, _Scope> key_ref(__address->first);
::cuda::atomic_ref<mapped_type, _Scope> payload_ref(__address->second);

const auto key_cas_success =
key_ref.compare_exchange_strong(__expected_key, __key_type{__desired.first}, ::cuda::memory_order_relaxed);
auto payload_cas_success =
payload_ref.compare_exchange_strong(__expected_payload, __desired.second, ::cuda::memory_order_relaxed);
const auto key_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>(
&__address->first, __expected_key, __key_type{__desired.first});
auto payload_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>(
&__address->second, __expected_payload, mapped_type{__desired.second});

// if __key success
if (key_cas_success)
{
while (!payload_cas_success)
{
payload_cas_success = payload_ref.compare_exchange_strong(
__expected_payload = empty_value_sentinel(), __desired.second, ::cuda::memory_order_relaxed);
__expected_payload = empty_value_sentinel();
payload_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>(
&__address->second, __expected_payload, mapped_type{__desired.second});
}
return __insert_result::__success;
}
else if (payload_cas_success)
{
// This is insert-specific, cannot for `erase` operations
payload_ref.store(empty_value_sentinel(), ::cuda::memory_order_relaxed);
::cuda::experimental::cuco::detail::__atomic_store<_Scope>(&__address->second, empty_value_sentinel());
}

// Our __key was already present in the slot, so our __key is a duplicate
Expand Down Expand Up @@ -795,16 +791,14 @@ public:
{
using mapped_type = ::cuda::std::decay_t<decltype(empty_value_sentinel())>;

::cuda::atomic_ref<__key_type, _Scope> key_ref(__address->first);
auto __expected_key = __expected.first;
const auto success =
key_ref.compare_exchange_strong(__expected_key, __key_type{__desired.first}, ::cuda::memory_order_relaxed);
const auto success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>(
&__address->first, __expected_key, __key_type{__desired.first});

// if __key success
if (success)
{
::cuda::atomic_ref<mapped_type, _Scope> payload_ref(__address->second);
payload_ref.store(__desired.second, ::cuda::memory_order_relaxed);
::cuda::experimental::cuco::detail::__atomic_store<_Scope>(&__address->second, mapped_type{__desired.second});
return __insert_result::__success;
}

Expand Down Expand Up @@ -892,12 +886,11 @@ public:
template <class _Value>
_CCCL_DEVICE_API void __wait_for_payload(_Value& __slot, _Value __sentinel) const noexcept
{
auto __ref = ::cuda::atomic_ref<_Value, _Scope>{__slot};
_Value __current;
// TODO exponential backoff strategy
do
{
__current = __ref.load(::cuda::std::memory_order_relaxed);
__current = ::cuda::experimental::cuco::detail::__atomic_load<_Scope>(&__slot);
} while (detail::__bitwise_compare(__current, __sentinel));
}
#endif // _CCCL_CUDA_COMPILATION()
Expand Down
Loading
Loading