From bb52d9ed62e8953dc9673b7e4d0fd13608a4dbf6 Mon Sep 17 00:00:00 2001 From: Jatin Chaudhary Date: Mon, 19 Aug 2024 18:19:04 +0100 Subject: [PATCH] SWDEV-480489 - fix unsafeAtomicAdd Integration into pytorch pointed out some issues, value narrowing, to fix this we are now using unions. Also removed check for -munsafe* compiler flag. The check is now just on builtin detection. Change-Id: I49364503fa429bd862952f9b29879072afa6d553 --- hipamd/include/hip/amd_detail/amd_hip_bf16.h | 13 ++++++++----- hipamd/include/hip/amd_detail/amd_hip_fp16.h | 10 +++++----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/hipamd/include/hip/amd_detail/amd_hip_bf16.h b/hipamd/include/hip/amd_detail/amd_hip_bf16.h index 6deaab77fa..0a5d1497fd 100644 --- a/hipamd/include/hip/amd_detail/amd_hip_bf16.h +++ b/hipamd/include/hip/amd_detail/amd_hip_bf16.h @@ -1835,12 +1835,15 @@ __BF16_DEVICE_STATIC__ __hip_bfloat162 h2trunc(const __hip_bfloat162 h) { */ __BF16_DEVICE_STATIC__ __hip_bfloat162 unsafeAtomicAdd(__hip_bfloat162* address, __hip_bfloat162 value) { -#if defined(__AMDGCN_UNSAFE_FP_ATOMICS__) && __has_builtin(__builtin_amdgcn_flat_atomic_fadd_v2bf16) +#if __has_builtin(__builtin_amdgcn_flat_atomic_fadd_v2bf16) typedef short __attribute__((ext_vector_type(2))) vec_short2; - __hip_bfloat162_raw bf2_v = value; - vec_short2 s2_in{bf2_v.x, bf2_v.y}; - vec_short2 s2_ret = __builtin_amdgcn_flat_atomic_fadd_v2bf16((vec_short2*)address, s2_in); - return __hip_bfloat162_raw{s2_ret[0], s2_ret[1]}; + static_assert(sizeof(vec_short2) == sizeof(__hip_bfloat162_raw)); + union { + __hip_bfloat162_raw bf162_raw; + vec_short2 vs2; + } u{static_cast<__hip_bfloat162_raw>(value)}; + u.vs2 = __builtin_amdgcn_flat_atomic_fadd_v2bf16((vec_short2*)address, u.vs2); + return static_cast<__hip_bfloat162>(u.bf162_raw); #else static_assert(sizeof(unsigned int) == sizeof(__hip_bfloat162_raw)); union u_hold { diff --git a/hipamd/include/hip/amd_detail/amd_hip_fp16.h b/hipamd/include/hip/amd_detail/amd_hip_fp16.h index 39a522a503..cfae082468 100644 --- a/hipamd/include/hip/amd_detail/amd_hip_fp16.h +++ b/hipamd/include/hip/amd_detail/amd_hip_fp16.h @@ -1509,17 +1509,17 @@ THE SOFTWARE. // Atomic #if defined(__clang__) && defined(__HIP__) inline __device__ __half2 unsafeAtomicAdd(__half2* address, __half2 value) { - #if defined(__AMDGCN_UNSAFE_FP_ATOMICS__) && __has_builtin(__builtin_amdgcn_flat_atomic_fadd_v2f16) + #if __has_builtin(__builtin_amdgcn_flat_atomic_fadd_v2f16) // The api expects an ext_vector_type of half - typedef __fp16 __attribute__((ext_vector_type(2))) vec_fp162; + typedef _Float16 __attribute__((ext_vector_type(2))) vec_fp162; static_assert(sizeof(vec_fp162) == sizeof(__half2_raw)); union { __half2_raw h2r; vec_fp162 fp16; - } u {value}; - vec_fp162 ret = + } u {static_cast<__half2_raw>(value)}; + u.fp16 = __builtin_amdgcn_flat_atomic_fadd_v2f16((vec_fp162*)address, u.fp16); - return __half2{ret[0], ret[1]}; + return static_cast<__half2>(u.h2r); #else static_assert(sizeof(__half2_raw) == sizeof(unsigned int)); union u_hold {