GPU/TPC: Forward MC labels for HIP clusters

Author: fweig

GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx

@@ -1059,8 +1059,12 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
           return;
         }
 
-        if (propagateMCLabels && fragment.index == 0) {
-          clusterer.PrepareMC();
+        if (propagateMCLabels) {
+          if (fragment.index == 0) {
+            // Must be only called on the first fragment as some buffers are used across the whole timeframe
+            clusterer.AllocMCBuffers();
+          }
+          clusterer.InitMCBuffersForFragment();
           clusterer.mPinputLabels = digitsMC->v[iSector];
           if (clusterer.mPinputLabels == nullptr) {
             GPUFatal("MC label container missing, sector %d", iSector);
@@ -1358,9 +1362,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
 
           if (doGPU && propagateMCLabels) {
             TransferMemoryResourceLinkToHost(RecoStep::TPCClusterFinding, clusterer.mScratchId, lane);
-            if (doGPU) {
-              SynchronizeStream(lane);
-            }
+            SynchronizeStream(lane);
             runKernel<GPUTPCCFClusterizer>({GetGrid(clusterer.mPmemory->counters.nClusters, lane, GPUReconstruction::krnlDeviceType::CPU), {iSector}}, 1); // Computes MC labels
           }
         } // if (nRegularClusters != 0) {
@@ -1369,11 +1371,16 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
         // TODO: Move this right after CheckPadBaseline once tail zeroing is moved into this kernel.
         if (rec()->GetParam().rec.tpc.hipTailFilter) {
           runKernel<GPUTPCCFHIPTailConnector>({GetGridBlk(GPUTPCGeometry::NROWS, lane), {iSector}});
-          runKernel<GPUTPCCFHIPClusterizer>({GetGridBlk(GPUTPCGeometry::NROWS, lane), {iSector}});
+          runKernel<GPUTPCCFHIPClusterizer>({GetGridBlk(GPUTPCGeometry::NROWS, lane), {iSector}}, 0);
           if (doGPU && (nRegularClusters == 0 || GetProcessingSettings().debugLevel >= 3)) {
             TransferMemoryResourceLinkToHost(RecoStep::TPCClusterFinding, clusterer.mMemoryId, lane);
             SynchronizeStream(lane);
           }
+          if (doGPU && propagateMCLabels) {
+            TransferMemoryResourceLinkToHost(RecoStep::TPCClusterFinding, clusterer.mScratchId, lane);
+            SynchronizeStream(lane);
+            runKernel<GPUTPCCFHIPClusterizer>({GetGrid(GPUTPCGeometry::NROWS, lane, GPUReconstruction::krnlDeviceType::CPU), {iSector}}, 1); // Computes MC labels
+          }
         }
 
         bool hasClusters = nRegularClusters != 0;
@@ -1445,20 +1452,16 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
       }
 
       if (not propagateMCLabels || not laneHasData[lane]) {
-        assert(propagateMCLabels ? mcLinearLabels.header.size() == nClsTotal : true);
         continue;
       }
 
       runKernel<GPUTPCCFMCLabelFlattener, GPUTPCCFMCLabelFlattener::setRowOffsets>({GetGrid(GPUTPCGeometry::NROWS, lane, GPUReconstruction::krnlDeviceType::CPU), {iSector}});
       GPUTPCCFMCLabelFlattener::setGlobalOffsetsAndAllocate(clusterer, mcLinearLabels);
       runKernel<GPUTPCCFMCLabelFlattener, GPUTPCCFMCLabelFlattener::flatten>({GetGrid(GPUTPCGeometry::NROWS, lane, GPUReconstruction::krnlDeviceType::CPU), {iSector}}, &mcLinearLabels);
-      clusterer.clearMCMemory();
       assert(propagateMCLabels ? mcLinearLabels.header.size() == nClsTotal : true);
     }
