Multi-threaded material budget LUT creation

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