faster mat LUT

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

Author: f3sch

Detectors/Base/include/DetectorsBase/MatLayerCyl.h

@@ -96,6 +96,7 @@ class MatLayerCyl : public o2::gpu::FlatObject
   // obtain material cell, cell ID must be valid
   GPUd() const MatCell& getCellPhiBin(int iphi, int iz) const { return mCells[getCellIDPhiBin(iphi, iz)]; }
   GPUd() const MatCell& getCell(int iphiSlice, int iz) const { return mCells[getCellID(iphiSlice, iz)]; }
+  GPUd() const MatCell* getCellRow(int iphiSlice) const { return mCells + iphiSlice * getNZBins(); }
 
 #ifndef GPUCA_ALIGPUCODE // this part is unvisible on GPU version
   MatCell& getCellPhiBin(int iphi, int iz)

Detectors/Base/include/DetectorsBase/MatLayerCylSet.h

@@ -41,10 +41,10 @@ struct MatLayerCylSetLayout {
   float mRMin2;         ///< precalculater rmin^2
   float mRMax2;         ///< precalculater rmax^2
   int mNLayers;         ///< number of layers
-  int mNRIntervals;     ///< number of R interval boundaries (gaps are possible)
+  int mNRIntervals;     ///< number of R interval boundaries, one more than the number of intervals (gaps are possible)
   MatLayerCyl* mLayers; //[mNLayers] set of cylinrical layers
-  float* mR2Intervals;  //[mNRIntervals+1] limits of layers
-  int* mInterval2LrID;  //[mNRIntervals] mapping from r2 interval to layer ID
+  float* mR2Intervals;  //[mNRIntervals] limits of layers
+  int* mInterval2LrID;  //[mNRIntervals-1] mapping from r2 interval to layer ID
 };
 
 class MatLayerCylSet : public o2::gpu::FlatObject
@@ -115,6 +115,16 @@ class MatLayerCylSet : public o2::gpu::FlatObject
   /// searches a layer based on r2 input, using a lookup table
   GPUd() int searchLayerFast(float r2, int low = -1, int high = -1) const;
 
+  /// resolves a layer from an already loaded lookup-table entry
+  GPUd() int resolveLayerRange(float r2, int voxel, uint16_t entry) const;
+
+  /// voxel holding this radius; the caller must have checked that r2 is inside the LUT
+  GPUd() int voxelIndex(float r2) const { return int(o2::gpu::CAMath::Sqrt(r2) * InvVoxelRDelta); }
+
+  /// Radial boundaries of a voxel.
+  GPUd() static constexpr float voxelRMin(int voxel) { return voxel * VoxelRDelta; }
+  GPUd() static constexpr float voxelRMax(int voxel) { return (voxel + 1) * VoxelRDelta; }
+
 #ifndef GPUCA_GPUCODE
   //-----------------------------------------------------------
   std::size_t estimateFlatBufferSize() const;
@@ -139,8 +149,10 @@ class MatLayerCylSet : public o2::gpu::FlatObject
   static constexpr float VoxelRDelta = 0.05; // voxel spacing for layer lookup; seems a natural choice - corresponding ~ to smallest spacing
   static constexpr float InvVoxelRDelta = 1.f / VoxelRDelta;
   static constexpr int NumVoxels = int(LayerRMax / VoxelRDelta);
+  static constexpr uint16_t VoxelAmbiguousBit = 0x8000u;
+  static constexpr uint16_t VoxelSegmentMask = 0x7fffu;
 
-  uint16_t mLayerVoxelLU[2 * NumVoxels]; //! helper structure to lookup a layer based on known radius (static dimension for easy copy to GPU)
+  uint16_t mLayerVoxelLU[NumVoxels];     //! first interval based on known radius, plus the ambiguity flag (static dimension for easy copy to GPU)
   bool mInitializedLayerVoxelLU = false; //! if the voxels have been initialized
 
   ClassDefNV(MatLayerCylSet, 1);

Detectors/Base/src/MatLayerCylSet.cxx

