Flatten cell neighbours structure (#12121)

Author: mpuccio

Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h

@@ -156,7 +156,8 @@ class TimeFrame
   std::vector<std::vector<float>>& getCellSeedsChi2() { return mCellSeedsChi2; }
 
   std::vector<std::vector<int>>& getCellsLookupTable();
-  std::vector<std::vector<std::vector<int>>>& getCellsNeighbours();
+  std::vector<std::vector<int>>& getCellsNeighbours();
+  std::vector<std::vector<int>>& getCellsNeighboursLUT();
   std::vector<Road<5>>& getRoads();
   std::vector<TrackITSExt>& getTracks(int rof) { return mTracks[rof]; }
   std::vector<MCCompLabel>& getTracksLabel(const int rof) { return mTracksLabel[rof]; }
@@ -256,7 +257,8 @@ class TimeFrame
   std::vector<std::vector<o2::track::TrackParCovF>> mCellSeeds;
   std::vector<std::vector<float>> mCellSeedsChi2;
   std::vector<std::vector<int>> mCellsLookupTable;
-  std::vector<std::vector<std::vector<int>>> mCellsNeighbours;
+  std::vector<std::vector<int>> mCellsNeighbours;
+  std::vector<std::vector<int>> mCellsNeighboursLUT;
   std::vector<Road<5>> mRoads;
   std::vector<std::vector<MCCompLabel>> mTracksLabel;
   std::vector<std::vector<TrackITSExt>> mTracks;
@@ -550,10 +552,8 @@ inline std::vector<std::vector<int>>& TimeFrame::getCellsLookupTable()
   return mCellsLookupTable;
 }
 
-inline std::vector<std::vector<std::vector<int>>>& TimeFrame::getCellsNeighbours()
-{
-  return mCellsNeighbours;
-}
+inline std::vector<std::vector<int>>& TimeFrame::getCellsNeighbours() { return mCellsNeighbours; }
+inline std::vector<std::vector<int>>& TimeFrame::getCellsNeighboursLUT() { return mCellsNeighboursLUT; }
 
 inline std::vector<Road<5>>& TimeFrame::getRoads() { return mRoads; }
 

Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx

@@ -266,6 +266,7 @@ void TimeFrame::initialise(const int iteration, const TrackingParameters& trkPar
     mCellSeedsChi2.resize(trkParam.CellsPerRoad());
     mCellsLookupTable.resize(trkParam.CellsPerRoad() - 1);
     mCellsNeighbours.resize(trkParam.CellsPerRoad() - 1);
+    mCellsNeighboursLUT.resize(trkParam.CellsPerRoad() - 1);
     mCellLabels.resize(trkParam.CellsPerRoad());
     mTracklets.resize(std::min(trkParam.TrackletsPerRoad(), maxLayers - 1));
     mTrackletLabels.resize(trkParam.TrackletsPerRoad());
@@ -406,6 +407,7 @@ void TimeFrame::initialise(const int iteration, const TrackingParameters& trkPar
     if (iLayer < (int)mCells.size() - 1) {
       mCellsLookupTable[iLayer].clear();
       mCellsNeighbours[iLayer].clear();
+      mCellsNeighboursLUT[iLayer].clear();
     }
   }
 }
@@ -420,9 +422,7 @@ unsigned long TimeFrame::getArtefactsMemory()
     size += sizeof(Cell) * cells.size();
   }
   for (auto& cellsN : mCellsNeighbours) {
-    for (auto& vec : cellsN) {
-      size += sizeof(int) * vec.size();
-    }
+    size += sizeof(int) * cellsN.size();
   }
   return size + sizeof(Road<5>) * mRoads.size();
 }

Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx

@@ -15,6 +15,7 @@
 
 #include "ITStracking/TrackerTraits.h"
 
+#include <algorithm>
 #include <cassert>
 #include <iostream>
 
@@ -394,8 +395,13 @@ void TrackerTraits::findCellsNeighbours(const int iteration)
 
     int layerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer].size())};
     const int nextLayerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer + 1].size())};
-    mTimeFrame->getCellsNeighbours()[iLayer].resize(nextLayerCellsNum);
 
+    mTimeFrame->getCellsNeighboursLUT()[iLayer].clear();
+    mTimeFrame->getCellsNeighboursLUT()[iLayer].resize(nextLayerCellsNum, 0);
+    std::vector<std::pair<int, int>> cellsNeighbours;
+    cellsNeighbours.reserve(nextLayerCellsNum);
+
+    std::vector<std::vector<int>> easyWay(nextLayerCellsNum);
     for (int iCell{0}; iCell < layerCellsNum; ++iCell) {
 
       const Cell& currentCell{mTimeFrame->getCells()[iLayer][iCell]};
@@ -426,15 +432,25 @@ void TrackerTraits::findCellsNeighbours(const int iteration)
           continue;
         }
 
