Implement new kind of multiplicity mask

Author: mpuccio

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h

@@ -31,6 +31,7 @@ class TimeFrameGPU final : public TimeFrame<NLayers>
   using typename TimeFrame<NLayers>::IndexTableUtilsN;
   using typename TimeFrame<NLayers>::ROFOverlapTableN;
   using typename TimeFrame<NLayers>::ROFVertexLookupTableN;
+  using typename TimeFrame<NLayers>::ROFMaskTableN;
 
  public:
   TimeFrameGPU() = default;
@@ -107,6 +108,7 @@ class TimeFrameGPU final : public TimeFrame<NLayers>
   IndexTableUtilsN* getDeviceIndexTableUtils() { return mIndexTableUtilsDevice; }
   const auto getDeviceROFOverlapTableView() { return mDeviceROFOverlapTableView; }
   const auto getDeviceROFVertexLookupTableView() { return mDeviceROFVertexLookupTableView; }
+  const auto getDeviceROFMaskTableView() { return mDeviceROFMaskTableView; }
   int* getDeviceROFramesClusters(const int layer) { return mROFramesClustersDevice[layer]; }
   auto& getTrackITSExt() { return mTrackITSExt; }
   Vertex* getDeviceVertices() { return mPrimaryVerticesDevice; }
@@ -177,9 +179,11 @@ class TimeFrameGPU final : public TimeFrame<NLayers>
   // device navigation views
   ROFOverlapTableN::View mDeviceROFOverlapTableView;
   ROFVertexLookupTableN::View mDeviceROFVertexLookupTableView;
+  ROFMaskTableN::View mDeviceROFMaskTableView;
 
   // Hybrid pref
   uint8_t* mMultMaskDevice;
+  int32_t* mMultMaskOffsetsDevice;
   Vertex* mPrimaryVerticesDevice;
   int* mROFramesPVDevice;
   std::array<Cluster*, NLayers> mClustersDevice;

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu

