ITS: new CPU + GPU seeding vertexer

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

@@ -15,6 +15,8 @@
 #include "GPUCommonDef.h"
 #ifndef GPUCA_GPUCODE_DEVICE
 #include <type_traits>
+#include <unordered_map>
+#include <utility>
 #endif
 #include "ReconstructionDataFormats/Vertex.h"
 #include "SimulationDataFormat/MCCompLabel.h"
@@ -25,6 +27,36 @@ namespace o2::its
 // NOTE: this uses the internal asymmetrical time reprenstation!
 using Vertex = o2::dataformats::Vertex<o2::its::TimeEstBC>;
 using VertexLabel = std::pair<o2::MCCompLabel, float>;
+
+#ifndef GPUCA_GPUCODE_DEVICE
+/// Majority-vote MC label of a vertex: the most frequent (source, event) among its
+/// contributors, flagged fake when no label reaches more than half of them. Templated
+/// on the container so both bounded_vector and std::vector callers share one copy
+template <typename Container>
+VertexLabel computeMainVertexLabel(const Container& elements)
+{
+  // we only care about the source&event of the tracks, not the trackId
+  auto composeVtxLabel = [](const o2::MCCompLabel& lbl) -> o2::MCCompLabel {
+    return {o2::MCCompLabel::maxTrackID(), lbl.getEventID(), lbl.getSourceID(), lbl.isFake()};
+  };
+  std::unordered_map<o2::MCCompLabel, size_t> frequency;
+  for (const auto& element : elements) {
+    ++frequency[composeVtxLabel(element)];
+  }
+  o2::MCCompLabel elem{};
+  size_t maxCount = 0;
+  for (const auto& [key, count] : frequency) {
+    if (count > maxCount) {
+      maxCount = count;
+      elem = key;
+    }
+  }
+  if (maxCount <= 1) { // need >50%
+    elem.setFakeFlag();
+  }
+  return std::make_pair(elem, static_cast<float>(maxCount) / static_cast<float>(elements.size()));
+}
+#endif
 } // namespace o2::its
 
 #ifndef GPUCA_GPUCODE_DEVICE

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

@@ -21,6 +21,8 @@
 #include "ITStracking/Configuration.h"
 #include "ITStracking/TrackExtensionHypothesis.h"
 #include "ITStrackingGPU/Utils.h"
