ITS: re-enable the possibility of extending tracks

Author: mpuccio

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/TrackExtensionCandidate.h"
 #include "ITStrackingGPU/Utils.h"
 
 namespace o2::its::gpu
@@ -90,8 +91,14 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   void createNeighboursDevice(const unsigned int layer);
   void createNeighboursLUTDevice(const int, const unsigned int);
   void createTrackITSExtDevice(const size_t);
+  void loadTrackExtensionStartStatesDevice();
+  void createTrackExtensionCandidatesDevice(const size_t);
+  void createTrackExtensionScratchDevice(const int nThreads, const int beamWidth);
+  void createTrackExtensionResultsDevice(const size_t);
   void downloadTrackITSExtDevice();
   void downloadCellsNeighboursDevice(std::vector<bounded_vector<CellNeighbour>>&, const int);
+  void downloadTrackExtensionCandidatesDevice();
+  void downloadTrackExtensionResultsDevice();
   void downloadNeighboursLUTDevice(bounded_vector<int>&, const int);
   void downloadCellsDevice();
   void downloadCellsLUTDevice();
@@ -118,13 +125,21 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   const auto getDeviceTrackingTopologyView() const { return mDeviceTrackingTopologyView; }
   int* getDeviceROFramesClusters(const int layer) { return mROFramesClustersDevice[layer]; }
   auto& getTrackITSExt() { return mTrackITSExt; }
+  auto& getTrackExtensionCandidates() { return mTrackExtensionCandidates; }
+  auto& getTrackExtensionResults() { return mTrackExtensionResults; }
   Vertex* getDeviceVertices() { return mPrimaryVerticesDevice; }
   int* getDeviceROFramesPV() { return mROFramesPVDevice; }
   unsigned char* getDeviceUsedClusters(const int);
   const o2::base::Propagator* getChainPropagator();
 
   // Hybrid
   TrackITSExt* getDeviceTrackITSExt() { return mTrackITSExtDevice; }
+  TrackExtensionStartState<NLayers>* getDeviceTrackExtensionStartStates() { return mTrackExtensionStartStatesDevice; }
+  TrackExtensionCandidate<NLayers>* getDeviceTrackExtensionCandidates() { return mTrackExtensionCandidatesDevice; }
+  int* getDeviceTrackExtensionCandidateOffsets() { return mTrackExtensionCandidateOffsetsDevice; }
+  TrackExtensionHypothesis<NLayers>* getDeviceActiveTrackExtensionHypotheses() { return mActiveTrackExtensionHypothesesDevice; }
+  TrackExtensionHypothesis<NLayers>* getDeviceNextTrackExtensionHypotheses() { return mNextTrackExtensionHypothesesDevice; }
+  TrackExtensionResult<NLayers>* getDeviceTrackExtensionResults() { return mTrackExtensionResultsDevice; }
   int* getDeviceNeighboursLUT(const int layer) { return mNeighboursLUTDevice[layer]; }
   gsl::span<int*> getDeviceNeighboursLUTs() { return mNeighboursLUTDevice; }
   CellNeighbour** getDeviceArrayNeighbours() { return mNeighboursDeviceArray; }
@@ -222,6 +237,13 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   float** mCellSeedsChi2DeviceArray;
 
   TrackITSExt* mTrackITSExtDevice;
+  TrackExtensionStartState<NLayers>* mTrackExtensionStartStatesDevice{nullptr};
+  TrackExtensionCandidate<NLayers>* mTrackExtensionCandidatesDevice{nullptr};
+  int* mTrackExtensionCandidateOffsetsDevice{nullptr};
+  TrackExtensionHypothesis<NLayers>* mActiveTrackExtensionHypothesesDevice{nullptr};
+  TrackExtensionHypothesis<NLayers>* mNextTrackExtensionHypothesesDevice{nullptr};
+  TrackExtensionResult<NLayers>* mTrackExtensionResultsDevice{nullptr};
+  unsigned int mNTrackExtensionResults{0};
   std::array<CellNeighbour*, MaxCells> mNeighboursDevice{};
   CellNeighbour** mNeighboursDeviceArray{nullptr};
   std::array<TrackingFrameInfo*, NLayers> mTrackingFrameInfoDevice;
@@ -238,6 +260,12 @@ class TimeFrameGPU : public TimeFrame<NLayers>
 
   // Temporary buffer for storing output tracks from GPU tracking
   bounded_vector<TrackITSExt> mTrackITSExt;
+  // Temporary buffer for compact track states used by GPU track extension
+  bounded_vector<TrackExtensionStartState<NLayers>> mTrackExtensionStartStates;
+  // Temporary buffer for compact track extension proposals from GPU tracking
+  bounded_vector<TrackExtensionCandidate<NLayers>> mTrackExtensionCandidates;
+  // Temporary buffer for fitted track extension proposals from GPU tracking
+  bounded_vector<TrackExtensionResult<NLayers>> mTrackExtensionResults;
 };
 
 template <int NLayers>

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

@@ -35,6 +35,7 @@ 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;
 
