ITS: allow sharing of arena in Tracker & Vertexer

Author: f3sch

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;
+  tbb::task_arena mTaskArena;
 };
 
 } // namespace o2::its

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

@@ -1266,9 +1266,16 @@ void TrackerTraits<nLayers>::setNThreads(int n)
   mNThreads = n > 0 ? n : 1;
 #if defined(OPTIMISATION_OUTPUT) || defined(CA_DEBUG)
   mNThreads = 1; // only works while serial
-#endif
   mTaskArena.initialize(mNThreads);
-  LOGP(info, "Setting tracker with {} threads.", mNThreads);
+#else
+  if (tbb::this_task_arena::current_thread_index() == tbb::task_arena::not_initialized) {
+    mTaskArena.initialize(mNThreads);
+    LOGP(info, "Setting tracker with {} threads.", mNThreads);
+  } else {
+    mTaskArena.initialize(tbb::task_arena::attach{});
+    LOGP(info, "Attaching tracker to calling threads arena");
+  }
+#endif
 }
 
 template class TrackerTraits<7>;

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

@@ -148,6 +148,18 @@ void ITSTrackingInterface::initialise()
   }
   mTracker->setParameters(trackParams);
   mVertexer->setParameters(vertParams);
+  if (trackConf.nThreads == vertConf.nThreads && tbb::this_task_arena::current_thread_index() == tbb::task_arena::not_initialized) {
+    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 {} threads{}", nThreads, (clamped) ? " (clamped)" : "");
+    mTaskArena.initialize(std::abs(nThreads));
+  }
 }
 
 void ITSTrackingInterface::run(framework::ProcessingContext& pc)

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

@@ -695,7 +695,14 @@ void VertexerTraits::setNThreads(int n)
   mNThreads = n > 0 ? n : 1;
 #if defined(VTX_DEBUG)
   mNThreads = 1;
-#endif
   mTaskArena.initialize(mNThreads);
-  LOGP(info, "Setting seeding vertexer with {} threads.", mNThreads);
+#else
+  if (tbb::this_task_arena::current_thread_index() == tbb::task_arena::not_initialized) {
+    mTaskArena.initialize(mNThreads);
+    LOGP(info, "Setting seeding vertexer with {} threads.", mNThreads);
+  } else {
+    mTaskArena.initialize(tbb::task_arena::attach{});
+    LOGP(info, "Attaching vertexer to calling thread's arena");
+  }
+#endif
 }