Skip forward fitting pass by keeping CKF states

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

@@ -0,0 +1,75 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2023-2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Local include(s).
+#include "traccc/device/global_index.hpp"
+
+// Project include(s).
+#include "traccc/definitions/qualifiers.hpp"
+#include "traccc/edm/track_candidate_container.hpp"
+#include "traccc/edm/track_parameters.hpp"
+#include "traccc/edm/track_state.hpp"
+#include "traccc/finding/candidate_link.hpp"
+#include "traccc/finding/finding_config.hpp"
+
+// VecMem include(s).
+#include <vecmem/containers/data/jagged_vector_view.hpp>
+#include <vecmem/containers/data/vector_view.hpp>
+
+namespace traccc::device {
+
+/// (Event Data) Payload for the @c traccc::device::build_tracks function
+struct build_fitted_tracks_payload {
+    /**
+     * @brief View objects to the vector of measurements
+     */
+    const measurement_collection_types::const_view measurements_view;
+
+    /**
+     * @brief View object to the vector of seeds
+     */
+    bound_track_parameters_collection_types::const_view seeds_view;
+
+    /**
+     * @brief View object to the track parameters
+     */
+    bound_track_parameters_collection_types::const_view track_param_view;
+
+    /**
+     * @brief View object to the vector of candidate links
+     */
+    vecmem::data::vector_view<const candidate_link> links_view;
+
+    /**
+     * @brief View object to the vector of tips
+     */
+    vecmem::data::vector_view<const unsigned int> tips_view;
+
+    /**
+     * @brief View object to the vector of track candidates
+     */
+    track_state_container_types::view track_states_view;
+};
+
+/// Function for building full tracks from the link container:
+/// The full tracks are built using the link container and tip link container.
+/// Since every link holds an information of the link from the previous step,
+/// we can build a full track by iterating from a tip link backwardly.
+///
+/// @param[in] globalIndex         The index of the current thread
+/// @param[in] cfg                    Track finding config object
+/// @param[inout] payload      The function call payload
+///
+TRACCC_HOST_DEVICE inline void build_fitted_tracks(
+    global_index_t globalIndex, const build_fitted_tracks_payload& payload);
+
+}  // namespace traccc::device
+
+// Include the implementation.
+#include "./impl/build_fitted_tracks.ipp"

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

