ITS: improve track follower

Author: mpuccio

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

@@ -87,7 +87,7 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   void createTrackITSExtDevice(const size_t);
   void loadTrackExtensionStartStatesDevice();
   void createTrackExtensionCandidatesDevice(const size_t);
-  void createTrackExtensionScratchDevice(const size_t, const int);
+  void createTrackExtensionScratchDevice(const int nThreads, const int beamWidth);
   void createTrackExtensionResultsDevice(const size_t);
   void downloadTrackITSExtDevice();
   void downloadTrackExtensionCandidatesDevice();

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

@@ -13,6 +13,7 @@
 #ifndef ITSTRACKINGGPU_TRACKINGKERNELS_H_
 #define ITSTRACKINGGPU_TRACKINGKERNELS_H_
 
+#include <array>
 #include <gsl/gsl>
 
 #include "ITStracking/BoundedAllocator.h"
@@ -36,6 +37,10 @@ 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,
@@ -50,8 +55,8 @@ void computeTrackExtensionCandidatesHandler(const TrackExtensionStartState<NLaye
                                             int* candidateOffsets,
                                             TrackExtensionHypothesis<NLayers>* activeHypotheses,
                                             TrackExtensionHypothesis<NLayers>* nextHypotheses,
-                                            const std::vector<float>& layerRadiiHost,
-                                            const std::vector<float>& layerxX0Host,
+                                            const std::array<float, NLayers> layerRadii,
+                                            const std::array<float, NLayers> layerxX0,
                                             const int nTracks,
                                             const int nLayers,
                                             const int phiBins,
@@ -73,7 +78,7 @@ void computeTrackExtensionResultsHandler(const TrackExtensionStartState<NLayers>
                                          const int* candidateOffsets,
                                          TrackExtensionResult<NLayers>* results,
                                          const TrackingFrameInfo** trackingFrameInfo,
-                                         const std::vector<float>& layerxX0Host,
+                                         const std::array<float, NLayers> layerxX0,
                                          const int nTracks,
                                          const int nLayers,
                                          const float bz,

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

@@ -578,10 +578,10 @@ void TimeFrameGPU<NLayers>::createTrackExtensionCandidatesDevice(const size_t nT
 }
 
 template <int NLayers>
-void TimeFrameGPU<NLayers>::createTrackExtensionScratchDevice(const size_t nTracks, const int beamWidth)
+void TimeFrameGPU<NLayers>::createTrackExtensionScratchDevice(const int nThreads, const int beamWidth)
 {
   GPUTimer timer("reserving track extension scratch");
-  const size_t nHypotheses = nTracks * std::max(1, beamWidth);
+  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;

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

@@ -13,6 +13,7 @@
 #include <unistd.h>
 
 #include <algorithm>
+#include <array>
 
 #include "ITStrackingGPU/TrackerTraitsGPU.h"
 #include "ITStrackingGPU/TrackingKernels.h"
@@ -398,7 +399,13 @@ void TrackerTraitsGPU<NLayers>::buildTrackExtensionCandidates(const int iteratio
   mTimeFrameGPU->syncStreams();
   mTimeFrameGPU->loadTrackExtensionStartStatesDevice();
   mTimeFrameGPU->createTrackExtensionCandidatesDevice(nTracks);
-  mTimeFrameGPU->createTrackExtensionScratchDevice(nTracks, beamWidth);
+  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->getDeviceTrackExtensionStartStates(),
                                                   mTimeFrameGPU->getDeviceIndexTableUtils(),
                                                   mTimeFrameGPU->getDeviceROFMaskTableView(),
@@ -412,8 +419,8 @@ void TrackerTraitsGPU<NLayers>::buildTrackExtensionCandidates(const int iteratio
                                                   mTimeFrameGPU->getDeviceTrackExtensionCandidateOffsets(),
                                                   mTimeFrameGPU->getDeviceActiveTrackExtensionHypotheses(),
                                                   mTimeFrameGPU->getDeviceNextTrackExtensionHypotheses(),
-                                                  this->mTrkParams[iteration].LayerRadii,
-                                                  this->mTrkParams[iteration].LayerxX0,
+                                                  layerRadii,
+                                                  layerxX0,
                                                   static_cast<int>(nTracks),
                                                   this->mTrkParams[iteration].NLayers,
                                                   this->mTrkParams[iteration].PhiBins,
@@ -434,7 +441,7 @@ void TrackerTraitsGPU<NLayers>::buildTrackExtensionCandidates(const int iteratio
                                                mTimeFrameGPU->getDeviceTrackExtensionCandidateOffsets(),
                                                mTimeFrameGPU->getDeviceTrackExtensionResults(),
                                                mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(),
-                                               this->mTrkParams[iteration].LayerxX0,
+                                               layerxX0,
                                                static_cast<int>(nTracks),
                                                this->mTrkParams[iteration].NLayers,
                                                this->mBz,
@@ -477,7 +484,7 @@ bool TrackerTraitsGPU<NLayers>::materializeTrackExtensionCandidate(TrackITSExt&
       diff |= (0x1u << iLayer);
     }
   }
-  track.setPattern(track.getPattern() | (diff << 24));
+  applyExtendedClustersPattern<NLayers>(track, diff);
   return true;
 }