Access the track parameters in a sorted order based on the theta

Author: beomki-yeo

device/common/CMakeLists.txt

@@ -26,6 +26,7 @@ traccc_add_library( traccc_device_common device_common TYPE SHARED
    "include/traccc/edm/device/triplet_counter.hpp"
    "include/traccc/edm/device/device_doublet.hpp"
    "include/traccc/edm/device/device_triplet.hpp"
+   "include/traccc/edm/device/theta_id.hpp"
    # Clusterization function(s).
    "include/traccc/clusterization/device/form_spacepoints.hpp"
    "include/traccc/clusterization/device/impl/form_spacepoints.ipp"

device/common/include/traccc/edm/device/theta_id.hpp

@@ -0,0 +1,30 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/definitions/primitives.hpp"
+
+namespace traccc::device {
+
+/// Theta of track parameters and its index
+struct theta_id {
+    traccc::scalar theta;
+    unsigned int id;
+};  // struct theta_id
+
+/// Comparator based on theta
+struct theta_id_sort_comp {
+    TRACCC_HOST_DEVICE
+    bool operator()(const theta_id& lhs, const theta_id& rhs) {
+        return math::abs(lhs.theta - constant<traccc::scalar>::pi_2) <
+               math::abs(rhs.theta - constant<traccc::scalar>::pi_2);
+    }
+};
+
+}  // namespace traccc::device

device/common/include/traccc/finding/device/add_links_for_holes.hpp

@@ -22,6 +22,7 @@ TRACCC_DEVICE inline void add_links_for_holes(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int& n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_total_candidates);
 

device/common/include/traccc/finding/device/find_tracks.hpp

@@ -38,6 +38,7 @@ namespace traccc::device {
 /// @param[in] step               Step index
 /// @param[in] n_max_candidates   Number of maximum candidates
 /// @param[out] out_params_view   Output parameters
+/// @param[out] theta_id_view     Pair of theta and id of output parameters
 /// @param[out] n_candidates_view  Number of candidates per input parameter
 /// @param[out] links_view        link container for the current step
 /// @param[out] n_total_candidates  The number of total candidates for the
@@ -56,6 +57,7 @@ TRACCC_DEVICE inline void find_tracks(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int& n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<unsigned int> n_candidates_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_total_candidates);

device/common/include/traccc/finding/device/impl/add_links_for_holes.ipp

@@ -19,6 +19,7 @@ TRACCC_DEVICE inline void add_links_for_holes(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int& n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_total_candidates) {
 
@@ -43,6 +44,9 @@ TRACCC_DEVICE inline void add_links_for_holes(
     // Output parameters
     bound_track_parameters_collection_types::device out_params(out_params_view);
 
+    // Theta id
+    vecmem::device_vector<device::theta_id> theta_id(theta_id_view);
+
     // Links
     vecmem::device_vector<candidate_link> links(links_view);
 
@@ -84,6 +88,7 @@ TRACCC_DEVICE inline void add_links_for_holes(
             skip_counter + 1};
 
         out_params.at(l_pos) = in_params.at(globalIndex);
+        theta_id.at(l_pos) = {out_params.at(l_pos).theta(), l_pos};
     }
 }
 

device/common/include/traccc/finding/device/impl/find_tracks.ipp

@@ -28,6 +28,7 @@ TRACCC_DEVICE inline void find_tracks(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int& n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<unsigned int> n_candidates_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_total_candidates) {
@@ -52,6 +53,9 @@ TRACCC_DEVICE inline void find_tracks(
     // Output parameters
     bound_track_parameters_collection_types::device out_params(out_params_view);
 
+    // Theta id
+    vecmem::device_vector<device::theta_id> theta_id(theta_id_view);
+
     // Number of candidates per parameter
     vecmem::device_vector<unsigned int> n_candidates(n_candidates_view);
 
@@ -145,7 +149,8 @@ TRACCC_DEVICE inline void find_tracks(
                 n_candidates[in_param_id]);
             num_candidates.fetch_add(1);
 
-            out_params[l_pos] = trk_state.filtered();
+            out_params.at(l_pos) = trk_state.filtered();
+            theta_id.at(l_pos) = {out_params.at(l_pos).theta(), l_pos};
         }
     }
 }

device/common/include/traccc/finding/device/impl/propagate_to_next_surface.ipp

@@ -15,6 +15,7 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     typename propagator_t::detector_type::view_type det_data,
     bfield_t field_data,
     bound_track_parameters_collection_types::const_view in_params_view,
+    vecmem::data::vector_view<const device::theta_id> theta_id_view,
     vecmem::data::vector_view<const candidate_link> links_view,
     const unsigned int step, const unsigned int& n_in_params,
     bound_track_parameters_collection_types::view out_params_view,
@@ -28,6 +29,9 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
         return;
     }
 