-    if (propagateMCLabels) {
-      for (int32_t lane = 0; lane < maxLane; lane++) {
-        processors()->tpcClusterer[iSectorBase + lane].clearMCMemory();
-      }
+    for (int32_t lane = 0; lane < maxLane; lane++) {
+      processors()->tpcClusterer[iSectorBase + lane].FreeMCBuffers();
     }
     if (buildNativeHost && buildNativeGPU && anyLaneHasData) {
       if (GetProcessingSettings().delayedOutput) {

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFCheckPadBaseline.cxx

@@ -20,6 +20,7 @@
 #include "GPUCommonAlgorithm.h"
 
 #ifndef GPUCA_GPUCODE
+#include "MCLabelAccumulator.h"
 #include "utils/VcShim.h"
 #endif
 
@@ -504,7 +505,7 @@ GPUd() void GPUTPCCFHIPTailConnector::Thread<0>(int32_t nBlocks, int32_t nThread
 // ======== HIP Clusterizer Kernel ========
 
 template <>
-GPUd() void GPUTPCCFHIPClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer)
+GPUd() void GPUTPCCFHIPClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer, uint8_t onlyMC)
 {
   if (iBlock >= (int32_t)GPUTPCGeometry::NROWS) {
     return;
@@ -514,33 +515,32 @@ GPUd() void GPUTPCCFHIPClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads,
   uint32_t nTails = clusterer.mPnHIPTails[row];
   nTails = CAMath::Min(nTails, (uint32_t)MaxHIPTailsPerRow - 1);
 
-  HIPTailDescriptor* tails = GetHIPTails(clusterer, row);
+  const auto* tails = GetHIPTails(clusterer, row);
   const auto& fragment = clusterer.mPmemory->fragment;
 
-  for (uint32_t iTail = iThread + 1; iTail <= nTails; iTail += nThreads) {
+  auto* clusterPosInRow = clusterer.mPhipClusterPosInRow
+                            ? clusterer.mPhipClusterPosInRow + row * MaxHIPTailsPerRow
+                            : nullptr;
 
-    auto* tail = &tails[iTail];
+  for (uint32_t iTail = iThread + 1; iTail <= nTails; iTail += nThreads) {
 
+    const auto* tail = &tails[iTail];
     if (tail->iPrev != 0) {
       continue;
     }
 
-    float qTot = tail->qTot;
-    float qMax = tail->qMax;
-    const float firstWeight = tail->qTot;
-    const float firstPad = tail->pad;
-    const float firstTime = HIPTailTimeMean(*tail);
-    float padSum = firstWeight * firstPad;
-    float padSqSum = firstWeight * firstPad * firstPad;
-    float timeSum = firstWeight * firstTime;
+    CPU_ONLY(auto labelAcc = MCLabelAccumulator{clusterer});
 
-    uint32_t tailStart = tail->tailStart;
-    uint32_t tailEnd = tail->tailEnd;
-
-    while (tail->iNext != 0) {
-
-      tail = &tails[tail->iNext];
+    float qTot = 0;
+    float qMax = 0;
+    float padSum = 0;
+    float padSqSum = 0;
+    float timeSum = 0;
+    uint32_t tailStart = (uint32_t)-1;
+    uint32_t tailEnd = 0;
 
+    // Zero-th element is empty tail
+    for (; tail != tails; tail = &tails[tail->iNext]) {
       const float tailWeight = tail->qTot;
       const float tailPad = tail->pad;
       const float tailTime = HIPTailTimeMean(*tail);
@@ -551,12 +551,14 @@ GPUd() void GPUTPCCFHIPClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads,
       timeSum += tailWeight * tailTime;
       tailStart = CAMath::Min<uint32_t>(tailStart, tail->tailStart);
       tailEnd = CAMath::Max<uint32_t>(tailEnd, tail->tailEnd);
+
+      CPU_ONLY(labelAcc.collectTail(row, tail->pad, tail->tailStart, tail->tailEnd));
     }
 
     const float weightSum = CAMath::Max(qTot, 1.f);
-    float padMean = padSum / weightSum;
-    float timeMean = timeSum / weightSum; // TODO: Use timebin of saturated signal instead! Time mean is biased for long tails.
-    float padSigma = CAMath::Sqrt(CAMath::Max(0.f, padSqSum / weightSum - padMean * padMean));
+    const float padMean = padSum / weightSum;
+    const float timeMean = timeSum / weightSum; // TODO: Use timebin of saturated signal instead! Time mean is biased for long tails.
+    const float padSigma = CAMath::Sqrt(CAMath::Max(0.f, padSqSum / weightSum - padMean * padMean));
 
     tpc::ClusterNative cn;
     cn.qMax = qMax;
@@ -568,13 +570,26 @@ GPUd() void GPUTPCCFHIPClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads,
     cn.setSigmaPad(padSigma);
 
     if (cn.qMax >= 1023) {
-      // Cut off clusters where the tail connection failed for some reason
-      // TODO: Deduplicate with GPUTPCCFClusterizer::sortIntoBuckets (can't call cross-kernel).
-      // TODO: Add error reporting for row cluster overflow.
-      uint32_t index = CAMath::AtomicAdd(&clusterer.mPclusterInRow[row], 1u);
-      if (index < clusterer.mNMaxClusterPerRow) {
-        clusterer.mPclusterByRow[clusterer.mNMaxClusterPerRow * row + index] = cn;
+
+      uint32_t index;
+
+      if (!onlyMC) {
+        // Cut off clusters where the tail connection failed for some reason
+        // TODO: Deduplicate with GPUTPCCFClusterizer::sortIntoBuckets (can't call cross-kernel).
+        // TODO: Add error reporting for row cluster overflow.
+        index = CAMath::AtomicAdd(&clusterer.mPclusterInRow[row], 1u);
+        if (index < clusterer.mNMaxClusterPerRow) {
+          clusterer.mPclusterByRow[clusterer.mNMaxClusterPerRow * row + index] = cn;
+        }
+        if (clusterPosInRow) {
+          clusterPosInRow[iTail] = index;
+        }
+      } else {
+        index = clusterPosInRow[iTail];
       }
+
+      CPU_ONLY(labelAcc.commit(row, index, clusterer.mNMaxClusterPerRow));
     }
-  }
+
+  } // for (uint32_t iTail = iThread + 1; iTail <= nTails; iTail += nThreads)
 }

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFCheckPadBaseline.h

@@ -193,7 +193,7 @@ class GPUTPCCFHIPClusterizer : public GPUKernelTemplate
   }
 
   template <int32_t iKernel = defaultKernel>
-  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer);
+  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& clusterer, uint8_t onlyMC);
 };
 
 } // namespace o2::gpu

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFClusterizer.cxx

