ITS: integrate track extension in road finding

mpuccio · mpuccio · commit c916cd34aa5a · 2026-05-20T18:54:00.000+02:00
diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackerTraitsGPU.h b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackerTraitsGPU.h
@@ -35,7 +35,6 @@ class TrackerTraitsGPU final : public TrackerTraits<NLayers>
   void computeLayerCells(const int iteration) final;
   void findCellsNeighbours(const int iteration) final;
   void findRoads(const int iteration) final;
-  void extendTracks(const int iteration) final;
 
   void setBz(float) final;
 
@@ -48,11 +47,6 @@ class TrackerTraitsGPU final : public TrackerTraits<NLayers>
   int getTFNumberOfCells() const override;
 
  private:
-  bool hasTrackFollower(const int iteration) const;
-
-  void buildTrackExtensionCandidates(const int iteration, typename TrackerTraits<NLayers>::TrackExtensionCandidates& candidatesPerTrack) final;
-  bool materializeTrackExtensionCandidate(TrackITSExt& track, const typename TrackerTraits<NLayers>::TrackExtensionCandidateN& candidate, const int iteration) final;
-
   IndexTableUtilsN* mDeviceIndexTableUtils;
   gpu::TimeFrameGPU<NLayers>* mTimeFrameGPU;
 };
diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx
@@ -394,126 +394,9 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
     mTimeFrameGPU->loadUsedClustersDevice();
   }
   this->markTracks(iteration, sharedFirstClusters);
