ITS: allow sharing of arena in Tracker & Vertexer

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

Author: f3sch

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

@@ -27,6 +27,8 @@
 #include <utility>
 #include <sstream>
 
+#include <oneapi/tbb/task_arena.h>
+
 #include "ITStracking/Configuration.h"
 #include "CommonConstants/MathConstants.h"
 #include "ITStracking/Definitions.h"
@@ -73,8 +75,7 @@ class Tracker
   void setBz(float bz) { mTraits->setBz(bz); }
   void setCorrType(const o2::base::PropagatorImpl<float>::MatCorrType type) { mTraits->setCorrType(type); }
   bool isMatLUT() const { return mTraits->isMatLUT(); }
-  void setNThreads(int n) { mTraits->setNThreads(n); }
-  int getNThreads() const { return mTraits->getNThreads(); }
+  void setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena) { mTraits->setNThreads(n, arena); }
   void printSummary() const;
 
  private:

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

@@ -80,8 +80,8 @@ class TrackerTraits
   void SetRecoChain(o2::gpu::GPUChainITS* chain) { mChain = chain; }
   void setSmoothing(bool v) { mApplySmoothing = v; }
   bool getSmoothing() const { return mApplySmoothing; }
-  void setNThreads(int n);
-  int getNThreads() const { return mNThreads; }
+  void setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena);
+  int getNThreads() { return mTaskArena->max_concurrency(); }
 
   o2::gpu::GPUChainITS* getChain() const { return mChain; }
 
@@ -94,10 +94,9 @@ class TrackerTraits
   track::TrackParCov buildTrackSeed(const Cluster& cluster1, const Cluster& cluster2, const TrackingFrameInfo& tf3);
   bool fitTrack(TrackITSExt& track, int start, int end, int step, float chi2clcut = o2::constants::math::VeryBig, float chi2ndfcut = o2::constants::math::VeryBig, float maxQoverPt = o2::constants::math::VeryBig, int nCl = 0);
 
-  int mNThreads = 1;
   bool mApplySmoothing = false;
   std::shared_ptr<BoundedMemoryResource> mMemoryPool;
-  tbb::task_arena mTaskArena;
+  std::shared_ptr<tbb::task_arena> mTaskArena;
 
  protected:
   o2::base::PropagatorImpl<float>::MatCorrType mCorrType = o2::base::PropagatorImpl<float>::MatCorrType::USEMatCorrNONE;

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

@@ -28,6 +28,8 @@
 #include "GPUO2Interface.h"
 #include "GPUChainITS.h"
 
+#include <oneapi/tbb/task_arena.h>
+
 namespace o2::its
 {
 class ITSTrackingInterface
@@ -97,6 +99,7 @@ class ITSTrackingInterface
   std::unique_ptr<Vertexer> mVertexer = nullptr;
   const o2::dataformats::MeanVertexObject* mMeanVertex;
   std::shared_ptr<BoundedMemoryResource> mMemoryPool;
+  std::shared_ptr<tbb::task_arena> mTaskArena;
 };
 
 } // namespace o2::its

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

@@ -21,6 +21,9 @@
 #include <iomanip>
 #include <array>
 #include <iosfwd>
+#include <memory>
+
+#include <oneapi/tbb/task_arena.h>
 
 #include "ITStracking/ROframe.h"
 #include "ITStracking/Constants.h"
@@ -90,6 +93,8 @@ class Vertexer
                    const unsigned selectedN, const unsigned int vertexN, const float initT,
                    const float trackletT, const float selecT, const float vertexT);
 
+  void setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena) { mTraits->setNThreads(n, arena); }
+
  private:
   std::uint32_t mTimeFrameCounter = 0;
 

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

@@ -17,6 +17,7 @@
 #define O2_ITS_TRACKING_VERTEXER_TRAITS_H_
 
 #include <array>
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -93,8 +94,8 @@ class VertexerTraits
   auto getVertexingParameters() const { return mVrtParams; }
   void setVertexingParameters(std::vector<VertexingParameters>& vertParams) { mVrtParams = vertParams; }
   void dumpVertexerTraits();
-  void setNThreads(int n);
-  int getNThreads() const { return mNThreads; }
+  void setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena);
+  int getNThreads() { return mTaskArena->max_concurrency(); }
   virtual bool isGPU() const noexcept { return false; }
   virtual const char* getName() const noexcept { return "CPU"; }
   virtual bool usesMemoryPool() const noexcept { return true; }
@@ -116,16 +117,14 @@ class VertexerTraits
   }
 
  protected:
-  int mNThreads = 1;
-
   std::vector<VertexingParameters> mVrtParams;
   IndexTableUtils mIndexTableUtils;
 
   // Frame related quantities
   TimeFrame7* mTimeFrame = nullptr; // observer ptr
  private:
   std::shared_ptr<BoundedMemoryResource> mMemoryPool;