@@ -33,7 +33,7 @@ GPUdii() void GPUTPCCFClusterizer::Thread<0>(int32_t nBlocks, int32_t nThreads,
   CfArray2D<PackedCharge> chargeMap(reinterpret_cast<PackedCharge*>(clusterer.mPchargeMap));
   CPU_ONLY(MCLabelAccumulator labelAcc(clusterer));
 
-  tpc::ClusterNative* clusterOut = (onlyMC) ? nullptr : clusterer.mPclusterByRow;
+  tpc::ClusterNative* clusterOut = onlyMC ? nullptr : clusterer.mPclusterByRow;
 
   GPUTPCCFClusterizer::computeClustersImpl(get_num_groups(0), get_local_size(0), get_group_id(0), get_local_id(0), clusterer, clusterer.mPmemory->fragment, smem, chargeMap, clusterer.mPfilteredPeakPositions, clusterer.Param().rec, CPU_PTR(&labelAcc), clusterer.mPmemory->counters.nClusters, clusterer.mNMaxClusterPerRow, clusterer.mPclusterInRow, clusterOut, clusterer.mPclusterPosInRow, true);
 }

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFClusterizer.inc

@@ -112,7 +112,6 @@ GPUdii() void GPUTPCCFClusterizer::updateClusterInner(
     PackedCharge p = buf[N * lid + i];
 
     Charge q = cluster->updateInner(p, d);
-
     CPU_ONLY(labelAcc->collect(pos.delta(d), q));
 
     aboveThreshold |= (uint8_t(q > calib.tpc.cfInnerThreshold) << i);
@@ -139,9 +138,7 @@ GPUdii() void GPUTPCCFClusterizer::updateClusterOuter(
 
     Delta2 d = cfconsts::OuterNeighbors[i];
 
-    Charge q = cluster->updateOuter(p, d);
-    static_cast<void>(q); // Avoid unused varible warning on GPU.
-
+    [[maybe_unused]] Charge q = cluster->updateOuter(p, d);
     CPU_ONLY(labelAcc->collect(pos.delta(d), q));
   }
 }

