ITS: allow sharing of arena in Tracker & Vertexer (#14394)

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
@@ -98,6 +100,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