GPU: implement the work-group and warp collectives for Metal

GPU/Common/GPUCommonAlgorithm.h

@@ -500,6 +500,111 @@ GPUdi() T warp_reduce_max(T v)
   return __syncthreads_count(pred);
 }
 
+#elif defined(__METAL__)
+// MSL has SIMD-group collectives only, so the work-group ones scan within each
+// SIMD group and then across the group totals, which go through smem as with
+// cub. The thread index is a kernel parameter in MSL, so like smem it is taken
+// from the caller's scope. A 1D threadgroup is split into SIMD groups of
+// consecutive threads, of GPUCA_WARP_SIZE each. The SIMD scans take 32-bit
+// lanes; the narrower types are counts and are widened. Sums are returned
+// promoted, as v + 0 would be, since the callers pass !!pred: a bool would
+// truncate them to 0 / 1.
+
+template <class T>
+GPUdi() auto warp_scan_inclusive_add(T v) -> decltype(+v)
+{
+  return simd_prefix_inclusive_sum((int32_t)v);
+}
+
+template <class T>
+GPUdi() T warp_broadcast(T v, int32_t i)
+{
+  return (T)simd_shuffle((int32_t)v, (ushort)i);
+}
+
+template <class T>
+GPUdi() T warp_shuffle_xor(T v, int32_t laneMask)
+{
+  return simd_shuffle_xor(v, (ushort)laneMask);
+}
+
+template <int32_t Width, class T>
+GPUdi() T warp_reduce_add(T v)
+{
+  for (int32_t m = Width / 2; m > 0; m >>= 1) {
+    v += warp_shuffle_xor(v, m);
+  }
+  return v;
+}
+
+template <int32_t Width, class T>
+GPUdi() T warp_reduce_min(T v)
+{
+  for (int32_t m = Width / 2; m > 0; m >>= 1) {
+    const T o = warp_shuffle_xor(v, m);
+    v = o < v ? o : v;
+  }
+  return v;
+}
+
+template <int32_t Width, class T>
+GPUdi() T warp_reduce_max(T v)
+{
+  for (int32_t m = Width / 2; m > 0; m >>= 1) {
+    const T o = warp_shuffle_xor(v, m);
+    v = o > v ? o : v;
+  }
+  return v;
+}
+
+#define work_group_scan_inclusive_add(v) work_group_scan_inclusive_add_FUNC(v, smem, iThread, nThreads)
+template <class T, class S>
+GPUdi() auto work_group_scan_inclusive_add_FUNC(T v, S& smem, int32_t iThread, int32_t nThreads) -> decltype(+v)
+{
+  const int32_t iWarp = iThread / GPUCA_WARP_SIZE;
+  const int32_t nWarps = (nThreads + GPUCA_WARP_SIZE - 1) / GPUCA_WARP_SIZE;
+  int32_t x = simd_prefix_inclusive_sum((int32_t)v);
+  if (iThread % GPUCA_WARP_SIZE == GPUCA_WARP_SIZE - 1 || iThread == nThreads - 1) {
+    smem.warpPredicateSum[iWarp] = x;
+  }
+  GPUbarrier();
+  if (iWarp == 0) {
+    const int32_t s = simd_prefix_inclusive_sum(iThread < nWarps ? smem.warpPredicateSum[iThread] : 0);
+    if (iThread < nWarps) {
+      smem.warpPredicateSum[iThread] = s;
+    }
+  }
+  GPUbarrier();
+  if (iWarp > 0) {
+    x += smem.warpPredicateSum[iWarp - 1];
+  }
+  GPUbarrier();
+  return x;
+}
+
+#define work_group_broadcast(v, i) work_group_broadcast_FUNC(v, i, smem, iThread)
+template <class T, class S>
+GPUdi() T work_group_broadcast_FUNC(T v, int32_t i, S& smem, int32_t iThread)
+{
+  if (iThread == i) {
+    smem.tmpBroadcast = v;
+  }
+  GPUbarrier();
+  T retVal = smem.tmpBroadcast;
+  GPUbarrier();
+  return retVal;
+}
+
+#define work_group_reduce_add(v) work_group_reduce_add_FUNC(v, smem, iThread, nThreads)
+template <class T, class S>
+GPUdi() auto work_group_reduce_add_FUNC(T v, S& smem, int32_t iThread, int32_t nThreads) -> decltype(+v)
+{
+  return work_group_broadcast_FUNC(work_group_scan_inclusive_add_FUNC(v, smem, iThread, nThreads), nThreads - 1, smem, iThread);
+}
+
+#define work_group_count(pred) ((uint32_t)work_group_reduce_add_FUNC((int32_t)(pred), smem, iThread, nThreads))
+#define work_group_any(pred) (work_group_count(pred) != 0)
+
 #else
 // Trivial implementation for the CPU
 

GPU/GPUTracking/Base/GPUGeneralKernels.h

@@ -76,6 +76,12 @@ class GPUKernelTemplate
       int32_t tmpBroadcast;
       int32_t warpPredicateSum[I / GPUCA_WARP_SIZE];
     };
+#elif defined(__METAL__) && defined(GPUCA_GPUCODE)
+    // What the Metal work-group collectives in GPUCommonAlgorithm.h go through
+    union {
+      int32_t tmpBroadcast;
+      int32_t warpPredicateSum[I / GPUCA_WARP_SIZE];
+    };
 #endif
   };
 