@@ -47,6 +48,11 @@ 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;
 };

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/TrackExtensionCandidate.h"
 #include "ITStrackingGPU/Utils.h"
 #include "DetectorsBase/Propagator.h"
 
@@ -35,6 +37,58 @@ class Cluster;
 class TrackITSExt;
 class ExternalAllocator;
 
+inline constexpr int kTrackExtensionLaunchBlocks = 60;
+inline constexpr int kTrackExtensionLaunchThreadsPerBlock = 256;
+inline constexpr int kTrackExtensionLaunchThreads = kTrackExtensionLaunchBlocks * kTrackExtensionLaunchThreadsPerBlock;
+
+template <int NLayers>
+void computeTrackExtensionCandidatesHandler(const TrackExtensionStartState<NLayers>* tracks,
+                                            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,
+                                            const TrackingFrameInfo** trackingFrameInfo,
+                                            TrackExtensionCandidate<NLayers>* candidates,
+                                            int* candidateOffsets,
+                                            TrackExtensionHypothesis<NLayers>* activeHypotheses,
+                                            TrackExtensionHypothesis<NLayers>* nextHypotheses,
+                                            const std::array<float, NLayers> layerRadii,
+                                            const std::array<float, NLayers> layerxX0,
+                                            const int nTracks,
+                                            const int nLayers,
+                                            const int phiBins,
+                                            const int beamWidth,
+                                            const bool extendTop,
+                                            const bool extendBot,
+                                            const float bz,
+                                            const float maxChi2ClusterAttachment,
+                                            const float maxChi2NDF,
+                                            const float nSigmaCutPhi,
+                                            const float nSigmaCutZ,
+                                            const o2::base::Propagator* propagator,
+                                            const o2::base::PropagatorF::MatCorrType matCorrType,
+                                            gpu::Stream& stream);
+
+template <int NLayers>
+void computeTrackExtensionResultsHandler(const TrackExtensionStartState<NLayers>* tracks,
+                                         const TrackExtensionCandidate<NLayers>* candidates,
+                                         const int* candidateOffsets,
+                                         TrackExtensionResult<NLayers>* results,
+                                         const TrackingFrameInfo** trackingFrameInfo,
+                                         const std::array<float, NLayers> layerxX0,
+                                         const int nTracks,
+                                         const int nLayers,
+                                         const float bz,
+                                         const float maxChi2ClusterAttachment,
+                                         const float maxChi2NDF,
+                                         const o2::base::Propagator* propagator,
+                                         const o2::base::PropagatorF::MatCorrType matCorrType,
+                                         const bool shiftRefToCluster,
+                                         gpu::Stream& stream);
+
 template <int NLayers>
 void countTrackletsInROFsHandler(const IndexTableUtils<NLayers>* utils,
                                  const typename ROFMaskTable<NLayers>::View& rofMask,

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,87 @@ void TimeFrameGPU<NLayers>::createTrackITSExtDevice(const size_t nSeeds)
   GPUChkErrS(cudaMemset(mTrackITSExtDevice, 0, mNTracks * sizeof(o2::its::TrackITSExt)));
 }
 