@@ -234,7 +234,7 @@ void MatLayerCylSet::finalizeStructures()
   o2::gpu::FlatObject::resizeArray(get()->mR2Intervals, 0, nR2Int);
   o2::gpu::FlatObject::resizeArray(get()->mInterval2LrID, 0, nR2Int);
   get()->mR2Intervals[0] = get()->mRMin2;
-  get()->mR2Intervals[1] = get()->mRMax2;
+  get()->mR2Intervals[1] = getLayer(0).getRMax2();
   get()->mInterval2LrID[0] = 0;
   auto& nRIntervals = get()->mNRIntervals;
   nRIntervals = 1;
@@ -319,14 +319,17 @@ void MatLayerCylSet::initLayerVoxelLU()
   if (LayerRMax < get()->mRMax) {
     LOG(fatal) << "Cannot initialized layer voxel lookup due to dimension problem (fix constants in MatLayerCylSet.h)";
   }
+  // the top bit of an entry carries the ambiguity flag, so the interval index has one bit less
+  if (get()->mNRIntervals > VoxelSegmentMask) {
+    LOG(fatal) << "Too many R intervals (" << get()->mNRIntervals << ") to pack into a layer voxel lookup entry";
+  }
   for (int voxel = 0; voxel < NumVoxels; ++voxel) {
     // check the 2 extremes of this voxel "covering"
-    const auto lowerR = voxel * VoxelRDelta;
-    const auto upperR = lowerR + VoxelRDelta;
+    const auto lowerR = voxelRMin(voxel);
+    const auto upperR = voxelRMax(voxel);
     const auto lowerSegment = searchSegment(lowerR * lowerR);
     const auto upperSegment = searchSegment(upperR * upperR);
-    mLayerVoxelLU[2 * voxel] = lowerSegment;
-    mLayerVoxelLU[2 * voxel + 1] = upperSegment;
+    mLayerVoxelLU[voxel] = uint16_t(lowerSegment) | (lowerSegment != upperSegment ? VoxelAmbiguousBit : uint16_t{0});
   }
   mInitializedLayerVoxelLU = true;
 }
@@ -477,7 +480,10 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
           tEndPhi = cross2;
           checkMorePhi = false;
         } else { // last phi slice still not reached
-          tEndPhi = ray.crossRadial(lr, (stepPhiID > 0 ? phiID + 1 : phiID) % nphiSlices);
+          const int boundaryPhiID = stepPhiID > 0 ? phiID + 1 : phiID;
+          // phiID may be offset by one revolution to handle wrapping, but never by more.
+          const int wrappedBoundaryPhiID = boundaryPhiID < nphiSlices ? boundaryPhiID : boundaryPhiID - nphiSlices;
+          tEndPhi = ray.crossRadial(lr, wrappedBoundaryPhiID);
           if (tEndPhi == Ray::InvalidT) {
             break; // ray parallel to radial line, abandon check for phi bin change
           }
@@ -490,6 +496,8 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
         }
         auto zID = lr.getZBinID(ray.getZ(tStartPhi));
         auto zIDLast = lr.getZBinID(ray.getZ(tEndPhi));
+        const int wrappedPhiID = phiID < nphiSlices ? phiID : phiID - nphiSlices;
+        const auto* cellRow = lr.getCellRow(wrappedPhiID);
         // check if Zbins are crossed
 
 #ifdef _DBG_LOC_
@@ -512,7 +520,7 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
             }
             // account materials of this step
             float step = tEndZ > tStartZ ? tEndZ - tStartZ : tStartZ - tEndZ; // the real step is ray.getDist(tEnd-tStart), will rescale all later
-            const auto& cell = lr.getCell(phiID % nphiSlices, zID);
+            const auto& cell = cellRow[zID];
             rval.meanRho += cell.meanRho * step;
             rval.meanX2X0 += cell.meanX2X0 * step;
             rval.length += step;