+    // Theta id
+    vecmem::device_vector<const device::theta_id> theta_id(theta_id_view);
+
     // Number of tracks per seed
     vecmem::device_vector<unsigned int> n_tracks_per_seed(
         n_tracks_per_seed_view);
@@ -71,7 +75,8 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     vecmem::device_vector<unsigned int> param_to_link(param_to_link_view);
 
     // Input bound track parameter
-    const bound_track_parameters in_par = in_params.at(globalIndex);
+    const unsigned int in_par_id = theta_id.at(globalIndex).id;
+    const bound_track_parameters in_par = in_params.at(in_par_id);
 
     // Create propagator
     propagator_t propagator(cfg.propagation);

device/common/include/traccc/finding/device/propagate_to_next_surface.hpp

@@ -26,6 +26,7 @@ namespace traccc::device {
 /// @param[in] cfg                Track finding config object
 /// @param[in] det_data           Detector view object
 /// @param[in] in_params_view     Input parameters
+/// @param[in] theta_id_view      Theta and id of input parameters
 /// @param[in] links_view         Link container for the current step
 /// @param[in] step               Step index
 /// @param[in] n_in_params        The number of input parameters
@@ -41,6 +42,7 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     typename propagator_t::detector_type::view_type det_data,
     bfield_t field_data,
     bound_track_parameters_collection_types::const_view in_params_view,
+    vecmem::data::vector_view<const device::theta_id> theta_id_view,
     vecmem::data::vector_view<const candidate_link> links_view,
     const unsigned int step, const unsigned int& n_in_params,
     bound_track_parameters_collection_types::view out_params_view,

device/cuda/src/finding/finding_algorithm.cu

@@ -13,6 +13,7 @@
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/device/finding_global_counter.hpp"
+#include "traccc/edm/device/theta_id.hpp"
 #include "traccc/finding/candidate_link.hpp"
 #include "traccc/finding/device/add_links_for_holes.hpp"
 #include "traccc/finding/device/apply_interaction.hpp"
@@ -105,6 +106,7 @@ __global__ void find_tracks(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<unsigned int> n_candidates_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_candidates) {
@@ -115,7 +117,7 @@ __global__ void find_tracks(
         gid, cfg, det_data, measurements_view, in_params_view,
         n_measurements_prefix_sum_view, ref_meas_idx_view, prev_links_view,
         prev_param_to_link_view, step, n_max_candidates, out_params_view,
-        n_candidates_view, links_view, n_candidates);
+        theta_id_view, n_candidates_view, links_view, n_candidates);
 }
 
 /// CUDA kernel for running @c traccc::device::add_links_for_holes
@@ -126,15 +128,16 @@ __global__ void add_links_for_holes(
     vecmem::data::vector_view<const unsigned int> prev_param_to_link_view,
     const unsigned int step, const unsigned int n_max_candidates,
     bound_track_parameters_collection_types::view out_params_view,
+    vecmem::data::vector_view<device::theta_id> theta_id_view,
     vecmem::data::vector_view<candidate_link> links_view,
     unsigned int& n_total_candidates) {
 
     int gid = threadIdx.x + blockIdx.x * blockDim.x;
 
     device::add_links_for_holes(gid, n_candidates_view, in_params_view,
                                 prev_links_view, prev_param_to_link_view, step,
-                                n_max_candidates, out_params_view, links_view,
-                                n_total_candidates);
+                                n_max_candidates, out_params_view,
+                                theta_id_view, links_view, n_total_candidates);
 }
 
 /// CUDA kernel for running @c traccc::device::propagate_to_next_surface
