Merge pull request #48698 from SegmentLinking/full_dyn_t5

cmsbuild · web-flow · commit 78cc7ca98f8e · 2025-08-27T20:13:19.000+02:00
Dynamic Memory Caps for LS, T3, T5 Objects in LST
diff --git a/RecoTracker/LSTCore/interface/MiniDoubletsSoA.h b/RecoTracker/LSTCore/interface/MiniDoubletsSoA.h
@@ -39,7 +39,8 @@ namespace lst {
                       SOA_COLUMN(float, outerHighEdgeX),
                       SOA_COLUMN(float, outerHighEdgeY),
                       SOA_COLUMN(float, outerLowEdgeX),
-                      SOA_COLUMN(float, outerLowEdgeY))
+                      SOA_COLUMN(float, outerLowEdgeY),
+                      SOA_COLUMN(unsigned int, connectedMax))
 
   GENERATE_SOA_LAYOUT(MiniDoubletsOccupancySoALayout,
                       SOA_COLUMN(unsigned int, nMDs),
diff --git a/RecoTracker/LSTCore/interface/SegmentsSoA.h b/RecoTracker/LSTCore/interface/SegmentsSoA.h
@@ -19,7 +19,8 @@ namespace lst {
                       SOA_COLUMN(uint16_t, outerLowerModuleIndices),
                       SOA_COLUMN(Params_LS::ArrayUxLayers, mdIndices),
                       SOA_COLUMN(unsigned int, innerMiniDoubletAnchorHitIndices),
-                      SOA_COLUMN(unsigned int, outerMiniDoubletAnchorHitIndices))
+                      SOA_COLUMN(unsigned int, outerMiniDoubletAnchorHitIndices),
+                      SOA_COLUMN(unsigned int, connectedMax))
 
   GENERATE_SOA_LAYOUT(SegmentsOccupancySoALayout,
                       SOA_COLUMN(unsigned int, nSegments),  //number of segments per inner lower module
diff --git a/RecoTracker/LSTCore/interface/TripletsSoA.h b/RecoTracker/LSTCore/interface/TripletsSoA.h
@@ -13,13 +13,14 @@ namespace lst {
                       SOA_COLUMN(Params_T3::ArrayU16xLayers, lowerModuleIndices),  // lower module index in each layer
                       SOA_COLUMN(Params_T3::ArrayU8xLayers, logicalLayers),        // layer ID
                       SOA_COLUMN(Params_T3::ArrayUxHits, hitIndices),              // hit indices
-                      SOA_COLUMN(FPX, betaIn),            // beta/chord angle of the inner segment
-                      SOA_COLUMN(float, centerX),         // lower/anchor-hit based circle center x
-                      SOA_COLUMN(float, centerY),         // lower/anchor-hit based circle center y
-                      SOA_COLUMN(float, radius),          // lower/anchor-hit based circle radius
-                      SOA_COLUMN(float, fakeScore),       // DNN confidence score for fake t3
-                      SOA_COLUMN(float, promptScore),     // DNN confidence score for real (prompt) t3
-                      SOA_COLUMN(float, displacedScore),  // DNN confidence score for real (displaced) t3
+                      SOA_COLUMN(FPX, betaIn),                 // beta/chord angle of the inner segment
+                      SOA_COLUMN(float, centerX),              // lower/anchor-hit based circle center x
+                      SOA_COLUMN(float, centerY),              // lower/anchor-hit based circle center y
+                      SOA_COLUMN(float, radius),               // lower/anchor-hit based circle radius
+                      SOA_COLUMN(float, fakeScore),            // DNN confidence score for fake t3
+                      SOA_COLUMN(float, promptScore),          // DNN confidence score for real (prompt) t3
+                      SOA_COLUMN(float, displacedScore),       // DNN confidence score for real (displaced) t3
+                      SOA_COLUMN(unsigned int, connectedMax),  // number of outer-triplets that pass the MD-equality cut
 #ifdef CUT_VALUE_DEBUG
                       SOA_COLUMN(float, zOut),
                       SOA_COLUMN(float, rtOut),
diff --git a/RecoTracker/LSTCore/src/alpaka/LSTEvent.dev.cc b/RecoTracker/LSTCore/src/alpaka/LSTEvent.dev.cc
@@ -209,6 +209,10 @@ void LSTEvent::createMiniDoublets() {
     alpaka::memset(queue_, totOccupancyMDs_view, 0u);
   }
 
+  auto mdView = miniDoubletsDC_->view<MiniDoubletsSoA>();
+  auto connView = cms::alpakatools::make_device_view(queue_, mdView.connectedMax(), mdView.metadata().size());
+  alpaka::memset(queue_, connView, 0u);
+
   unsigned int mdSize = pixelSize_ * 2;
   auto src_view_mdSize = cms::alpakatools::make_host_view(mdSize);
 
@@ -255,6 +259,17 @@ void LSTEvent::createMiniDoublets() {
 
 void LSTEvent::createSegmentsWithModuleMap() {
   if (!segmentsDC_) {
+    auto const countMDConn_wd = cms::alpakatools::make_workdiv<Acc3D>({nLowerModules_, 1, 1}, {1, 8, 32});
+
+    alpaka::exec<Acc3D>(queue_,
+                        countMDConn_wd,
+                        CountMiniDoubletConnections{},
+                        modules_.const_view<ModulesSoA>(),
+                        miniDoubletsDC_->view<MiniDoubletsSoA>(),
+                        miniDoubletsDC_->const_view<MiniDoubletsOccupancySoA>(),
+                        rangesDC_->const_view(),
+                        ptCut_);
+
     auto const createSegmentArrayRanges_workDiv = cms::alpakatools::make_workdiv<Acc1D>(1, 1024);
 
     alpaka::exec<Acc1D>(queue_,
@@ -263,8 +278,7 @@ void LSTEvent::createSegmentsWithModuleMap() {
                         modules_.const_view<ModulesSoA>(),
                         rangesDC_->view(),
                         miniDoubletsDC_->const_view<MiniDoubletsSoA>(),
-                        miniDoubletsDC_->const_view<MiniDoubletsOccupancySoA>(),
-                        ptCut_);
+                        miniDoubletsDC_->const_view<MiniDoubletsOccupancySoA>());
 
     auto rangesOccupancy = rangesDC_->view();
     auto nTotalSegments_view_h = cms::alpakatools::make_host_view(nTotalSegments_);
@@ -278,12 +292,15 @@ void LSTEvent::createSegmentsWithModuleMap() {
     segmentsDC_.emplace(segments_sizes, queue_);
 
     auto segmentsOccupancy = segmentsDC_->view<SegmentsOccupancySoA>();
+    auto segments = segmentsDC_->view<SegmentsSoA>();
     auto nSegments_view =
         cms::alpakatools::make_device_view(queue_, segmentsOccupancy.nSegments(), segmentsOccupancy.metadata().size());
     auto totOccupancySegments_view = cms::alpakatools::make_device_view(
         queue_, segmentsOccupancy.totOccupancySegments(), segmentsOccupancy.metadata().size());
     alpaka::memset(queue_, nSegments_view, 0u);
     alpaka::memset(queue_, totOccupancySegments_view, 0u);
+    auto conn_view = cms::alpakatools::make_device_view(queue_, segments.connectedMax(), segments.metadata().size());
+    alpaka::memset(queue_, conn_view, 0u);
 
     auto src_view_size = cms::alpakatools::make_host_view(pixelSize_);
 
@@ -297,7 +314,7 @@ void LSTEvent::createSegmentsWithModuleMap() {
     alpaka::wait(queue_);
   }
 
-  auto const createSegments_workDiv = cms::alpakatools::make_workdiv<Acc3D>({nLowerModules_, 1, 1}, {1, 1, 64});
+  auto const createSegments_workDiv = cms::alpakatools::make_workdiv<Acc3D>({nLowerModules_, 1, 1}, {1, 8, 32});
 
   alpaka::exec<Acc3D>(queue_,
                       createSegments_workDiv,
@@ -326,6 +343,16 @@ void LSTEvent::createSegmentsWithModuleMap() {
 
 void LSTEvent::createTriplets() {
   if (!tripletsDC_) {
+    auto const countSegConn_wd = cms::alpakatools::make_workdiv<Acc3D>({nLowerModules_, 1, 1}, {1, 16, 16});
+
+    alpaka::exec<Acc3D>(queue_,
+                        countSegConn_wd,
+                        CountSegmentConnections{},
+                        modules_.const_view<ModulesSoA>(),
+                        segmentsDC_->view<SegmentsSoA>(),
+                        segmentsDC_->const_view<SegmentsOccupancySoA>(),
+                        rangesDC_->const_view());
+
     auto const createTripletArrayRanges_workDiv = cms::alpakatools::make_workdiv<Acc1D>(1, 1024);
 
     alpaka::exec<Acc1D>(queue_,
@@ -334,8 +361,7 @@ void LSTEvent::createTriplets() {
                         modules_.const_view<ModulesSoA>(),
                         rangesDC_->view(),
                         segmentsDC_->const_view<SegmentsSoA>(),
-                        segmentsDC_->const_view<SegmentsOccupancySoA>(),
-                        ptCut_);
+                        segmentsDC_->const_view<SegmentsOccupancySoA>());
 
     // TODO: Why are we pulling this back down only to put it back on the device in a new struct?
     auto rangesOccupancy = rangesDC_->view();
@@ -362,6 +388,9 @@ void LSTEvent::createTriplets() {
     alpaka::memset(queue_, partOfT5_view, 0u);
     auto partOfPT3_view = cms::alpakatools::make_device_view(queue_, triplets.partOfPT3(), triplets.metadata().size());
     alpaka::memset(queue_, partOfPT3_view, 0u);
+    auto connectedMax_view =
+        cms::alpakatools::make_device_view(queue_, triplets.connectedMax(), triplets.metadata().size());
+    alpaka::memset(queue_, connectedMax_view, 0u);
   }
 
   uint16_t nonZeroModules = 0;
@@ -715,6 +744,17 @@ void LSTEvent::createPixelTriplets() {
 }
 
 void LSTEvent::createQuintuplets() {
+  auto const countConn_workDiv = cms::alpakatools::make_workdiv<Acc3D>({nLowerModules_, 1, 1}, {1, 8, 32});
+
+  alpaka::exec<Acc3D>(queue_,
+                      countConn_workDiv,
+                      CountTripletConnections{},
+                      modules_.const_view<ModulesSoA>(),
+                      segmentsDC_->const_view<SegmentsSoA>(),
+                      tripletsDC_->view<TripletsSoA>(),
+                      tripletsDC_->const_view<TripletsOccupancySoA>(),
+                      rangesDC_->const_view());
+
   auto const createEligibleModulesListForQuintuplets_workDiv = cms::alpakatools::make_workdiv<Acc1D>(1, 1024);
 
   alpaka::exec<Acc1D>(queue_,
@@ -723,8 +763,7 @@ void LSTEvent::createQuintuplets() {
                       modules_.const_view<ModulesSoA>(),
                       tripletsDC_->const_view<TripletsOccupancySoA>(),
                       rangesDC_->view(),
-                      tripletsDC_->view<TripletsSoA>(),
-                      ptCut_);
+                      tripletsDC_->view<TripletsSoA>());
 
   auto nEligibleT5Modules_buf = cms::alpakatools::make_host_buffer<uint16_t>(queue_);
   auto nTotalQuintuplets_buf = cms::alpakatools::make_host_buffer<unsigned int>(queue_);
diff --git a/RecoTracker/LSTCore/src/alpaka/Quintuplet.h b/RecoTracker/LSTCore/src/alpaka/Quintuplet.h
@@ -1697,7 +1697,9 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
                                   ObjectRangesConst ranges,
                                   uint16_t nEligibleT5Modules,
                                   const float ptCut) const {
-      for (int iter : cms::alpakatools::uniform_elements_z(acc, nEligibleT5Modules)) {
+      ALPAKA_ASSERT_ACC((alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[1] == 1) &&
+                        (alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[2] == 1));
+      for (int iter : cms::alpakatools::uniform_groups_z(acc, nEligibleT5Modules)) {
         uint16_t lowerModule1 = ranges.indicesOfEligibleT5Modules()[iter];
         short layer2_adjustment;
         int layer = modules.layers()[lowerModule1];
@@ -1819,17 +1821,70 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
     }
   };
 
+  ALPAKA_FN_ACC ALPAKA_FN_INLINE bool isValidQuintRegion(ModulesConst modules, uint16_t lowerModule) {
+    const short layer = modules.layers()[lowerModule];
+    const short subdet = modules.subdets()[lowerModule];
+    // Quintuplets starting outside these regions are not built.
+    return (subdet == Barrel && layer < 3) || (subdet == Endcap && layer <= 1);
+  }
+
+  struct CountTripletConnections {
+    ALPAKA_FN_ACC void operator()(Acc3D const& acc,
+                                  ModulesConst modules,
+                                  SegmentsConst segments,
+                                  Triplets triplets,
+                                  TripletsOccupancyConst tripletsOcc,
+                                  ObjectRangesConst ranges) const {
+      // The atomicAdd below with hierarchy::Threads{} requires one block in x, y dimensions.
+      ALPAKA_ASSERT_ACC((alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[1] == 1) &&
+                        (alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[2] == 1));
+      const auto& mdIndices = segments.mdIndices();
+      const auto& segIdx = triplets.segmentIndices();
+      const auto& lmIdx = triplets.lowerModuleIndices();
+      const auto& tripIdx = ranges.tripletModuleIndices();
+
+      for (uint16_t lowerModule1 : cms::alpakatools::uniform_groups_z(acc, modules.nLowerModules())) {
+        if (!isValidQuintRegion(modules, lowerModule1))
+          continue;
+
+        const unsigned int nInnerTriplets = tripletsOcc.nTriplets()[lowerModule1];
+        if (nInnerTriplets == 0)
+          continue;
+
+        for (unsigned int innerTripletArrayIndex : cms::alpakatools::uniform_elements_y(acc, nInnerTriplets)) {
+          const unsigned int innerTripletIndex = tripIdx[lowerModule1] + innerTripletArrayIndex;
+
+          const uint16_t lowerModule3 = lmIdx[innerTripletIndex][2];
+          const unsigned int nOuterTriplets = tripletsOcc.nTriplets()[lowerModule3];
+          if (nOuterTriplets == 0)
+            continue;
+
+          const unsigned int secondSegIdx = segIdx[innerTripletIndex][1];
+          const unsigned int secondMDOuter = mdIndices[secondSegIdx][1];
+
+          for (unsigned int outerTripletArrayIndex : cms::alpakatools::uniform_elements_x(acc, nOuterTriplets)) {
+            const unsigned int outerTripletIndex = tripIdx[lowerModule3] + outerTripletArrayIndex;
+            const unsigned int thirdSegIdx = segIdx[outerTripletIndex][0];
+            const unsigned int thirdMDInner = mdIndices[thirdSegIdx][0];
+
+            if (secondMDOuter == thirdMDInner) {
+              alpaka::atomicAdd(acc, &triplets.connectedMax()[innerTripletIndex], 1u, alpaka::hierarchy::Threads{});
+            }
+          }
+        }
+      }
+    }
+  };
+
   struct CreateEligibleModulesListForQuintuplets {
     ALPAKA_FN_ACC void operator()(Acc1D const& acc,
                                   ModulesConst modules,
-                                  TripletsOccupancyConst tripletsOccupancy,
+                                  TripletsOccupancyConst tripletsOcc,
                                   ObjectRanges ranges,
-                                  Triplets triplets,
-                                  const float ptCut) const {
-      // implementation is 1D with a single block
+                                  TripletsConst triplets) const {
+      // Single-block kernel
       ALPAKA_ASSERT_ACC((alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[0] == 1));
 
-      // Initialize variables in shared memory and set to 0
       int& nEligibleT5Modulesx = alpaka::declareSharedVar<int, __COUNTER__>(acc);
       int& nTotalQuintupletsx = alpaka::declareSharedVar<int, __COUNTER__>(acc);
       if (cms::alpakatools::once_per_block(acc)) {
@@ -1838,75 +1893,31 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
       }
       alpaka::syncBlockThreads(acc);
 
-      // Occupancy matrix for 0.8 GeV pT Cut
-      constexpr int p08_occupancy_matrix[4][4] = {
-          {336, 414, 231, 146},  // category 0
-          {0, 0, 0, 0},          // category 1
-          {0, 0, 0, 0},          // category 2
-          {0, 0, 191, 106}       // category 3
-      };
-
-      // Occupancy matrix for 0.6 GeV pT Cut, 99.99%
-      constexpr int p06_occupancy_matrix[4][4] = {
-          {325, 237, 217, 176},  // category 0
-          {0, 0, 0, 0},          // category 1
-          {0, 0, 0, 0},          // category 2
-          {0, 0, 129, 180}       // category 3
-      };
-
-      // Select the appropriate occupancy matrix based on ptCut
-      const auto& occupancy_matrix = (ptCut < 0.8f) ? p06_occupancy_matrix : p08_occupancy_matrix;
-
-      for (int i : cms::alpakatools::uniform_elements(acc, modules.nLowerModules())) {
-        // Condition for a quintuple to exist for a module
-        // TCs don't exist for layers 5 and 6 barrel, and layers 2,3,4,5 endcap
-        short module_rings = modules.rings()[i];
-        short module_layers = modules.layers()[i];
-        short module_subdets = modules.subdets()[i];
-        float module_eta = alpaka::math::abs(acc, modules.eta()[i]);
-
-        if (tripletsOccupancy.nTriplets()[i] == 0)
+      for (uint16_t lowerModule : cms::alpakatools::uniform_elements(acc, modules.nLowerModules())) {
+        if (!isValidQuintRegion(modules, lowerModule))
           continue;
-        if (module_subdets == Barrel && module_layers >= 3)
-          continue;
-        if (module_subdets == Endcap && module_layers > 1)
+
+        unsigned int nInnerTriplets = tripletsOcc.nTriplets()[lowerModule];
+        if (nInnerTriplets == 0)
           continue;
 
+        // Sum the real connectivity for triplets in this module
         int dynamic_count = 0;
-
-        // How many triplets are in module i?
-        int nTriplets_i = tripletsOccupancy.nTriplets()[i];
-        int firstTripletIdx = ranges.tripletModuleIndices()[i];
-
-        // Loop over all triplets that live in module i
-        for (int t = 0; t < nTriplets_i; t++) {
-          int tripletIndex = firstTripletIdx + t;
-          uint16_t outerModule = triplets.lowerModuleIndices()[tripletIndex][2];
-          dynamic_count += tripletsOccupancy.nTriplets()[outerModule];
+        const unsigned int firstTripletIdx = ranges.tripletModuleIndices()[lowerModule];
+        for (unsigned int t = 0; t < nInnerTriplets; ++t) {
+          unsigned int tripletIndex = firstTripletIdx + t;
+          dynamic_count += triplets.connectedMax()[tripletIndex];
         }
 
-        int category_number = getCategoryNumber(module_layers, module_subdets, module_rings);
-        int eta_number = getEtaBin(module_eta);
-
-#ifdef WARNINGS
-        if (category_number == -1 || eta_number == -1) {
-          printf("Unhandled case in createEligibleModulesListForQuintupletsGPU! Module index = %i\n", i);
-        }
-#endif
-
-        // Get matrix-based cap (use dynamic_count as fallback)
-        int matrix_cap =
-            (category_number != -1 && eta_number != -1) ? occupancy_matrix[category_number][eta_number] : 0;
-
-        // Cap occupancy at minimum of dynamic count and matrix value
-        int occupancy = alpaka::math::min(acc, dynamic_count, matrix_cap);
+        if (dynamic_count == 0)
+          continue;
 
         int nEligibleT5Modules = alpaka::atomicAdd(acc, &nEligibleT5Modulesx, 1, alpaka::hierarchy::Threads{});
-        int nTotQ = alpaka::atomicAdd(acc, &nTotalQuintupletsx, occupancy, alpaka::hierarchy::Threads{});
+        int nTotQ = alpaka::atomicAdd(acc, &nTotalQuintupletsx, dynamic_count, alpaka::hierarchy::Threads{});
 
-        ranges.quintupletModuleIndices()[i] = nTotQ;
-        ranges.indicesOfEligibleT5Modules()[nEligibleT5Modules] = i;
-        ranges.quintupletModuleOccupancy()[i] = occupancy;
+        ranges.quintupletModuleIndices()[lowerModule] = nTotQ;
+        ranges.indicesOfEligibleT5Modules()[nEligibleT5Modules] = lowerModule;
+        ranges.quintupletModuleOccupancy()[lowerModule] = dynamic_count;
       }
 
       // Wait for all threads to finish before reporting final values
diff --git a/RecoTracker/LSTCore/src/alpaka/Segment.h b/RecoTracker/LSTCore/src/alpaka/Segment.h
diff --git a/RecoTracker/LSTCore/src/alpaka/Triplet.h b/RecoTracker/LSTCore/src/alpaka/Triplet.h