GPU/GPUTracking/TPCClusterFinder/GPUTPCCFMCLabelFlattener.cxx

@@ -49,12 +49,7 @@ GPUd() void GPUTPCCFMCLabelFlattener::Thread<GPUTPCCFMCLabelFlattener::setRowOff
   const Row row = get_global_id(0);
   const size_t clusterInRow = clusterer.mPclusterInRow[row];
 
-  // Label Flattener assumes 1 label container per cluster,
-  // but HIP clusters don't support MC labels yet and containers are missing for those clusters.
-  // So append empty label container for each HIP cluster.
-  // Note: This assumes that HIP cluster are store behind regular clusters!
   auto& labels = clusterer.mPlabelsByRow[row].data;
-  labels.resize(std::max(labels.size(), clusterInRow));
 
   uint32_t labelCount = 0;
 

GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.cxx

@@ -27,6 +27,8 @@
 #include "CfArray2D.h"
 #include "GPUTPCCFCheckPadBaseline.h"
 
+#include <algorithm>
+
 using namespace o2::gpu;
 using namespace o2::tpc;
 
@@ -38,7 +40,7 @@ void GPUTPCClusterFinder::InitializeProcessor()
 GPUTPCClusterFinder::~GPUTPCClusterFinder()
 {
   delete[] mMinMaxCN;
-  clearMCMemory();
+  FreeMCBuffers();
 }
 
 void* GPUTPCClusterFinder::SetPointersMemory(void* mem)
@@ -86,8 +88,10 @@ void* GPUTPCClusterFinder::SetPointersScratch(void* mem)
   computePointerWithAlignment(mem, mPfilteredPeakPositions, mNMaxClusters);
   if (mRec->GetProcessingSettings().runMC) {
     computePointerWithAlignment(mem, mPclusterPosInRow, mNMaxClusters);
+    computePointerWithAlignment(mem, mPhipClusterPosInRow, GPUTPCGeometry::NROWS * GPUTPCCFHIPClusterizer::MaxHIPTailsPerRow);
   } else {
     mPclusterPosInRow = nullptr;
+    mPhipClusterPosInRow = nullptr;
   }
   computePointerWithAlignment(mem, mPisPeak, mNMaxDigitsFragment);
   computePointerWithAlignment(mem, mPchargeMap, TPCMapMemoryLayout<decltype(*mPchargeMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen));
@@ -165,17 +169,24 @@ uint32_t GPUTPCClusterFinder::getNSteps(size_t items) const
   return c;
 }
 
-void GPUTPCClusterFinder::PrepareMC()
+void GPUTPCClusterFinder::AllocMCBuffers()
 {
   assert(mNMaxClusterPerRow > 0);
 
-  clearMCMemory();
-  mPindexMap = new uint32_t[TPCMapMemoryLayout<decltype(*mPindexMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen)];
+  FreeMCBuffers();
+  const size_t nItems = TPCMapMemoryLayout<decltype(*mPindexMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen);
+  mPindexMap = new uint32_t[nItems];
   mPlabelsByRow = new GPUTPCClusterMCInterimArray[GPUTPCGeometry::NROWS];
   mPlabelsInRow = new uint32_t[GPUTPCGeometry::NROWS];
 }
 
