Multi-threaded material budget LUT creation

Introduces multi-threaded creation of the material budget lookup table, reducing it
from hours to minutes. Creating the LUT walks every cell of every layer through TGeo
and the cells are independent, so they are spread over TBB tasks with one
TGeoNavigator per thread. The innermost 20 layers at 60 trials/cell drop from 28 min
to 72 s on 28 cores.

Effective only with ROOT >= v6-36-10-alice3, which removes a per-query thread-id
lookup and the false sharing between per-thread scratch buffers of TGeo shapes. On
older ROOT the parallel path is correct, just slower -- it saturates near 12x.

All layers map onto a single flat cell index so the load stays balanced despite very
different cell counts per layer; a binary search maps a flat index back to
(layer, iz, iphi). The worker navigators are given back at the end.

meanMaterialBudget() takes an optional navigator: a caller passing its own runs
lock-free, a caller passing none shares gGeoManager's and still takes the mutex.
Deciding from the argument keeps it local, so process-global state cannot break it.

Thread count comes from the new populateFromTGeo() argument, falling back to
NTHREADS_MATBUD; the default is the previous serial path. Results are independent of
the thread count -- compareMatBudLUT.C checks a parallel LUT against a serial one cell
by cell, and they match exactly over all 129523 cells.

Supervised-by: Sandro Wenzel <sandro.wenzel@cern.ch>

Author: Tristan Wenzel

Detectors/Base/CMakeLists.txt

@@ -92,6 +92,7 @@ if(BUILD_SIMULATION)
 endif()
 
 install(FILES test/buildMatBudLUT.C
+              test/compareMatBudLUT.C
               test/extractLUTLayers.C
               test/rescaleLUT.C
               DESTINATION share/macro/)
@@ -100,6 +101,10 @@ o2_add_test_root_macro(test/buildMatBudLUT.C
                        PUBLIC_LINK_LIBRARIES O2::DetectorsBase
                        LABELS detectorsbase)
 
+o2_add_test_root_macro(test/compareMatBudLUT.C
+                       PUBLIC_LINK_LIBRARIES O2::DetectorsBase
+                       LABELS detectorsbase)
+
 o2_add_test_root_macro(test/extractLUTLayers.C
                        PUBLIC_LINK_LIBRARIES O2::DetectorsBase
                        LABELS detectorsbase)

Detectors/Base/include/DetectorsBase/GeometryManager.h

@@ -29,6 +29,7 @@
 #include <mutex>
 class TGeoHMatrix; // lines 11-11
 class TGeoManager; // lines 9-9
+class TGeoNavigator;
 
 namespace o2
 {
@@ -96,14 +97,18 @@ class GeometryManager : public TObject
     ClassDefNV(MatBudgetExt, 1);
   };
 
-  static o2::base::MatBudget meanMaterialBudget(float x0, float y0, float z0, float x1, float y1, float z1);
-  static o2::base::MatBudget meanMaterialBudget(const math_utils::Point3D<float>& start, const math_utils::Point3D<float>& end)
+  /// Mean material budget between two points. Pass a navigator owned by the calling thread to
+  /// run lock-free from several threads; with nav = nullptr the shared navigator is used under
+  /// a mutex, as before.
+  static o2::base::MatBudget meanMaterialBudget(float x0, float y0, float z0, float x1, float y1, float z1,
+                                                TGeoNavigator* nav = nullptr);
+  static o2::base::MatBudget meanMaterialBudget(const math_utils::Point3D<float>& start, const math_utils::Point3D<float>& end, TGeoNavigator* nav = nullptr)
   {
-    return meanMaterialBudget(start.X(), start.Y(), start.Z(), end.X(), end.Y(), end.Z());
+    return meanMaterialBudget(start.X(), start.Y(), start.Z(), end.X(), end.Y(), end.Z(), nav);
   }
-  static o2::base::MatBudget meanMaterialBudget(const math_utils::Point3D<double>& start, const math_utils::Point3D<double>& end)
+  static o2::base::MatBudget meanMaterialBudget(const math_utils::Point3D<double>& start, const math_utils::Point3D<double>& end, TGeoNavigator* nav = nullptr)
   {
-    return meanMaterialBudget(start.X(), start.Y(), start.Z(), end.X(), end.Y(), end.Z());
+    return meanMaterialBudget(start.X(), start.Y(), start.Z(), end.X(), end.Y(), end.Z(), nav);
   }
 
   static MatBudgetExt meanMaterialBudgetExt(float x0, float y0, float z0, float x1, float y1, float z1);