+template <int NLayers>
+void TimeFrameGPU<NLayers>::loadTrackExtensionStartStatesDevice()
+{
+  GPUTimer timer("loading track extension start states");
+  GPULog("gpu-transfer: loading {} track extension start states, for {:.2f} MB.", this->mTracks.size(), this->mTracks.size() * sizeof(o2::its::TrackExtensionStartState<NLayers>) / constants::MB);
+  mTrackExtensionStartStatesDevice = nullptr;
+  mTrackExtensionStartStates = bounded_vector<TrackExtensionStartState<NLayers>>(this->mTracks.size(), {}, this->getMemoryPool().get());
+  if (this->mTracks.empty()) {
+    return;
+  }
+  for (size_t iTrack{0}; iTrack < this->mTracks.size(); ++iTrack) {
+    const auto& track = this->mTracks[iTrack];
+    auto& state = mTrackExtensionStartStates[iTrack];
+    state.paramIn = track.getParamIn();
+    state.paramOut = track.getParamOut();
+    state.time = track.getTimeStamp();
+    state.chi2 = track.getChi2();
+    state.nClusters = track.getNClusters();
+    state.firstClusterLayer = static_cast<int>(track.getFirstClusterLayer());
+    state.lastClusterLayer = static_cast<int>(track.getLastClusterLayer());
+    for (int iLayer{0}; iLayer < NLayers; ++iLayer) {
+      state.clusters[iLayer] = track.getClusterIndex(iLayer);
+    }
+  }
+  allocMem(reinterpret_cast<void**>(&mTrackExtensionStartStatesDevice), mTrackExtensionStartStates.size() * sizeof(o2::its::TrackExtensionStartState<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+  GPUChkErrS(cudaMemcpy(mTrackExtensionStartStatesDevice, mTrackExtensionStartStates.data(), mTrackExtensionStartStates.size() * sizeof(o2::its::TrackExtensionStartState<NLayers>), cudaMemcpyHostToDevice));
+}
+
+template <int NLayers>
+void TimeFrameGPU<NLayers>::createTrackExtensionCandidatesDevice(const size_t nTracks)
+{
+  GPUTimer timer("reserving track extension candidates");
+  const size_t nCandidates = nTracks * MaxTrackExtensionCandidatesPerTrack;
+  GPULog("gpu-allocation: reserving {} track extension candidates, for {:.2f} MB.", nCandidates, nCandidates * sizeof(o2::its::TrackExtensionCandidate<NLayers>) / constants::MB);
+  mTrackExtensionCandidates = bounded_vector<TrackExtensionCandidate<NLayers>>(nCandidates, {}, this->getMemoryPool().get());
+  mTrackExtensionCandidatesDevice = nullptr;
+  mTrackExtensionCandidateOffsetsDevice = nullptr;
+  if (mTrackExtensionCandidates.empty()) {
+    return;
+  }
+  allocMem(reinterpret_cast<void**>(&mTrackExtensionCandidatesDevice), nCandidates * sizeof(o2::its::TrackExtensionCandidate<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+  allocMem(reinterpret_cast<void**>(&mTrackExtensionCandidateOffsetsDevice), (nTracks + 1) * sizeof(int), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+}
+
+template <int NLayers>
+void TimeFrameGPU<NLayers>::createTrackExtensionScratchDevice(const int nThreads, const int beamWidth)
+{
+  GPUTimer timer("reserving track extension scratch");
+  const size_t nHypotheses = static_cast<size_t>(std::max(1, nThreads)) * std::max(1, beamWidth);
+  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>::createTrackExtensionResultsDevice(const size_t nTracks)
+{
+  GPUTimer timer("reserving fitted track extension results");
+  mNTrackExtensionResults = 0;
+  if (nTracks == 0 || mTrackExtensionCandidateOffsetsDevice == nullptr) {
+    mTrackExtensionResults = bounded_vector<TrackExtensionResult<NLayers>>(0, {}, this->getMemoryPool().get());
+    mTrackExtensionResultsDevice = nullptr;
+    return;
+  }
+  int nResults{0};
+  GPUChkErrS(cudaMemcpy(&nResults, mTrackExtensionCandidateOffsetsDevice + nTracks, sizeof(int), cudaMemcpyDeviceToHost));
+  mNTrackExtensionResults = nResults;
+  GPULog("gpu-allocation: reserving {} fitted track extension results, for {:.2f} MB.", mNTrackExtensionResults, mNTrackExtensionResults * sizeof(o2::its::TrackExtensionResult<NLayers>) / constants::MB);
+  mTrackExtensionResults = bounded_vector<TrackExtensionResult<NLayers>>(mNTrackExtensionResults, {}, this->getMemoryPool().get());
+  mTrackExtensionResultsDevice = nullptr;
+  if (mTrackExtensionResults.empty()) {
+    return;
+  }
+  allocMem(reinterpret_cast<void**>(&mTrackExtensionResultsDevice), mNTrackExtensionResults * sizeof(o2::its::TrackExtensionResult<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
+}
+
 template <int NLayers>
 void TimeFrameGPU<NLayers>::downloadCellsDevice()
 {
@@ -627,6 +709,28 @@ void TimeFrameGPU<NLayers>::downloadTrackITSExtDevice()
   GPUChkErrS(cudaMemcpy(mTrackITSExt.data(), mTrackITSExtDevice, mTrackITSExt.size() * sizeof(o2::its::TrackITSExt), cudaMemcpyDeviceToHost));
 }
 
+template <int NLayers>
+void TimeFrameGPU<NLayers>::downloadTrackExtensionCandidatesDevice()
+{
+  GPUTimer timer("downloading track extension candidates");
+  GPULog("gpu-transfer: downloading {} track extension candidates, for {:.2f} MB.", mTrackExtensionCandidates.size(), mTrackExtensionCandidates.size() * sizeof(o2::its::TrackExtensionCandidate<NLayers>) / constants::MB);
+  if (mTrackExtensionCandidates.empty()) {
+    return;
+  }
+  GPUChkErrS(cudaMemcpy(mTrackExtensionCandidates.data(), mTrackExtensionCandidatesDevice, mTrackExtensionCandidates.size() * sizeof(o2::its::TrackExtensionCandidate<NLayers>), cudaMemcpyDeviceToHost));
+}
+
+template <int NLayers>
+void TimeFrameGPU<NLayers>::downloadTrackExtensionResultsDevice()
+{
+  GPUTimer timer("downloading fitted track extension results");
+  GPULog("gpu-transfer: downloading {} fitted track extension results, for {:.2f} MB.", mTrackExtensionResults.size(), mTrackExtensionResults.size() * sizeof(o2::its::TrackExtensionResult<NLayers>) / constants::MB);
+  if (mTrackExtensionResults.empty()) {
+    return;
+  }
+  GPUChkErrS(cudaMemcpy(mTrackExtensionResults.data(), mTrackExtensionResultsDevice, mTrackExtensionResults.size() * sizeof(o2::its::TrackExtensionResult<NLayers>), cudaMemcpyDeviceToHost));
+}
+
 template <int NLayers>
 void TimeFrameGPU<NLayers>::unregisterHostMemory(const int maxLayers)
 {