GPU TPC clusterizer: decode the dense-link ZS channels after the link headers on Metal

GPUTPCCFDecodeZSDenseLink decodes a page with one warp, walking it
timebin by timebin and link by link. For every link header, the warp
also collected the channels with data, a thread per channel with three
warp scans, inside a serial chain of dependent loads that is the
latency limit of the kernel.

On Metal the header loop now only walks the headers, and thread 0
stores each link's level-2 bitmask and the offset of its channel
bitmask in shared memory (+64 bytes). The channels are then collected
with one thread per byte of a link's channel bitmask: one warp scan per
32 bytes instead of about three per link. Same samples in the same
order, same pointers into the page, including into the next one, same
error paths.

On an M1 Max the decode takes 23 to 39% less time per launch in normal
multi-stream runs (pbpb30 4.9 -> 3.7 ms, 50 kHz TF0 6.7 -> 4.2 ms),
about 1 to 2% of the full time frame. The cluster dumps are
byte-identical: sim5pp against the CPU, 30 kHz (1 of 4) and 50 kHz TF0
before against after. pbpb30 in deterministic mode gives 22,183 tracks
as before, and a 36 against 36 pass comparison at 50 kHz shows no
difference in the track count.

It is Metal-only for now, since it was measured only there: CUDA, HIP,
OpenCL and the CPU compile the existing code unchanged. The chain it
shortens is there on every backend, so it could be enabled for them
once timed.

Author: ktf

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDecodeZS.cxx

@@ -670,6 +670,9 @@ GPUd() int16_t GPUTPCCFDecodeZSDenseLink::DecodeTB(
   int32_t timeBin = (linkBC + (uint64_t)(raw::RDHUtils::getHeartBeatOrbit(*rawDataHeader) - ctx.firstHBF) * o2::constants::lhc::LHCMaxBunches) / LHCBCPERTIMEBIN;
 
   int16_t nSamplesInTB = 0;
+#ifdef __METAL__
+  const uint8_t* const tbStart = ctx.page;
+#endif
 
   // Read timebin link headers
   for (uint8_t iLink = 0; iLink < nLinksInTimebin; iLink++) {
@@ -687,6 +690,15 @@ GPUd() int16_t GPUTPCCFDecodeZSDenseLink::DecodeTB(
       MAYBE_PAGE_OVERFLOW(ctx.page);
     }
 
+#ifdef __METAL__ // Channels collected after this loop, see below. Measured on Metal only, not yet timed on CUDA and HIP
+    if (ctx.iThread == 0) {
+      smem.linkBitmaskL2[iLink] = bitmaskL2;
+      smem.linkBitmaskOffset[iLink] = ctx.page - tbStart;
+    }
+
+    ConsumeBytes(ctx.page, CAMath::Popcount(bitmaskL2));
+    MAYBE_PAGE_OVERFLOW(ctx.page);
+#else
     int32_t nBytesBitmask = CAMath::Popcount(bitmaskL2);
 
     for (int32_t chan = ctx.iThread; chan < CAMath::nextMultipleOf<NTHREADS>(80); chan += NTHREADS) {
@@ -713,9 +725,39 @@ GPUd() int16_t GPUTPCCFDecodeZSDenseLink::DecodeTB(
     if (ctx.iThread == 0) {
       smem.samplesPerLinkEnd[iLink] = nSamplesInTB;
     }
+#endif
 
   } // for (uint8_t iLink = 0; iLink < nLinksInTimebin; iLink++)
 
+#ifdef __METAL__
+  GPUbarrierWarp(); // Ensure the link headers are in shared memory before reading them
+
+  // Collect the channels with data, one byte of a link's channel bitmask per thread. A thread per
+  // channel would need 3 warp-wide scans for each link, in a chain that is long enough as it is.
+  for (int32_t slot = ctx.iThread; slot < CAMath::nextMultipleOf<NTHREADS>(nLinksInTimebin * 10); slot += NTHREADS) {
+    const int32_t iLink = slot / 10;
+    const int32_t chanL2Idx = slot % 10;
+    uint8_t bitmask = 0;
+    if (iLink < nLinksInTimebin) {
+      const uint16_t bitmaskL2 = smem.linkBitmaskL2[iLink];
+      if (TEST_BIT(bitmaskL2, chanL2Idx)) {
+        const int32_t chanByteOffset = CAMath::Popcount(bitmaskL2) - 1 - CAMath::Popcount(bitmaskL2 >> (chanL2Idx + 1));
+        bitmask = PEEK_OVERFLOW(tbStart + smem.linkBitmaskOffset[iLink], chanByteOffset);
+      }
+    }
+    const int32_t nSamplesThread = CAMath::Popcount(bitmask);
+    const int32_t threadSampleEnd = warp_scan_inclusive_add(nSamplesThread);
+    int32_t sampleOffset = nSamplesInTB + threadSampleEnd - nSamplesThread;
+    for (uint32_t bits = bitmask; bits; bits &= bits - 1) {
+      smem.rawFECChannels[sampleOffset++] = chanL2Idx * 8 + CAMath::Ctz(bits);
+    }
+    if (chanL2Idx == 9 && iLink < nLinksInTimebin) {
+      smem.samplesPerLinkEnd[iLink] = nSamplesInTB + threadSampleEnd;
+    }
+    nSamplesInTB += warp_broadcast(threadSampleEnd, NTHREADS - 1);
+  }
+#endif
+
   GPUbarrierWarp(); // Ensure all writes to shared memory are finished, before reading it
 
   if (nSamplesInTB > nSamplesLeftInPage) {

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFDecodeZS.h

@@ -165,6 +165,10 @@ class GPUTPCCFDecodeZSDenseLink : public GPUTPCCFDecodeZSLinkBase
     // GPUCA_SHARED_STORAGE(uint32_t ZSPage[o2::tpc::TPCZSHDR::TPC_ZS_PAGE_SIZE / sizeof(uint32_t)]);
     uint16_t samplesPerLinkEnd[MaxNLinksPerTimebin]; // Offset from end of TB link header to first sample not in this link
     uint8_t linkIds[MaxNLinksPerTimebin];
+#ifdef __METAL__
+    uint16_t linkBitmaskL2[MaxNLinksPerTimebin];
+    uint16_t linkBitmaskOffset[MaxNLinksPerTimebin]; // Offset of the link's channel bitmask from the timebin header
+#endif
     uint8_t rawFECChannels[MaxNLinksPerTimebin * 80];
   };