GPU TPC: Add possibility to run last way of TPC fit in separate kernel to rebuild track ion between

Author: davidrohr

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -138,6 +138,7 @@ AddOptionRTC(cfNoiseSuppressionEpsilon, uint8_t, 10, "", 0, "Cluster Finder: Dif
 AddOptionRTC(cfNoiseSuppressionEpsilonRelative, uint8_t, 76, "", 0, "Cluster Finder: Difference between peak and charge for the charge to count as a minima during noise suppression, relative as fraction of 255")
 AddOptionRTC(cfEdgeTwoPads, uint8_t, 0, "", 0, "Flag clusters with peak on the 2 pads closes to the sector edge as edge cluster")
 AddOptionRTC(nWays, uint8_t, 3, "", 0, "Do N fit passes in final fit of merger (must be odd to end with inward fit)")
+AddOptionRTC(rebuildTrackInFit, uint8_t, 1, "", 0, "Rebuild track completely during fit based on clusters closed to interpolated track positions")
 AddOptionRTC(trackFitRejectMode, int8_t, 5, "", 0, "0: no limit on rejection or missed hits, >0: break after n rejected hits, <0: reject at max -n hits")
 AddOptionRTC(rejectIFCLowRadiusCluster, uint8_t, 1, "", 0, "Reject clusters that get the IFC mask error during refit")
 AddOptionRTC(dEdxTruncLow, uint8_t, 2, "", 0, "Low truncation threshold, fraction of 128")

GPU/GPUTracking/Global/GPUChainTracking.cxx

@@ -261,6 +261,10 @@ bool GPUChainTracking::ValidateSettings()
     GPUError("Cannot do error interpolation with NWays < 3!");
     return false;
   }
+  if (param().rec.tpc.rebuildTrackInFit && !param().rec.tpc.mergerInterpolateErrors) {
+    GPUError("Need error interpolation to rebuild tracks during fit");
+    return false;
+  }
   if (param().continuousMaxTimeBin > (int32_t)GPUSettings::TPC_MAX_TF_TIME_BIN) {
     GPUError("configured max time bin exceeds 256 orbits");
     return false;

GPU/GPUTracking/Global/GPUChainTrackingMerger.cxx

@@ -234,7 +234,10 @@ int32_t GPUChainTracking::RunTPCTrackingMerger(bool synchronizeOutput)
     mOutputQueue.clear();
   }
 
-  runKernel<GPUTPCGMMergerTrackFit>(doGPU ? GetGrid(Merger.NMergedTracks(), 0) : GetGridAuto(0), mergerSortTracks ? 1 : 0);
+  runKernel<GPUTPCGMMergerTrackFit>(doGPU ? GetGrid(Merger.NMergedTracks(), 0) : GetGridAuto(0), mergerSortTracks ? 1 : 0, 0);
+  if (param().rec.tpc.rebuildTrackInFit) {
+    runKernel<GPUTPCGMMergerTrackFit>(doGPU ? GetGrid(Merger.NMergedTracks(), 0) : GetGridAuto(0), mergerSortTracks ? 1 : 0, 1);
+  }
   runKernel<GPUTPCGMMergerFollowLoopers>(GetGridAuto(0));
 
   DoDebugAndDump(RecoStep::TPCMerging, GPUChainTrackingDebugFlags::TPCMergingRefit, Merger, &GPUTPCGMMerger::DumpRefit, *mDebugFile);

GPU/GPUTracking/Merger/GPUTPCGMMergerGPU.cxx

@@ -19,11 +19,11 @@
 using namespace o2::gpu;
 
 template <>
-GPUdii() void GPUTPCGMMergerTrackFit::Thread<0>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& GPUrestrict() merger, int32_t mode)
+GPUdii() void GPUTPCGMMergerTrackFit::Thread<0>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& GPUrestrict() merger, int32_t mode, int32_t rebuilt)
 {
   GPUCA_TBB_KERNEL_LOOP(merger.GetRec(), int32_t, ii, merger.NMergedTracks(), {
     const int32_t i = mode ? merger.TrackOrderProcess()[ii] : ii;
-    GPUTPCGMTrackParam::RefitTrack(merger.MergedTracks()[i], i, &merger);
+    GPUTPCGMTrackParam::RefitTrack(merger.MergedTracks()[i], i, &merger, rebuilt);
   });
 }
 

GPU/GPUTracking/Merger/GPUTPCGMMergerGPU.h

