Access track parameters in sorted order for CKF propagation

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/sort_key.hpp"
    # Clusterization function(s).
    "include/traccc/clusterization/device/form_spacepoints.hpp"
    "include/traccc/clusterization/device/impl/form_spacepoints.ipp"
@@ -64,12 +65,14 @@ traccc_add_library( traccc_device_common device_common TYPE SHARED
    "include/traccc/finding/device/apply_interaction.hpp"
    "include/traccc/finding/device/build_tracks.hpp"
    "include/traccc/finding/device/find_tracks.hpp"
+   "include/traccc/finding/device/fill_sort_keys.hpp"
    "include/traccc/finding/device/make_barcode_sequence.hpp"
    "include/traccc/finding/device/propagate_to_next_surface.hpp"
    "include/traccc/finding/device/prune_tracks.hpp"
    "include/traccc/finding/device/impl/apply_interaction.ipp"
    "include/traccc/finding/device/impl/build_tracks.ipp"
    "include/traccc/finding/device/impl/find_tracks.ipp"
+   "include/traccc/finding/device/impl/fill_sort_keys.ipp"
    "include/traccc/finding/device/impl/make_barcode_sequence.ipp"
    "include/traccc/finding/device/impl/propagate_to_next_surface.ipp"
    "include/traccc/finding/device/impl/prune_tracks.ipp"

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

@@ -0,0 +1,24 @@
+/** 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"
+#include "traccc/edm/track_parameters.hpp"
+
+namespace traccc::device {
+
+using sort_key = traccc::scalar;
+
+TRACCC_HOST_DEVICE
+inline sort_key get_sort_key(const bound_track_parameters& params) {
+    // key = |theta - pi/2|
+    return math::abs(params.theta() - constant<traccc::scalar>::pi_2);
+}
+
+}  // namespace traccc::device

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

@@ -0,0 +1,32 @@
+/** 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/edm/device/sort_key.hpp"
+#include "traccc/edm/track_candidate.hpp"
+
+namespace traccc::device {
+
+/// Function used for fill key container
+///
+/// @param[in] globalIndex   The index of the current thread
+/// @param[in] params_view   The input parameters
+/// @param[out] keys_view    The key values
+/// @param[out] ids_view     The param ids
+///
+TRACCC_HOST_DEVICE inline void fill_sort_keys(
+    std::size_t globalIndex,
+    bound_track_parameters_collection_types::const_view params_view,
+    vecmem::data::vector_view<device::sort_key> keys_view,
+    vecmem::data::vector_view<unsigned int> ids_view);
+
+}  // namespace traccc::device
+
+// Include the implementation.
+#include "traccc/finding/device/impl/fill_sort_keys.ipp"

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

@@ -0,0 +1,34 @@
+/** 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
+
+namespace traccc::device {
+
+TRACCC_HOST_DEVICE inline void fill_sort_keys(
+    std::size_t globalIndex,
+    bound_track_parameters_collection_types::const_view params_view,
+    vecmem::data::vector_view<device::sort_key> keys_view,
+    vecmem::data::vector_view<unsigned int> ids_view) {
+
+    bound_track_parameters_collection_types::const_device params(params_view);
+
+    // Keys
+    vecmem::device_vector<device::sort_key> keys_device(keys_view);
+
+    // Param id
+    vecmem::device_vector<unsigned int> ids_device(ids_view);
+
+    if (globalIndex >= keys_device.size()) {
+        return;
+    }
+
+    keys_device.at(globalIndex) = device::get_sort_key(params.at(globalIndex));
+    ids_device.at(globalIndex) = globalIndex;
+}
+
+}  // namespace traccc::device

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,
+    const vecmem::data::vector_view<const unsigned int>& param_ids_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,11 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
         return;
     }
 
+    // Theta id
+    vecmem::device_vector<const unsigned int> param_ids(param_ids_view);
+
+    const unsigned int param_id = param_ids.at(globalIndex);
+
     // Number of tracks per seed
     vecmem::device_vector<unsigned int> n_tracks_per_seed(
         n_tracks_per_seed_view);
@@ -36,7 +42,7 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     vecmem::device_vector<const candidate_link> links(links_view);
 
     // Seed id
-    unsigned int orig_param_id = links.at(globalIndex).seed_idx;
+    unsigned int orig_param_id = links.at(param_id).seed_idx;
 
     // Count the number of tracks per seed
     vecmem::device_atomic_ref<unsigned int> num_tracks_per_seed(
@@ -52,8 +58,8 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     vecmem::device_vector<typename candidate_link::link_index_type> tips(
         tips_view);
 
-    if (links[globalIndex].n_skipped > cfg.max_num_skipping_per_cand) {
-        tips.push_back({step, globalIndex});
+    if (links.at(param_id).n_skipped > cfg.max_num_skipping_per_cand) {
+        tips.push_back({step, param_id});
         return;
     }
 
@@ -71,7 +77,7 @@ 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 bound_track_parameters in_par = in_params.at(param_id);
 
     // Create propagator
     propagator_t propagator(cfg.propagation);
@@ -115,17 +121,17 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
 
         out_params[out_param_id] = propagation._stepping._bound_params;
 
-        param_to_link[out_param_id] = static_cast<unsigned int>(globalIndex);
+        param_to_link[out_param_id] = param_id;
     }
     // Unless the track found a surface, it is considered a tip
     else if (!s4.success && step >= cfg.min_track_candidates_per_track - 1) {
-        tips.push_back({step, globalIndex});
+        tips.push_back({step, param_id});
     }
 
     // If no more CKF step is expected, current candidate is
     // kept as a tip
     if (s4.success && step == cfg.max_track_candidates_per_track - 1) {
-        tips.push_back({step, globalIndex});
+        tips.push_back({step, param_id});
     }
 }
 

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] param_ids_view     Sorted param ids
 /// @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,
+    const vecmem::data::vector_view<const unsigned int>& param_ids_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

@@ -15,9 +15,11 @@
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/device/finding_global_counter.hpp"
+#include "traccc/edm/device/sort_key.hpp"
 #include "traccc/finding/candidate_link.hpp"
 #include "traccc/finding/device/apply_interaction.hpp"
 #include "traccc/finding/device/build_tracks.hpp"
+#include "traccc/finding/device/fill_sort_keys.hpp"
 #include "traccc/finding/device/find_tracks.hpp"
 #include "traccc/finding/device/make_barcode_sequence.hpp"
 #include "traccc/finding/device/propagate_to_next_surface.hpp"
@@ -108,13 +110,24 @@ __global__ void find_tracks(
         shared_candidates_size);
 }
 
+/// CUDA kernel for running @c traccc::device::fill_sort_keys
+__global__ void fill_sort_keys(
+    bound_track_parameters_collection_types::const_view params_view,
+    vecmem::data::vector_view<device::sort_key> keys_view,
+    vecmem::data::vector_view<unsigned int> ids_view) {
+
+    device::fill_sort_keys(threadIdx.x + blockIdx.x * blockDim.x, params_view,
+                           keys_view, ids_view);
+}
+
 /// CUDA kernel for running @c traccc::device::propagate_to_next_surface
 template <typename propagator_t, typename bfield_t, typename config_t>
 __global__ void propagate_to_next_surface(
     const config_t cfg,
     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 unsigned int> param_ids_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,
@@ -127,9 +140,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, param_ids_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
@@ -365,7 +378,34 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
         m_stream.synchronize();
 
         /*****************************************************************
-         * Kernel4: Propagate to the next surface
+         * Kernel4: Get key and value for parameter sorting
+         *****************************************************************/
+
+        vecmem::data::vector_buffer<device::sort_key> keys_buffer(
+            global_counter_host.n_candidates, m_mr.main);
+        vecmem::data::vector_buffer<unsigned int> param_ids_buffer(
+            global_counter_host.n_candidates, m_mr.main);
+
+        nThreads = m_warp_size * 2;
+
+        if (global_counter_host.n_candidates > 0) {
+            nBlocks =
+                (global_counter_host.n_candidates + nThreads - 1) / nThreads;
+            kernels::fill_sort_keys<<<nBlocks, nThreads, 0, stream>>>(
+                updated_params_buffer, keys_buffer, param_ids_buffer);
+            TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+
+            // Sort the key and values
+            vecmem::device_vector<device::sort_key> keys_device(keys_buffer);
+            vecmem::device_vector<unsigned int> param_ids_device(
+                param_ids_buffer);
+            thrust::sort_by_key(thrust::cuda::par.on(stream),
+                                keys_device.begin(), keys_device.end(),
+                                param_ids_device.begin());
+        }
+
+        /*****************************************************************
+         * Kernel5: Propagate to the next surface
          *****************************************************************/
 
         // Buffer for out parameters for the next step
@@ -390,8 +430,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],
+                    param_ids_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());
@@ -477,7 +518,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
     }
 
     /*****************************************************************
-     * Kernel5: Build tracks
+     * Kernel6: Build tracks
      *****************************************************************/
 
     // Create track candidate buffer