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. The cast has to keep the address space, which a generic pointer does not carry into atomic_*_explicit, hence the two overloads: threadgroup stays threadgroup for the *Shared entry points, anything else resolves to device. GPUTPCCFDecodeZS's shared memory had to be declared GPUsharedref() for that to hold; without it the reference is generic by the time AtomicAddShared sees it.
ktf · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026 · Sep 25, 2026
commit ab2f08346036b94feef3a71c7f2ff9e4456c6458
@@ -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__

@@ -475,6 +475,17 @@ 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 overloads keep the address space, which a
+// generic pointer would not carry into atomic_*_explicit.
+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)
 {
@@ -485,7 +496,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);
@@ -505,7 +516,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
@@ -523,7 +534,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
@@ -541,7 +552,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)) {
@@ -559,7 +570,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)) {

@@ -37,12 +37,12 @@ using namespace o2::tpc::constants;
 // ===========================================================================
 
 template <>
-GPUdii() void GPUTPCCFDecodeZS::Thread<GPUTPCCFDecodeZS::decodeZS>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer, int32_t firstHBF, int32_t tpcTimeBinCut)
+GPUdii() void GPUTPCCFDecodeZS::Thread<GPUTPCCFDecodeZS::decodeZS>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& clusterer, int32_t firstHBF, int32_t tpcTimeBinCut)
 {
   GPUTPCCFDecodeZS::decode(clusterer, smem, nBlocks, nThreads, iBlock, iThread, firstHBF, tpcTimeBinCut);
 }
 
-GPUdii() void GPUTPCCFDecodeZS::decode(GPUTPCClusterFinder& clusterer, GPUSharedMemory& s, int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, int32_t firstHBF, int32_t tpcTimeBinCut)
+GPUdii() void GPUTPCCFDecodeZS::decode(GPUTPCClusterFinder& clusterer, GPUsharedref() GPUSharedMemory& s, int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, int32_t firstHBF, int32_t tpcTimeBinCut)
 {
   const uint32_t sector = clusterer.mISector;
 #ifdef GPUCA_GPUCODE

@@ -45,7 +45,7 @@ class GPUTPCCFDecodeZS : public GPUKernelTemplate
     decodeZS,
   };
 
-  static GPUd() void decode(GPUTPCClusterFinder& clusterer, GPUSharedMemory& s, int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, int32_t firstHBF, int32_t tpcTimeBinCut);
+  static GPUd() void decode(GPUTPCClusterFinder& clusterer, GPUsharedref() GPUSharedMemory& s, int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, int32_t firstHBF, int32_t tpcTimeBinCut);
 
   typedef GPUTPCClusterFinder processorType;
   GPUhdi() static processorType* Processor(GPUConstantMem& processors)
@@ -59,7 +59,7 @@ class GPUTPCCFDecodeZS : public GPUKernelTemplate
   }
 
   template <int32_t iKernel = defaultKernel, typename... Args>
-  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer, Args... args);
+  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& clusterer, Args... args);
 };
 
 class GPUTPCCFDecodeZSLinkBase : public GPUKernelTemplate