From e112c935057434897bb12d9ab3910380a8bd5f58 Mon Sep 17 00:00:00 2001
From: Alex Voicu <alexandru.voicu@amd.com>
Date: Tue, 19 Mar 2024 18:38:52 +0200
Subject: [PATCH] Fix_msvc_atomics (#61)

* Fix MSVC build.
---
 include/hip/hip_atomic.h               | 29 ++++++++++++++++++--------
 src/include/hip/detail/atomic_msvc.hpp | 20 ++++++++++--------
 2 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/include/hip/hip_atomic.h b/include/hip/hip_atomic.h
index dc8a91c..d854e95 100644
--- a/include/hip/hip_atomic.h
+++ b/include/hip/hip_atomic.h
@@ -43,7 +43,8 @@ unsigned long atomicAdd(unsigned long* address, unsigned long val) noexcept
 }
 
 inline
-unsigned long long atomicAdd(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicAdd(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_add(address, val);
 }
@@ -91,7 +92,8 @@ long long atomicAnd(long long* address, long long val) noexcept
 }
 
 inline
-unsigned long long atomicAnd(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicAnd(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_and(address, val);
 }
@@ -101,7 +103,9 @@ unsigned long long atomicAnd(unsigned long long* address, unsigned long long val
 // FIXME: not available in CUDA or in HIP/ROCm
 inline
 unsigned short atomicCAS(
-    unsigned short* address, unsigned short compare, unsigned short val) noexcept
+    unsigned short* address,
+    unsigned short compare,
+    unsigned short val) noexcept
 {
     return hip::detail::atomic_cas(address, compare, val);
 }
@@ -129,7 +133,9 @@ unsigned long atomicCAS(
 
 inline
 unsigned long long atomicCAS(
-    unsigned long long* address, unsigned long long compare, unsigned long long val) noexcept
+    unsigned long long* address,
+    unsigned long long compare,
+    unsigned long long val) noexcept
 {
     return hip::detail::atomic_cas(address, compare, val);
 }
@@ -165,7 +171,8 @@ unsigned long atomicExch(unsigned long* address, unsigned long val) noexcept
 }
 
 inline
-unsigned long long atomicExch(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicExch(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_exchange(address, val);
 }
@@ -213,7 +220,8 @@ long long atomicMax(long long* address, long long val) noexcept
 }
 
 inline
-unsigned long long atomicMax(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicMax(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_max(address, val);
 }
@@ -246,7 +254,8 @@ long long atomicMin(long long* address, long long val) noexcept
 }
 
 inline
-unsigned long long atomicMin(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicMin(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_min(address, val);
 }
@@ -306,7 +315,8 @@ unsigned long atomicSub(unsigned long* address, unsigned long val) noexcept
 }
 
 inline
-unsigned long long atomicSub(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicSub(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_sub(address, val);
 }
@@ -340,7 +350,8 @@ long long atomicXor(long long* address, long long val) noexcept
 }
 
 inline
-unsigned long long atomicXor(unsigned long long* address, unsigned long long val) noexcept
+unsigned long long atomicXor(
+    unsigned long long* address, unsigned long long val) noexcept
 {
     return hip::detail::atomic_xor(address, val);
 }
diff --git a/src/include/hip/detail/atomic_msvc.hpp b/src/include/hip/detail/atomic_msvc.hpp
index d23146f..d46b011 100644
--- a/src/include/hip/detail/atomic_msvc.hpp
+++ b/src/include/hip/detail/atomic_msvc.hpp
@@ -26,13 +26,15 @@ namespace hip
         T atomic_add(T* p, T v) noexcept
         {
             using U = std::conditional_t<
-                std::is_same_v<T, long> || std::is_same_v<T, unsigned long> ||
-                    std::is_same_v<T, long long>,
+                std::is_same_v<T, long> || std::is_same_v<T, long long>,
                 T,
                 std::conditional_t<sizeof(T) == sizeof(long), long, long long>>;
 
             if constexpr (std::is_same_v<T, U>) {
-                return _InterlockedExchangeAdd(p, v);
+                if constexpr (sizeof(T) == sizeof(long)) {
+                    return _InterlockedExchangeAdd(p, v);
+                }
+                else return _InterlockedExchangeAdd64(p, v);
             }
             else if constexpr (sizeof(U) == sizeof(long)) {
                 return bit_cast<T>(_InterlockedExchangeAdd(
@@ -49,13 +51,15 @@ namespace hip
         T atomic_cas(T* p, T cmp, T v) noexcept
         {
             using U = std::conditional_t<
-                std::is_same_v<T, long long> ||
-                    std::is_same_v<T, unsigned long> || std::is_same_v<T, long>,
+                std::is_same_v<T, long> ||  std::is_same_v<T, long long>,
                 T,
                 std::conditional_t<sizeof(T) == sizeof(long), long, long long>>;
 
             if constexpr (std::is_same_v<T, U>) {
-                return _InterlockedCompareExchange(p, v, cmp);
+                if constexpr(sizeof(T) == sizeof(long)) {
+                    return _InterlockedCompareExchange(p, v, cmp);
+                }
+                else return _InterlockedCompareExchange64(p, v, cmp);
             }
             else if constexpr (sizeof(T) == sizeof(long)) {
                 return bit_cast<T>(_InterlockedCompareExchange(
@@ -157,9 +161,7 @@ namespace hip
         T atomic_exchange(T* p, T v) noexcept
         {
             using U = std::conditional_t<
-                std::is_same_v<T, long> ||
-                    std::is_same_v<T, unsigned long> ||
-                    std::is_same_v<T, long long>,
+                std::is_same_v<T, long> || std::is_same_v<T, long long>,
                 T,
                 std::conditional_t<sizeof(T) == sizeof(long), long, long long>>;