@@ -36,7 +36,7 @@ class GPUTPCGMMergerTrackFit : public GPUTPCGMMergerGeneral
 {
  public:
   template <int32_t iKernel = defaultKernel>
-  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& merger, int32_t mode);
+  GPUd() static void Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& merger, int32_t mode, int32_t rebuilt);
 };
 
 class GPUTPCGMMergerFollowLoopers : public GPUTPCGMMergerGeneral

GPU/GPUTracking/Merger/GPUTPCGMTrackParam.cxx

@@ -47,7 +47,7 @@
 using namespace o2::gpu;
 using namespace o2::tpc;
 
-GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_t iTrk, GPUTPCGMMergedTrackHit* GPUrestrict() clusters, int32_t& GPUrestrict() N, int32_t& GPUrestrict() NTolerated, float& GPUrestrict() Alpha, float maxSinPhi, GPUTPCGMMergedTrack& GPUrestrict() track)
+GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_t iTrk, GPUTPCGMMergedTrackHit* GPUrestrict() clusters, int32_t& GPUrestrict() N, int32_t& GPUrestrict() NTolerated, float& GPUrestrict() Alpha, float maxSinPhi, GPUTPCGMMergedTrack& GPUrestrict() track, bool rebuilt)
 {
   static constexpr float kDeg2Rad = M_PI / 180.f;
   CADEBUG(static constexpr float kSectAngle = 2 * M_PI / 18.f);
@@ -60,7 +60,7 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_
   prop.SetMaterialTPC();
   prop.SetPolynomialField(&param.polynomialField);
   prop.SetMaxSinPhi(maxSinPhi);
-  if (param.rec.tpc.mergerInterpolateErrors) {
+  if (param.rec.tpc.mergerInterpolateErrors && !rebuilt) {
     for (int32_t i = 0; i < N; i++) {
       interpolation.hit[i].errorY = -1;
     }
@@ -70,12 +70,9 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_
   const int32_t maxN = N;
   int32_t ihitStart = 0;
   float covYYUpd = 0.f;
-  float lastUpdateX = -1.f;
-  uint8_t lastRow = 255;
-  uint8_t lastSector = 255;
   float deltaZ = 0.f;
 
-  for (int32_t iWay = 0; iWay < nWays; iWay++) {
+  for (int32_t iWay = rebuilt ? nWays - 1 : 0; iWay < nWays; iWay++) { // DR: Unrolling has no performance improvement on GPU, why?
     int32_t nMissed = 0, nMissed2 = 0;
     float sumInvSqrtCharge = 0.f;
     int32_t nAvgCharge = 0;
@@ -95,12 +92,14 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_
     prop.SetFitInProjections(true); // param.rec.fitInProjections == -1 ? (iWay == 0) : param.rec.fitInProjections); // TODO: Reenable once fixed
     prop.SetPropagateBzOnly(param.rec.fitPropagateBzOnly == -1 ? !finalFit : param.rec.fitPropagateBzOnly);
     prop.SetMatLUT((param.rec.useMatLUT && finalFit) ? merger->GetConstantMem()->calibObjects.matLUT : nullptr);
-    prop.SetTrack(this, iWay ? prop.GetAlpha() : Alpha);
+    prop.SetTrack(this, iWay && !rebuilt ? prop.GetAlpha() : Alpha);
     ConstrainSinPhi(iWay == 0 ? 0.95f : constants::MAX_SIN_PHI_LOW);
     CADEBUG(printf("Fitting track %d way %d (sector %d, alpha %f) !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n", iTrk, iWay, CAMath::Float2IntRn(prop.GetAlpha() / kSectAngle) + (mP[1] < 0 ? 18 : 0), prop.GetAlpha()));
 
     N = 0;
-    lastUpdateX = -1;
+    uint8_t lastRow = 255;
+    uint8_t lastSector = 255;
+    float lastUpdateX = -1;
     const bool inFlyDirection = iWay & 1;
     const int32_t wayDirection = (iWay & 1) ? -1 : 1;
 
@@ -110,9 +109,12 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_
         lastSector = clusters[ihit].sector;
       }
 
-      if ((param.rec.tpc.trackFitRejectMode > 0 && nMissed >= param.rec.tpc.trackFitRejectMode) || nMissed2 >= param.rec.tpc.trackFitMaxRowMissedHard || clusters[ihit].state & GPUTPCGMMergedTrackHit::flagReject) {
+      if ((param.rec.tpc.trackFitRejectMode > 0 && nMissed >= param.rec.tpc.trackFitRejectMode) || nMissed2 >= param.rec.tpc.trackFitMaxRowMissedHard || (clusters[ihit].state & GPUTPCGMMergedTrackHit::flagReject) || (rebuilt && (clusters[ihit].state & GPUTPCGMMergedTrackHit::flagHighIncl))) {
         CADEBUG(printf("\tSkipping hit %d, %d hits rejected, flag %X\n", ihit, nMissed, (int32_t)clusters[ihit].state));
-        if (finalOutInFit && !(clusters[ihit].state & GPUTPCGMMergedTrackHit::flagReject)) {
+        if (rebuilt && (clusters[ihit].state & GPUTPCGMMergedTrackHit::flagHighIncl)) {
+          NTolerated++;
+        }
+        if (finalOutInFit && !(clusters[ihit].state & (GPUTPCGMMergedTrackHit::flagReject | GPUTPCGMMergedTrackHit::flagHighIncl))) {
           clusters[ihit].state |= GPUTPCGMMergedTrackHit::flagRejectErr;
         }
         continue;
@@ -327,6 +329,14 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger* GPUrestrict() merger, int32_
     } else {
       deltaZ = 0.f;
     }
+
+    if (param.rec.tpc.rebuildTrackInFit && iWay == nWays - 2) {
+      Alpha = prop.GetAlpha();
+      if (ihitStart != 0) {
+        MarkClusters(clusters, 0, ihitStart - 1, 1, GPUTPCGMMergedTrackHit::flagHighIncl);
+      }
+      return true;
+    }
   }
   ConstrainSinPhi();
 
@@ -894,7 +904,7 @@ GPUd() bool GPUTPCGMTrackParam::CheckNumericalQuality(float overrideCovYY) const
   return ok;
 }
 
-GPUdii() void GPUTPCGMTrackParam::RefitTrack(GPUTPCGMMergedTrack& GPUrestrict() track, int32_t iTrk, GPUTPCGMMerger* GPUrestrict() merger) // VS: GPUd changed to GPUdii. No change in output and no performance penalty.
+GPUdii() void GPUTPCGMTrackParam::RefitTrack(GPUTPCGMMergedTrack& GPUrestrict() track, int32_t iTrk, GPUTPCGMMerger* GPUrestrict() merger, bool rebuilt) // VS: GPUd changed to GPUdii. No change in output and no performance penalty.
 {
   if (!track.OK()) {
     return;
@@ -908,21 +918,21 @@ GPUdii() void GPUTPCGMTrackParam::RefitTrack(GPUTPCGMMergedTrack& GPUrestrict()
   int32_t NTolerated = 0; // Clusters not fit but tollerated for track length cut
   GPUTPCGMTrackParam t = track.Param();
   float Alpha = track.Alpha();
-  CADEBUG(int32_t nTrackHitsOld = nTrackHits; float ptOld = t.QPt());
-  bool ok = t.Fit(merger, iTrk, merger->Clusters() + track.FirstClusterRef(), nTrackHits, NTolerated, Alpha, constants::MAX_SIN_PHI, track);
-  CADEBUG(printf("Finished Fit Track %d\n", iTrk));
-  CADEBUG(printf("OUTPUT hits %d -> %d+%d = %d, QPt %f -> %f, SP %f, OK %d chi2 %f chi2ndf %f\n", nTrackHitsOld, nTrackHits, NTolerated, nTrackHits + NTolerated, ptOld, t.QPt(), t.SinPhi(), (int32_t)ok, t.Chi2(), t.Chi2() / CAMath::Max(1, nTrackHits)));
+  bool ok = t.Fit(merger, iTrk, merger->Clusters() + track.FirstClusterRef(), nTrackHits, NTolerated, Alpha, constants::MAX_SIN_PHI, track, rebuilt);
+  CADEBUG(if (!merger->Param().rec.tpc.rebuildTrackInFit || rebuilt) printf("Finished Fit Track %7d --- OUTPUT hits %d -> %d+%d = %d, QPt %f -> %f, SP %f, OK %d chi2 %f chi2ndf %f\n", iTrk, track.NClusters(), nTrackHits, NTolerated, nTrackHits + NTolerated, track.GetParam().GetQPt(), t.QPt(), t.SinPhi(), (int32_t)ok, t.Chi2(), t.Chi2() / CAMath::Max(1, nTrackHits)));
 
   if (CAMath::Abs(t.QPt()) < 1.e-4f) {
-    t.QPt() = 1.e-4f;
+    t.QPt() = CAMath::Copysign(1.e-4f, t.QPt());
   }
 
   CADEBUG(if (t.GetX() > 250) { printf("ERROR, Track %d at impossible X %f, Pt %f, Looper %d\n", iTrk, t.GetX(), CAMath::Abs(1.f / t.QPt()), (int32_t)merger->MergedTracks()[iTrk].Looper()); });
 
   track.SetOK(ok);
-  track.SetNClustersFitted(nTrackHits);
   track.Param() = t;
   track.Alpha() = Alpha;
+  if (!merger->Param().rec.tpc.rebuildTrackInFit || rebuilt) {
+    track.SetNClustersFitted(nTrackHits);
+  }
 
   // if (track.OK()) merger->DebugRefitMergedTrack(track);
 }

GPU/GPUTracking/Merger/GPUTPCGMTrackParam.h

@@ -141,7 +141,8 @@ class GPUTPCGMTrackParam
   GPUd() bool CheckNumericalQuality(float overrideCovYY = -1.f) const;
   GPUd() bool CheckCov() const;
 
-  GPUd() bool Fit(GPUTPCGMMerger* merger, int32_t iTrk, GPUTPCGMMergedTrackHit* clusters, int32_t& N, int32_t& NTolerated, float& Alpha, float maxSinPhi, GPUTPCGMMergedTrack& track);
+  GPUd() bool Fit(GPUTPCGMMerger* merger, int32_t iTrk, GPUTPCGMMergedTrackHit* clusters, int32_t& N, int32_t& NTolerated, float& Alpha, float maxSinPhi, GPUTPCGMMergedTrack& track, bool rebuilt);
+  GPUd() static void RefitTrack(GPUTPCGMMergedTrack& track, int32_t iTrk, GPUTPCGMMerger* merger, bool rebuilt);
   GPUd() void MoveToReference(GPUTPCGMPropagator& prop, const GPUParam& param, float& alpha);
   GPUd() void MirrorTo(GPUTPCGMPropagator& prop, float toY, float toZ, bool inFlyDirection, const GPUParam& param, uint8_t row, uint8_t clusterState, bool mirrorParameters, int8_t sector);
   GPUd() int32_t MergeDoubleRowClusters(int32_t& ihit, int32_t wayDirection, GPUTPCGMMergedTrackHit* clusters, const GPUTPCGMMerger* merger, GPUTPCGMPropagator& prop, float& xx, float& yy, float& zz, int32_t maxN, float clAlpha, uint8_t& clusterState, bool rejectChi2);
@@ -200,8 +201,6 @@ class GPUTPCGMTrackParam
     }
   }
 
-  GPUd() static void RefitTrack(GPUTPCGMMergedTrack& track, int32_t iTrk, GPUTPCGMMerger* merger);
-
   GPUdi() void ConstrainSinPhi(float limit = constants::MAX_SIN_PHI)
   {
     if (mP[2] > limit) {

GPU/GPUTracking/kernels.cmake

@@ -51,7 +51,7 @@ o2_gpu_add_kernel("GPUTPCGlobalDebugSortKernels, mergedTracks2"           "= TPC
 o2_gpu_add_kernel("GPUTPCGlobalDebugSortKernels, borderTracks"            "= TPCMERGER"                                           NO int8_t parameter)
 o2_gpu_add_kernel("GPUTPCCreateOccupancyMap, fill"                        "= TPCOCCUPANCY"                                        LB GPUTPCClusterOccupancyMapBin* map)
 o2_gpu_add_kernel("GPUTPCCreateOccupancyMap, fold"                        "= TPCOCCUPANCY"                                        LB GPUTPCClusterOccupancyMapBin* map uint32_t* output)
-o2_gpu_add_kernel("GPUTPCGMMergerTrackFit"                                "GPUTPCGMMergerGPU TPCMERGER TPCTRACKER MATLUT TPCDEDX" LB int32_t mode)
+o2_gpu_add_kernel("GPUTPCGMMergerTrackFit"                                "GPUTPCGMMergerGPU TPCMERGER TPCTRACKER MATLUT TPCDEDX" LB int32_t mode int32_t rebuilt)
 o2_gpu_add_kernel("GPUTPCGMMergerFollowLoopers"                           "GPUTPCGMMergerGPU TPCMERGER TPCTRACKER MATLUT"         LB)
 o2_gpu_add_kernel("GPUTPCGMMergerUnpackResetIds"                          "GPUTPCGMMergerGPU TPCMERGER"                           LB int32_t iSector)
 o2_gpu_add_kernel("GPUTPCGMMergerSectorRefit"                             "GPUTPCGMMergerGPU TPCMERGER MATLUT"                    LB int32_t iSector)