-  if (!hasTrackFollower(iteration)) {
-    // wipe the artefact memory
-    mTimeFrameGPU->popMemoryStack(iteration);
-  }
-};
-
-template <int NLayers>
-void TrackerTraitsGPU<NLayers>::extendTracks(const int iteration)
-{
-  TrackerTraits<NLayers>::extendTracks(iteration);
-  mTimeFrameGPU->loadUsedClustersDevice();
-  // wipe the artefact memory kept alive for the track follower
+  // wipe the artefact memory
   mTimeFrameGPU->popMemoryStack(iteration);
-}
-
-template <int NLayers>
-bool TrackerTraitsGPU<NLayers>::hasTrackFollower(const int iteration) const
-{
-  return this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerTop] ||
-         this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerBot];
-}
-
-template <int NLayers>
-void TrackerTraitsGPU<NLayers>::buildTrackExtensionCandidates(const int iteration, typename TrackerTraits<NLayers>::TrackExtensionCandidates& candidatesPerTrack)
-{
-  const auto nTracks = this->mTimeFrame->getTracks().size();
-  const int beamWidth = std::max(1, this->mTrkParams[iteration].TrackFollowerBeamWidth);
-  mTimeFrameGPU->syncStreams();
-  mTimeFrameGPU->loadTrackExtensionStartTracksDevice();
-  mTimeFrameGPU->createTrackExtensionCandidatesDevice(nTracks);
-  mTimeFrameGPU->createTrackExtensionScratchDevice(kTrackExtensionLaunchThreads, beamWidth);
-  std::array<float, NLayers> layerRadii{};
-  std::array<float, NLayers> layerxX0{};
-  for (int iLayer{0}; iLayer < this->mTrkParams[iteration].NLayers; ++iLayer) {
-    layerRadii[iLayer] = this->mTrkParams[iteration].LayerRadii[iLayer];
-    layerxX0[iLayer] = this->mTrkParams[iteration].LayerxX0[iLayer];
-  }
-  computeTrackExtensionCandidatesHandler<NLayers>(mTimeFrameGPU->getDeviceTrackExtensionStartTracks(),
-                                                  mTimeFrameGPU->getDeviceIndexTableUtils(),
-                                                  mTimeFrameGPU->getDeviceROFMaskTableView(),
-                                                  mTimeFrameGPU->getDeviceROFOverlapTableView(),
-                                                  mTimeFrameGPU->getDeviceArrayClusters(),
-                                                  (const uint8_t**)mTimeFrameGPU->getDeviceArrayUsedClusters(),
-                                                  mTimeFrameGPU->getDeviceArrayClustersIndexTables(),
-                                                  mTimeFrameGPU->getDeviceROFrameClusters(),
-                                                  mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(),
-                                                  mTimeFrameGPU->getDeviceTrackExtensionCandidates(),
-                                                  mTimeFrameGPU->getDeviceTrackExtensionCandidateOffsets(),
-                                                  mTimeFrameGPU->getDeviceActiveTrackExtensionHypotheses(),
-                                                  mTimeFrameGPU->getDeviceNextTrackExtensionHypotheses(),
-                                                  layerRadii,
-                                                  layerxX0,
-                                                  static_cast<int>(nTracks),
-                                                  this->mTrkParams[iteration].NLayers,
-                                                  this->mTrkParams[iteration].PhiBins,
-                                                  beamWidth,
-                                                  this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerTop],
-                                                  this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerBot],
-                                                  this->mBz,
-                                                  this->mTrkParams[iteration].MaxChi2ClusterAttachment,
-                                                  this->mTrkParams[iteration].MaxChi2NDF,
-                                                  this->mTrkParams[iteration].TrackFollowerNSigmaCutPhi,
-                                                  this->mTrkParams[iteration].TrackFollowerNSigmaCutZ,
-                                                  mTimeFrameGPU->getDevicePropagator(),
-                                                  this->mTrkParams[iteration].CorrType,
-                                                  mTimeFrameGPU->getStream(0));
-  mTimeFrameGPU->createTrackExtensionResultsDevice(nTracks);
-  computeTrackExtensionResultsHandler<NLayers>(mTimeFrameGPU->getDeviceTrackExtensionStartTracks(),
-                                               mTimeFrameGPU->getDeviceTrackExtensionCandidates(),
-                                               mTimeFrameGPU->getDeviceTrackExtensionCandidateOffsets(),
-                                               mTimeFrameGPU->getDeviceTrackExtensionResults(),
-                                               mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(),
-                                               layerxX0,
-                                               static_cast<int>(nTracks),
-                                               this->mTrkParams[iteration].NLayers,
-                                               this->mBz,
-                                               this->mTrkParams[iteration].MaxChi2ClusterAttachment,
-                                               this->mTrkParams[iteration].MaxChi2NDF,
-                                               mTimeFrameGPU->getDevicePropagator(),
-                                               this->mTrkParams[iteration].CorrType,
-                                               this->mTrkParams[iteration].ShiftRefToCluster,
-                                               mTimeFrameGPU->getStream(0));
-  mTimeFrameGPU->downloadTrackExtensionResultsDevice();
-
-  const auto& results = mTimeFrameGPU->getTrackExtensionResults();
-  for (int iResult{0}; iResult < static_cast<int>(results.size()); ++iResult) {
-    const auto& result = results[iResult];
-    if (!result.isValid()) {
-      continue;
-    }
-    auto candidate = result.candidate;
-    candidate.resultIndex = iResult;
-    if (candidatesPerTrack.add(candidate.trackIndex, candidate) < 0) {
-      continue;
-    }
-  }
-}
-
-template <int NLayers>
-bool TrackerTraitsGPU<NLayers>::materializeTrackExtensionCandidate(TrackITSExt& track, const typename TrackerTraits<NLayers>::TrackExtensionCandidateN& candidate, const int iteration)
-{
-  const auto& results = mTimeFrameGPU->getTrackExtensionResults();
-  if (candidate.resultIndex < 0 || candidate.resultIndex >= static_cast<int>(results.size())) {
-    return TrackerTraits<NLayers>::materializeTrackExtensionCandidate(track, candidate, iteration);
-  }
-  const auto& result = results[candidate.resultIndex];
-  if (!result.isValid() || result.candidate.trackIndex != candidate.trackIndex) {
-    return false;
-  }
-  track = result.track;
-  this->updateExtendedTrackTimeStamp(track, iteration);
-  uint32_t diff{0};
-  for (int iLayer{0}; iLayer < this->mTrkParams[iteration].NLayers; ++iLayer) {
-    if (candidate.addedClusters[iLayer] != constants::UnusedIndex) {
-      diff |= (0x1u << iLayer);
-    }
-  }
-  applyExtendedClustersPattern<NLayers>(track, diff);
-  return true;
-}
+};
 
 template <int NLayers>
 int TrackerTraitsGPU<NLayers>::getTFNumberOfClusters() const
diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h
@@ -279,7 +279,6 @@ struct TimeFrame {
   std::vector<bounded_vector<Tracklet>> mTracklets;
   std::vector<bounded_vector<CellSeed>> mCells;
   bounded_vector<TrackITSExt> mTracks;
-  std::vector<TrackITSExt> mFittedExtensionTracks;
   bounded_vector<MCCompLabel> mTracksLabel;
   std::vector<bounded_vector<int>> mCellsNeighbours;
   std::vector<bounded_vector<int>> mCellsNeighboursTopology;
diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/Tracker.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/Tracker.h
@@ -74,7 +74,6 @@ class Tracker
   void computeCells(int iteration) { mTraits->computeLayerCells(iteration); }
   void findCellsNeighbours(int iteration) { mTraits->findCellsNeighbours(iteration); }
   void findRoads(int iteration) { mTraits->findRoads(iteration); }
-  void extendTracks(int iteration) { mTraits->extendTracks(iteration); }
 
   void rectifyClusterIndices();
   void sortTracks();
diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackerTraits.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackerTraits.h
@@ -20,6 +20,7 @@
 #include <oneapi/tbb.h>
 #include <vector>
 
+#include "DetectorsBase/Propagator.h"
 #include "ITStracking/Configuration.h"
 #include "ITStracking/Constants.h"
 #include "ITStracking/IndexTableUtils.h"
@@ -57,8 +58,6 @@ class TrackerTraits
   virtual void computeLayerCells(const int iteration);
   virtual void findCellsNeighbours(const int iteration);
   virtual void findRoads(const int iteration);
-  virtual bool supportsExtendTracks() const noexcept { return true; }
-  virtual void extendTracks(const int iteration);
 
   template <typename InputSeed>
   void processNeighbours(int iteration, int defaultCellTopologyId, int iLevel, const bounded_vector<InputSeed>& currentCellSeed, const bounded_vector<int>& currentCellId, const bounded_vector<int>& currentCellTopologyId, bounded_vector<TrackSeedN>& updatedCellSeed, bounded_vector<int>& updatedCellId, bounded_vector<int>& updatedCellTopologyId);
@@ -93,54 +92,19 @@ class TrackerTraits
   std::shared_ptr<tbb::task_arena> mTaskArena;
 
  protected:
-  using TrackExtensionCandidateN = TrackExtensionCandidate<NLayers>;
-  struct TrackExtensionCandidates {
-    TrackExtensionCandidates() = default;
-    explicit TrackExtensionCandidates(size_t nTracks)
-      : candidates(nTracks * MaxTrackExtensionCandidatesPerTrack), counts(nTracks, 0)
-    {
-    }
-
-    int add(int trackIndex, const TrackExtensionCandidateN& candidate)
-    {
-      auto& count = counts[trackIndex];
-      if (count >= MaxTrackExtensionCandidatesPerTrack) {
-        return -1;
-      }
-      const int flatIndex = static_cast<int>(getFlatTrackExtensionCandidateIndex(trackIndex, count));
-      candidates[flatIndex] = candidate;
-      ++count;
-      return flatIndex;
-    }
-
-    void pop_back(int trackIndex)
-    {
-      --counts[trackIndex];
-    }
-
-    bool empty(int trackIndex) const { return counts[trackIndex] == 0; }
-    int size(int trackIndex) const { return counts[trackIndex]; }
-    TrackExtensionCandidateN* begin(int trackIndex) { return candidates.data() + getFlatTrackExtensionCandidateIndex(trackIndex, 0); }
-    TrackExtensionCandidateN* end(int trackIndex) { return begin(trackIndex) + counts[trackIndex]; }
-    TrackExtensionCandidateN& get(int trackIndex, int candidateIndex) { return candidates[getFlatTrackExtensionCandidateIndex(trackIndex, candidateIndex)]; }
-    const TrackExtensionCandidateN& get(int trackIndex, int candidateIndex) const { return candidates[getFlatTrackExtensionCandidateIndex(trackIndex, candidateIndex)]; }
-    TrackExtensionCandidateN& getFlat(int flatIndex) { return candidates[flatIndex]; }
-
-    std::vector<TrackExtensionCandidateN> candidates;
-    std::vector<int> counts;
-  };
-
   struct TrackFollowerScratch {
     std::vector<TrackExtensionHypothesis<NLayers>> activeHypotheses;
     std::vector<TrackExtensionHypothesis<NLayers>> nextHypotheses;
   };
 
+  bool finaliseTrackSeed(const TrackSeedN& seed,
+                         TrackITSExt& track,
+                         const int iteration,
+                         const TrackingFrameInfo* const* tfInfos,
+                         const Cluster* const* unsortedClusters,
+                         const o2::base::Propagator* propagator);
   bool trackFollowing(TrackITSExt* track, bool outward, const int iteration, TrackFollowerScratch& scratch);
   bool refitExtendedTrack(TrackITSExt& track, const int iteration);
-  void updateExtendedTrackTimeStamp(TrackITSExt& track, const int iteration);
-  virtual bool materializeTrackExtensionCandidate(TrackITSExt& track, const TrackExtensionCandidateN& candidate, const int iteration);
-  virtual void buildTrackExtensionCandidates(const int iteration, TrackExtensionCandidates& candidatesPerTrack);
-  void applyTrackExtensionCandidates(const int iteration, TrackExtensionCandidates& candidatesPerTrack);
 
   o2::gpu::GPUChainITS* mChain = nullptr;
   TimeFrame<NLayers>* mTimeFrame;
diff --git a/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx b/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx
@@ -475,7 +475,6 @@ template <int NLayers>
 void TimeFrame<NLayers>::wipe()
 {
   deepVectorClear(mTracks);
-  deepVectorClear(mFittedExtensionTracks);
   deepVectorClear(mTracklets);
   deepVectorClear(mCells);
   deepVectorClear(mCellsNeighbours);
diff --git a/Detectors/ITSMFT/ITS/tracking/src/Tracker.cxx b/Detectors/ITSMFT/ITS/tracking/src/Tracker.cxx
@@ -93,18 +93,6 @@ float Tracker<NLayers>::clustersToTracks(const LogFunc& logger, const LogFunc& e
       logger(std::format(" - Neighbours finding: {} neighbours found in {:.2f} ms", nNeighbours, timeNeighbours));
       logger(std::format(" - Track finding: {} tracks found in {:.2f} ms", nTracks + mTimeFrame->getNumberOfTracks(), timeRoads));
       total += timeTracklets + timeCells + timeNeighbours + timeRoads;
-      if (mTraits->supportsExtendTracks() && (mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerTop] || mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerBot])) {
-        const int nClustersBefore = mTimeFrame->getNumberOfUsedClusters();
-        const int nTracksBefore = std::count_if(mTimeFrame->getTracks().begin(), mTimeFrame->getTracks().end(), [](const auto& track) {
-          return track.getPattern() & 0xff000000;
-        });
-        const auto timeExtending = evaluateTask(&Tracker::extendTracks, StateNames[mCurStep = Extending], iteration, evalLog, iteration);
-        const int nTracksAfter = std::count_if(mTimeFrame->getTracks().begin(), mTimeFrame->getTracks().end(), [](const auto& track) {
-          return track.getPattern() & 0xff000000;
-        });
-        total += timeExtending;
-        logger(std::format(" - Extending tracks: {} tracks using {} clusters in {:.2f} ms", nTracksAfter - nTracksBefore, mTimeFrame->getNumberOfUsedClusters() - nClustersBefore, timeExtending));
-      }
     }
   } catch (const BoundedMemoryResource::MemoryLimitExceeded& err) {
     handleException(err);
diff --git a/Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx b/Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx

Original file line number	Diff line number	Diff line change
`@@ -475,7 +475,6 @@ template <int NLayers>`
`475`	`475`	`void TimeFrame<NLayers>::wipe()`
`476`	`476`	`{`
`477`	`477`	`deepVectorClear(mTracks);`
`478`		`- deepVectorClear(mFittedExtensionTracks);`
`479`	`478`	`deepVectorClear(mTracklets);`
`480`	`479`	`deepVectorClear(mCells);`
`481`	`480`	`deepVectorClear(mCellsNeighbours);`