Detectors/Base/include/DetectorsBase/MatLayerCyl.h

@@ -25,6 +25,8 @@
 #include "GPUCommonMath.h"
 #include "DetectorsBase/MatCell.h"
 
+class TGeoNavigator;
+
 namespace o2
 {
 namespace base
@@ -66,7 +68,7 @@ class MatLayerCyl : public o2::gpu::FlatObject
   void initSegmentation(float rMin, float rMax, float zHalfSpan, int nz, int nphi);
   void initSegmentation(float rMin, float rMax, float zHalfSpan, float dzMin, float drphiMin);
   void populateFromTGeo(int ntrPerCell = 10);
-  void populateFromTGeo(int ip, int iz, int ntrPerCell);
+  void populateFromTGeo(int ip, int iz, int ntrPerCell, TGeoNavigator* nav = nullptr);
   void print(bool data = false) const;
 #endif // !GPUCA_ALIGPUCODE
 

Detectors/Base/include/DetectorsBase/MatLayerCylSet.h

@@ -73,7 +73,10 @@ class MatLayerCylSet : public o2::gpu::FlatObject
 #ifndef GPUCA_ALIGPUCODE // this part is unvisible on GPU version
   void print(bool data = false) const;
   void addLayer(float rmin, float rmax, float zmax, float dz, float drphi);
-  void populateFromTGeo(int ntrPerCel = 10);
+  /// Populate the LUT from TGeo. nThreads > 1 fills the cells in parallel (one TGeoNavigator
+  /// per thread); nThreads < 0 takes the count from the NTHREADS_MATBUD environment variable.
+  void populateFromTGeo(int ntrPerCel = 10, int nThreads = -1);
+  static int getNThreadsFromEnv();
   void optimizePhiSlices(float maxRelDiff = 0.05);
 
   void dumpToTree(const std::string& outName = "matbudTree.root") const;

Detectors/Base/src/GeometryManager.cxx

@@ -16,6 +16,7 @@
 #include <TCollection.h>       // for TIter
 #include <TFile.h>
 #include <TGeoMatrix.h>       // for TGeoHMatrix
+#include <TGeoNavigator.h>    // for TGeoNavigator
 #include <TGeoNode.h>         // for TGeoNode
 #include <TGeoPhysicalNode.h> // for TGeoPhysicalNode, TGeoPNEntry
 #include <string>
@@ -398,7 +399,8 @@ GeometryManager::MatBudgetExt GeometryManager::meanMaterialBudgetExt(float x0, f
 }
 
 //_____________________________________________________________________________________
-o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, float z0, float x1, float y1, float z1)
+o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, float z0, float x1, float y1, float z1,
+                                                        TGeoNavigator* nav)
 {
   //
   // Calculate mean material budget and material properties between
@@ -414,6 +416,8 @@ o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, floa
   //
   //  Ported to O2: ruben.shahoyan@cern.ch
   //
+  //  Multi-threaded execution: pass a navigator owned by the calling thread.
+  //
 
   double length, startD[3] = {x0, y0, z0};
   double dir[3] = {x1 - x0, y1 - y0, z1 - z0};
@@ -425,9 +429,17 @@ o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, floa
   for (int i = 3; i--;) {
     dir[i] *= invlen;
   }
-  std::lock_guard<std::mutex> guard(sTGMutex);
+  // A caller that passes its own navigator owns it exclusively, so no locking is needed. A caller
+  // that passes none shares gGeoManager's current navigator and must still serialize. Deciding
+  // this from the argument keeps the choice local: it does not depend on -- and cannot be broken
+  // by -- process-global state such as TGeoManager::GetMaxThreads().
+  std::unique_lock<std::mutex> guard(sTGMutex, std::defer_lock);
+  if (!nav) {
+    guard.lock();
+    nav = gGeoManager->GetCurrentNavigator();
+  }
   // Initialize start point and direction
-  TGeoNode* currentnode = gGeoManager->InitTrack(startD, dir);
+  TGeoNode* currentnode = nav->InitTrack(startD, dir);
   if (!currentnode) {
     LOG(error) << "start point out of geometry: " << x0 << ':' << y0 << ':' << z0;
     return o2::base::MatBudget(); // return empty struct
@@ -439,11 +451,11 @@ o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, floa
 
   // Locate next boundary within length without computing safety.
   // Propagate either with length (if no boundary found) or just cross boundary
-  gGeoManager->FindNextBoundaryAndStep(length, kFALSE);
+  nav->FindNextBoundaryAndStep(length, kFALSE);
   Double_t stepTot = 0.0; // Step made
-  Double_t step = gGeoManager->GetStep();
+  Double_t step = nav->GetStep();
   // If no boundary within proposed length, return current step data
-  if (!gGeoManager->IsOnBoundary()) {
+  if (!nav->IsOnBoundary()) {
     budStep.meanX2X0 = budStep.length / budStep.meanX2X0;
     return o2::base::MatBudget(budStep);
   }
@@ -458,7 +470,7 @@ o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, floa
     if (nzero > 3) {
       // This means navigation has problems on one boundary
       // Try to cross by making a small step
-      const double* curPos = gGeoManager->GetCurrentPoint();
+      const double* curPos = nav->GetCurrentPoint();
       LOG(warning) << "Cannot cross boundary at (" << curPos[0] << ',' << curPos[1] << ',' << curPos[2] << ')';
       budTotal.meanRho /= stepTot;
       budTotal.length = stepTot;
@@ -472,14 +484,14 @@ o2::base::MatBudget GeometryManager::meanMaterialBudget(float x0, float y0, floa
     if (step >= length) {
       break;
     }
-    currentnode = gGeoManager->GetCurrentNode();
+    currentnode = nav->GetCurrentNode();
     if (!currentnode) {
       break;
     }
     length -= step;
     accountMaterial(currentnode->GetVolume()->GetMedium()->GetMaterial(), budStep);
-    gGeoManager->FindNextBoundaryAndStep(length, kFALSE);
-    step = gGeoManager->GetStep();
+    nav->FindNextBoundaryAndStep(length, kFALSE);
+    step = nav->GetStep();
   }
   budTotal.meanRho /= stepTot;
   budTotal.length = stepTot;

Detectors/Base/src/MatLayerCyl.cxx

@@ -123,7 +123,7 @@ void MatLayerCyl::populateFromTGeo(int ntrPerCell)
 }
 
 //________________________________________________________________________________
-void MatLayerCyl::populateFromTGeo(int ip, int iz, int ntrPerCell)
+void MatLayerCyl::populateFromTGeo(int ip, int iz, int ntrPerCell, TGeoNavigator* nav)
 {
   /// populate cell with info extracted from TGeometry, using ntrPerCell test tracks per cell
 
@@ -136,7 +136,7 @@ void MatLayerCyl::populateFromTGeo(int ip, int iz, int ntrPerCell)
     float dzt = zs > 0.f ? 0.25 * dz : -0.25 * dz; // to avoid 90 degree polar angle
     for (int isp = ntrPerCell; isp--;) {
       o2::math_utils::sincos(phmn + (isp + 0.5) * getDPhi() / ntrPerCell, sn, cs);
-      auto bud = o2::base::GeometryManager::meanMaterialBudget(rMin * cs, rMin * sn, zs - dzt, rMax * cs, rMax * sn, zs + dzt);
+      auto bud = o2::base::GeometryManager::meanMaterialBudget(rMin * cs, rMin * sn, zs - dzt, rMax * cs, rMax * sn, zs + dzt, nav);
       if (bud.length > 0.) {
         meanRho += bud.length * bud.meanRho;
         meanX2X0 += bud.meanX2X0; // we store actually not X2X0 but 1./X0

Detectors/Base/src/MatLayerCylSet.cxx

@@ -17,7 +17,16 @@
 #ifndef GPUCA_ALIGPUCODE // this part is unvisible on GPU version
 #include "GPUCommonLogger.h"
 #include <TFile.h>
+#include <TGeoManager.h>
 #include "CommonUtils/TreeStreamRedirector.h"
+#include <tbb/blocked_range.h>
+#include <tbb/enumerable_thread_specific.h>
+#include <tbb/global_control.h>
+#include <tbb/parallel_for.h>
+#include <algorithm>
+#include <chrono>
+#include <cstdlib>
+#include <vector>
 //#define _DBG_LOC_ // for local debugging only
 
 #endif // !GPUCA_ALIGPUCODE
@@ -69,9 +78,26 @@ void MatLayerCylSet::addLayer(float rmin, float rmax, float zmax, float dz, floa
 }
 
 //________________________________________________________________________________
-void MatLayerCylSet::populateFromTGeo(int ntrPerCell)
+int MatLayerCylSet::getNThreadsFromEnv()
 {
-  ///< populate layers, using ntrPerCell test tracks per cell
+  ///< number of threads requested via NTHREADS_MATBUD, or 1 if unset/invalid
+  const char* env = std::getenv("NTHREADS_MATBUD");
+  if (!env) {
+    return 1;
+  }
+  int n = std::atoi(env);
+  if (n < 1) {
+    LOG(warning) << "Ignoring invalid NTHREADS_MATBUD=" << env;
+    return 1;
+  }
+  return n;
+}
+
+//________________________________________________________________________________
+void MatLayerCylSet::populateFromTGeo(int ntrPerCell, int nThreads)
+{
+  ///< populate layers, using ntrPerCell test tracks per cell.
+  ///< nThreads < 0 takes the number of threads from the NTHREADS_MATBUD environment variable.
   assert(mConstructionMask == InProgress);
 
   int nlr = getNLayers();
@@ -83,12 +109,86 @@ void MatLayerCylSet::populateFromTGeo(int ntrPerCell)
     LOG(error) << "The LUT is already populated";
     return;
   }
+
+  if (nThreads < 0) {
+    nThreads = getNThreadsFromEnv();
+  }
+
+  using Clock = std::chrono::steady_clock;
+  auto seconds = [](Clock::time_point a, Clock::time_point b) {
+    return std::chrono::duration<double>(b - a).count();
+  };
+
+  if (nThreads <= 1) {
+    for (int i = 0; i < nlr; i++) {
+      LOG(info) << "Populating with " << ntrPerCell << " trials Lr " << i;
+      get()->mLayers[i].print();
+    }
+    const auto tFillStart = Clock::now();
+    for (int i = 0; i < nlr; i++) {
+      get()->mLayers[i].populateFromTGeo(ntrPerCell);
+    }
+    const auto tFillEnd = Clock::now();
+    finalizeStructures();
+    LOG(info) << "LUT fill: 1 thread, cells " << seconds(tFillStart, tFillEnd) << " s";
+    return;
+  }
+
+  // Cells of all layers form one flat index range so that the load is balanced across
+  // threads even though layers differ a lot in cell count. layerOffsets[i] is the first
+  // flat index of layer i; a binary search maps a flat index back to (layer, iz, iphi).
+  std::vector<size_t> layerOffsets(nlr + 1, 0);
   for (int i = 0; i < nlr; i++) {
-    printf("Populating with %d trials Lr  %3d ", ntrPerCell, i);
+    LOG(info) << "Queuing " << ntrPerCell << " trials Lr " << i;
     get()->mLayers[i].print();
-    get()->mLayers[i].populateFromTGeo(ntrPerCell);
+    const auto& lr = get()->mLayers[i];
+    layerOffsets[i + 1] = layerOffsets[i] + size_t(lr.getNZBins()) * lr.getNPhiBins();
   }
+  const size_t totalCells = layerOffsets[nlr];
+
+  const auto tSetupStart = Clock::now();
+
+  // TGeo has to be told that several threads will navigate it, and each thread needs its own
+  // navigator. SetMaxThreads() is one-way -- TGeoManager has no API to return to
+  // single-threaded mode -- so we do not pretend to restore it; that is harmless because
+  // meanMaterialBudget() decides whether to lock from its own argument, not from this global.
+  // The navigators we book are ours, though, so those we do give back.
+  gGeoManager->SetMaxThreads(nThreads);
+
+  tbb::enumerable_thread_specific<TGeoNavigator*> threadNavigators(
+    []() { return gGeoManager->AddNavigator(); });
+
+  const auto tFillStart = Clock::now();
+  {
+    tbb::global_control threadControl(tbb::global_control::max_allowed_parallelism, nThreads);
+    tbb::parallel_for(tbb::blocked_range<size_t>(0, totalCells),
+                      [this, ntrPerCell, &layerOffsets, &threadNavigators](const tbb::blocked_range<size_t>& range) {
+                        TGeoNavigator* nav = threadNavigators.local();
+                        for (size_t idx = range.begin(); idx != range.end(); ++idx) {
+                          auto it = std::upper_bound(layerOffsets.begin(), layerOffsets.end(), idx);
+                          const int layerIdx = int(std::distance(layerOffsets.begin(), it)) - 1;
+                          const size_t cellInLayer = idx - layerOffsets[layerIdx];
+                          auto& layer = this->get()->mLayers[layerIdx];
+                          const int nphi = layer.getNPhiBins();
+                          layer.populateFromTGeo(int(cellInLayer % nphi), int(cellInLayer / nphi), ntrPerCell, nav);
+                        }
+                      });
+  }
+
+  const auto tFillEnd = Clock::now();
+
+  for (TGeoNavigator* nav : threadNavigators) {
+    gGeoManager->RemoveNavigator(nav);
+  }
+
   finalizeStructures();
+  const auto tEnd = Clock::now();
+
+  // Reported separately because only the middle term scales: the setup walks every volume
+  // in the geometry (TGeoManager::SetMaxThreads) and the teardown is serial by nature.
+  LOG(info) << "LUT fill: " << nThreads << " threads, setup " << seconds(tSetupStart, tFillStart)
+            << " s, cells " << seconds(tFillStart, tFillEnd)
+            << " s, finalize " << seconds(tFillEnd, tEnd) << " s";
 }
 
 //________________________________________________________________________________

Detectors/Base/test/README.md

@@ -13,6 +13,16 @@ root -b -q O2/Detectors/Base/test/buildMatBudLUT.C+
 
 The generation is quite time consuming (may take ~30 min).
 
+It can be filled in parallel, one `TGeoNavigator` per thread, by passing a thread count as the
+5th argument of `buildMatBudLUT` or by setting the environment variable:
+```
+export NTHREADS_MATBUD=16
+```
+The result does not depend on the number of threads. Scaling beyond a few threads needs
+ROOT >= v6-36-10-alice3, which removes the per-query thread-id lookup and the false sharing
+between the per-thread scratch buffers of TGeo shapes; with older ROOT the parallel path is
+still correct, just slower.
+
 The optimized LUT will be stored in the matbud.root file.
 
 Load it as:

Detectors/Base/test/buildMatBudLUT.C

@@ -27,7 +27,10 @@ o2::base::MatLayerCylSet mbLUT;
 
 bool testMBLUT(const std::string& lutFile = "matbud.root");
 
-bool buildMatBudLUT(int nTst = 60, int maxLr = -1, const std::string& outFile = "matbud.root", const std::string& geomName = "o2sim_geometry-aligned.root");
+/// Build the material budget LUT. nThreads < 0 takes the thread count from NTHREADS_MATBUD.
+bool buildMatBudLUT(int nTst = 60, int maxLr = -1, const std::string& outFile = "matbud.root",
+                    const std::string& geomNamePrefix = "o2sim", const std::string& opts = "",
+                    int nThreads = -1);
 
 struct LrData {
   float rMin = 0.f;
@@ -42,7 +45,8 @@ struct LrData {
 std::vector<LrData> lrData;
 void configLayers();
 
-bool buildMatBudLUT(int nTst, int maxLr, const std::string& outFile, const std::string& geomNamePrefix, const std::string& opts)
+bool buildMatBudLUT(int nTst, int maxLr, const std::string& outFile, const std::string& geomNamePrefix,
+                    const std::string& opts, int nThreads)
 {
   auto geomName = o2::base::NameConf::getGeomFileName(geomNamePrefix);
   if (gSystem->AccessPathName(geomName.c_str())) { // if needed, create geometry
@@ -67,7 +71,7 @@ bool buildMatBudLUT(int nTst, int maxLr, const std::string& outFile, const std::
   }
 
   TStopwatch sw;
-  mbLUT.populateFromTGeo(nTst);
+  mbLUT.populateFromTGeo(nTst, nThreads);
   mbLUT.optimizePhiSlices(); // move to populateFromTGeo
   mbLUT.flatten();           // move to populateFromTGeo