-void GPUTPCClusterFinder::clearMCMemory()
+void GPUTPCClusterFinder::InitMCBuffersForFragment()
+{
+  const size_t nItems = TPCMapMemoryLayout<decltype(*mPindexMap)>::items(mRec->GetProcessingSettings().overrideClusterizerFragmentLen);
+  std::fill_n(mPindexMap, nItems, uint32_t(-1));
+}
+
+void GPUTPCClusterFinder::FreeMCBuffers()
 {
   delete[] mPindexMap;
   mPindexMap = nullptr;

GPU/GPUTracking/TPCClusterFinder/GPUTPCClusterFinder.h

@@ -92,8 +92,9 @@ class GPUTPCClusterFinder : public GPUProcessor
   uint32_t getNSteps(size_t items) const;
   void SetNMaxDigits(size_t nDigits, size_t nPages, size_t nDigitsFragment, size_t nDigitsEndpointMax);
 
-  void PrepareMC();
-  void clearMCMemory();
+  void AllocMCBuffers();
+  void InitMCBuffersForFragment();
+  void FreeMCBuffers();
 #endif
   uint8_t* mPzs = nullptr;
   ZSOffset* mPzsOffsets = nullptr;
@@ -107,6 +108,7 @@ class GPUTPCClusterFinder : public GPUProcessor
   uint32_t* mPclusterPosInRow = nullptr; // store the index where the corresponding cluster is stored in a bucket.
                                          // Required when MC are enabled to write the mc data to the correct position.
                                          // Set to >= mNMaxClusterPerRow if cluster was discarded.
+  uint32_t* mPhipClusterPosInRow = nullptr; // Identical to mPclusterPosInRow. Need a seperate array for HIP cluster because tail index is used to identify clusters across GPU and CPU
   uint16_t* mPchargeMap = nullptr;
   uint8_t* mPpeakMap = nullptr;
   uint32_t* mPindexMap = nullptr;

GPU/GPUTracking/TPCClusterFinder/MCLabelAccumulator.cxx

@@ -26,13 +26,20 @@ MCLabelAccumulator::MCLabelAccumulator(GPUTPCClusterFinder& clusterer)
 {
 }
 
-void MCLabelAccumulator::collect(const CfChargePos& pos, Charge q)
+MCLabelAccumulator::~MCLabelAccumulator() = default;
+
+void MCLabelAccumulator::collect(const CfChargePos& pos, float q)
 {
   if (q == 0 || !engaged()) {
     return;
   }
 
+  // Use -1 as sentinel to indicate a missing label.
+  // Can't use zero charge, as HIP filter will zero existing digits.
   uint32_t index = mIndexMap[pos];
+  if (index == uint32_t(-1)) {
+    return;
+  }
 
   const auto& labels = mLabels->getLabels(index);
 
@@ -51,6 +58,22 @@ void MCLabelAccumulator::collect(const CfChargePos& pos, Charge q)
   }
 }
 
+void MCLabelAccumulator::collectTail(tpccf::Row row, tpccf::Pad pad, uint16_t tailStart, uint16_t tailEnd)
+{
+  if (!engaged()) {
+    return;
+  }
+
+  const auto basePos = CfChargePos{row, pad, 0};
+
+  for (uint16_t t = tailStart; t < tailEnd; t++) {
+    const auto pos = basePos.delta({0, (int16_t)t});
+    // Charge passed to collect() doesn't matter, collect() skips zero charges
+    // But we know there's an interesting value, but it was zeroed in chargeMap by tail filter
+    collect(pos, 1023.f);
+  }
+}
+
 void MCLabelAccumulator::commit(Row row, uint32_t indexInRow, uint32_t maxElemsPerBucket)
 {
   if (indexInRow >= maxElemsPerBucket || !engaged()) {

GPU/GPUTracking/TPCClusterFinder/MCLabelAccumulator.h

@@ -43,12 +43,15 @@ class MCLabelAccumulator
 
  public:
   MCLabelAccumulator(GPUTPCClusterFinder&);
+  ~MCLabelAccumulator(); // Explicit destructor to allow forward declaring MCCompLabel with std::vector
 
-  void collect(const CfChargePos&, tpccf::Charge);
+  void collect(const CfChargePos& pos, float q);
+
+  void collectTail(tpccf::Row row, tpccf::Pad pad, uint16_t tailStart, uint16_t tailEnd);
 
   bool engaged() const { return mLabels != nullptr && mOutput != nullptr; }
 
-  void commit(tpccf::Row, uint32_t, uint32_t);
+  void commit(tpccf::Row row, uint32_t indexInRow, uint32_t maxElemsPerBucket);
 
  private:
   CfArray2D<const uint32_t> mIndexMap;