@@ -144,6 +147,7 @@ __global__ void propagate_to_next_surface(
     typename propagator_t::detector_type::view_type det_data,
     bfield_t field_data,
     bound_track_parameters_collection_types::const_view in_params_view,
+    vecmem::data::vector_view<const device::theta_id> theta_id_view,
     vecmem::data::vector_view<const candidate_link> links_view,
     const unsigned int step, const unsigned int& n_candidates,
     bound_track_parameters_collection_types::view out_params_view,
@@ -156,9 +160,9 @@ __global__ void propagate_to_next_surface(
     int gid = threadIdx.x + blockIdx.x * blockDim.x;
 
     device::propagate_to_next_surface<propagator_t, bfield_t, config_t>(
-        gid, cfg, det_data, field_data, in_params_view, links_view, step,
-        n_candidates, out_params_view, param_to_link_view, tips_view,
-        n_tracks_per_seed_view, n_out_params);
+        gid, cfg, det_data, field_data, in_params_view, theta_id_view,
+        links_view, step, n_candidates, out_params_view, param_to_link_view,
+        tips_view, n_tracks_per_seed_view, n_out_params);
 }
 
 /// CUDA kernel for running @c traccc::device::build_tracks
@@ -412,6 +416,8 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
         bound_track_parameters_collection_types::buffer updated_params_buffer(
             n_in_params * m_cfg.max_num_branches_per_surface, m_mr.main);
+        vecmem::data::vector_buffer<device::theta_id> theta_and_id_buffer(
+            n_in_params * m_cfg.max_num_branches_per_surface, m_mr.main);
 
         // Create the link map
         link_map[step] = {n_in_params * m_cfg.max_num_branches_per_surface,
@@ -428,7 +434,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
                     n_measurements_prefix_sum_buffer, ref_meas_idx_buffer,
                     link_map[prev_step], param_to_link_map[prev_step], step,
                     n_max_candidates, updated_params_buffer,
-                    n_candidates_buffer, link_map[step],
+                    theta_and_id_buffer, n_candidates_buffer, link_map[step],
                     (*global_counter_device).n_candidates);
             TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
         }
@@ -443,7 +449,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
             kernels::add_links_for_holes<<<nBlocks, nThreads, 0, stream>>>(
                 n_candidates_buffer, in_params_buffer, link_map[prev_step],
                 param_to_link_map[prev_step], step, n_max_candidates,
-                updated_params_buffer, link_map[step],
+                updated_params_buffer, theta_and_id_buffer, link_map[step],
                 (*global_counter_device).n_candidates);
             TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
         }
@@ -456,6 +462,15 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
         m_stream.synchronize();
 
+        // Sort the theta_and_id vector with the link vector
+        vecmem::device_vector<device::theta_id> theta_and_id_device(
+            theta_and_id_buffer);
+        vecmem::device_vector<candidate_link> link_device(link_map[step]);
+        thrust::sort_by_key(
+            thrust::cuda::par.on(stream), theta_and_id_device.begin(),
+            theta_and_id_device.begin() + global_counter_host.n_candidates,
+            link_device.begin(), device::theta_id_sort_comp());
+
         /*****************************************************************
          * Kernel6: Propagate to the next surface
          *****************************************************************/
@@ -482,8 +497,9 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
                                                config_type>
                 <<<nBlocks, nThreads, 0, stream>>>(
                     m_cfg, det_view, field_view, updated_params_buffer,
-                    link_map[step], step, (*global_counter_device).n_candidates,
-                    out_params_buffer, param_to_link_map[step], tips_map[step],
+                    theta_and_id_buffer, link_map[step], step,
+                    (*global_counter_device).n_candidates, out_params_buffer,
+                    param_to_link_map[step], tips_map[step],
                     n_tracks_per_seed_buffer,
                     (*global_counter_device).n_out_params);
             TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());