-  tbb::task_arena mTaskArena;
+  std::shared_ptr<tbb::task_arena> mTaskArena;
 };
 
 inline void VertexerTraits::initialise(const TrackingParameters& trackingParams, const int iteration)

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

@@ -342,7 +342,6 @@ void Tracker::getGlobalConfiguration()
   } else {
     mTraits->setCorrType(o2::base::PropagatorImpl<float>::MatCorrType::USEMatCorrLUT);
   }
-  setNThreads(tc.nThreads);
   int nROFsPerIterations = tc.nROFsPerIterations > 0 ? tc.nROFsPerIterations : -1;
   if (tc.nOrbitsPerIterations > 0) {
     /// code to be used when the number of ROFs per orbit is known, this gets priority over the number of ROFs per iteration

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

@@ -71,7 +71,7 @@ void TrackerTraits<nLayers>::computeLayerTracklets(const int iteration, int iROF
     int minRof = o2::gpu::CAMath::Max(startROF, rof0 - mTrkParams[iteration].DeltaROF);
     int maxRof = o2::gpu::CAMath::Min(endROF - 1, rof0 + mTrkParams[iteration].DeltaROF);
 
-    mTaskArena.execute([&] {
+    mTaskArena->execute([&] {
       tbb::parallel_for(
         tbb::blocked_range<int>(0, mTrkParams[iteration].TrackletsPerRoad()),
         [&](const tbb::blocked_range<int>& Layers) {
@@ -200,7 +200,7 @@ void TrackerTraits<nLayers>::computeLayerTracklets(const int iteration, int iROF
     return a.firstClusterIndex == b.firstClusterIndex && a.secondClusterIndex == b.secondClusterIndex;
   };
 
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     tbb::parallel_for(
       tbb::blocked_range<int>(0, mTrkParams[iteration].CellsPerRoad()),
       [&](const tbb::blocked_range<int>& Layers) {
@@ -229,7 +229,7 @@ void TrackerTraits<nLayers>::computeLayerTracklets(const int iteration, int iROF
   /// Layer 0 is done outside the loop
   // in-place deduplication
   auto& trklt0 = mTimeFrame->getTracklets()[0];
-  mTaskArena.execute([&] { tbb::parallel_sort(trklt0.begin(), trklt0.end(), sortTracklets); });
+  mTaskArena->execute([&] { tbb::parallel_sort(trklt0.begin(), trklt0.end(), sortTracklets); });
   trklt0.erase(std::unique(trklt0.begin(), trklt0.end(), equalTracklets), trklt0.end());
   trklt0.shrink_to_fit();
 
@@ -275,7 +275,7 @@ void TrackerTraits<nLayers>::computeLayerCells(const int iteration)
     }
   }
 
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     tbb::parallel_for(
       tbb::blocked_range<int>(0, mTrkParams[iteration].CellsPerRoad()),
       [&](const tbb::blocked_range<int>& Layers) {
@@ -496,7 +496,7 @@ void TrackerTraits<nLayers>::findCellsNeighbours(const int iteration)
       continue;
     }
 
-    mTaskArena.execute([&] {
+    mTaskArena->execute([&] {
       int layerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer].size())};
 
       bounded_vector<int> perCellCount(layerCellsNum + 1, 0, mMemoryPool.get());
@@ -621,7 +621,7 @@ void TrackerTraits<nLayers>::processNeighbours(int iLayer, int iLevel, const bou
   int failed[5]{0, 0, 0, 0, 0}, attempts{0}, failedByMismatch{0};
 #endif
 
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     bounded_vector<int> perCellCount(currentCellSeed.size() + 1, 0, mMemoryPool.get());
     tbb::parallel_for(
       tbb::blocked_range<int>(0, (int)currentCellSeed.size()),
@@ -812,7 +812,7 @@ void TrackerTraits<nLayers>::findRoads(const int iteration)
     }
 
     bounded_vector<TrackITSExt> tracks(mMemoryPool.get());
-    mTaskArena.execute([&] {
+    mTaskArena->execute([&] {
       bounded_vector<int> perSeedCount(trackSeeds.size() + 1, 0, mMemoryPool.get());
       tbb::parallel_for(
         tbb::blocked_range<int>(0, (int)trackSeeds.size()),
@@ -1258,17 +1258,19 @@ bool TrackerTraits<nLayers>::isMatLUT() const
 }
 
 template <int nLayers>
-void TrackerTraits<nLayers>::setNThreads(int n)
+void TrackerTraits<nLayers>::setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena)
 {
-  if (mNThreads == n && mTaskArena.is_active()) {
-    return;
-  }
-  mNThreads = n > 0 ? n : 1;
 #if defined(OPTIMISATION_OUTPUT) || defined(CA_DEBUG)
-  mNThreads = 1; // only works while serial
+  mTaskArena = std::make_shared<tbb::task_arena>(1);
+#else
+  if (arena == nullptr) {
+    mTaskArena = std::make_shared<tbb::task_arena>(std::abs(n));
+    LOGP(info, "Setting tracker with {} threads.", n);
+  } else {
+    mTaskArena = arena;
+    LOGP(info, "Attaching tracker to calling thread's arena");
+  }
 #endif
-  mTaskArena.initialize(mNThreads);
-  LOGP(info, "Setting tracker with {} threads.", mNThreads);
 }
 
 template class TrackerTraits<7>;

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

@@ -16,6 +16,7 @@
 #include "ITSReconstruction/FastMultEst.h"
 
 #include "ITStracking/TrackingInterface.h"
+#include <oneapi/tbb/task_arena.h>
 #include <memory>
 
 #include "DataFormatsITSMFT/ROFRecord.h"
@@ -148,6 +149,20 @@ void ITSTrackingInterface::initialise()
   }
   mTracker->setParameters(trackParams);
   mVertexer->setParameters(vertParams);
+  if (trackConf.nThreads == vertConf.nThreads) {
+    bool clamped{false};
+    int nThreads = trackConf.nThreads;
+    if (nThreads > 0) {
+      const int hw = std::thread::hardware_concurrency();
+      const int maxThreads = (hw == 0 ? 1 : hw);
+      nThreads = std::clamp(nThreads, 1, maxThreads);
+      clamped = trackConf.nThreads > maxThreads;
+    }
+    LOGP(info, "Tracker and Vertexer will share the task arena with {} thread(s){}", nThreads, (clamped) ? " (clamped)" : "");
+    mTaskArena = std::make_shared<tbb::task_arena>(std::abs(nThreads));
+  }
+  mVertexer->setNThreads(vertConf.nThreads, mTaskArena);
+  mTracker->setNThreads(trackConf.nThreads, mTaskArena);
 }
 
 void ITSTrackingInterface::run(framework::ProcessingContext& pc)

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

@@ -11,6 +11,7 @@
 ///
 
 #include <iostream>
+#include <memory>
 #include <string>
 #include <chrono>
 
@@ -168,13 +169,12 @@ void VertexerTraits::updateVertexingParameters(const std::vector<VertexingParame
     par.phiSpan = static_cast<int>(std::ceil(mIndexTableUtils.getNphiBins() * par.phiCut / o2::constants::math::TwoPI));
     par.zSpan = static_cast<int>(std::ceil(par.zCut * mIndexTableUtils.getInverseZCoordinate(0)));
   }
-  setNThreads(vrtPar[0].nThreads);
 }
 
 // Main functions
 void VertexerTraits::computeTracklets(const int iteration)
 {
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     tbb::parallel_for(
       tbb::blocked_range<short>(0, (short)mTimeFrame->getNrof()),
       [&](const tbb::blocked_range<short>& Rofs) {
@@ -220,7 +220,7 @@ void VertexerTraits::computeTracklets(const int iteration)
   mTimeFrame->getTracklets()[0].resize(mTimeFrame->getTotalTrackletsTF(0));
   mTimeFrame->getTracklets()[1].resize(mTimeFrame->getTotalTrackletsTF(1));
 
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     tbb::parallel_for(
       tbb::blocked_range<short>(0, (short)mTimeFrame->getNrof()),
       [&](const tbb::blocked_range<short>& Rofs) {
@@ -329,7 +329,7 @@ void VertexerTraits::computeTracklets(const int iteration)
 
 void VertexerTraits::computeTrackletMatching(const int iteration)
 {
-  mTaskArena.execute([&] {
+  mTaskArena->execute([&] {
     tbb::parallel_for(
       tbb::blocked_range<short>(0, (short)mTimeFrame->getNrof()),
       [&](const tbb::blocked_range<short>& Rofs) {
@@ -687,15 +687,17 @@ void VertexerTraits::computeVerticesInRof(int rofId,
   verticesInRof.push_back(foundVertices);
 }
 
-void VertexerTraits::setNThreads(int n)
+void VertexerTraits::setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena)
 {
-  if (mNThreads == n && mTaskArena.is_active()) {
-    return;
-  }
-  mNThreads = n > 0 ? n : 1;
 #if defined(VTX_DEBUG)
-  mNThreads = 1;
+  mTaskArena = std::make_shared<tbb::task_arena>(1);
+#else
+  if (arena == nullptr) {
+    mTaskArena = std::make_shared<tbb::task_arena>(std::abs(n));
+    LOGP(info, "Setting seeding vertexer with {} threads.", n);
+  } else {
+    mTaskArena = arena;
+    LOGP(info, "Attaching vertexer to calling thread's arena");
+  }
 #endif
-  mTaskArena.initialize(mNThreads);
-  LOGP(info, "Setting seeding vertexer with {} threads.", mNThreads);
 }