@@ -23,8 +23,9 @@
 
 namespace traccc::device {
 
-/// (Event Data) Payload for the @c traccc::device::build_tracks function
-struct build_tracks_payload {
+/// (Event Data) Payload for the @c traccc::device::build_unfitted_tracks
+/// function
+struct build_unfitted_tracks_payload {
     /**
      * @brief View object to the vector of measurements
      */
@@ -55,10 +56,10 @@ struct build_tracks_payload {
 /// @param[in] cfg                    Track finding config object
 /// @param[inout] payload      The function call payload
 ///
-TRACCC_HOST_DEVICE inline void build_tracks(
-    global_index_t globalIndex, const build_tracks_payload& payload);
+TRACCC_HOST_DEVICE inline void build_unfitted_tracks(
+    global_index_t globalIndex, const build_unfitted_tracks_payload& payload);
 
 }  // namespace traccc::device
 
 // Include the implementation.
-#include "./impl/build_tracks.ipp"
+#include "./impl/build_unfitted_tracks.ipp"

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

@@ -125,6 +125,18 @@ struct find_tracks_payload {
      * @brief View object to the temporary link vector
      */
     vecmem::data::vector_view<candidate_link> tmp_links_view;
+
+    /**
+     * @brief View object to the persistent track parameters, estabilishing
+     * a direct map between links and parameters at the same index.
+     */
+    bound_track_parameters_collection_types::view persistent_parameters_view;
+
+    /**
+     * @brief Flag that, if true, indicates whether holes should be counted
+     * in the length of the tip in the tip length output.
+     */
+    bool count_holes = false;
 };
 
 /// (Shared Event Data) Payload for the @c traccc::device::find_tracks function

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

@@ -0,0 +1,87 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2023-2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/utils/prob.hpp"
+
+namespace traccc::device {
+
+TRACCC_HOST_DEVICE inline void build_fitted_tracks(
+    const global_index_t globalIndex,
+    const build_fitted_tracks_payload& payload) {
+
+    const measurement_collection_types::const_device measurements(
+        payload.measurements_view);
+
+    const bound_track_parameters_collection_types::const_device track_params(
+        payload.track_param_view);
+    const bound_track_parameters_collection_types::const_device seeds(
+        payload.seeds_view);
+
+    const vecmem::device_vector<const candidate_link> links(payload.links_view);
+
+    const vecmem::device_vector<const unsigned int> tips(payload.tips_view);
+
+    track_state_container_types::device track_states(payload.track_states_view);
+
+    if (globalIndex >= tips.size()) {
+        return;
+    }
+
+    const auto tip = tips.at(globalIndex);
+
+    auto track = track_states.at(globalIndex).items;
+    auto header = track_states.at(globalIndex).header;
+
+    // Get the link corresponding to tip
+    unsigned int link_idx = tip;
+    auto L = links.at(link_idx);
+    const unsigned int n_meas = measurements.size();
+
+    // Track summary variables
+    scalar ndf_sum = 0.f;
+    scalar chi2_sum = 0.f;
+
+    // Reversely iterate to fill the track candidates
+    for (auto it = track.rbegin(); it != track.rend(); it++) {
+        if (L.meas_idx >= n_meas) {
+            it->is_hole = true;
+        } else {
+            *it = track_state<default_algebra>(measurements.at(L.meas_idx));
+            it->is_hole = false;
+            it->filtered_chi2() = L.chi2;
+            it->filtered() = track_params.at(link_idx);
+
+            // Sanity check on chi2
+            assert(L.chi2 < std::numeric_limits<traccc::scalar>::max());
+            assert(L.chi2 >= 0.f);
+
+            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 != track.rend() - 1) {
+            link_idx = L.previous_candidate_idx;
+            L = links.at(link_idx);
+        }
+    }
+
+    header.fit_outcome = fitter_outcome::SUCCESS;
+    header.fit_params = seeds.at(L.seed_idx);
+    header.trk_quality.ndf = ndf_sum - 5.f;
+    header.trk_quality.chi2 = chi2_sum;
+    header.trk_quality.pval =
+        prob(header.trk_quality.chi2, header.trk_quality.ndf);
+    header.trk_quality.n_holes = L.n_skipped;
+}
+
+}  // namespace traccc::device

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

@@ -12,8 +12,9 @@
 
 namespace traccc::device {
 
-TRACCC_HOST_DEVICE inline void build_tracks(
-    const global_index_t globalIndex, const build_tracks_payload& payload) {
+TRACCC_HOST_DEVICE inline void build_unfitted_tracks(
+    const global_index_t globalIndex,
+    const build_unfitted_tracks_payload& payload) {
 
     const measurement_collection_types::const_device measurements(
         payload.track_candidates_view.measurements);

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

@@ -88,6 +88,8 @@ TRACCC_HOST_DEVICE inline void find_tracks(
     vecmem::device_vector<candidate_link> tmp_links(payload.tmp_links_view);
     bound_track_parameters_collection_types::device tmp_params(
         payload.tmp_params_view);
+    bound_track_parameters_collection_types::device persistent_params(
+        payload.persistent_parameters_view);
     vecmem::device_vector<const detray::geometry::barcode> barcodes(
         payload.barcodes_view);
     vecmem::device_vector<const unsigned int> upper_bounds(
@@ -630,6 +632,10 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                     .chi2 = std::numeric_limits<traccc::scalar>::max(),
                     .chi2_sum = prev_chi2_sum,
                     .ndf_sum = prev_ndf_sum};
+                if (persistent_params.capacity() > 0) {
+                    persistent_params.at(out_offset) =
+                        in_params.at(in_param_id);
+                }
 
                 unsigned int param_pos = out_offset - payload.curr_links_idx;
 
@@ -641,7 +647,12 @@ TRACCC_HOST_DEVICE inline void find_tracks(
 
                 if (n_cands >= cfg.min_track_candidates_per_track) {
                     auto tip_pos = tips.push_back(prev_link_idx);
-                    tip_lengths.at(tip_pos) = n_cands;
+
+                    if (payload.count_holes) {
+                        tip_lengths.at(tip_pos) = payload.step + 1u;
+                    } else {
+                        tip_lengths.at(tip_pos) = n_cands;
+                    }
                 }
             }
         } else {
@@ -659,13 +670,21 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                 out_params_liveness.at(param_pos) =
                     static_cast<unsigned int>(!last_step);
                 links.at(out_offset) = tmp_links.at(in_offset);
+                if (persistent_params.capacity() > 0) {
+                    persistent_params.at(out_offset) = tmp_params.at(in_offset);
+                }
 
                 const unsigned int n_cands = payload.step + 1 - n_skipped;
 
                 if (last_step &&
                     n_cands >= cfg.min_track_candidates_per_track) {
                     auto tip_pos = tips.push_back(param_pos);
-                    tip_lengths.at(tip_pos) = n_cands;
+
+                    if (payload.count_holes) {
+                        tip_lengths.at(tip_pos) = payload.step + 1u;
+                    } else {
+                        tip_lengths.at(tip_pos) = n_cands;
+                    }
                 }
             }
         }

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

@@ -109,7 +109,12 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(
 
         if (n_cands >= cfg.min_track_candidates_per_track) {
             auto tip_pos = tips.push_back(link_idx);
-            tip_lengths.at(tip_pos) = n_cands;
+
+            if (payload.count_holes) {
+                tip_lengths.at(tip_pos) = link.step + 1;
+            } else {
+                tip_lengths.at(tip_pos) = n_cands;
+            }
         }
     }
 }

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

@@ -79,6 +79,12 @@ struct propagate_to_next_surface_payload {
      * @brief Vector to hold the number of track states per tip
      */
     vecmem::data::vector_view<unsigned int> tip_lengths_view;
+
+    /**
+     * @brief Flag that, if true, indicates whether holes should be counted
+     * in the length of the tip in the tip length output.
+     */
+    bool count_holes = false;
 };
 
 /// Function for propagating the kalman-updated tracks to the next surface

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

@@ -0,0 +1,28 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+namespace traccc::device {
+/**
+ * @defgroup Track finding return type specifiers
+ *
+ * Optional parameters to track finding algorithms which instruct the
+ * algorithm to return either unfitted tracks or fitted tracks directly.
+ *
+ * @{
+ * @brief Return tracks with fitted track states.
+ */
+struct finding_return_fitted {};
+/*
+ * @brief Return tracks with unfitted track states.
+ */
+struct finding_return_unfitted {};
+/*
+ * @}
+ */
+}  // namespace traccc::device

device/common/include/traccc/fitting/device/fill_fitting_state_sort_keys.hpp

@@ -0,0 +1,50 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024-2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Local include(s).
+#include "traccc/device/global_index.hpp"
+#include "traccc/edm/device/sort_key.hpp"
+
+// Project include(s).
+#include "traccc/edm/track_state.hpp"
+
+namespace traccc::device {
+
+/// Function used to fill key container
+///
+/// @param[in] globalIndex   The index of the current thread
+/// @param[in] track_candidates_view The input track states
+/// @param[out] keys_view    The key values
+/// @param[out] ids_view     The param ids
+///
+TRACCC_HOST_DEVICE inline void fill_fitting_state_sort_keys(
+    global_index_t globalIndex,
+    track_state_container_types::const_view& track_states_view,
+    vecmem::data::vector_view<device::sort_key> keys_view,
+    vecmem::data::vector_view<unsigned int> ids_view) {
+    const track_state_container_types::const_device track_states(
+        track_states_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;
+    }
+
+    // Key = The number of measurements
+    keys_device.at(globalIndex) =
+        static_cast<traccc::scalar>(track_states.at(globalIndex).items.size());
+    ids_device.at(globalIndex) = globalIndex;
+}
+
+}  // namespace traccc::device