+#include "ITStracking/ClusterLines.h"
+#include "ITStracking/LineProjection.h"
 
 namespace o2::its::gpu
 {
@@ -54,10 +56,14 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   void createTrackingFrameInfoDeviceArray(const int = NLayers);
   void loadUnsortedClustersDevice(const int);
   void createUnsortedClustersDeviceArray(const int = NLayers);
-  void loadClustersDevice(const int);
   void createClustersDeviceArray(const int = NLayers);
   void loadClustersIndexTables(const int);
   void createClustersIndexTablesArray(const int = NLayers);
+  void createClustersDevice(const int);
+  void createClustersIndexTables(const int);
+  void createClusterRadiiDevice();
+  void uploadClusterRadii();
+  void sortClustersDevice(const int layer, const TrackingParameters& trkParam);
   void createUsedClustersDevice(const int);
   void createUsedClustersDeviceArray(const int = NLayers);
   void loadUsedClustersDevice();
@@ -87,6 +93,20 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   void createTrackExtensionScratchDevice(const int nThreads, const int maxHypotheses);
   void downloadTrackITSExtDevice();
 
+  // Seeding-vertexer
+  void createClusterOwnersDeviceArray();
+  void createClusterOwnersDevice();
+  void resetClusterOwnersDevice();
+  void createClusterSortScratchDevice(const int layer);
+
+  void createLinesDevice(const int nCells);
+  void createDiamondDevice(const Vertex& diamond);
+  unsigned int downloadLinesDevice();
+  unsigned int getNLines();
+  const auto& getHostLines() const { return mLinesHost; }
+  const auto& getHostLineRof() const { return mLineRofHost; }
+  const auto& getHostLineClusters() const { return mLineClustersHost; }
+
   /// synchronization
   auto& getStream(const size_t stream) { return mGpuStreams[stream]; }
   auto& getStreams() { return mGpuStreams; }
@@ -111,6 +131,17 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   auto& getTrackITSExt() { return mTrackITSExt; }
   auto& getTrackIndices() { return mTrackIndices; }
   Vertex* getDeviceVertices() { return mPrimaryVerticesDevice; }
+  int* getDeviceROFramesClusters(const int layer) { return mROFramesClustersDevice[layer]; }
+  int* getDeviceClusterSortKeys(const int layer) { return mClusterSortKeysDevice[layer]; }
+  int* getDeviceClusterSortPerm(const int layer) { return mClusterSortPermDevice[layer]; }
+  Cluster* getDeviceUnsortedClusters(const int layer) { return mUnsortedClustersDevice[layer]; }
+  Cluster* getDeviceClusters(const int layer) { return mClustersDevice[layer]; }
+  int* getDeviceClustersIndexTable(const int layer) { return mClustersIndexTablesDevice[layer]; }
+  const float* getDeviceMinRs() const { return mClusterMinRDevice; }
+  const float* getDeviceMaxRs() const { return mClusterMaxRDevice; }
+  int* getDeviceROFramesPV() { return mROFramesPVDevice; }
+  unsigned char* getDeviceUsedClusters(const int);
+  const o2::base::Propagator* getChainPropagator();
 
   // Hybrid
   TrackITSExt* getDeviceTrackITSExt() { return mTrackITSExtDevice; }
@@ -119,6 +150,39 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   TrackExtensionHypothesis<NLayers>* getDeviceNextTrackExtensionHypotheses() { return mNextTrackExtensionHypothesesDevice; }
   int* getDeviceNeighboursLUT(const int layer) { return mNeighboursLUTDevice[layer]; }
   CellNeighbour** getDeviceArrayNeighbours() { return mNeighboursDeviceArray; }
+  unsigned long long** getDeviceArrayClusterOwners() { return mClusterOwnersDeviceArray; }
+  o2::its::Line* getDeviceLines() { return mLinesDevice; }
+  int* getDeviceLineSlots() { return mLineSlotsDevice; }
+  int* getDeviceLineRof() { return mLineRofDevice; }
+  int* getDeviceLineClusters() { return mLineClustersDevice; }
+  float* getDeviceLineChi2() { return mLineChi2Device; }
+  float* getDeviceLinePt() { return mLinePtDevice; }
+  float* getDeviceLineZs() { return mLineZsDevice; }
+  o2::its::TimeEstBC* getDeviceLineTimes() { return mLineTimesDevice; }
+  int* getDeviceLineSortedIdx() { return mLinesSortedIdx; }
+  LineProjSoA getLineProjSoA() { return {mLineZsDevice, mLineTimesDevice, mLinesSortedIdx, mLineRofDevice}; }
+  LineProjSoA getLineProjSortedSoA() { return {mLineZsSortedDevice, mLineTimesSortedDevice, mLinesSortedIdx, mLineRofSortedDevice}; }
+  int* getDeviceRofLineOffsets() { return mRofLineOffsetsDevice; }
+  int* getDeviceLineDensity() { return mLineDensityDevice; }
+  gpu::LineWindow* getDeviceLineWin() { return mLineWinDevice; }
+  uint8_t* getDeviceLineIsPeak() { return mLineIsPeakDevice; }
+  int* getDeviceLineDensityFine() { return mLineDensityFineDevice; }
+  gpu::LineWindow* getDeviceLineWinFine() { return mLineWinFineDevice; }
+  uint8_t* getDeviceLineIsPeakFine() { return mLineIsPeakFineDevice; }
+  int* getDevicePeakScan() { return mPeakScanDevice; }
+  int* getDevicePeakLineIdx() { return mPeakLineIdxDevice; }
+  int* getDevicePeakOffsets() { return mPeakOffsetsDevice; }
+  const int* getDeviceNPeaks() { return mPeakOffsetsDevice + this->getNrof(1); }
+  VertexCand* getDeviceVertexCands() { return mVertexCandsDevice; }
+  int downloadVertexCandsDevice();
+  void downloadPeakMembershipInputs(); // MC-only: peak indices, z-windows and the sorted time/idx columns
+  const auto& getHostVertexCands() const { return mVertexCandsHost; }
+  const auto& getHostPeakOffsets() const { return mPeakOffsetsHost; }
+  const auto& getHostPeakMembership() const { return mPeakMembershipHost; }
+  std::vector<o2::MCCompLabel>& getLineLabelFlat() { return mLineLabelFlatHost; }
+  const std::vector<o2::MCCompLabel>& getLineLabelFlat() const { return mLineLabelFlatHost; }
+  Vertex* getDeviceDiamond() { return mDiamondDevice; }
+  std::array<CellNeighbour*, MaxCells>& getDeviceNeighboursAll() { return mNeighboursDevice; }
   CellNeighbour* getDeviceNeighbours(const int layer) { return mNeighboursDevice[layer]; }
   const TrackingFrameInfo** getDeviceArrayTrackingFrameInfo() const { return mTrackingFrameInfoDeviceArray; }
   const Cluster** getDeviceArrayClusters() const { return mClustersDeviceArray; }
@@ -156,6 +220,10 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   size_t getNumberOfCells() const final;
   size_t getNumberOfNeighbours() const final;
 
+ protected:
+  void prepareClusters(const TrackingParameters& trkParam, const int maxLayers) override;
+  void allocateClusterSortStorage(const TrackingParameters& trkParam, const int maxLayers) override;
+
  private:
   enum class SlotInit {
     Raw, ///< whatever the allocator handed back
@@ -215,6 +283,11 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   const int** mClustersIndexTablesDeviceArray{nullptr};
   uint8_t** mUsedClustersDeviceArray{nullptr};
   const int** mROFramesClustersDeviceArray{nullptr};
+  int* mROFramesPVDevice;
+  std::array<int*, NLayers> mClusterSortKeysDevice{};
+  std::array<int*, NLayers> mClusterSortPermDevice{};
+  float* mClusterMinRDevice{nullptr};
+  float* mClusterMaxRDevice{nullptr};
   std::array<Tracklet*, MaxLinks> mTrackletsDevice{};
   std::array<int*, MaxLinks> mTrackletsLUTDevice{};
   std::array<int*, MaxCells> mCellsLUTDevice{};
@@ -239,6 +312,40 @@ class TimeFrameGPU : public TimeFrame<NLayers>
   CellNeighbour** mNeighboursDeviceArray{nullptr};
   std::array<TrackingFrameInfo*, NLayers> mTrackingFrameInfoDevice{};
   const TrackingFrameInfo** mTrackingFrameInfoDeviceArray{nullptr};
+  std::array<unsigned long long*, 3> mClusterOwnersDevice{};
+  unsigned long long** mClusterOwnersDeviceArray{nullptr};
+  int* mLineSlotsDevice{nullptr};
+  o2::its::Line* mLinesDevice{nullptr};
+  int* mLineRofDevice{nullptr};
+  int* mLineClustersDevice{nullptr};
+  float* mLineChi2Device{nullptr};
+  float* mLinePtDevice{nullptr};
+  float* mLineZsDevice{nullptr};
+  o2::its::TimeEstBC* mLineTimesDevice{nullptr};
+  float* mLineZsSortedDevice{nullptr};
+  o2::its::TimeEstBC* mLineTimesSortedDevice{nullptr};
+  int* mLinesSortedIdx{nullptr};
+  int* mLineRofSortedDevice{nullptr};       // per (sorted) line's ROF
+  int* mRofLineOffsetsDevice{nullptr};      // CSR offsets into the (rof,z)-sorted lines, size nRofs+1
+  int* mLineDensityDevice{nullptr};         // per (sorted) line: count of time-compatible neighbours in its z-window
+  gpu::LineWindow* mLineWinDevice{nullptr}; // per (sorted) line: [lo,hi) bounds of its z-window (sorted coords)
+  uint8_t* mLineIsPeakDevice{nullptr};      // per (sorted) line: 1 if it is a local density peak (vertex candidate)
+  int* mLineDensityFineDevice{nullptr};
+  gpu::LineWindow* mLineWinFineDevice{nullptr};
+  uint8_t* mLineIsPeakFineDevice{nullptr};
+  int* mPeakScanDevice{nullptr};    // per (sorted) line: number of peaks strictly before it
+  int* mPeakLineIdxDevice{nullptr}; // per peak slot: the sorted line index it came from
+  int* mPeakOffsetsDevice{nullptr}; // CSR offsets into the compacted peaks
+  VertexCand* mVertexCandsDevice{nullptr};
+  int mNLinesCapacity{0}; // = nCells the line buffers were sized for
+  std::vector<o2::its::Line> mLinesHost;
+  std::vector<int> mLineRofHost;
+  std::vector<int> mLineClustersHost;
+  std::vector<VertexCand> mVertexCandsHost;
+  std::vector<int> mPeakOffsetsHost;
+  PeakMembershipHost mPeakMembershipHost;
+  std::vector<o2::MCCompLabel> mLineLabelFlatHost;
+  Vertex* mDiamondDevice{nullptr};
 
   // State
   Streams mGpuStreams;

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

@@ -29,6 +29,9 @@ class TrackerTraitsGPU final : public TrackerTraits<NLayers>
   void adoptTimeFrame(TimeFrame<NLayers>* tf) final;
   void initialiseTimeFrame(const int iteration) final;
 
+  void computeVertexCandidates(const int iteration) final;
+  void computeVertices(const int iteration) final;
+
   void computeLayerTracklets(const int iteration, int) final;
   void computeLayerCells(const int iteration) final;
   void findCellsNeighbours(const int iteration) final;

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

@@ -22,6 +22,8 @@
 #include "ITStracking/TrackingTopology.h"
 #include "ITStracking/TrackExtensionHypothesis.h"
 #include "ITStrackingGPU/Utils.h"
+#include "ITStracking/ClusterLines.h"
+#include "ITStracking/LineProjection.h"
 #include "DetectorsBase/Propagator.h"
 
 namespace o2::its
@@ -52,6 +54,7 @@ struct TrackingKernels {
                                            const typename ROFVertexLookupTable<NLayers>::View& vertexLUT,
                                            const int vertexId,
                                            const Vertex* vertices,
+                                           const bool vtxMode,
                                            const Cluster** clusters,
                                            const std::vector<unsigned int>& nClusters,
                                            const int** ROFClusters,
@@ -67,8 +70,8 @@ struct TrackingKernels {
                                            const typename TrackingTopology<NLayers>::View topology,
                                            bounded_vector<float>& linkPhiCuts,
                                            const float resolutionPV,
-                                           std::array<float, NLayers>& minR,
-                                           std::array<float, NLayers>& maxR,
+                                           const float* minRs,
+                                           const float* maxRs,
                                            bounded_vector<float>& resolutions,
                                            std::vector<float>& radii,
                                            bounded_vector<float>& linkMSAngles,
@@ -89,6 +92,7 @@ struct TrackingKernels {
                                  const float bz,
                                  const float maxChi2ClusterAttachment,
                                  const float cellDeltaTanLambdaSigma,
+                                 const float cellDeltaPhiCut,
                                  const float nSigmaCut,
                                  const float* layerxX0,
                                  o2::its::ExternalAllocator* alloc,
@@ -172,6 +176,108 @@ struct TrackingKernels {
                                      const o2::base::Propagator* propagator,
                                      const o2::base::PropagatorF::MatCorrType matCorrType,
                                      o2::its::ExternalAllocator* alloc);
+
+  static void sortClustersHandler(const Cluster* unsorted,
+                                  Cluster* sorted,
+                                  const int* clusterOffsets,
+                                  int* indexTable,
+                                  const IndexTableUtils<NLayers>* utils,
+                                  const typename ROFMaskTable<NLayers>::View& rofMask,
+                                  float beamX, float beamY,
+                                  int zBins, int phiBins, int nRofs, int nClustersLayer, int iLayer,
+                                  float* minRadiusLayer, float* maxRadiusLayer,
+                                  int* keys,
+                                  int* perm,
+                                  o2::its::ExternalAllocator* alloc,
+                                  gpu::Stream& stream);
+
+  static void registerClusterOwnershipHandler(const CellSeed* cellsLayersDevice,
+                                              const int nCells,
+                                              unsigned long long** clusterOwnersDeviceArray,
+                                              gpu::Stream& stream);
+
+  static void linearizeCellsToLinesHandler(const int nCells,
+                                           const CellSeed* cells,
+                                           const unsigned long long* const* clusterOwners,
+                                           const int* rofFramesClustersL1,
+                                           const int nRofsL1,
+                                           const int ownedClustersCut,
+                                           o2::its::Line* lines,
+                                           int* lineRof,
+                                           int* lineClusters,
+                                           int* lineSlots,
+                                           const float beamX,
+                                           const float beamY,
+                                           const float maxZ,
+                                           const float minPt,
+                                           float* linesZs,
+                                           o2::its::TimeEstBC* lineTimes,
+                                           float* lineChi2,
+                                           float* linePt,
+                                           o2::its::ExternalAllocator* alloc,
+                                           gpu::Stream& stream);
+
+  static void sortLinesHandler(const int nLines,
+                               const int nRofs,
+                               const gpu::LineProjSoA soa,
+                               const gpu::LineProjSoA sortedSoa,
+                               const int* lineRof,
+                               int* rofOffsets,
+                               o2::its::ExternalAllocator* alloc,
+                               gpu::Stream& stream);
+
+  static void scanDensityHandler(const int nLines,
+                                 const gpu::LineProjSoA sortedSoa,
+                                 const int* rofOffsets,
+                                 int* density,
+                                 gpu::LineWindow* win,
+                                 const float zWindow,
+                                 gpu::Stream& stream);
+
+  static void findPeaksHandler(const int nLines,
+                               const int nRofs,
+                               const gpu::LineProjSoA sortedSoa,
+                               const int* rofOffsets,
+                               const int* density,
+                               const gpu::LineWindow* win,
+                               uint8_t* isPeak,
+                               const int* densityFine,
+                               const gpu::LineWindow* winFine,
+                               const int fineMinDensity,
+                               uint8_t* isPeakFine,
+                               int* peakScan,
+                               int* peakLineIdx,
+                               int* peakOffsets,
+                               o2::its::ExternalAllocator* alloc,
+                               gpu::Stream& stream);
+
+  static void fitPeaksHandler(const int* nPeaksDevice,
+                              const int* peakLineIdx,
+                              const gpu::LineWindow* win,
+                              const gpu::LineProjSoA sortedSoa,
+                              const o2::its::Line* lines,
+                              const float* lineChi2,
+                              const float* linePt,
+                              const float goodLineChi2Cut,
+                              const float goodLinePtCut,
+                              const float pairCut2,
+                              const float nSigmaCut,
+                              const int minContributors,
+                              const float beamX,
+                              const float beamY,
+                              const uint8_t* isPeakFine,
+                              const float fineMaxDrift,
+                              gpu::VertexCand* cands,
+                              gpu::Stream& stream);
+
+  static void dedupVertexCandidatesHandler(const int* nPeaksDevice,
+                                           const int* peakLineIdx,
+                                           const int* peakOffsets,
+                                           const gpu::LineProjSoA sortedSoa,
+                                           const float duplicateZCut,
+                                           const float duplicateZScale,
+                                           gpu::VertexCand* cands,
+                                           gpu::Stream& stream);
 };
 
 void resetOutputCounterHandler(int* outputCounter, gpu::Stream& stream);