SWDEV-492165 - Add support for 16 bit atomicAdd and atomicCAS

Change-Id: I7f1e4876fe2960dff1775d27cb6443a89e146c86


[ROCm/clr commit: 93cff75928]

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<bool B, typename T, typename F> struct Cond_t;
@@ -113,6 +114,25 @@ T hip_cas_extrema_expander(T* p, T x, Cmp cmp, F f) noexcept
   return reinterpret_cast<const T&>(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

projects/clr/hipamd/include/hip/amd_detail/amd_hip_bf16.h

@@ -111,13 +111,10 @@
 #if !defined(__HIPCC_RTC__)
 #include <hip/amd_detail/amd_hip_common.h>
 #include <hip/amd_detail/amd_warp_functions.h>  // 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<unsigned short int *>(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

projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h

@@ -30,9 +30,6 @@ THE SOFTWARE.
   #define __HOST_DEVICE__ __host__ __device__
   #include <hip/amd_detail/amd_hip_common.h>
   #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 <assert.h>
   #if defined(__cplusplus)
     #include <algorithm>
@@ -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<unsigned short int *>(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