ITS: add back the code for track following

DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TrackITS.h

@@ -35,6 +35,11 @@ namespace its
 
 class TrackITS : public o2::track::TrackParCov
 {
+ public:
+  static constexpr unsigned int ExtendedPatternShift = 24;
+  static constexpr int MaxLayersInTrackPattern = 8;
+
+ private:
   enum UserBits {
     kSharedClusters = 1 << 28
   };
@@ -106,16 +111,47 @@ class TrackITS : public o2::track::TrackParCov
   GPUhdi() uint32_t getPattern() const { return mPattern; }
   bool hasHitOnLayer(uint32_t i) const { return mPattern & (0x1 << i); }
   bool isFakeOnLayer(uint32_t i) const { return !(mPattern & (0x1 << (16 + i))); }
-  bool isExtendedOnLayer(uint32_t i) const { return (mPattern & (0x1 << (24 + i))); } // only correct if getNClusters <= 8 on layers <= 8
-  uint32_t getLastClusterLayer() const
+  bool isExtendedOnLayer(uint32_t i) const { return (mPattern & (0x1 << (ExtendedPatternShift + i))); } // only correct if getNClusters <= 8 on layers <= 8
+  template <int NLayers>
+  GPUhdi() static constexpr uint32_t getLayerPatternMask()
+  {
+    return (NLayers >= 32) ? 0xffffffffu : ((1u << NLayers) - 1u);
+  }
+  template <int NLayers>
+  GPUhdi() void setExtendedLayerPattern(uint32_t pattern)
+  {
+    pattern &= getLayerPatternMask<NLayers>();
+    setUserField(static_cast<uint16_t>(pattern));
+    if constexpr (NLayers <= MaxLayersInTrackPattern) {
+      setPattern(getPattern() | (pattern << ExtendedPatternShift));
+    }
+  }
+  template <int NLayers>
+  GPUhdi() uint32_t getExtendedLayerPattern() const
+  {
+    const auto mask = getLayerPatternMask<NLayers>();
+    if constexpr (NLayers <= MaxLayersInTrackPattern) {
+      const auto pattern = (getPattern() >> ExtendedPatternShift) & mask;
+      if (pattern) {
+        return pattern;
+      }
+    }
+    return getUserField() & mask;
+  }
+  GPUhdi() void clearExtendedLayerPattern()
+  {
+    setUserField(0);
+    getParamOut().setUserField(0);
+  }
+  GPUhdi() uint32_t getLastClusterLayer() const
   {
     uint32_t r{0}, v{mPattern & ((1 << 16) - 1)};
     while (v >>= 1) {
       r++;
     }
     return r;
   }
-  uint32_t getFirstClusterLayer() const
+  GPUhdi() uint32_t getFirstClusterLayer() const
   {
     int s{0};
     while (!(mPattern & (1 << s))) {

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h

@@ -19,6 +19,7 @@
 #include "ITStracking/BoundedAllocator.h"
 #include "ITStracking/TimeFrame.h"
 #include "ITStracking/Configuration.h"
+#include "ITStracking/TrackExtensionHypothesis.h"
 #include "ITStrackingGPU/Utils.h"
 
 namespace o2::its::gpu
@@ -90,6 +91,7 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   void createNeighboursDevice(const unsigned int layer);
   void createNeighboursLUTDevice(const int, const unsigned int);
   void createTrackITSExtDevice(const size_t);
+  void createTrackExtensionScratchDevice(const int nThreads, const int maxHypotheses);
   void downloadTrackITSExtDevice();
   void downloadCellsNeighboursDevice(std::vector<bounded_vector<CellNeighbour>>&, const int);
   void downloadNeighboursLUTDevice(bounded_vector<int>&, const int);
@@ -125,6 +127,8 @@ class TimeFrameGPU : public TimeFrame<NLayers>
 
   // Hybrid
   TrackITSExt* getDeviceTrackITSExt() { return mTrackITSExtDevice; }
+  TrackExtensionHypothesis<NLayers>* getDeviceActiveTrackExtensionHypotheses() { return mActiveTrackExtensionHypothesesDevice; }
+  TrackExtensionHypothesis<NLayers>* getDeviceNextTrackExtensionHypotheses() { return mNextTrackExtensionHypothesesDevice; }
   int* getDeviceNeighboursLUT(const int layer) { return mNeighboursLUTDevice[layer]; }
   gsl::span<int*> getDeviceNeighboursLUTs() { return mNeighboursLUTDevice; }
   CellNeighbour** getDeviceArrayNeighbours() { return mNeighboursDeviceArray; }
@@ -222,6 +226,8 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   float** mCellSeedsChi2DeviceArray;
 
   TrackITSExt* mTrackITSExtDevice;
+  TrackExtensionHypothesis<NLayers>* mActiveTrackExtensionHypothesesDevice{nullptr};
+  TrackExtensionHypothesis<NLayers>* mNextTrackExtensionHypothesesDevice{nullptr};
   std::array<CellNeighbour*, MaxCells> mNeighboursDevice{};
   CellNeighbour** mNeighboursDeviceArray{nullptr};
   std::array<TrackingFrameInfo*, NLayers> mTrackingFrameInfoDevice;

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h

@@ -13,11 +13,13 @@
 #ifndef ITSTRACKINGGPU_TRACKINGKERNELS_H_
 #define ITSTRACKINGGPU_TRACKINGKERNELS_H_
 
+#include <array>
 #include <gsl/gsl>
 
 #include "ITStracking/BoundedAllocator.h"
 #include "ITStracking/ROFLookupTables.h"
 #include "ITStracking/TrackingTopology.h"
+#include "ITStracking/TrackExtensionHypothesis.h"
 #include "ITStrackingGPU/Utils.h"
 #include "DetectorsBase/Propagator.h"
 
@@ -208,7 +210,6 @@ void countTrackSeedHandler(TrackSeed<NLayers>* trackSeeds,
                            const std::vector<float>& layerxX0Host,
                            const unsigned int nSeeds,
                            const float Bz,
-                           const int startLevel,
                            const float maxChi2ClusterAttachment,
                            const float maxChi2NDF,
                            const int reseedIfShorter,
@@ -222,20 +223,35 @@ template <int NLayers>
 void computeTrackSeedHandler(TrackSeed<NLayers>* trackSeeds,
                              const TrackingFrameInfo** foundTrackingFrameInfo,
                              const Cluster** unsortedClusters,
+                             const IndexTableUtils<NLayers>* utils,
+                             const typename ROFMaskTable<NLayers>::View& rofMask,
+                             const typename ROFOverlapTable<NLayers>::View& rofOverlaps,
+                             const Cluster** clusters,
+                             const unsigned char** usedClusters,
+                             const int** clustersIndexTables,
+                             const int** ROFClusters,
                              o2::its::TrackITSExt* tracks,
                              const int* seedLUT,
+                             TrackExtensionHypothesis<NLayers>* activeHypotheses,
+                             TrackExtensionHypothesis<NLayers>* nextHypotheses,
                              const std::vector<float>& layerRadiiHost,
                              const std::vector<float>& minPtsHost,
                              const std::vector<float>& layerxX0Host,
                              const unsigned int nSeeds,
                              const unsigned int nTracks,
                              const float Bz,
-                             const int startLevel,
                              const float maxChi2ClusterAttachment,
                              const float maxChi2NDF,
                              const int reseedIfShorter,
                              const bool repeatRefitOut,
                              const bool shiftRefToCluster,
+                             const int nLayers,
+                             const int phiBins,
+                             const int maxHypotheses,
+                             const bool extendTop,
+                             const bool extendBot,
+                             const float nSigmaCutPhi,
+                             const float nSigmaCutZ,
                              const o2::base::Propagator* propagator,
                              const o2::base::PropagatorF::MatCorrType matCorrType,
                              o2::its::ExternalAllocator* alloc);

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu

@@ -12,6 +12,7 @@
 
 #include <cuda_runtime.h>
 
+#include <algorithm>
 #include <unistd.h>
 #include <vector>
 
@@ -581,6 +582,21 @@ void TimeFrameGPU<NLayers>::createTrackITSExtDevice(const size_t nSeeds)
   GPUChkErrS(cudaMemset(mTrackITSExtDevice, 0, mNTracks * sizeof(o2::its::TrackITSExt)));
 }
 
+template <int NLayers>
+void TimeFrameGPU<NLayers>::createTrackExtensionScratchDevice(const int nThreads, const int maxHypotheses)
+{
+  GPUTimer timer("reserving track extension scratch");
+  const size_t nHypotheses = static_cast<size_t>(std::max(1, nThreads)) * std::max(1, maxHypotheses);
+  GPULog("gpu-allocation: reserving {} track extension hypotheses per scratch buffer, for {:.2f} MB each.", nHypotheses, nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>) / constants::MB);
+  mActiveTrackExtensionHypothesesDevice = nullptr;
+  mNextTrackExtensionHypothesesDevice = nullptr;
+  if (nHypotheses == 0) {
+    return;
+  }
+  allocMem(reinterpret_cast<void**>(&mActiveTrackExtensionHypothesesDevice), nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+  allocMem(reinterpret_cast<void**>(&mNextTrackExtensionHypothesesDevice), nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+}
+
 template <int NLayers>
 void TimeFrameGPU<NLayers>::downloadCellsDevice()
 {

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx

@@ -12,12 +12,19 @@
 
 #include <unistd.h>
 
+#include <algorithm>
+#include <array>
+
 #include "ITStrackingGPU/TrackerTraitsGPU.h"
 #include "ITStrackingGPU/TrackingKernels.h"
 #include "ITStracking/Configuration.h"
 
 namespace o2::its
 {
+namespace
+{
+constexpr int trackExtensionLaunchThreads = 60 * 256;
+}
 
 template <int NLayers>
 void TrackerTraitsGPU<NLayers>::initialiseTimeFrame(const int iteration)
@@ -301,10 +308,11 @@ template <int NLayers>
 void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
 {
   bounded_vector<bounded_vector<int>> firstClusters(this->mTrkParams[iteration].NLayers, bounded_vector<int>(this->getMemoryPool().get()), this->getMemoryPool().get());
-  bounded_vector<bounded_vector<int>> sharedFirstClusters(this->mTrkParams[iteration].NLayers, bounded_vector<int>(this->getMemoryPool().get()), this->getMemoryPool().get());
   firstClusters.resize(this->mTrkParams[iteration].NLayers);
-  sharedFirstClusters.resize(this->mTrkParams[iteration].NLayers);
   const auto hostTopology = mTimeFrameGPU->getTrackingTopologyView();
+  const bool extendTop = this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerTop];
+  const bool extendBot = this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerBot];
+  const bool extendTracks = extendTop || extendBot;
   for (int startLevel{this->mTrkParams[iteration].CellsPerRoad()}; startLevel >= this->mTrkParams[iteration].CellMinimumLevel(); --startLevel) {
     bounded_vector<TrackSeed<NLayers>> trackSeeds(this->getMemoryPool().get());
     for (int startCellTopologyId{0}; startCellTopologyId < hostTopology.nCells; ++startCellTopologyId) {
@@ -353,7 +361,6 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
                           this->mTrkParams[iteration].LayerxX0,
                           trackSeeds.size(),
                           this->mBz,
-                          startLevel,
                           this->mTrkParams[iteration].MaxChi2ClusterAttachment,
                           this->mTrkParams[iteration].MaxChi2NDF,
                           this->mTrkParams[iteration].ReseedIfShorter,
@@ -363,23 +370,41 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
                           this->mTrkParams[iteration].CorrType,
                           mTimeFrameGPU->getFrameworkAllocator());
     mTimeFrameGPU->createTrackITSExtDevice(trackSeeds.size());
+    if (extendTracks) {
+      mTimeFrameGPU->createTrackExtensionScratchDevice(trackExtensionLaunchThreads, this->mTrkParams[iteration].TrackFollowerMaxHypotheses);
+    }
     computeTrackSeedHandler(mTimeFrameGPU->getDeviceTrackSeeds(),
                             mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(),
                             mTimeFrameGPU->getDeviceArrayUnsortedClusters(),
+                            mTimeFrameGPU->getDeviceIndexTableUtils(),
+                            mTimeFrameGPU->getDeviceROFMaskTableView(),
+                            mTimeFrameGPU->getDeviceROFOverlapTableView(),
+                            mTimeFrameGPU->getDeviceArrayClusters(),
+                            (const unsigned char**)mTimeFrameGPU->getDeviceArrayUsedClusters(),
+                            mTimeFrameGPU->getDeviceArrayClustersIndexTables(),
+                            mTimeFrameGPU->getDeviceROFrameClusters(),
                             mTimeFrameGPU->getDeviceTrackITSExt(),
                             mTimeFrameGPU->getDeviceTrackSeedsLUT(),
+                            extendTracks ? mTimeFrameGPU->getDeviceActiveTrackExtensionHypotheses() : nullptr,
+                            extendTracks ? mTimeFrameGPU->getDeviceNextTrackExtensionHypotheses() : nullptr,
                             this->mTrkParams[iteration].LayerRadii,
                             this->mTrkParams[iteration].MinPt,
                             this->mTrkParams[iteration].LayerxX0,
                             trackSeeds.size(),
                             mTimeFrameGPU->getNTrackSeeds(),
                             this->mBz,
-                            startLevel,
                             this->mTrkParams[iteration].MaxChi2ClusterAttachment,
                             this->mTrkParams[iteration].MaxChi2NDF,
                             this->mTrkParams[iteration].ReseedIfShorter,
                             this->mTrkParams[iteration].RepeatRefitOut,
                             this->mTrkParams[iteration].ShiftRefToCluster,
+                            this->mTrkParams[iteration].NLayers,
+                            this->mTrkParams[iteration].PhiBins,
+                            this->mTrkParams[iteration].TrackFollowerMaxHypotheses,
+                            extendTop,
+                            extendBot,
+                            this->mTrkParams[iteration].TrackFollowerNSigmaCutPhi,
+                            this->mTrkParams[iteration].TrackFollowerNSigmaCutZ,
                             mTimeFrameGPU->getDevicePropagator(),
                             this->mTrkParams[iteration].CorrType,
                             mTimeFrameGPU->getFrameworkAllocator());