ITS: chunking

Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>

Author: f3sch

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

@@ -17,6 +17,7 @@
 #define TRACKINGITSU_INCLUDE_TRACKERTRAITS_H_
 
 #include <oneapi/tbb.h>
+#include <span>
 
 #include "DetectorsBase/Propagator.h"
 #include "ITStracking/Configuration.h"
@@ -55,7 +56,7 @@ class TrackerTraits
   virtual void findRoads(const int iteration);
 
   template <typename InputSeed>
-  void processNeighbours(int iLayer, int iLevel, const bounded_vector<InputSeed>& currentCellSeed, const bounded_vector<int>& currentCellId, bounded_vector<TrackSeedN>& updatedCellSeed, bounded_vector<int>& updatedCellId);
+  void processNeighbours(int iLayer, int iLevel, std::span<const InputSeed> currentCellSeed, std::span<const int> currentCellId, int cellOffset, bounded_vector<TrackSeedN>& updatedCellSeed, bounded_vector<int>& updatedCellId);
 
   void updateTrackingParameters(const std::vector<TrackingParameters>& trkPars) { mTrkParams = trkPars; }
   TimeFrame<NLayers>* getTimeFrame() { return mTimeFrame; }

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

@@ -524,7 +524,7 @@ void TrackerTraits<NLayers>::findCellsNeighbours(const int iteration)
 
 template <int NLayers>
 template <typename InputSeed>
-void TrackerTraits<NLayers>::processNeighbours(int iLayer, int iLevel, const bounded_vector<InputSeed>& currentCellSeed, const bounded_vector<int>& currentCellId, bounded_vector<TrackSeedN>& updatedCellSeeds, bounded_vector<int>& updatedCellsIds)
+void TrackerTraits<NLayers>::processNeighbours(int iLayer, int iLevel, std::span<const InputSeed> currentCellSeed, std::span<const int> currentCellId, int cellOffset, bounded_vector<TrackSeedN>& updatedCellSeeds, bounded_vector<int>& updatedCellsIds)
 {
   auto propagator = o2::base::Propagator::Instance();
 
@@ -543,7 +543,7 @@ void TrackerTraits<NLayers>::processNeighbours(int iLayer, int iLevel, const bou
         }
       }
 
-      const int cellId = currentCellId.empty() ? iCell : currentCellId[iCell];
+      const int cellId = currentCellId.empty() ? (iCell + cellOffset) : currentCellId[iCell];
       const int startNeighbourId{cellId ? mTimeFrame->getCellsNeighboursLUT()[iLayer - 1][cellId - 1] : 0};
       const int endNeighbourId{mTimeFrame->getCellsNeighboursLUT()[iLayer - 1][cellId]};
       int foundSeeds{0};
@@ -658,30 +658,42 @@ void TrackerTraits<NLayers>::findRoads(const int iteration)
     };
 
     bounded_vector<TrackSeedN> trackSeeds(mMemoryPool.get());
