From 17ad1834d4b2fffe1c4d4f8027a8867d4d0ce2a7 Mon Sep 17 00:00:00 2001 From: Rahul Manocha Date: Mon, 16 Dec 2024 14:27:58 -0800 Subject: [PATCH] SWDEV-492165 - Add support for 16 bit atomicAdd and atomicCAS Change-Id: I7f1e4876fe2960dff1775d27cb6443a89e146c86 [ROCm/clr commit: 93cff75928e68fe5e6327085c82855ec98de34e6] --- .../include/hip/amd_detail/amd_hip_atomic.h | 30 +++++++++++++++++ .../include/hip/amd_detail/amd_hip_bf16.h | 33 +++++++++++++++---- .../include/hip/amd_detail/amd_hip_fp16.h | 25 ++++++++++++-- 3 files changed, 79 insertions(+), 9 deletions(-) diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_atomic.h b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_atomic.h index a5c06f0321..08c08d58b0 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_atomic.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_atomic.h @@ -26,6 +26,7 @@ THE SOFTWARE. #include "amd_device_functions.h" #endif + #if __has_builtin(__hip_atomic_compare_exchange_strong) template struct Cond_t; @@ -113,6 +114,25 @@ T hip_cas_extrema_expander(T* p, T x, Cmp cmp, F f) noexcept return reinterpret_cast(tmp); } +__device__ +inline +unsigned short int atomicCAS(unsigned short int* address, unsigned short int compare, + unsigned short int val) { + __hip_atomic_compare_exchange_strong(address, &compare, val, __ATOMIC_RELAXED, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + return compare; +} + +__device__ +inline +unsigned short int atomicCAS_system(unsigned short int* address, unsigned short int compare, + unsigned short int val) { + __hip_atomic_compare_exchange_strong(address, &compare, val, __ATOMIC_RELAXED, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_SYSTEM); + return compare; +} + + __device__ inline int atomicCAS(int* address, int compare, int val) { @@ -1328,6 +1348,16 @@ unsigned long long atomicXor_system(unsigned long long* address, unsigned long l } #else // __hip_atomic_compare_exchange_strong +__device__ +inline +unsigned short int atomicCAS(unsigned short int* address, unsigned short int compare, + unsigned short int val) +{ + __atomic_compare_exchange_n( + address, &compare, val, false, __ATOMIC_RELAXED, __ATOMIC_RELAXED); + + return compare; +} __device__ inline diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_bf16.h b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_bf16.h index 1b433ba531..87c41177b8 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_bf16.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_bf16.h @@ -111,13 +111,10 @@ #if !defined(__HIPCC_RTC__) #include #include // Sync functions -#include "amd_hip_vector_types.h" // float2 etc -#include "device_library_decls.h" // ocml conversion functions -#if defined(__clang__) && defined(__HIP__) -#include "amd_hip_atomic.h" -#endif // defined(__clang__) && defined(__HIP__) +#include "amd_hip_vector_types.h" // float2 etc +#include "device_library_decls.h" // ocml conversion functions #include "math_fwd.h" // ocml device functions -#endif // !defined(__HIPCC_RTC__) +#endif // !defined(__HIPCC_RTC__) #define __BF16_DEVICE__ __device__ #if defined(__HIPCC_RTC__) @@ -1860,5 +1857,29 @@ __BF16_DEVICE_STATIC__ __hip_bfloat162 unsafeAtomicAdd(__hip_bfloat162* address, return old_val.h2r; #endif } +__BF16_DEVICE_STATIC__ __hip_bfloat16 unsafeAtomicAdd(__hip_bfloat16 *address, + __hip_bfloat16 value) { + static_assert(sizeof(unsigned short int) == sizeof(__hip_bfloat16_raw)); + unsigned short int* address_as_short = reinterpret_cast(address); + // Align to 4 bytes + unsigned int* aligned_addr = __builtin_bit_cast(unsigned int*, + __builtin_bit_cast(unsigned long long int, address_as_short) & + (unsigned long long int)(~0x3)); + + bool is_lower = __builtin_bit_cast(unsigned long long int, aligned_addr) == + __builtin_bit_cast(unsigned long long int, address); + + __hip_bfloat162 fval; + if (is_lower) + fval = __halves2bfloat162(value, 0); + else + fval = __halves2bfloat162(0, value); + + __hip_bfloat162 *in = (__hip_bfloat162 *)(aligned_addr); + __hip_bfloat162 out = unsafeAtomicAdd(in , fval); + if (is_lower) + return __low2bfloat16(out); + return __high2bfloat16(out); +} #endif // defined(__clang__) && defined(__HIP__) #endif diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h index 1a08bb8df1..797f33f25c 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h @@ -30,9 +30,6 @@ THE SOFTWARE. #define __HOST_DEVICE__ __host__ __device__ #include #include "hip/amd_detail/host_defines.h" -#if defined(__clang__) && defined(__HIP__) - #include "hip/amd_detail/amd_hip_atomic.h" -#endif // defined(__clang__) && defined(__HIP__) #include #if defined(__cplusplus) #include @@ -1558,6 +1555,28 @@ THE SOFTWARE. return old_val.h2r; #endif } + inline __device__ __half unsafeAtomicAdd(__half* address, __half value) { + static_assert(sizeof(unsigned short int) == sizeof(__half_raw)); + unsigned short int* address_as_short = reinterpret_cast(address); + // Align to 4 bytes + unsigned int* aligned_addr = __builtin_bit_cast(unsigned int*, + __builtin_bit_cast(unsigned long long int, address_as_short) + & (unsigned long long int)(~0x3)); + + bool is_lower = __builtin_bit_cast(unsigned long long int, aligned_addr) == + __builtin_bit_cast(unsigned long long int, address); + __half2 fval; + if (is_lower) + fval = __halves2half2(value, 0); + else + fval = __halves2half2(0, value); + + __half2 *in = (__half2 *)(aligned_addr); + __half2 out = unsafeAtomicAdd(in , fval); + if (is_lower) + return __low2half(out); + return __high2half(out); + } #endif // defined(__clang__) && defined(__HIP__) // Math functions