@@ -523,7 +531,7 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
             printf(
               "Lr#%3d / cross#%d : account %f<t<%f at phiSlice %d | Zbin: %3d (%3d) |[%+e %+e +%e]:[%+e %+e %+e] "
               "Step: %.3e StrpCor: %.3e\n",
-              lrID, ic, tEndZ, tStartZ, phiID % nphiSlices, zID, zIDLast,
+              lrID, ic, tEndZ, tStartZ, wrappedPhiID, zID, zIDLast,
               pos0[0], pos0[1], pos0[2], pos1[0], pos1[1], pos1[2], step, ray.getDist(step));
 #endif
 
@@ -532,7 +540,7 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
           } while (checkMoreZ);
         } else {
           float step = tEndPhi > tStartPhi ? tEndPhi - tStartPhi : tStartPhi - tEndPhi; // the real step is |ray.getDist(tEnd-tStart)|, will rescale all later
-          const auto& cell = lr.getCell(phiID % nphiSlices, zID);
+          const auto& cell = cellRow[zID];
           rval.meanRho += cell.meanRho * step;
           rval.meanX2X0 += cell.meanX2X0 * step;
           rval.length += step;
@@ -543,7 +551,7 @@ GPUd() MatBudget MatLayerCylSet::getMatBudget(float x0, float y0, float z0, floa
           printf(
             "Lr#%3d / cross#%d : account %f<t<%f at phiSlice %d | Zbin: %3d ----- |[%+e %+e +%e]:[%+e %+e %+e]"
             "Step: %.3e StrpCor: %.3e\n",
-            lrID, ic, tEndPhi, tStartPhi, phiID % nphiSlices, zID,
+            lrID, ic, tEndPhi, tStartPhi, wrappedPhiID, zID,
             pos0[0], pos0[1], pos0[2], pos1[0], pos1[1], pos1[2], step, ray.getDist(step));
 #endif
         }
@@ -585,8 +593,15 @@ GPUd() bool MatLayerCylSet::getLayersRange(const Ray& ray, short& lmin, short& l
     lmxInt = rmax2 < getRMax2() ? searchSegment(rmax2, 0) : get()->mNRIntervals - 2;
     lmnInt = rmin2 >= getRMin2() ? searchSegment(rmin2, 0, lmxInt + 1) : 0;
   } else {
-    lmxInt = rmax2 < getRMax2() ? searchLayerFast(rmax2, 0) : get()->mNRIntervals - 2;
-    lmnInt = rmin2 >= getRMin2() ? searchLayerFast(rmin2, 0, lmxInt + 1) : 0;
+    // The two lookups are independent so overlapping the pair is worth the clumsier shape.
+    const bool useMax = rmax2 < getRMax2();
+    const bool useMin = rmin2 >= getRMin2();
+    const int ixMax = useMax ? voxelIndex(rmax2) : NumVoxels - 1;
+    const int ixMin = useMin ? voxelIndex(rmin2) : 0;
+    const uint16_t eMax = mLayerVoxelLU[ixMax];
+    const uint16_t eMin = mLayerVoxelLU[ixMin];
+    lmxInt = useMax ? resolveLayerRange(rmax2, ixMax, eMax) : get()->mNRIntervals - 2;
+    lmnInt = useMin ? resolveLayerRange(rmin2, ixMin, eMin) : 0;
   }
 
   const auto* interval2LrID = get()->mInterval2LrID;
@@ -605,11 +620,17 @@ GPUd() bool MatLayerCylSet::getLayersRange(const Ray& ray, short& lmin, short& l
 GPUd() int MatLayerCylSet::searchLayerFast(float r2, int low, int high) const
 {
   // we can avoid the sqrt .. at the cost of more memory in the lookup
-  const auto index = 2 * int(o2::gpu::CAMath::Sqrt(r2) * InvVoxelRDelta);
-  const auto layersfirst = mLayerVoxelLU[index];
-  const auto layerslast = mLayerVoxelLU[index + 1];
-  if (layersfirst != layerslast) {
-    // this means the voxel is undecided and we revert to search
+  const auto index = voxelIndex(r2);
+  return resolveLayerRange(r2, index, mLayerVoxelLU[index]);
+}
+
+GPUd() int MatLayerCylSet::resolveLayerRange(float r2, int voxel, uint16_t entry) const
+{
+  const int layersfirst = entry & VoxelSegmentMask;
+  if (entry & VoxelAmbiguousBit) {
+    // Recreate the upper candidate only for the small fraction of undecided voxels
+    const auto upperR = voxelRMax(voxel);
+    const auto layerslast = searchSegment(upperR * upperR);
     return searchSegment(r2, layersfirst, layerslast + 1);
   }
   return layersfirst;
@@ -663,12 +684,13 @@ void MatLayerCylSet::flatten()
   offs = alignSize(offs + nLr * sizeof(MatLayerCyl), MatLayerCyl::getClassAlignmentBytes()); // account for the alignment
 
   // move array of R2 boundaries to the flat array
-  delete[] o2::gpu::FlatObject::resizeArray(get()->mR2Intervals, nLr + 1, nLr + 1, (float*)(mFlatBufferPtr + offs));
-  offs = alignSize(offs + (nLr + 1) * sizeof(float), getBufferAlignmentBytes()); // account for the alignment
+  const int nRBound = get()->mNRIntervals;
+  delete[] o2::gpu::FlatObject::resizeArray(get()->mR2Intervals, nRBound, nRBound, (float*)(mFlatBufferPtr + offs));
+  offs = alignSize(offs + nRBound * sizeof(float), getBufferAlignmentBytes()); // account for the alignment
 
-  // move array of R2 boundaries to the flat array
-  delete[] o2::gpu::FlatObject::resizeArray(get()->mInterval2LrID, nLr, nLr, (int*)(mFlatBufferPtr + offs));
-  offs = alignSize(offs + nLr * sizeof(int), getBufferAlignmentBytes()); // account for the alignment
+  // move array of interval -> layer ID to the flat array
+  delete[] o2::gpu::FlatObject::resizeArray(get()->mInterval2LrID, nRBound - 1, nRBound - 1, (int*)(mFlatBufferPtr + offs));
+  offs = alignSize(offs + (nRBound - 1) * sizeof(int), getBufferAlignmentBytes()); // account for the alignment
 
   for (int il = 0; il < nLr; il++) {
     MatLayerCyl& lr = get()->mLayers[il];
@@ -710,6 +732,11 @@ void MatLayerCylSet::cloneFromObject(const MatLayerCylSet& obj, char* newFlatBuf
   /// Initializes from another object, copies data to newBufferPtr
   flatObject::cloneFromObject(obj, newFlatBufferPtr);
   fixPointers(mFlatBufferPtr);
+  // the voxel lookup lives outside the flat buffer
+  if (obj.mInitializedLayerVoxelLU) {
+    std::copy(obj.mLayerVoxelLU, obj.mLayerVoxelLU + NumVoxels, mLayerVoxelLU);
+    mInitializedLayerVoxelLU = true;
+  }
 }
 
 //______________________________________________

Detectors/Base/test/buildMatBudLUT.C

@@ -21,6 +21,7 @@
 #include <TFile.h>
 #include <TSystem.h>
 #include <TStopwatch.h>
+#include <TRandom.h>
 #endif
 
 using MatbudGeomBackend = o2::base::MatbudGeomBackend;
@@ -30,6 +31,11 @@ o2::base::MatLayerCylSet mbLUT;
 bool testMBLUT(const std::string& lutFile = "matbud.root");
 MatbudGeomBackend parseBackend(const std::string& s);
 
+/// mR2Intervals must be non-decreasing
+bool testMBLUTIntervalsSorted(const o2::base::MatLayerCylSet* lut);
+/// getLayersRange() must agree with and without the voxel lookup
+bool testMBLUTVoxelConsistency(o2::base::MatLayerCylSet* lut, int nRays = 5000);
+
 /// Build the material budget LUT. nThreads < 0 takes the thread count from NTHREADS_MATBUD.
 /// geomBackend is "ROOT" (default) or "VECGEOM" (requires O2 built against TGeo2VecGeom).
 bool buildMatBudLUT(int nTst = 60, int maxLr = -1, const std::string& outFile = "matbud.root",
@@ -186,6 +192,61 @@ bool testMBLUT(const std::string& lutFile)
   return true;
 }
 
+//_______________________________________________________________________
+bool testMBLUTIntervalsSorted(const o2::base::MatLayerCylSet* lut)
+{
+  // searchSegment() is a binary search over mR2Intervals, enfore order
+  const auto* layout = lut->get();
+  for (int i = 1; i < layout->mNRIntervals; i++) { // mNRIntervals counts boundaries, last index is mNRIntervals-1
+    if (layout->mR2Intervals[i] < layout->mR2Intervals[i - 1]) {
+      LOGP(error, "mR2Intervals not monotonic at {}: {} > {}", i, layout->mR2Intervals[i - 1], layout->mR2Intervals[i]);
+      return false;
+    }
+  }
+  return true;
+}
+
+//_______________________________________________________________________
+bool testMBLUTVoxelConsistency(o2::base::MatLayerCylSet* lut, int nRays)
+{
+  // The voxel lookup is only a shortcut into searchSegment(), so it must not change the answer.
+  if (!lut->mInitializedLayerVoxelLU) {
+    LOG(error) << "voxel lookup is not initialized, nothing to compare against";
+    return false;
+  }
+  const float rMax = lut->getRMax(), zMax = lut->getZMax();
+  TRandom rnd(20260825);
+  int nBad = 0, nInside = 0;
+  for (int i = 0; i < nRays; i++) {
+    float x0 = rnd.Uniform(-rMax, rMax), y0 = rnd.Uniform(-rMax, rMax), z0 = rnd.Uniform(-zMax, zMax);
+    float x1 = rnd.Uniform(-rMax, rMax), y1 = rnd.Uniform(-rMax, rMax), z1 = rnd.Uniform(-zMax, zMax);
+    o2::base::Ray ray(x0, y0, z0, x1, y1, z1);
+    short lmin = -1, lmax = -1, lminRef = -1, lmaxRef = -1;
+    const bool ok = lut->getLayersRange(ray, lmin, lmax);
+    lut->mInitializedLayerVoxelLU = false; // force the plain binary search
+    const bool okRef = lut->getLayersRange(ray, lminRef, lmaxRef);
+    lut->mInitializedLayerVoxelLU = true;
+    if (ok) {
+      nInside++;
+    }
+    if (ok != okRef || (ok && (lmin != lminRef || lmax != lmaxRef))) {
+      if (++nBad < 10) {
+        LOGP(error, "ray {} ({:.3f},{:.3f},{:.3f})->({:.3f},{:.3f},{:.3f}): voxel LU gives {} [{},{}], search gives {} [{},{}]",
+             i, x0, y0, z0, x1, y1, z1, ok, lmin, lmax, okRef, lminRef, lmaxRef);
+      }
+    }
+  }
+  if (nInside < nRays / 10) {
+    LOGP(error, "only {} of {} test rays crossed the LUT, the comparison is not meaningful", nInside, nRays);
+    return false;
+  }
+  if (nBad) {
+    LOGP(error, "{} of {} rays disagree between the voxel lookup and searchSegment()", nBad, nRays);
+    return false;
+  }
+  return true;
+}
+
 //_______________________________________________________________________
 void configLayers()
 {

Detectors/Base/test/testMatBudLUT.cxx

@@ -28,5 +28,10 @@ BOOST_AUTO_TEST_CASE(MatBudLUT)
   matBudFile += std::to_string(getpid()) + ".root";
   BOOST_CHECK(buildMatBudLUT(2, 20, matBudFile, geomPrefix + std::to_string(getpid()), "align-geom.mDetectors=none")); // generate LUT
   BOOST_CHECK(testMBLUT(matBudFile));                                                                                  // test LUT manipulations
+
+  o2::base::MatLayerCylSet* lut = o2::base::MatLayerCylSet::loadFromFile(matBudFile);
+  BOOST_REQUIRE(lut != nullptr);
+  BOOST_CHECK(testMBLUTIntervalsSorted(lut));  // mR2Intervals is monotonic
+  BOOST_CHECK(testMBLUTVoxelConsistency(lut)); // voxel lookup agrees with the plain search
 }
 } // namespace o2