Improve memory usage in track finding postamble

device/common/CMakeLists.txt

@@ -64,14 +64,12 @@ traccc_add_library( traccc_device_common device_common TYPE INTERFACE
    "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"
    # Track fitting funtions(s).
    "include/traccc/fitting/device/fit.hpp"
    "include/traccc/fitting/device/impl/fit.ipp"

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

@@ -49,20 +49,9 @@ struct build_tracks_payload {
     vecmem::data::vector_view<const unsigned int> tips_view;
 
     /**
-     * @brief View object to the vector of track candidates
+     * @brief View object to the vector of pruned track candidates
      */
-    track_candidate_container_types::view track_candidates_view;
-
-    /**
-     * @brief View object to the vector of indices meeting the selection
-     * criteria
-     */
-    vecmem::data::vector_view<unsigned int> valid_indices_view;
-
-    /**
-     * @brief The number of valid tracks meeting criteria
-     */
-    unsigned int* n_valid_tracks;
+    track_candidate_container_types::view final_candidates_view;
 };
 
 /// Function for building full tracks from the link container:

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

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <vecmem/containers/device_vector.hpp>
 namespace traccc::device {
 
 TRACCC_DEVICE inline void build_tracks(const global_index_t globalIndex,
@@ -23,106 +24,70 @@ TRACCC_DEVICE inline void build_tracks(const global_index_t globalIndex,
 
     const vecmem::device_vector<const unsigned int> tips(payload.tips_view);
 
-    track_candidate_container_types::device track_candidates(
-        payload.track_candidates_view);
-
-    vecmem::device_vector<unsigned int> valid_indices(
-        payload.valid_indices_view);
+    track_candidate_container_types::device final_candidates(
+        payload.final_candidates_view);
 
     if (globalIndex >= tips.size()) {
         return;
     }
 
     const auto tip = tips.at(globalIndex);
-    auto& seed = track_candidates[globalIndex].header.seed_params;
-    auto& trk_quality = track_candidates[globalIndex].header.trk_quality;
-    auto cands_per_track = track_candidates[globalIndex].items;
+
+    const auto& L_origin = links.at(tip);
 
     // Get the link corresponding to tip
-    auto L = links.at(tip);
+    auto L = L_origin;
+    const unsigned int num_meas = final_candidates.at(globalIndex).items.size();
     const unsigned int n_meas = measurements.size();
 
-    const unsigned int n_cands = L.step + 1 - L.n_skipped;
-
-    // Resize the candidates with the exact size
-    cands_per_track.resize(n_cands);
-
-    bool success = true;
+    assert(num_meas >= cfg.min_track_candidates_per_track &&
+           num_meas <= cfg.max_track_candidates_per_track);
 
     // Track summary variables
     scalar ndf_sum = 0.f;
     scalar chi2_sum = 0.f;
 
-    [[maybe_unused]] std::size_t num_inserted = 0;
-
     // Reversely iterate to fill the track candidates
-    for (auto it = cands_per_track.rbegin(); it != cands_per_track.rend();
-         it++) {
-
+    for (unsigned int i = num_meas - 1; i < num_meas; --i) {
         while (L.meas_idx >= n_meas && L.step != 0u) {
 
             L = links.at(L.previous_candidate_idx);
         }
 
-        // Break if the measurement is still invalid
-        if (L.meas_idx >= measurements.size()) {
-            success = false;
-            break;
-        }
+        assert(L.meas_idx < measurements.size());
 
-        *it = {measurements.at(L.meas_idx)};
-        num_inserted++;
+        final_candidates.at(globalIndex).items.at(i) =
+            measurements.at(L.meas_idx);
 
         // Sanity check on chi2
         assert(L.chi2 < std::numeric_limits<traccc::scalar>::max());
         assert(L.chi2 >= 0.f);
 
-        ndf_sum += static_cast<scalar>(it->meas_dim);
+        ndf_sum += static_cast<scalar>(measurements.at(L.meas_idx).meas_dim);
         chi2_sum += L.chi2;
 
-        // Break the loop if the iterator is at the first candidate and fill the
-        // seed and track quality
-        if (it == cands_per_track.rend() - 1) {
-            seed = seeds.at(L.seed_idx);
-            trk_quality.ndf = ndf_sum - 5.f;
-            trk_quality.chi2 = chi2_sum;
-            trk_quality.n_holes = L.n_skipped;
-        } else {
+        if (i != 0) {
             L = links.at(L.previous_candidate_idx);
         }
     }
 
+    final_candidates.at(globalIndex).header.seed_params =
+        seeds.at(L_origin.seed_idx);
+    final_candidates.at(globalIndex).header.trk_quality = {
+        .ndf = ndf_sum - 5.f, .chi2 = chi2_sum, .n_holes = L_origin.n_skipped};
+
 #ifndef NDEBUG
-    if (success) {
-        // Assert that we inserted exactly as many elements as we reserved
-        // space for.
-        assert(num_inserted == cands_per_track.size());
-
-        // Assert that we did not make any duplicate track states.
-        for (unsigned int i = 0; i < cands_per_track.size(); ++i) {
-            for (unsigned int j = 0; j < cands_per_track.size(); ++j) {
-                if (i != j) {
-                    // TODO: Re-enable me!
-                    // assert(cands_per_track.at(i).measurement_id !=
-                    //       cands_per_track.at(j).measurement_id);
-                }
+    // Assert that we did not make any duplicate track states.
+    for (unsigned int i = 0; i < num_meas; ++i) {
+        for (unsigned int j = 0; j < num_meas; ++j) {
+            if (i != j) {
+                // TODO: Re-enable me!
+                // assert(cands_per_track.at(i).measurement_id !=
+                //       cands_per_track.at(j).measurement_id);
             }
         }
     }
 #endif
-
-    // NOTE: We may at some point want to assert that `success` is true
-
-    // Criteria for valid tracks
-    if (n_cands >= cfg.min_track_candidates_per_track &&
-        n_cands <= cfg.max_track_candidates_per_track && success) {
-
-        vecmem::device_atomic_ref<unsigned int> num_valid_tracks(
-            *payload.n_valid_tracks);
-
-        const unsigned int pos = num_valid_tracks.fetch_add(1);
-        valid_indices[pos] = globalIndex;
-    }
 }
 
 }  // namespace traccc::device

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

@@ -48,79 +48,91 @@ TRACCC_DEVICE inline void propagate_to_next_surface(
     const unsigned int s_pos = num_tracks_per_seed.fetch_add(1);
     vecmem::device_vector<unsigned int> params_liveness(
         payload.params_liveness_view);
+    vecmem::device_vector<unsigned int> tips(payload.tips_view);
+    vecmem::device_vector<unsigned int> tip_lengths(payload.tip_lengths_view);
+
+    bool create_tip = false;
 
     if (s_pos >= cfg.max_num_branches_per_seed) {
-        params_liveness[param_id] = 0u;
-        return;
+        params_liveness.at(param_id) = 0u;
+    } else if (links.at(payload.prev_links_idx + param_id).n_skipped >
+               cfg.max_num_skipping_per_cand) {
+        params_liveness.at(param_id) = 0u;
+        create_tip = true;
     }
 
-    // tips
-    vecmem::device_vector<unsigned int> tips(payload.tips_view);
+    if (params_liveness.at(param_id) != 0u) {
+        // Detector
+        typename propagator_t::detector_type det(payload.det_data);
+
+        // Parameters
+        bound_track_parameters_collection_types::device params(
+            payload.params_view);
+
+        // Input bound track parameter
+        const bound_track_parameters<> in_par = params.at(param_id);
+
+        // Create propagator
+        propagator_t propagator(cfg.propagation);
+
+        // Create propagator state
+        typename propagator_t::state propagation(in_par, payload.field_data,
+                                                 det);
+        propagation.set_particle(
+            detail::correct_particle_hypothesis(cfg.ptc_hypothesis, in_par));
+        propagation._stepping
+            .template set_constraint<detray::step::constraint::e_accuracy>(
+                cfg.propagation.stepping.step_constraint);
+
+        // Actor state
+        // @TODO: simplify the syntax here
+        // @NOTE: Post material interaction might be required here
+        using actor_tuple_type =
+            typename propagator_t::actor_chain_type::actor_tuple;
+        typename detray::detail::tuple_element<0, actor_tuple_type>::type::state
+            s0{};
+        typename detray::detail::tuple_element<3, actor_tuple_type>::type::state
+            s3{};
+        typename detray::detail::tuple_element<2, actor_tuple_type>::type::state
+            s2{s3};
+        typename detray::detail::tuple_element<4, actor_tuple_type>::type::state
+            s4;
+        s4.min_step_length = cfg.min_step_length_for_next_surface;
+        s4.max_count = cfg.max_step_counts_for_next_surface;
+
+        // Propagate to the next surface
+        propagator.propagate_sync(propagation, detray::tie(s0, s2, s3, s4));
+
+        // If a surface found, add the parameter for the next step
+        if (s4.success) {
+            params[param_id] = propagation._stepping.bound_params();
+
+            if (payload.step == cfg.max_track_candidates_per_track - 1) {
+                create_tip = true;
+                params_liveness[param_id] = 0u;
+            } else {
+                params_liveness[param_id] = 1u;
+            }
+        } else {
+            params_liveness[param_id] = 0u;
 
-    if (links.at(payload.prev_links_idx + param_id).n_skipped >
-        cfg.max_num_skipping_per_cand) {
-        params_liveness[param_id] = 0u;
-        tips.push_back(payload.prev_links_idx + param_id);
-        return;
+            if (payload.step >= cfg.min_track_candidates_per_track - 1) {
+                create_tip = true;
+            }
+        }
     }
 
-    // Detector
-    typename propagator_t::detector_type det(payload.det_data);
+    if (create_tip) {
+        const auto& L = links.at(payload.prev_links_idx + param_id);
 
-    // Parameters
-    bound_track_parameters_collection_types::device params(payload.params_view);
-
-    if (params_liveness.at(param_id) == 0u) {
-        return;
-    }
-
-    // Input bound track parameter
-    const bound_track_parameters<> in_par = params.at(param_id);
-
-    // Create propagator
-    propagator_t propagator(cfg.propagation);
-
-    // Create propagator state
-    typename propagator_t::state propagation(in_par, payload.field_data, det);
-    propagation.set_particle(
-        detail::correct_particle_hypothesis(cfg.ptc_hypothesis, in_par));
-    propagation._stepping
-        .template set_constraint<detray::step::constraint::e_accuracy>(
-            cfg.propagation.stepping.step_constraint);
-
-    // Actor state
-    // @TODO: simplify the syntax here
-    // @NOTE: Post material interaction might be required here
-    using actor_tuple_type =
-        typename propagator_t::actor_chain_type::actor_tuple;
-    typename detray::detail::tuple_element<0, actor_tuple_type>::type::state
-        s0{};
-    typename detray::detail::tuple_element<3, actor_tuple_type>::type::state
-        s3{};
-    typename detray::detail::tuple_element<2, actor_tuple_type>::type::state s2{
-        s3};
-    typename detray::detail::tuple_element<4, actor_tuple_type>::type::state s4;
-    s4.min_step_length = cfg.min_step_length_for_next_surface;
-    s4.max_count = cfg.max_step_counts_for_next_surface;
-
-    // Propagate to the next surface
-    propagator.propagate_sync(propagation, detray::tie(s0, s2, s3, s4));
-
-    // If a surface found, add the parameter for the next step
-    if (s4.success) {
-        params[param_id] = propagation._stepping.bound_params();
-
-        if (payload.step == cfg.max_track_candidates_per_track - 1) {
-            tips.push_back(payload.prev_links_idx + param_id);
-            params_liveness[param_id] = 0u;
-        } else {
-            params_liveness[param_id] = 1u;
-        }
-    } else {
-        params_liveness[param_id] = 0u;
+        const unsigned int num_meas = payload.step + 1 - L.n_skipped;
 
-        if (payload.step >= cfg.min_track_candidates_per_track - 1) {
-            tips.push_back(payload.prev_links_idx + param_id);
+        // Criteria for valid tracks
+        if (num_meas >= cfg.min_track_candidates_per_track &&
+            num_meas <= cfg.max_track_candidates_per_track) {
+            const unsigned int tip_pos =
+                tips.push_back(payload.prev_links_idx + param_id);
+            tip_lengths.at(tip_pos) = num_meas;
         }
     }
 }

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

@@ -75,6 +75,8 @@ struct propagate_to_next_surface_payload {
      */
     vecmem::data::vector_view<unsigned int> tips_view;
 
+    vecmem::data::vector_view<unsigned int> tip_lengths_view;
+
     /**
      * @brief View object to the vector of the number of tracks per initial
      * input seed