-        mTimeFrame->getCellsNeighbours()[iLayer][iNextCell].push_back(iCell);
-
+        mTimeFrame->getCellsNeighboursLUT()[iLayer][iNextCell]++;
+        cellsNeighbours.push_back(std::make_pair(iCell, iNextCell));
+        easyWay[iNextCell].push_back(iCell);
         const int currentCellLevel{currentCell.getLevel()};
 
         if (currentCellLevel >= nextCell.getLevel()) {
           nextCell.setLevel(currentCellLevel + 1);
         }
       }
     }
+    std::sort(cellsNeighbours.begin(), cellsNeighbours.end(), [](const std::pair<int, int>& a, const std::pair<int, int>& b) {
+      return a.second < b.second;
+    });
+    mTimeFrame->getCellsNeighbours()[iLayer].clear();
+    mTimeFrame->getCellsNeighbours()[iLayer].reserve(cellsNeighbours.size());
+    for (auto& cellNeighboursIndex : cellsNeighbours) {
+      mTimeFrame->getCellsNeighbours()[iLayer].push_back(cellNeighboursIndex.first);
+    }
+    std::inclusive_scan(mTimeFrame->getCellsNeighboursLUT()[iLayer].begin(), mTimeFrame->getCellsNeighboursLUT()[iLayer].end(), mTimeFrame->getCellsNeighboursLUT()[iLayer].begin());
   }
 }
 
@@ -457,11 +473,11 @@ void TrackerTraits::findRoads(const int iteration)
         if (iLevel == 1) {
           continue;
         }
-        const int cellNeighboursNum{static_cast<int>(
-          mTimeFrame->getCellsNeighbours()[iLayer - 1][iCell].size())};
+        const int startNeighbourId{iCell ? mTimeFrame->getCellsNeighboursLUT()[iLayer - 1][iCell - 1] : 0};
+        const int endNeighbourId{mTimeFrame->getCellsNeighboursLUT()[iLayer - 1][iCell]};
         bool isFirstValidNeighbour = true;
-        for (int iNeighbourCell{0}; iNeighbourCell < cellNeighboursNum; ++iNeighbourCell) {
-          const int neighbourCellId = mTimeFrame->getCellsNeighbours()[iLayer - 1][iCell][iNeighbourCell];
+        for (int iNeighbourCell{startNeighbourId}; iNeighbourCell < endNeighbourId; ++iNeighbourCell) {
+          const int neighbourCellId = mTimeFrame->getCellsNeighbours()[iLayer - 1][iNeighbourCell];
           const Cell& neighbourCell = mTimeFrame->getCells()[iLayer - 1][neighbourCellId];
           if (iLevel - 1 != neighbourCell.getLevel()) {
             continue;
@@ -973,14 +989,11 @@ void TrackerTraits::traverseCellsTree(const int currentCellId, const int current
   mTimeFrame->getRoads().back().addCell(currentLayerId, currentCellId);
 
   if (currentLayerId > 0 && currentCellLevel > 1) {
-    const int cellNeighboursNum{static_cast<int>(
-      mTimeFrame->getCellsNeighbours()[currentLayerId - 1][currentCellId].size())};
     bool isFirstValidNeighbour = true;
-
-    for (int iNeighbourCell{0}; iNeighbourCell < cellNeighboursNum; ++iNeighbourCell) {
-
-      const int neighbourCellId =
-        mTimeFrame->getCellsNeighbours()[currentLayerId - 1][currentCellId][iNeighbourCell];
+    const int startNeighbourId = currentCellId ? mTimeFrame->getCellsNeighboursLUT()[currentLayerId - 1][currentCellId - 1] : 0;
+    const int endNeighbourId = mTimeFrame->getCellsNeighboursLUT()[currentLayerId - 1][currentCellId];
+    for (int iNeighbourCell{startNeighbourId}; iNeighbourCell < endNeighbourId; ++iNeighbourCell) {
+      const int neighbourCellId = mTimeFrame->getCellsNeighbours()[currentLayerId - 1][iNeighbourCell];
       const Cell& neighbourCell = mTimeFrame->getCells()[currentLayerId - 1][neighbourCellId];
 
       if (currentCellLevel - 1 != neighbourCell.getLevel()) {