From 2335200dc0ab65a39607069ae1796e21f088962e Mon Sep 17 00:00:00 2001 From: Giulio Eulisse <10544+ktf@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:59:43 +0100 Subject: [PATCH] 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. --- GPU/Common/GPUCommonAlgorithm.h | 105 ++++++++++++++++++ GPU/GPUTracking/Base/GPUGeneralKernels.h | 6 + GPU/GPUTracking/TPCClusterFinder/CfUtils.h | 8 +- .../GPUTPCCFDeconvolution.cxx | 4 +- .../TPCClusterFinder/GPUTPCCFPeakFinder.cxx | 2 +- .../GPUTPCCFStreamCompaction.cxx | 4 +- 6 files changed, 120 insertions(+), 9 deletions(-) diff --git a/GPU/Common/GPUCommonAlgorithm.h b/GPU/Common/GPUCommonAlgorithm.h index a96eb05a7aa65..fddfdd2f889da 100644 --- a/GPU/Common/GPUCommonAlgorithm.h +++ b/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 +GPUdi() auto warp_scan_inclusive_add(T v) -> decltype(+v) +{ + return simd_prefix_inclusive_sum((int32_t)v); +} + +template +GPUdi() T warp_broadcast(T v, int32_t i) +{ + return (T)simd_shuffle((int32_t)v, (ushort)i); +} + +template +GPUdi() T warp_shuffle_xor(T v, int32_t laneMask) +{ + return simd_shuffle_xor(v, (ushort)laneMask); +} + +template +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 +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 +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 +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 +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 +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 diff --git a/GPU/GPUTracking/Base/GPUGeneralKernels.h b/GPU/GPUTracking/Base/GPUGeneralKernels.h index 871cc21ee2bfa..904d66a79d6ee 100644 --- a/GPU/GPUTracking/Base/GPUGeneralKernels.h +++ b/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 }; diff --git a/GPU/GPUTracking/TPCClusterFinder/CfUtils.h b/GPU/GPUTracking/TPCClusterFinder/CfUtils.h index 3bc426522deff..7a3166452315d 100644 --- a/GPU/GPUTracking/TPCClusterFinder/CfUtils.h +++ b/GPU/GPUTracking/TPCClusterFinder/CfUtils.h @@ -65,7 +65,7 @@ class CfUtils } template - 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 - 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 - 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(smem, int32_t(!pred && participates), &part); + int32_t lpos = blockPredicateScan(smem, int32_t(!pred && participates), &part GPUCA_THREAD_INFO_PROVIDE); uint16_t pos = (participates && !pred) ? lpos : part; diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDeconvolution.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDeconvolution.cxx index d6b8703a9b35d..60c3acaa50c04 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDeconvolution.cxx +++ b/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(smem, ll, exclude3x3, SCRATCH_PAD_WORK_GROUP_SIZE, &in3x3); + partId = CfUtils::partition(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(smem, partId, peakCount > 0 && !exclude3x3, in3x3, &in5x5); + partId = CfUtils::partition(smem, partId, peakCount > 0 && !exclude3x3, in3x3, &in5x5 GPUCA_THREAD_INFO_PROVIDE); if (partId < in5x5) { smem.posBcast1[partId] = pos; diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFPeakFinder.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFPeakFinder.cxx index d0b134badc49a..7e02d68ccc44d 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFPeakFinder.cxx +++ b/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; diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFStreamCompaction.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFStreamCompaction.cxx index 0f2fd235dc0d0..8ea2b209907ed 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFStreamCompaction.cxx +++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCCFStreamCompaction.cxx @@ -36,7 +36,7 @@ GPUdii() void GPUTPCCFStreamCompaction::Thread(smem, pred); + int32_t nElemsInBlock = CfUtils::blockPredicateSum(smem, pred GPUCA_THREAD_INFO_PROVIDE); int32_t lastThread = nThreads - 1; if (iThread == lastThread) { @@ -115,7 +115,7 @@ GPUdii() void GPUTPCCFStreamCompaction::Thread= nElems); int32_t pred = (iAmDummy) ? 0 : predicate[iThreadGlobal]; - int32_t offsetInBlock = CfUtils::blockPredicateScan(smem, pred); + int32_t offsetInBlock = CfUtils::blockPredicateScan(smem, pred, nullptr GPUCA_THREAD_INFO_PROVIDE); SizeT globalOffsetOut = offsetInBlock; if (iBlock > 0) {