atomic.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_TESTS_UTILS_ATOMIC_HPP
#define KOKKOS_EXECUTION_TESTS_UTILS_ATOMIC_HPP
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#    include "Kokkos_Core.hpp"
#endif
 
#if defined(KOKKOS_ENABLE_CUDA)
#    include <cuda/atomic>
#elif defined(KOKKOS_ENABLE_HIP)
#    include "desul/atomics/Adapt_HIP.hpp"
#elif defined(KOKKOS_ENABLE_SYCL)
#    include "desul/atomics/Adapt_SYCL.hpp"
#endif

 
#if defined(KOKKOS_ENABLE_CUDA)
namespace desul::Impl {
template <class MemoryScope>
struct CudaMemoryScope;
 
template <>
struct CudaMemoryScope<MemoryScopeDevice> {
    static constexpr auto value = cuda::thread_scope_device;
};
 
template <>
struct CudaMemoryScope<MemoryScopeSystem> {
    static constexpr auto value = cuda::thread_scope_system;
};
 
template <class MemoryOrder>
struct CudaMemoryOrder;
 
template <>
struct CudaMemoryOrder<MemoryOrderRelaxed> {
    static constexpr auto value = cuda::memory_order_relaxed;
};
} // namespace desul::Impl
#endif
 
namespace Tests::Utils {

template <typename Scope, typename Order, typename T>
KOKKOS_FUNCTION auto atomic_fetch_add(T* ptr, const T val) {
#if defined(KOKKOS_ENABLE_CUDA)
    return cuda::atomic_ref<T, desul::Impl::CudaMemoryScope<Scope>::value>(*ptr)
        .fetch_add(val, desul::Impl::CudaMemoryOrder<Order>::value);
#elif defined(KOKKOS_ENABLE_HIP)
    return __hip_atomic_fetch_add(
        ptr, val, desul::Impl::HIPMemoryOrder<Order>::value, desul::Impl::HIPMemoryScope<Scope>::value);
#elif defined(KOKKOS_ENABLE_SYCL)
    return desul::Impl::sycl_atomic_ref<T, Order, Scope>{*ptr}.fetch_add(val);
#else
    static_assert(std::same_as<Order, desul::MemoryOrderRelaxed>);
    return std::atomic_ref<T>(*ptr).fetch_add(val, std::memory_order_relaxed);
#endif
}
 
} // namespace Tests::Utils
 
#endif // KOKKOS_EXECUTION_TESTS_UTILS_ATOMIC_HPP