GPU: make the Metal atomics operate on plain counters

GPUAtomic() expanded to metal::atomic, but the counters it guards are plain
integers in structures that the host allocates and transfers, and the code reads
them directly outside the atomic operations. That is exactly the CUDA and HIP
situation, where GPUAtomic() is a no-op and atomicAdd() takes a plain pointer,
so Metal now does the same.

MSL's atomic_*_explicit do want metal::atomic, which has the same size and
alignment, so the operand is cast at the one point where the operation is
issued. That cast has to be made in the address space the pointer is really in,
since casting a threadgroup pointer to a device atomic is a hard error, hence
the two overloads: threadgroup stays threadgroup for the *Shared entry points,
anything else resolves to device.

Author: ktf

GPU/Common/GPUCommonDefAPI.h

@@ -173,7 +173,7 @@
   #define GPUdouble() float
   #define GPUbarrier() threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup)
   #define GPUbarrierWarp() simdgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup)
-  #define GPUAtomic(type) atomic<type>                      // atomic variable type
+  #define GPUAtomic(type) type                      // atomic variable type
 #elif defined(__HIPCC__) //Defines for HIP
   #define GPUd() __device__
   #define GPUdDefault() __device__

GPU/Common/GPUCommonMath.h

@@ -489,6 +489,25 @@ GPUhdi() constexpr int32_t GPUCommonMath::Abs<int32_t>(int32_t x)
   return GPUCA_CHOICE(abs(x), abs(x), abs(x));
 }
 
+#ifdef __METAL__
+// The counters these operate on are plain integers in the transferred structures,
+// as they are for CUDA and HIP. MSL's atomic operations want metal::atomic, which
+// has the same size and alignment. The overload is picked by address space, and a
+// threadgroup pointer has to keep its own: casting one to a device atomic is an
+// error while the qualifier is visible, but a generic pointer resolves here to
+// device and the cast then goes through silently.
+template <class T>
+GPUdi() threadgroup metal::atomic<T>* GPUCommonMathMetalAtomic(threadgroup T* p)
+{
+  return reinterpret_cast<threadgroup metal::atomic<T>*>(p);
+}
+template <class T>
+GPUdi() device metal::atomic<T>* GPUCommonMathMetalAtomic(T* p)
+{
+  return (device metal::atomic<T>*)p;
+}
+#endif
+
 template <class S, class T>
 GPUdi() uint32_t GPUCommonMath::AtomicExchInternal(S* addr, T val)
 {
@@ -499,7 +518,7 @@ GPUdi() uint32_t GPUCommonMath::AtomicExchInternal(S* addr, T val)
 #elif defined(GPUCA_GPUCODE) && (defined(__CUDACC__) || defined(__HIPCC__))
   return ::atomicExch(addr, val);
 #elif defined(GPUCA_GPUCODE) && defined(__METAL__)
-  return atomic_exchange_explicit(addr, val, memory_order_relaxed);
+  return atomic_exchange_explicit(GPUCommonMathMetalAtomic(addr), val, memory_order_relaxed);
 #elif defined(WITH_OPENMP)
   uint32_t old;
   __atomic_exchange(addr, &val, &old, __ATOMIC_SEQ_CST);
@@ -519,7 +538,7 @@ GPUdi() bool GPUCommonMath::AtomicCASInternal(S* addr, T cmp, T val)
 #elif defined(GPUCA_GPUCODE) && (defined(__CUDACC__) || defined(__HIPCC__))
   return ::atomicCAS(addr, cmp, val) == cmp;
 #elif defined(GPUCA_GPUCODE) && defined(__METAL__)
-  return atomic_compare_exchange_weak_explicit(addr, &cmp, val, memory_order_relaxed, memory_order_relaxed);
+  return atomic_compare_exchange_weak_explicit(GPUCommonMathMetalAtomic(addr), &cmp, val, memory_order_relaxed, memory_order_relaxed);
 #elif defined(WITH_OPENMP)
   return __atomic_compare_exchange(addr, &cmp, &val, true, __ATOMIC_SEQ_CST, __ATOMIC_SEQ_CST);
 #else
@@ -537,7 +556,7 @@ GPUdi() uint32_t GPUCommonMath::AtomicAddInternal(S* addr, T val)
 #elif defined(GPUCA_GPUCODE) && (defined(__CUDACC__) || defined(__HIPCC__))
   return ::atomicAdd(addr, val);
 #elif defined(GPUCA_GPUCODE) && defined(__METAL__)
-  return atomic_fetch_add_explicit(addr, val, memory_order_relaxed);
+  return atomic_fetch_add_explicit(GPUCommonMathMetalAtomic(addr), val, memory_order_relaxed);
 #elif defined(WITH_OPENMP)
   return __atomic_add_fetch(addr, val, __ATOMIC_SEQ_CST) - val;
 #else
@@ -555,7 +574,7 @@ GPUdi() void GPUCommonMath::AtomicMaxInternal(S* addr, T val)
 #elif defined(GPUCA_GPUCODE) && (defined(__CUDACC__) || defined(__HIPCC__))
   ::atomicMax(addr, val);
 #elif defined(GPUCA_GPUCODE) && defined(__METAL__)
-  atomic_fetch_max_explicit(addr, val, memory_order_relaxed);
+  atomic_fetch_max_explicit(GPUCommonMathMetalAtomic(addr), val, memory_order_relaxed);
 #else
   S current;
   while ((current = *(volatile S*)addr) < val && !AtomicCASInternal(addr, current, val)) {
@@ -573,7 +592,7 @@ GPUdi() void GPUCommonMath::AtomicMinInternal(S* addr, T val)
 #elif defined(GPUCA_GPUCODE) && (defined(__CUDACC__) || defined(__HIPCC__))
   ::atomicMin(addr, val);
 #elif defined(GPUCA_GPUCODE) && defined(__METAL__)
-  atomic_fetch_min_explicit(addr, val, memory_order_relaxed);
+  atomic_fetch_min_explicit(GPUCommonMathMetalAtomic(addr), val, memory_order_relaxed);
 #else
   S current;
   while ((current = *(volatile S*)addr) > val && !AtomicCASInternal(addr, current, val)) {