Merge pull request #1097 from stephenswat/feat/spbm_in_triplet

stephenswat · web-flow · commit 9b5e684a5287 · 2025-07-30T14:42:50.000+02:00
Embed bottom and middle spacepoint in triplet
diff --git a/core/include/traccc/seeding/seed_selecting_helper.hpp b/core/include/traccc/seeding/seed_selecting_helper.hpp
@@ -70,16 +70,12 @@ struct seed_selecting_helper {
     /// @param seed             current seed to possibly cut
     ///
     /// @return boolean value
+    template <typename spacepoint_type>
     static TRACCC_HOST_DEVICE bool cut_per_middle_sp(
-        const seedfilter_config& filter_config,
-        const edm::spacepoint_collection::const_device& spacepoints,
-        const details::spacepoint_grid_types::const_device& grid,
-        const triplet& seed) {
+        const seedfilter_config& filter_config, const spacepoint_type& spB,
+        const scalar weight) {
 
-        const edm::spacepoint_collection::const_device::const_proxy_type spB =
-            spacepoints.at(grid.bin(seed.sp1.bin_idx)[seed.sp1.sp_idx]);
-
-        return (seed.weight > filter_config.seed_min_weight ||
+        return (weight > filter_config.seed_min_weight ||
                 spB.radius() > filter_config.spB_min_radius);
     }
 };
diff --git a/core/src/seeding/seed_filtering.cpp b/core/src/seeding/seed_filtering.cpp
@@ -95,9 +95,14 @@ void seed_filtering::operator()(
             std::min(triplets_passing_single_seed_cuts.size(),
                      static_cast<std::size_t>(m_finder_config.maxSeedsPerSpM));
         for (std::size_t i = 1; i < itLength; ++i) {
+            const traccc::details::spacepoint_grid_types::const_device
+                sp_grid_accessor(sp_grid_data);
+            const auto& this_seed = triplets_passing_single_seed_cuts[i].get();
             if (seed_selecting_helper::cut_per_middle_sp(
-                    m_filter_config, spacepoints, sp_grid_data,
-                    triplets_passing_single_seed_cuts[i])) {
+                    m_filter_config,
+                    spacepoints.at(sp_grid_accessor.bin(
+                        this_seed.sp1.bin_idx)[this_seed.sp1.sp_idx]),
+                    this_seed.weight)) {
                 triplets_passing_final_cuts.push_back(
                     triplets_passing_single_seed_cuts[i]);
             }
diff --git a/device/alpaka/src/seeding/seed_finding.cpp b/device/alpaka/src/seeding/seed_finding.cpp
@@ -134,7 +134,6 @@ struct UpdateTripletWeights {
     ALPAKA_FN_ACC void operator()(
         TAcc const& acc, seedfilter_config filter_config,
         edm::spacepoint_collection::const_view spacepoints,
-        traccc::details::spacepoint_grid_types::const_view sp_grid,
         device::triplet_counter_spM_collection_types::const_view spM_tc,
         device::triplet_counter_collection_types::const_view midBot_tc,
         device::device_triplet_collection_types::view triplet_view) const {
@@ -151,8 +150,8 @@ struct UpdateTripletWeights {
         scalar* dataPos = &data[localThreadIdx * filter_config.compatSeedLimit];
 
         device::update_triplet_weights(globalThreadIdx, filter_config,
-                                       spacepoints, sp_grid, spM_tc, midBot_tc,
-                                       dataPos, triplet_view);
+                                       spacepoints, spM_tc, midBot_tc, dataPos,
+                                       triplet_view);
     }
 };
 
@@ -175,10 +174,12 @@ struct SelectSeeds {
 
         // Array for temporary storage of quality parameters for comparing
         // triplets within weight updating kernel
-        triplet* const data = ::alpaka::getDynSharedMem<triplet>(acc);
+        device::device_triplet* const data =
+            ::alpaka::getDynSharedMem<device::device_triplet>(acc);
 
-        // Each thread uses maxSeedsPerSpM elements of the array
-        triplet* dataPos = &data[localThreadIdx * finder_config.maxSeedsPerSpM];
+        // Each thread uses max_triplets_per_spM elements of the array
+        device::device_triplet* dataPos =
+            &data[localThreadIdx * finder_config.maxSeedsPerSpM];
 
         device::select_seeds(globalThreadIdx, finder_config, filter_config,
                              spacepoints, sp_view, spM_tc, midBot_tc,
@@ -360,7 +361,7 @@ edm::seed_collection::buffer seed_finding::operator()(
 
     // Update the weights of all spacepoint triplets.
     ::alpaka::exec<Acc>(queue, workDiv, kernels::UpdateTripletWeights{},
-                        m_seedfilter_config, spacepoints_view, g2_view,
+                        m_seedfilter_config, spacepoints_view,
                         vecmem::get_data(triplet_counter_spM_buffer),
                         vecmem::get_data(triplet_counter_midBot_buffer),
                         vecmem::get_data(triplet_buffer));
@@ -419,7 +420,7 @@ struct BlockSharedMemDynSizeBytes<traccc::alpaka::kernels::SelectSeeds, TAcc> {
         ) -> std::size_t {
         return static_cast<std::size_t>(finder_config.maxSeedsPerSpM *
                                         blockThreadExtent.prod()) *
-               sizeof(traccc::triplet);
+               sizeof(traccc::device::device_triplet);
     }
 };
 
diff --git a/device/common/include/traccc/edm/device/device_triplet.hpp b/device/common/include/traccc/edm/device/device_triplet.hpp
@@ -17,7 +17,7 @@ namespace traccc::device {
 /// Triplets of bottom, middle and top spacepoints
 struct device_triplet {
     // top spacepoint location in internal spacepoint container
-    sp_location spT;
+    unsigned int spB, spM, spT;
 
     using link_type = device::triplet_counter_collection_types::host::size_type;
     /// Link to triplet counter where the middle and bottom spacepoints are
diff --git a/device/common/include/traccc/seeding/device/impl/find_triplets.ipp b/device/common/include/traccc/seeding/device/impl/find_triplets.ipp
@@ -56,12 +56,14 @@ inline void find_triplets(
     const sp_location spB_loc = mid_bot_counter.spB;
 
     // middle spacepoint
+    const unsigned int spM_idx = sp_grid.bin(spM_loc.bin_idx)[spM_loc.sp_idx];
     const edm::spacepoint_collection::const_device::const_proxy_type spM =
-        spacepoints.at(sp_grid.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
+        spacepoints.at(spM_idx);
 
     // bottom spacepoint
+    const unsigned int spB_idx = sp_grid.bin(spB_loc.bin_idx)[spB_loc.sp_idx];
     const edm::spacepoint_collection::const_device::const_proxy_type spB =
-        spacepoints.at(sp_grid.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
+        spacepoints.at(spB_idx);
 
     // Set up the device result collection
     device_triplet_collection_types::device triplets(triplet_view);
@@ -95,8 +97,10 @@ inline void find_triplets(
     for (unsigned int i = mt_start_idx; i < mt_end_idx; ++i) {
         const sp_location spT_loc = mid_top_doublet_device[i].sp2;
 
+        const unsigned int spT_idx =
+            sp_grid.bin(spT_loc.bin_idx)[spT_loc.sp_idx];
         const edm::spacepoint_collection::const_device::const_proxy_type spT =
-            spacepoints.at(sp_grid.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
+            spacepoints.at(spT_idx);
 
         // Apply the conformal transformation to middle-top doublet
         const traccc::lin_circle lt =
@@ -110,7 +114,7 @@ inline void find_triplets(
 
             // Add triplet to jagged vector
             triplets.at(posTriplets++) = device_triplet(
-                {spT_loc, globalIndex, curvature,
+                {spB_idx, spM_idx, spT_idx, globalIndex, curvature,
                  -impact_parameter * filter_config.impactWeightFactor,
                  lb.Zo()});
         }
diff --git a/device/common/include/traccc/seeding/device/impl/select_seeds.ipp b/device/common/include/traccc/seeding/device/impl/select_seeds.ipp
@@ -20,7 +20,7 @@ namespace traccc::device {
 namespace details {
 // Finding minimum element algorithm
 template <typename Comparator>
-TRACCC_HOST_DEVICE std::size_t min_elem(const triplet* arr,
+TRACCC_HOST_DEVICE std::size_t min_elem(const device_triplet* arr,
                                         const std::size_t begin_idx,
                                         const std::size_t end_idx,
                                         Comparator comp) {
@@ -38,11 +38,11 @@ TRACCC_HOST_DEVICE std::size_t min_elem(const triplet* arr,
 
 // Sorting algorithm for sorting seeds in the local memory
 template <typename Comparator>
-TRACCC_HOST_DEVICE void insertionSort(triplet* arr,
+TRACCC_HOST_DEVICE void insertionSort(device_triplet* arr,
                                       const unsigned int begin_idx,
                                       const unsigned int n, Comparator comp) {
     int j = 0;
-    triplet key = arr[begin_idx];
+    device_triplet key = arr[begin_idx];
     for (unsigned int i = 0; i < n; ++i) {
         key = arr[begin_idx + i];
         j = static_cast<int>(i) - 1;
@@ -67,7 +67,7 @@ inline void select_seeds(
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view,
     const device_triplet_collection_types::const_view& triplet_view,
-    triplet* data, edm::seed_collection::view seed_view) {
+    device_triplet* data, edm::seed_collection::view seed_view) {
 
     // Check if anything needs to be done.
     const triplet_counter_spM_collection_types::const_device triplet_counts_spM(
@@ -90,8 +90,9 @@ inline void select_seeds(
     // Current work item = middle spacepoint
     const triplet_counter_spM spM_counter = triplet_counts_spM.at(globalIndex);
     const sp_location spM_loc = spM_counter.spM;
+    const unsigned int spM_idx = sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx];
     const edm::spacepoint_collection::const_device::const_proxy_type spM =
-        spacepoints.at(sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
+        spacepoints.at(spM_idx);
 
     // Number of triplets added for this spM
     unsigned int n_triplets_per_spM = 0;
@@ -103,14 +104,12 @@ inline void select_seeds(
         device_triplet aTriplet = triplets[i];
 
         // spacepoints bottom and top for this triplet
-        const sp_location spB_loc =
-            triplet_counts.at(static_cast<unsigned int>(aTriplet.counter_link))
-                .spB;
-        const sp_location spT_loc = aTriplet.spT;
+        const unsigned int spB_idx = aTriplet.spB;
         const edm::spacepoint_collection::const_device::const_proxy_type spB =
-            spacepoints.at(sp_device.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
+            spacepoints.at(spB_idx);
+        const unsigned int spT_idx = aTriplet.spT;
         const edm::spacepoint_collection::const_device::const_proxy_type spT =
-            spacepoints.at(sp_device.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
+            spacepoints.at(spT_idx);
 
         // update weight of triplet
         seed_selecting_helper::seed_weight(filter_config, spM, spB, spT,
@@ -126,61 +125,54 @@ inline void select_seeds(
         // the triplet with the lowest weight is removed
         if (n_triplets_per_spM >= finder_config.maxSeedsPerSpM) {
 
-            const std::size_t min_index =
-                details::min_elem(data, 0, finder_config.maxSeedsPerSpM,
-                                  [](const triplet lhs, const triplet rhs) {
-                                      return lhs.weight > rhs.weight;
-                                  });
+            const std::size_t min_index = details::min_elem(
+                data, 0, finder_config.maxSeedsPerSpM,
+                [](const device_triplet& lhs, const device_triplet& rhs) {
+                    return lhs.weight > rhs.weight;
+                });
 
             const scalar& min_weight = data[min_index].weight;
 
             if (aTriplet.weight > min_weight) {
-                data[min_index] = {spB_loc,         spM_loc,
-                                   spT_loc,         aTriplet.curvature,
-                                   aTriplet.weight, aTriplet.z_vertex};
+                data[min_index] = aTriplet;
             }
         }
 
         // if the number of good triplets is below the threshold, add
         // the current triplet to the array
         else if (n_triplets_per_spM < finder_config.maxSeedsPerSpM) {
-            data[n_triplets_per_spM] = {spB_loc,         spM_loc,
-                                        spT_loc,         aTriplet.curvature,
-                                        aTriplet.weight, aTriplet.z_vertex};
+            data[n_triplets_per_spM] = aTriplet;
             n_triplets_per_spM++;
         }
     }
 
     // sort the triplets per spM
     details::insertionSort(
         data, 0, n_triplets_per_spM,
-        traccc::details::triplet_sorter{spacepoints, sp_device});
+        [](const device_triplet& lhs, const device_triplet& rhs) {
+            return lhs.weight > rhs.weight;
+        });
 
     // the number of good seed per compatible middle spacepoint
     unsigned int n_seeds_per_spM = 0;
 
     // iterate over the good triplets for final selection of seeds
     for (unsigned int i = 0; i < n_triplets_per_spM; ++i) {
-        const triplet& aTriplet = data[i];
-        const sp_location& spB_loc = aTriplet.sp1;
-        const sp_location& spT_loc = aTriplet.sp3;
+        const device_triplet& aTriplet = data[i];
 
         // if the number of seeds reaches the threshold, break
         if (n_seeds_per_spM >= finder_config.maxSeedsPerSpM + 1) {
             break;
         }
 
         // check if it is a good triplet
-        if (seed_selecting_helper::cut_per_middle_sp(filter_config, spacepoints,
-                                                     sp_device, aTriplet) ||
+        if (seed_selecting_helper::cut_per_middle_sp(
+                filter_config, spacepoints.at(aTriplet.spB), aTriplet.weight) ||
             n_seeds_per_spM == 0) {
 
             n_seeds_per_spM++;
 
-            seeds_device.push_back(
-                {sp_device.bin(spB_loc.bin_idx)[spB_loc.sp_idx],
-                 sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx],
-                 sp_device.bin(spT_loc.bin_idx)[spT_loc.sp_idx]});
+            seeds_device.push_back({aTriplet.spB, aTriplet.spM, aTriplet.spT});
         }
     }
 }
diff --git a/device/common/include/traccc/seeding/device/impl/update_triplet_weights.ipp b/device/common/include/traccc/seeding/device/impl/update_triplet_weights.ipp
@@ -19,7 +19,6 @@ TRACCC_HOST_DEVICE
 inline void update_triplet_weights(
     const global_index_t globalIndex, const seedfilter_config& filter_config,
     const edm::spacepoint_collection::const_view& spacepoints_view,
-    const traccc::details::spacepoint_grid_types::const_view& sp_view,
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view, scalar* data,
     device_triplet_collection_types::view triplet_view) {
@@ -33,7 +32,6 @@ inline void update_triplet_weights(
     // Set up the device containers
     const edm::spacepoint_collection::const_device spacepoints{
         spacepoints_view};
-    const traccc::details::spacepoint_grid_types::const_device sp_grid(sp_view);
     const triplet_counter_spM_collection_types::const_device triplet_counts_spM(
         spM_tc_view);
     const triplet_counter_collection_types::const_device triplet_counts(
@@ -42,11 +40,8 @@ inline void update_triplet_weights(
     // Current work item
     device_triplet this_triplet = triplets.at(globalIndex);
 
-    const sp_location& spT_idx = this_triplet.spT;
-
     const edm::spacepoint_collection::const_device::const_proxy_type
-        current_spT =
-            spacepoints.at(sp_grid.bin(spT_idx.bin_idx)[spT_idx.sp_idx]);
+        current_spT = spacepoints.at(this_triplet.spT);
 
     const scalar currentTop_r = current_spT.radius();
 
@@ -82,10 +77,8 @@ inline void update_triplet_weights(
         }
 
         const device_triplet other_triplet = triplets[i];
-        const sp_location other_spT_idx = other_triplet.spT;
         const edm::spacepoint_collection::const_device::const_proxy_type
-            other_spT = spacepoints.at(
-                sp_grid.bin(other_spT_idx.bin_idx)[other_spT_idx.sp_idx]);
+            other_spT = spacepoints.at(other_triplet.spT);
 
         // compared top SP should have at least deltaRMin distance
         const scalar otherTop_r = other_spT.radius();
diff --git a/device/common/include/traccc/seeding/device/update_triplet_weights.hpp b/device/common/include/traccc/seeding/device/update_triplet_weights.hpp
@@ -37,7 +37,6 @@ TRACCC_HOST_DEVICE
 inline void update_triplet_weights(
     global_index_t globalIndex, const seedfilter_config& filter_config,
     const edm::spacepoint_collection::const_view& spacepoints,
-    const traccc::details::spacepoint_grid_types::const_view& sp_view,
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view, scalar* data,
     device_triplet_collection_types::view triplet_view);
diff --git a/device/cuda/src/seeding/seed_finding.cu b/device/cuda/src/seeding/seed_finding.cu
@@ -111,7 +111,6 @@ __global__ void find_triplets(
 __global__ void update_triplet_weights(
     seedfilter_config filter_config,
     edm::spacepoint_collection::const_view spacepoints,
-    traccc::details::spacepoint_grid_types::const_view sp_grid,
     device::triplet_counter_spM_collection_types::const_view spM_tc,
     device::triplet_counter_collection_types::const_view midBot_tc,
     device::device_triplet_collection_types::view triplet_view) {
@@ -123,8 +122,8 @@ __global__ void update_triplet_weights(
     scalar* dataPos = &data[threadIdx.x * filter_config.compatSeedLimit];
 
     device::update_triplet_weights(details::global_index1(), filter_config,
-                                   spacepoints, sp_grid, spM_tc, midBot_tc,
-                                   dataPos, triplet_view);
+                                   spacepoints, spM_tc, midBot_tc, dataPos,
+                                   triplet_view);
 }
 
 /// CUDA kernel for running @c traccc::device::select_seeds
@@ -139,9 +138,10 @@ __global__ void select_seeds(
 
     // Array for temporary storage of triplets for comparing within seed
     // selecting kernel
-    extern __shared__ triplet data2[];
-    // Each thread uses maxSeedsPerSpM elements of the array
-    triplet* dataPos = &data2[threadIdx.x * finder_config.maxSeedsPerSpM];
+    extern __shared__ device::device_triplet data2[];
+    // Each thread uses max_triplets_per_spM elements of the array
+    device::device_triplet* dataPos =
+        &data2[threadIdx.x * finder_config.maxSeedsPerSpM];
 
     device::select_seeds(details::global_index1(), finder_config, filter_config,
                          spacepoints, sp_view, spM_tc, midBot_tc, triplet_view,
@@ -336,7 +336,7 @@ edm::seed_collection::buffer seed_finding::operator()(
         nWeightUpdatingBlocks, nWeightUpdatingThreads,
         sizeof(scalar) * m_seedfilter_config.compatSeedLimit *
             nWeightUpdatingThreads,
-        stream>>>(m_seedfilter_config, spacepoints_view, g2_view,
+        stream>>>(m_seedfilter_config, spacepoints_view,
                   triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
                   triplet_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
@@ -355,14 +355,14 @@ edm::seed_collection::buffer seed_finding::operator()(
         nSeedSelectingThreads;
 
     // Create seeds out of selected triplets
-    kernels::
-        select_seeds<<<nSeedSelectingBlocks, nSeedSelectingThreads,
-                       sizeof(triplet) * m_seedfinder_config.maxSeedsPerSpM *
-                           nSeedSelectingThreads,
-                       stream>>>(
-            m_seedfinder_config, m_seedfilter_config, spacepoints_view, g2_view,
-            triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
-            triplet_buffer, seed_buffer);
+    kernels::select_seeds<<<nSeedSelectingBlocks, nSeedSelectingThreads,
+                            sizeof(device::device_triplet) *
+                                m_seedfinder_config.maxSeedsPerSpM *
+                                nSeedSelectingThreads,
+                            stream>>>(
+        m_seedfinder_config, m_seedfilter_config, spacepoints_view, g2_view,
+        triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
+        triplet_buffer, seed_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 
     return seed_buffer;
diff --git a/device/sycl/src/seeding/seed_finding.sycl b/device/sycl/src/seeding/seed_finding.sycl