GPU/GPUTracking/TPCClusterFinder/CfUtils.h

@@ -65,7 +65,7 @@ class CfUtils
   }
 
   template <size_t BlockSize, typename SharedMemory>
-  static GPUdi() int32_t blockPredicateScan(SharedMemory& smem, int32_t pred, int32_t* sum = nullptr)
+  static GPUdi() int32_t blockPredicateScan(SharedMemory& smem, int32_t pred, int32_t* sum GPUCA_THREAD_INFO_DECL)
   {
 #if defined(__HIPCC__) || defined(__CUDACC__)
     int32_t iThread =
@@ -115,7 +115,7 @@ class CfUtils
   }
 
   template <size_t BlockSize, typename SharedMemory>
-  static GPUdi() int32_t blockPredicateSum(SharedMemory& smem, int32_t pred)
+  static GPUdi() int32_t blockPredicateSum(SharedMemory& smem, int32_t pred GPUCA_THREAD_INFO_DECL)
   {
 #if defined(__HIPCC__) || defined(__CUDACC__)
     int32_t iThread =
@@ -152,12 +152,12 @@ class CfUtils
   }
 
   template <size_t SCRATCH_PAD_WORK_GROUP_SIZE, typename SharedMemory>
-  static GPUdi() uint16_t partition(SharedMemory& smem, uint16_t ll, bool pred, uint16_t partSize, uint16_t* newPartSize)
+  static GPUdi() uint16_t partition(SharedMemory& smem, uint16_t ll, bool pred, uint16_t partSize, uint16_t* newPartSize GPUCA_THREAD_INFO_DECL)
   {
     bool participates = ll < partSize;
 
     int32_t part;
-    int32_t lpos = blockPredicateScan<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, int32_t(!pred && participates), &part);
+    int32_t lpos = blockPredicateScan<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, int32_t(!pred && participates), &part GPUCA_THREAD_INFO_PROVIDE);
 
     uint16_t pos = (participates && !pred) ? lpos : part;
 

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDeconvolution.cxx

@@ -52,7 +52,7 @@ GPUdii() void GPUTPCCFDeconvolution::deconvolutionImpl(int32_t nBlocks, int32_t
 
   uint16_t in3x3 = 0;
   bool exclude3x3 = iamPeak || !pos.valid();
-  partId = CfUtils::partition<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, ll, exclude3x3, SCRATCH_PAD_WORK_GROUP_SIZE, &in3x3);
+  partId = CfUtils::partition<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, ll, exclude3x3, SCRATCH_PAD_WORK_GROUP_SIZE, &in3x3 GPUCA_THREAD_INFO_PROVIDE);
 
   if (partId < in3x3) {
     smem.posBcast1[partId] = pos;
@@ -76,7 +76,7 @@ GPUdii() void GPUTPCCFDeconvolution::deconvolutionImpl(int32_t nBlocks, int32_t
   }
 
   uint16_t in5x5 = 0;
-  partId = CfUtils::partition<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, partId, peakCount > 0 && !exclude3x3, in3x3, &in5x5);
+  partId = CfUtils::partition<SCRATCH_PAD_WORK_GROUP_SIZE>(smem, partId, peakCount > 0 && !exclude3x3, in3x3, &in5x5 GPUCA_THREAD_INFO_PROVIDE);
 
   if (partId < in5x5) {
     smem.posBcast1[partId] = pos;

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFPeakFinder.cxx

@@ -50,7 +50,7 @@ GPUdii() bool GPUTPCCFPeakFinder::isPeak(
     ll,
     belowThreshold,
     SCRATCH_PAD_WORK_GROUP_SIZE,
-    &lookForPeaks);
+    &lookForPeaks GPUCA_THREAD_INFO_PROVIDE);
 
   if (partId < lookForPeaks) {
     posBcast[partId] = pos;

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFStreamCompaction.cxx

@@ -36,7 +36,7 @@ GPUdii() void GPUTPCCFStreamCompaction::Thread<GPUTPCCFStreamCompaction::scanSta
     pred = predicate[iThreadGlobal];
   }
 
-  int32_t nElemsInBlock = CfUtils::blockPredicateSum<GPUCA_PAR_CF_SCAN_WORKGROUP_SIZE>(smem, pred);
+  int32_t nElemsInBlock = CfUtils::blockPredicateSum<GPUCA_PAR_CF_SCAN_WORKGROUP_SIZE>(smem, pred GPUCA_THREAD_INFO_PROVIDE);
 
   int32_t lastThread = nThreads - 1;
   if (iThread == lastThread) {
@@ -115,7 +115,7 @@ GPUdii() void GPUTPCCFStreamCompaction::Thread<GPUTPCCFStreamCompaction::compact
   bool iAmDummy = (iThreadGlobal >= nElems);
 
   int32_t pred = (iAmDummy) ? 0 : predicate[iThreadGlobal];
-  int32_t offsetInBlock = CfUtils::blockPredicateScan<GPUCA_PAR_CF_SCAN_WORKGROUP_SIZE>(smem, pred);
+  int32_t offsetInBlock = CfUtils::blockPredicateScan<GPUCA_PAR_CF_SCAN_WORKGROUP_SIZE>(smem, pred, nullptr GPUCA_THREAD_INFO_PROVIDE);
 
   SizeT globalOffsetOut = offsetInBlock;
   if (iBlock > 0) {