diff --git a/Src/Base/AMReX_Functional.H b/Src/Base/AMReX_Functional.H index bed6e7a4267..0098365db4a 100644 --- a/Src/Base/AMReX_Functional.H +++ b/Src/Base/AMReX_Functional.H @@ -63,6 +63,24 @@ struct LogicalOr } }; +template +struct Multiplies +{ + constexpr T operator() (const T & lhs, const T & rhs) const + { + return lhs * rhs; + } +}; + +template +struct Divides +{ + constexpr T operator() (const T & lhs, const T & rhs) const + { + return lhs / rhs; + } +}; + } #endif diff --git a/Src/Base/AMReX_GpuAtomic.H b/Src/Base/AMReX_GpuAtomic.H index deea6ae932e..9adc655298e 100644 --- a/Src/Base/AMReX_GpuAtomic.H +++ b/Src/Base/AMReX_GpuAtomic.H @@ -13,11 +13,13 @@ namespace amrex { namespace Gpu::Atomic { // For Add, Min and Max, we support int, unsigned int, long, unsigned long long, float and double. +// For Multiply and Divide, we support generic types provided they are the same size as int or unsigned long long +// and have *= and /= operators. // For LogicalOr and LogicalAnd, the data type is int. // For Exch and CAS, the data type is generic. // All these functions are non-atomic in host code!!! // If one needs them to be atomic in host code, use HostDevice::Atomic::*. Currently only -// HostDevice::Atomic is supported. We could certainly add more. +// HostDevice::Atomic::Add is supported. We could certainly add more. namespace detail { @@ -526,6 +528,78 @@ namespace detail { )) #endif } + +//////////////////////////////////////////////////////////////////////// +// Multiply +//////////////////////////////////////////////////////////////////////// + +#ifdef AMREX_USE_GPU + + template = 0> + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + T Multiply_device (T* const prod, T const value) noexcept + { + return detail::atomic_op(prod,value,amrex::Multiplies()); + } + + template = 0> + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + T Multiply_device (T* const prod, T const value) noexcept + { + return detail::atomic_op(prod,value,amrex::Multiplies()); + } + +#endif + + template + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + T Multiply (T* const prod, T const value) noexcept + { + AMREX_IF_ON_DEVICE(( + return Multiply_device(prod, value); + )) + AMREX_IF_ON_HOST(( + auto const old = *prod; + *prod *= value; + return old; + )) + } + +//////////////////////////////////////////////////////////////////////// +// Divide +//////////////////////////////////////////////////////////////////////// + +#ifdef AMREX_USE_GPU + + template = 0> + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + T Divide_device (T* const quot, T const value) noexcept + { + return detail::atomic_op(quot,value,amrex::Divides()); + } + + template = 0> + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + T Divide_device (T* const quot, T const value) noexcept + { + return detail::atomic_op(quot,value,amrex::Divides()); + } + +#endif + + template + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + T Divide (T* const quot, T const value) noexcept + { + AMREX_IF_ON_DEVICE(( + return Divide_device(quot, value); + )) + AMREX_IF_ON_HOST(( + auto const old = *quot; + *quot /= value; + return old; + )) + } } namespace HostDevice::Atomic {