@@ -252,11 +252,18 @@ void TimeFrameGPU<NLayers>::loadMultiplicityCutMask(const int iteration)
 {
   if (!iteration || iteration == 3) { // we need to re-load the swapped mult-mask in upc iteration
     GPUTimer timer("loading multiplicity cut mask");
-    GPULog("gpu-transfer: iteration {} loading multiplicity cut mask with {} elements, for {:.2f} MB.", iteration, this->mMultiplicityCutMask.size(), this->mMultiplicityCutMask.size() * sizeof(uint8_t) / constants::MB);
-    if (!iteration) { // only allocate on first call
-      allocMem(reinterpret_cast<void**>(&mMultMaskDevice), this->mMultiplicityCutMask.size() * sizeof(uint8_t), this->hasFrameworkAllocator());
+    const auto& hostTable = this->mMultiplicityCutMask;
+    const auto hostView = hostTable.getView();
+    GPULog("gpu-transfer: iteration {} loading multiplicity cut mask with {} elements, for {:.2f} MB.",
+           iteration, hostTable.getFlatMaskSize(), hostTable.getFlatMaskSize() * sizeof(uint8_t) / constants::MB);
+    if (!iteration) { // only allocate on first call; offsets are stable across iterations
+      allocMem(reinterpret_cast<void**>(&mMultMaskDevice), hostTable.getFlatMaskSize() * sizeof(uint8_t), this->hasFrameworkAllocator());
+      allocMem(reinterpret_cast<void**>(&mMultMaskOffsetsDevice), NLayers * sizeof(int32_t), this->hasFrameworkAllocator());
+      GPUChkErrS(cudaMemcpy(mMultMaskOffsetsDevice, hostView.mLayerROFOffsets, NLayers * sizeof(int32_t), cudaMemcpyHostToDevice));
     }
-    GPUChkErrS(cudaMemcpy(mMultMaskDevice, this->mMultiplicityCutMask.data(), this->mMultiplicityCutMask.size() * sizeof(uint8_t), cudaMemcpyHostToDevice));
+    // Re-copy the flat mask on every qualifying iteration (e.g. after swapMasks() for UPC)
+    GPUChkErrS(cudaMemcpy(mMultMaskDevice, hostView.mFlatMask, hostTable.getFlatMaskSize() * sizeof(uint8_t), cudaMemcpyHostToDevice));
+    mDeviceROFMaskTableView = hostTable.getDeviceView(mMultMaskDevice, mMultMaskOffsetsDevice);
   }
 }
 

Detectors/ITSMFT/ITS/tracking/include/ITStracking/FastMultEst.h

@@ -17,10 +17,12 @@
 #define ALICEO2_ITS_FASTMULTEST_
 
 #include "ITSMFTReconstruction/ChipMappingITS.h"
+#include "DataFormatsITS/Vertex.h"
 #include "DataFormatsITSMFT/ROFRecord.h"
 #include "DataFormatsITSMFT/CompCluster.h"
-#include <DataFormatsITSMFT/PhysTrigger.h>
+#include "DataFormatsITSMFT/PhysTrigger.h"
 #include "ITStracking/FastMultEstConfig.h"
+#include "ITStracking/ROFLookupTables.h"
 #include <gsl/span>
 #include <array>
 
@@ -32,32 +34,56 @@ namespace its
 struct FastMultEst {
 
   static constexpr int NLayers = o2::itsmft::ChipMappingITS::NLayers;
+  using ROFOverlapTableN = ROFOverlapTable<NLayers>;
+  using ROFMaskTableN = ROFMaskTable<NLayers>;
 
-  float mult = 0.;                         /// estimated signal clusters multipliciy at reference (1st?) layer
-  float noisePerChip = 0.;                 /// estimated or imposed noise per chip
-  float cov[3] = {0.};                     /// covariance matrix of estimation
-  float chi2 = 0.;                         /// chi2
-  int nLayersUsed = 0;                     /// number of layers actually used
+  float mult = 0.;                         /// estimated signal clusters multiplicity on the selected multiplicity layer
+  float noisePerChip = 0.;                 /// imposed noise per chip (when enabled by configuration)
+  float cov[3] = {0.};                     /// retained for compatibility; set to zero in single-layer mode
+  float chi2 = 0.;                         /// retained for compatibility; set to zero in single-layer mode
+  int nLayersUsed = 0;                     /// number of layers used by estimator (0/1 in single-layer mode)
   uint32_t lastRandomSeed = 0;             /// state of the gRandom before
-
-  std::array<int, NLayers> nClPerLayer{0}; // measured N Cl per layer selectROFs
   FastMultEst();
 
   static uint32_t getCurrentRandomSeed();
-  int selectROFs(const gsl::span<const o2::itsmft::ROFRecord> rofs, const gsl::span<const o2::itsmft::CompClusterExt> clus,
-                 const gsl::span<const o2::itsmft::PhysTrigger> trig, std::vector<uint8_t>& sel);
+  int selectROFs(const std::array<gsl::span<const o2::itsmft::ROFRecord>, NLayers>& rofs,
+                 const std::array<gsl::span<const o2::itsmft::CompClusterExt>, NLayers>& clus,
+                 const gsl::span<const o2::itsmft::PhysTrigger> trig,
+                 uint32_t firstTForbit,
+                 bool doStaggering,
+                 const ROFOverlapTableN::View& overlapView,
+                 ROFMaskTableN& sel);
+  void selectROFsWithVertices(const auto& vertices, const ROFOverlapTableN::View& overlapView, ROFMaskTableN& sel)
+  {
+    const auto& multEstConf = FastMultEstConfig::Instance();
+    if (!multEstConf.isVtxMultCutRequested()) {
+      return;
+    }
+
+    for (const auto& vertex : vertices) {
+      if (!multEstConf.isPassingVtxMultCut(vertex.getNContributors())) {
+        const auto& timestamp{vertex.getTimeStamp()};
+        for (int layer = 0; layer < NLayers; ++layer) {
+          uint32_t startROF = sel.getLayer(layer).getROF(timestamp.lower());
+          uint32_t endROF = sel.getLayer(layer).getROF(timestamp.upper());
+          for (uint32_t rof = startROF; rof <= endROF; ++rof) {
+            sel.setROFsEnabled(layer, rof, 0);
+          }
+        }
+      }
+    }
+  }
 
-  void fillNClPerLayer(const gsl::span<const o2::itsmft::CompClusterExt>& clusters);
-  float process(const std::array<int, NLayers> ncl)
+  int countClustersOnLayer(const gsl::span<const o2::itsmft::CompClusterExt>& clusters) const;
+  float process(int nClusters)
   {
-    return FastMultEstConfig::Instance().imposeNoisePerChip > 0 ? processNoiseImposed(ncl) : processNoiseFree(ncl);
+    return FastMultEstConfig::Instance().imposeNoisePerChip > 0 ? processNoiseImposed(nClusters) : processNoiseFree(nClusters);
   }
-  float processNoiseFree(const std::array<int, NLayers> ncl);
-  float processNoiseImposed(const std::array<int, NLayers> ncl);
+  float processNoiseFree(int nClusters);
+  float processNoiseImposed(int nClusters);
   float process(const gsl::span<const o2::itsmft::CompClusterExt>& clusters)
   {
-    fillNClPerLayer(clusters);
-    return process(nClPerLayer);
+    return process(countClustersOnLayer(clusters));
   }
   static bool sSeedSet;
 

Detectors/ITSMFT/ITS/tracking/include/ITStracking/FastMultEstConfig.h

@@ -34,6 +34,7 @@ struct FastMultEstConfig : public o2::conf::ConfigurableParamHelper<FastMultEstC
   float imposeNoisePerChip = 1.e-9 * 1024 * 512; // assumed noise, free parameter if<0
 
   // cuts to reject to low or too high mult events
+  int cutMultClusLayer = NLayers - 1;   /// layer used for cluster multiplicity selection (by default the outermost one)
   float cutMultClusLow = 0;   /// reject ROF with estimated cluster mult. below this value (no cut if <0)
   float cutMultClusHigh = -1; /// reject ROF with estimated cluster mult. above this value (no cut if <0)
   float cutMultVtxLow = -1;   /// reject seed vertex if its multiplicity below this value (no cut if <0)

Detectors/ITSMFT/ITS/tracking/include/ITStracking/ROFLookupTables.h

@@ -683,6 +683,167 @@ class ROFVertexLookupTable : public LayerTimingBase<NLayers>
   std::vector<TableEntry> mFlatTable;
 };
 
+// GPU-friendly view of the ROF mask table
+template <int32_t NLayers>
+struct ROFMaskTableView {
+  const uint8_t* mFlatMask{nullptr};
+  const int32_t* mLayerROFOffsets{nullptr}; // size NLayers+1
+
+  GPUhdi() bool isROFEnabled(int32_t layer, int32_t rofId) const noexcept
+  {
+    assert(layer >= 0 && layer < NLayers);
+    return mFlatMask[mLayerROFOffsets[layer] + rofId] != 0;
+  }
+
+#ifndef GPUCA_GPUCODE
+  GPUh() void printAll() const
+  {
+    for (int32_t i = 0; i < NLayers; ++i) {
+      printLayer(i);
+    }
+  }
+
+  GPUh() void printLayer(int32_t layer) const
+  {
+    constexpr int w_rof = 10;
+    constexpr int w_active = 10;
+    int32_t nROFs = mLayerROFOffsets[layer + 1] - mLayerROFOffsets[layer];
+    LOGF(info, "Mask table: Layer %d", layer);
+    LOGF(info, "%*s | %*s", w_rof, "ROF", w_active, "Enabled");
+    LOGF(info, "%.*s-+-%.*s", w_rof, "----------", w_active, "----------");
+    for (int32_t i = 0; i < nROFs; ++i) {
+      LOGF(info, "%*d | %*d", w_rof, i, w_active, (int)isROFEnabled(layer, i));
+    }
+  }
+#endif
+};
+
+// Per-ROF per-layer boolean mask (uint8_t for GPU compatibility).
+template <int32_t NLayers>
+class ROFMaskTable : public LayerTimingBase<NLayers>
+{
+ public:
+  using BCRange = dataformats::RangeReference<LayerTiming::BCType, LayerTiming::BCType>;
+  using View = ROFMaskTableView<NLayers>;
+
+  GPUdDefault() ROFMaskTable() = default;
+  GPUh() explicit ROFMaskTable(const LayerTimingBase<NLayers>& timingBase) : LayerTimingBase<NLayers>(timingBase) { init(); }
+
+  GPUh() void init()
+  {
+    int32_t totalROFs = 0;
+    for (int32_t layer{0}; layer < NLayers; ++layer) {
+      mLayerROFOffsets[layer] = totalROFs;
+      totalROFs += this->getLayer(layer).mNROFsTF;
+    }
+    mLayerROFOffsets[NLayers] = totalROFs; // sentinel
+    mFlatMask.resize(totalROFs, 1);
+  }
+
+  GPUh() size_t getFlatMaskSize() const noexcept { return mFlatMask.size(); }
+
+  GPUh() bool isROFEnabled(int32_t layer, int32_t rofId) const noexcept { return mFlatMask[mLayerROFOffsets[layer] + rofId] != 0; }
+
+  GPUh() void setROFEnabled(int32_t layer, int32_t rofId, uint8_t state = 1) noexcept
+  {
+    assert(layer >= 0 && layer < NLayers);
+    assert(rofId >= 0 && rofId < mLayerROFOffsets[layer + 1] - mLayerROFOffsets[layer]);
+    mFlatMask[mLayerROFOffsets[layer] + rofId] = state;
+  }
+
+  GPUh() void setROFsEnabled(int32_t layer, int32_t firstRof, int32_t nRofs, uint8_t state = 1) noexcept
+  {
+    assert(layer >= 0 && layer < NLayers);
+    assert(firstRof >= 0);
+    assert(firstRof + nRofs <= mLayerROFOffsets[layer + 1] - mLayerROFOffsets[layer]);
+    std::memset(mFlatMask.data() + mLayerROFOffsets[layer] + firstRof, state, nRofs);
+  }
+
+  // Enable all ROFs in all layers that are time-compatible with the given BC range
+  GPUh() void selectROF(const BCRange& t)
+  {
+    const int32_t bcStart = t.getFirstEntry();
+    const int32_t bcEnd = t.getEntriesBound();
+    for (int32_t layer{0}; layer < NLayers; ++layer) {
+      const auto& lay = this->getLayer(layer);
+      const int32_t offset = mLayerROFOffsets[layer];
+      for (int32_t rofId{0}; rofId < lay.mNROFsTF; ++rofId) {
+        if (static_cast<int32_t>(lay.getROFStartInBC(rofId)) < bcEnd &&
+            static_cast<int32_t>(lay.getROFEndInBC(rofId)) > bcStart) {
+          mFlatMask[offset + rofId] = 1;
+        }
+      }
+    }
+  }
+
+  // Reset mask to 0, then enable all ROFs compatible with any of the given BC ranges
+  GPUh() void selectROFs(const std::vector<BCRange>& ts)
+  {
+    resetMask();
+    for (const auto& t : ts) {
+      selectROF(t);
+    }
+  }
+
+  GPUh() void resetMask(uint8_t s = 0)
+  {
+    std::memset(mFlatMask.data(), s, mFlatMask.size());
+  }
+
+  GPUh() void invertMask()
+  {
+    std::ranges::transform(mFlatMask, mFlatMask.begin(), [](uint8_t x) { return 1 - x; });
+  }
+
+  GPUh() void swap(ROFMaskTable& other) noexcept
+  {
+    std::swap(mFlatMask, other.mFlatMask);
+    std::swap(mLayerROFOffsets, other.mLayerROFOffsets);
+  }
+
+  GPUh() View getView() const
+  {
+    View view;
+    view.mFlatMask = mFlatMask.data();
+    view.mLayerROFOffsets = mLayerROFOffsets;
+    return view;
+  }
+
+  GPUh() View getDeviceView(const uint8_t* deviceFlatMaskPtr, const int32_t* deviceOffsetPtr) const
+  {
+    View view;
+    view.mFlatMask = deviceFlatMaskPtr;
+    view.mLayerROFOffsets = deviceOffsetPtr;
+    return view;
+  }
+#ifndef GPUCA_GPUCODE
+  GPUh() std::string asString() const
+  {
+    std::string mask_str;
+    for (int32_t i = 0; i < NLayers; ++i) {
+      int32_t nROFs = mLayerROFOffsets[i + 1] - mLayerROFOffsets[i];
+      int32_t enabledROFs = 0;
+      for (int32_t j = 0; j < nROFs; ++j) {
+        if (isROFEnabled(i, j)) {
+          ++enabledROFs;
+        }
+      }
+      mask_str += std::format("Layer {} ROFs enabled: {}/{} | ", i, enabledROFs, nROFs);
+    }
+    return mask_str;
+  }
+
+  GPUh() void print() const
+  {
+    LOG(info) << asString();
+  }
+#endif
+
+ private:
+  int32_t mLayerROFOffsets[NLayers + 1]{}; // NLayers entries + 1 sentinel
+  std::vector<uint8_t> mFlatMask;
+};
+
 } // namespace o2::its
 
 #endif

Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h