+    constexpr int kStartCellChunk = 32000; // TODO: expose as tunable parameter
     for (int startLayer{mTrkParams[iteration].NeighboursPerRoad()}; startLayer >= startLevel - 1; --startLayer) {
       if ((mTrkParams[iteration].StartLayerMask & (1 << (startLayer + 2))) == 0) {
         continue;
       }
 
-      bounded_vector<int> lastCellId(mMemoryPool.get()), updatedCellId(mMemoryPool.get());
-      bounded_vector<TrackSeedN> lastCellSeed(mMemoryPool.get()), updatedCellSeed(mMemoryPool.get());
+      const auto& startCells = mTimeFrame->getCells()[startLayer];
+      const int nStartCells = static_cast<int>(startCells.size());
 
-      processNeighbours(startLayer, startLevel, mTimeFrame->getCells()[startLayer], lastCellId, updatedCellSeed, updatedCellId);
+      for (int chunkStart{0}; chunkStart < nStartCells; chunkStart += kStartCellChunk) {
+        const int chunkSize = std::min(kStartCellChunk, nStartCells - chunkStart);
+        std::span<const CellSeed> chunkSeeds{startCells.data() + chunkStart, static_cast<size_t>(chunkSize)};
 
-      int level = startLevel;
-      for (int iLayer{startLayer - 1}; iLayer > 0 && level > 2; --iLayer) {
-        lastCellSeed.swap(updatedCellSeed);
-        lastCellId.swap(updatedCellId);
-        deepVectorClear(updatedCellSeed); /// tame the memory peaks
-        deepVectorClear(updatedCellId);   /// tame the memory peaks
-        processNeighbours(iLayer, --level, lastCellSeed, lastCellId, updatedCellSeed, updatedCellId);
-      }
-      deepVectorClear(lastCellId);   /// tame the memory peaks
-      deepVectorClear(lastCellSeed); /// tame the memory peaks
+        bounded_vector<int> lastCellId(mMemoryPool.get()), updatedCellId(mMemoryPool.get());
+        bounded_vector<TrackSeedN> lastCellSeed(mMemoryPool.get()), updatedCellSeed(mMemoryPool.get());
+
+        processNeighbours(startLayer, startLevel, chunkSeeds, std::span<const int>{}, chunkStart, updatedCellSeed, updatedCellId);
 
-      if (!updatedCellSeed.empty()) {
-        trackSeeds.reserve(trackSeeds.size() + std::count_if(updatedCellSeed.begin(), updatedCellSeed.end(), seedFilter));
-        std::copy_if(updatedCellSeed.begin(), updatedCellSeed.end(), std::back_inserter(trackSeeds), seedFilter);
+        int level = startLevel;
+        for (int iLayer{startLayer - 1}; iLayer > 0 && level > 2; --iLayer) {
+          lastCellSeed.swap(updatedCellSeed);
+          lastCellId.swap(updatedCellId);
+          deepVectorClear(updatedCellSeed); /// tame the memory peaks
+          deepVectorClear(updatedCellId);   /// tame the memory peaks
+          processNeighbours(iLayer, --level,
+                            std::span<const TrackSeedN>{lastCellSeed.data(), lastCellSeed.size()},
+                            std::span<const int>{lastCellId.data(), lastCellId.size()},
+                            0, updatedCellSeed, updatedCellId);
+        }
+        deepVectorClear(lastCellId);   /// tame the memory peaks
+        deepVectorClear(lastCellSeed); /// tame the memory peaks
+
+        if (!updatedCellSeed.empty()) {
+          trackSeeds.reserve(trackSeeds.size() + std::count_if(updatedCellSeed.begin(), updatedCellSeed.end(), seedFilter));
+          std::copy_if(updatedCellSeed.begin(), updatedCellSeed.end(), std::back_inserter(trackSeeds), seedFilter);
+        }
       }
     }
 
@@ -1001,13 +1013,13 @@ void TrackerTraits<NLayers>::setNThreads(int n, std::shared_ptr<tbb::task_arena>
 }
 
 template class TrackerTraits<7>;
-template void TrackerTraits<7>::processNeighbours<CellSeed>(int, int, const bounded_vector<CellSeed>&, const bounded_vector<int>&, bounded_vector<TrackSeed<7>>&, bounded_vector<int>&);
-template void TrackerTraits<7>::processNeighbours<TrackSeed<7>>(int, int, const bounded_vector<TrackSeed<7>>&, const bounded_vector<int>&, bounded_vector<TrackSeed<7>>&, bounded_vector<int>&);
+template void TrackerTraits<7>::processNeighbours<CellSeed>(int, int, std::span<const CellSeed>, std::span<const int>, int, bounded_vector<TrackSeed<7>>&, bounded_vector<int>&);
+template void TrackerTraits<7>::processNeighbours<TrackSeed<7>>(int, int, std::span<const TrackSeed<7>>, std::span<const int>, int, bounded_vector<TrackSeed<7>>&, bounded_vector<int>&);
 // ALICE3 upgrade
 #ifdef ENABLE_UPGRADES
 template class TrackerTraits<11>;
-template void TrackerTraits<11>::processNeighbours<CellSeed>(int, int, const bounded_vector<CellSeed>&, const bounded_vector<int>&, bounded_vector<TrackSeed<11>>&, bounded_vector<int>&);
-template void TrackerTraits<11>::processNeighbours<TrackSeed<11>>(int, int, const bounded_vector<TrackSeed<11>>&, const bounded_vector<int>&, bounded_vector<TrackSeed<11>>&, bounded_vector<int>&);
+template void TrackerTraits<11>::processNeighbours<CellSeed>(int, int, std::span<const CellSeed>, std::span<const int>, int, bounded_vector<TrackSeed<11>>&, bounded_vector<int>&);
+template void TrackerTraits<11>::processNeighbours<TrackSeed<11>>(int, int, std::span<const TrackSeed<11>>, std::span<const int>, int, bounded_vector<TrackSeed<11>>&, bounded_vector<int>&);
 #endif
 
 } // namespace o2::its