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

Metal fell through to the CPU branch of GPUCommonAlgorithm.h, where a
scan, a reduction or a broadcast returns the thread's own value: right for
one thread per block, wrong for the 32- to 1024-thread groups Metal runs.
The TPC ZS decoding and the stream compaction of the cluster finder got
wrong offsets and counts, and found no clusters at all.

Implement them with MSL's SIMD-group functions, and the work-group ones by
scanning within each SIMD group and then across the group totals, which go
through smem as with cub. MSL has the thread index only as a kernel
parameter, so like smem it is taken from the caller's scope, and the
CfUtils block helpers get it through GPUCA_THREAD_INFO_DECL, which is
empty for the other backends. blockPredicateScan loses the default of its
last parameter for that. Sums are returned promoted, since the callers
pass !!pred and a bool would truncate them.

The CUDA, HIP, OpenCL and CPU code is unchanged.

Author: ktf

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) {