@@ -67,6 +67,7 @@ struct TimeFrame {
   using IndexTableUtilsN = IndexTableUtils<NLayers>;
   using ROFOverlapTableN = ROFOverlapTable<NLayers>;
   using ROFVertexLookupTableN = ROFVertexLookupTable<NLayers>;
+  using ROFMaskTableN = ROFMaskTable<NLayers>;
   using CellSeedN = CellSeed<NLayers>;
   friend class gpu::TimeFrameGPU<NLayers>;
 
@@ -220,8 +221,8 @@ struct TimeFrame {
   std::array<float, 2>& getBeamXY() { return mBeamPos; }
   // \Vertexer
 
-  void setMultiplicityCutMask(const std::vector<uint8_t>& cutMask) { mMultiplicityCutMask = cutMask; }
-  void setROFMask(const std::vector<uint8_t>& rofMask) { mROFMask = rofMask; }
+  void setMultiplicityCutMask(const ROFMaskTableN& cutMask) { mMultiplicityCutMask = cutMask; }
+  void setROFMask(const ROFMaskTableN& rofMask) { mROFMask = rofMask; }
   void swapMasks() { mMultiplicityCutMask.swap(mROFMask); }
 
   int hasBogusClusters() const { return std::accumulate(mBogusClusters.begin(), mBogusClusters.end(), 0); }
@@ -266,7 +267,7 @@ struct TimeFrame {
   bounded_vector<MCCompLabel> mTracksLabel;
   std::vector<bounded_vector<int>> mCellsNeighbours;
   std::vector<bounded_vector<int>> mCellsLookupTable;
-  std::vector<uint8_t> mMultiplicityCutMask;
+  ROFMaskTableN mMultiplicityCutMask;
 
   const o2::base::PropagatorImpl<float>* mPropagatorDevice = nullptr; // Needed only for GPU
 
@@ -290,7 +291,7 @@ struct TimeFrame {
   bounded_vector<float> mPositionResolution;
   std::array<bounded_vector<uint8_t>, NLayers> mClusterSize;
 
-  std::vector<uint8_t> mROFMask;
+  ROFMaskTableN mROFMask;
   bounded_vector<std::array<float, 2>> mPValphaX; /// PV x and alpha for track propagation
   std::vector<bounded_vector<MCCompLabel>> mTrackletLabels;
   std::vector<bounded_vector<MCCompLabel>> mCellLabels;