Allow fitting algorithms to use prefitted tracks

This commit enables our fitting algorithms to use prefitted tracks, such
that we can reuse the fitting information from the CKF. Note that we
also preserve the ability to work with unfitted tracks.

Author: stephenswat

device/alpaka/include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp

@@ -34,6 +34,10 @@ class kalman_fitting_algorithm
     : public algorithm<edm::track_fit_container<default_algebra>::buffer(
           const detector_buffer&, const magnetic_field&,
           const edm::track_candidate_container<default_algebra>::const_view&)>,
+      public algorithm<edm::track_fit_container<default_algebra>::buffer(
+          const detector_buffer&, const magnetic_field&,
+          edm::track_fit_container<default_algebra>::buffer&&,
+          const measurement_collection_types::const_view&)>,
       public messaging {
 
     public:
@@ -55,7 +59,7 @@ class kalman_fitting_algorithm
         vecmem::copy& copy, queue& q,
         std::unique_ptr<const Logger> logger = getDummyLogger().clone());
 
-    /// Execute the algorithm
+    /// Execute the algorithm using unfitted tracks
     ///
     /// @param det             The detector object
     /// @param bfield          The magnetic field object
@@ -68,6 +72,20 @@ class kalman_fitting_algorithm
         const edm::track_candidate_container<default_algebra>::const_view&
             track_candidates) const override;
 
+    /// Execute the algorithm using fitted tracks
+    ///
+    /// @param det             The detector object
+    /// @param bfield          The magnetic field object
+    /// @param track_candidates All track candidates to fit
+    ///
+    /// @return A container of the fitted track states
+    ///
+    output_type operator()(
+        const detector_buffer& det, const magnetic_field& bfield,
+        edm::track_fit_container<default_algebra>::buffer&& track_states,
+        const measurement_collection_types::const_view& measurements)
+        const override;
+
     private:
     /// Algorithm configuration
     config_type m_config;

device/alpaka/src/fitting/kalman_fitting.hpp

@@ -35,14 +35,13 @@ struct fill_fitting_sort_keys {
     template <typename TAcc>
     ALPAKA_FN_ACC void operator()(
         TAcc const& acc,
-        edm::track_candidate_collection<default_algebra>::const_view
-            track_candidates_view,
+        edm::track_fit_collection<default_algebra>::const_view track_fit_view,
         vecmem::data::vector_view<device::sort_key> keys_view,
         vecmem::data::vector_view<unsigned int> ids_view) const {
 
         const device::global_index_t globalThreadIdx =
             ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0];
-        device::fill_fitting_sort_keys(globalThreadIdx, track_candidates_view,
+        device::fill_fitting_sort_keys(globalThreadIdx, track_fit_view,
                                        keys_view, ids_view);
     }
 };
@@ -52,17 +51,15 @@ struct fit_prelude {
     template <typename TAcc>
     ALPAKA_FN_ACC void operator()(
         TAcc const& acc,
-        vecmem::data::vector_view<const unsigned int> param_ids_view,
         edm::track_candidate_container<default_algebra>::const_view
             track_candidates_view,
-        edm::track_fit_container<default_algebra>::view track_states_view,
-        vecmem::data::vector_view<unsigned int> param_liveness_view) const {
+        edm::track_fit_container<default_algebra>::view track_states_view)
+        const {
 
         const device::global_index_t globalThreadIdx =
             ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0];
         device::fit_prelude<default_algebra>(
-            globalThreadIdx, param_ids_view, track_candidates_view,
-            track_states_view, param_liveness_view);
+            globalThreadIdx, track_candidates_view, track_states_view);
     }
 };
 
@@ -96,14 +93,15 @@ struct fit_backward {
 
 }  // namespace kernels
 
-/// Templated implementation of the Alpaka track fitting algorithm.
+/// Templated implementation of the Alpaka track fitting algorithm for
+/// fitted tracks.
 ///
 /// @tparam detector_t The (device) detector type to use
 /// @tparam bfield_t   The magnetic field type to use
 ///
 /// @param[in] det_view     A view of the detector geometry
 /// @param[in] field_view   A view of the magnetic field
-/// @param[in] track_candidates_view All track candidates to fit
+/// @param[in] track_fit_view All track candidates to fit
 /// @param[in] config       The fitting configuration
 /// @param[in] mr           Memory resource(s) to use
 /// @param[in] copy         The copy object to use for memory transfers
@@ -116,41 +114,32 @@ typename edm::track_fit_container<typename detector_t::algebra_type>::buffer
 kalman_fitting(
     const typename detector_t::const_view_type& det_view,
     const bfield_t& field_view,
-    const typename edm::track_candidate_container<
-        typename detector_t::algebra_type>::const_view& track_candidates_view,
+    typename edm::track_fit_container<
+        typename detector_t::algebra_type>::buffer&& track_fit_buffer,
+    const measurement_collection_types::const_view& measurements,
     const fitting_config& config, const memory_resource& mr, vecmem::copy& copy,
-    Queue& queue) {
+    Queue& queue, bool forward_on_first_iteration = false) {
 
     // Number of threads per block to use.
     const Idx threadsPerBlock = getWarpSize<Acc>() * 2;
 
+    typename edm::track_fit_container<
+        typename detector_t::algebra_type>::const_view track_fit_view{
+        vecmem::get_data(track_fit_buffer.tracks),
+        vecmem::get_data(track_fit_buffer.states), measurements};
+
     // Get the number of tracks.
     const edm::track_candidate_collection<
         default_algebra>::const_device::size_type n_tracks =
-        copy.get_size(track_candidates_view.tracks);
+        copy.get_size(track_fit_view.tracks);
 
     // Get the sizes of the track candidates in each track.
     const std::vector<unsigned int> candidate_sizes =
-        copy.get_sizes(track_candidates_view.tracks);
-    const unsigned int n_states =
-        std::accumulate(candidate_sizes.begin(), candidate_sizes.end(), 0u);
-
-    // Create the result buffer.
-    typename edm::track_fit_container<typename detector_t::algebra_type>::buffer
-        track_states_buffer{
-            {candidate_sizes, mr.main, mr.host,
-             vecmem::data::buffer_type::resizable},
-            {n_states, mr.main, vecmem::data::buffer_type::resizable}};
-    vecmem::copy::event_type tracks_setup_event =
-        copy.setup(track_states_buffer.tracks);
-    vecmem::copy::event_type track_states_setup_event =
-        copy.setup(track_states_buffer.states);
+        copy.get_sizes(track_fit_view.tracks);
 
     // Return early, if there are no tracks.
     if (n_tracks == 0) {
-        tracks_setup_event->wait();
-        track_states_setup_event->wait();
-        return track_states_buffer;
+        return track_fit_buffer;
     }
 
     std::vector<unsigned int> seqs_sizes(candidate_sizes.size());
@@ -186,8 +175,7 @@ kalman_fitting(
 
     // Fill the keys and param_ids buffers.
     ::alpaka::exec<Acc>(queue, workDiv, kernels::fill_fitting_sort_keys{},
-                        track_candidates_view.tracks,
-                        vecmem::get_data(keys_buffer),
+                        track_fit_view.tracks, vecmem::get_data(keys_buffer),
                         vecmem::get_data(param_ids_buffer));
     ::alpaka::wait(queue);
 
@@ -197,28 +185,15 @@ kalman_fitting(
     details::sort_by_key(queue, mr, keys_device.begin(), keys_device.end(),
                          param_ids_device.begin());
 
-    // Run the fitting, using the sorted parameter IDs.
-    typename edm::track_fit_container<typename detector_t::algebra_type>::view
-        track_states_view{track_states_buffer.tracks,
-                          track_states_buffer.states,
-                          track_candidates_view.measurements};
-    tracks_setup_event->wait();
-    track_states_setup_event->wait();
-
-    ::alpaka::exec<Acc>(queue, workDiv, kernels::fit_prelude{},
-                        vecmem::get_data(param_ids_buffer),
-                        track_candidates_view, track_states_view,
-                        vecmem::get_data(param_liveness_buffer));
-    ::alpaka::wait(queue);
-
     // Allocate the fitting kernels's payload in host memory.
     using fitter_t = traccc::details::kalman_fitter_t<detector_t, bfield_t>;
     device::fit_payload<fitter_t> host_payload{
         .det_data = det_view,
         .field_data = field_view,
         .param_ids_view = param_ids_buffer,
         .param_liveness_view = param_liveness_buffer,
-        .tracks_view = track_states_view,
+        .tracks_view = {track_fit_buffer.tracks, track_fit_buffer.states,
+                        measurements},
         .barcodes_view = seqs_buffer};
     // Now copy it to device memory.
     vecmem::data::vector_buffer<device::fit_payload<fitter_t>> device_payload(
@@ -231,16 +206,113 @@ kalman_fitting(
 
     for (std::size_t i = 0; i < config.n_iterations; ++i) {
         // Run the track fitting
-        ::alpaka::exec<Acc>(queue, workDiv, kernels::fit_forward<fitter_t>{},
-                            config, device_payload.ptr());
-        ::alpaka::wait(queue);
+        if (i > 0 || forward_on_first_iteration) {
+            ::alpaka::exec<Acc>(queue, workDiv,
+                                kernels::fit_forward<fitter_t>{}, config,
+                                device_payload.ptr());
+            ::alpaka::wait(queue);
+        }
         ::alpaka::exec<Acc>(queue, workDiv, kernels::fit_backward<fitter_t>{},
                             config, device_payload.ptr());
         ::alpaka::wait(queue);
     }
 
     // Return the fitted tracks.
-    return track_states_buffer;
+    return track_fit_buffer;
+}
+
+/// Templated implementation of the Alpaka track fitting algorithm for
+/// unfitted tracks.
+///
+/// @tparam detector_t The (device) detector type to use
+/// @tparam bfield_t   The magnetic field type to use
+///
+/// @param[in] det_view     A view of the detector geometry
+/// @param[in] field_view   A view of the magnetic field
+/// @param[in] track_candidates_view All track candidates to fit
+/// @param[in] config       The fitting configuration
+/// @param[in] mr           Memory resource(s) to use
+/// @param[in] copy         The copy object to use for memory transfers
+/// @param[in] queue        The Alpaka queue to use for execution
+///
+/// @return A container of the fitted track states
+///
+template <typename detector_t, typename bfield_t>
+typename edm::track_fit_container<typename detector_t::algebra_type>::buffer
+kalman_fitting(
+    const typename detector_t::const_view_type& det_view,
+    const bfield_t& field_view,
+    const typename edm::track_candidate_container<
+        typename detector_t::algebra_type>::const_view& track_candidates_view,
+    const fitting_config& config, const memory_resource& mr, vecmem::copy& copy,
+    Queue& queue) {
+
+    // Number of threads per block to use.
+    const Idx threadsPerBlock = getWarpSize<Acc>() * 2;
+
+    // Get the number of tracks.
+    const edm::track_candidate_collection<
+        default_algebra>::const_device::size_type n_tracks =
+        copy.get_size(track_candidates_view.tracks);
+
+    // Get the sizes of the track candidates in each track.
+    const std::vector<unsigned int> candidate_sizes =
+        copy.get_sizes(track_candidates_view.tracks);
+    const unsigned int n_states =
+        std::accumulate(candidate_sizes.begin(), candidate_sizes.end(), 0u);
+
+    // Create the result buffer.
+    typename edm::track_fit_container<typename detector_t::algebra_type>::buffer
+        track_states_buffer{
+            {candidate_sizes, mr.main, mr.host,
+             vecmem::data::buffer_type::resizable},
+            {n_states, mr.main, vecmem::data::buffer_type::resizable}};
+    vecmem::copy::event_type tracks_setup_event =
+        copy.setup(track_states_buffer.tracks);
+    vecmem::copy::event_type track_states_setup_event =
+        copy.setup(track_states_buffer.states);
+
+    // Return early, if there are no tracks.
+    if (n_tracks == 0) {
+        tracks_setup_event->wait();
+        track_states_setup_event->wait();
+        return track_states_buffer;
+    }
+
+    std::vector<unsigned int> seqs_sizes(candidate_sizes.size());
+    std::transform(candidate_sizes.begin(), candidate_sizes.end(),
+                   seqs_sizes.begin(), [&config](const unsigned int sz) {
+                       return std::max(sz * config.barcode_sequence_size_factor,
+                                       config.min_barcode_sequence_capacity);
+                   });
+    vecmem::data::jagged_vector_buffer<detray::geometry::barcode> seqs_buffer{
+        seqs_sizes, mr.main, mr.host, vecmem::data::buffer_type::resizable};
+    copy.setup(seqs_buffer)->wait();
+
+    // The execution range for the two kernels of the function.
+    const Idx blocksPerGrid =
+        (n_tracks + threadsPerBlock - 1) / threadsPerBlock;
+    const auto workDiv = makeWorkDiv<Acc>(blocksPerGrid, threadsPerBlock);
+
+    // Run the fitting, using the sorted parameter IDs.
+    typename edm::track_fit_container<typename detector_t::algebra_type>::view
+        track_states_view{track_states_buffer.tracks,
+                          track_states_buffer.states,
+                          track_candidates_view.measurements};
+    tracks_setup_event->wait();
+    track_states_setup_event->wait();
+
+    ::alpaka::exec<Acc>(queue, workDiv, kernels::fit_prelude{},
+                        track_candidates_view, track_states_view);
+    ::alpaka::wait(queue);
+
+    return kalman_fitting<detector_t, bfield_t>(
+        det_view, field_view,
+        typename edm::track_fit_container<
+            typename detector_t::algebra_type>::buffer{
+            std::move(track_states_buffer.tracks),
+            std::move(track_states_buffer.states)},
+        track_candidates_view.measurements, config, mr, copy, queue, true);
 }
 
 }  // namespace traccc::alpaka::details

device/alpaka/src/fitting/kalman_fitting_algorithm.cpp

@@ -44,4 +44,23 @@ kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()(
         });
 }
 
+kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()(
+    const detector_buffer& det, const magnetic_field& bfield,
+    edm::track_fit_container<default_algebra>::buffer&& track_states,
+    const measurement_collection_types::const_view& measurements) const {
+
+    // Run the track fitting.
+    return detector_buffer_magnetic_field_visitor<
+        detector_type_list, alpaka::bfield_type_list<scalar>>(
+        det, bfield,
+        [&]<typename detector_t, typename bfield_view_t>(
+            const typename detector_t::view& detector,
+            const bfield_view_t& field) {
+            return details::kalman_fitting<typename detector_t::device>(
+                detector, field, std::move(track_states), measurements,
+                m_config, m_mr, m_copy.get(),
+                details::get_queue(m_queue.get()));
+        });
+}
+
 }  // namespace traccc::alpaka

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

@@ -12,21 +12,21 @@
 #include "traccc/edm/device/sort_key.hpp"
 
 // Project include(s).
-#include "traccc/edm/track_candidate_collection.hpp"
+#include "traccc/edm/track_fit_collection.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 candidates
+/// @param[in] track_fit_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_sort_keys(
     global_index_t globalIndex,
-    const edm::track_candidate_collection<default_algebra>::const_view&
-        track_candidates_view,
+    const edm::track_fit_collection<default_algebra>::const_view&
+        track_fit_view,
     vecmem::data::vector_view<device::sort_key> keys_view,
     vecmem::data::vector_view<unsigned int> ids_view);
 

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

@@ -24,11 +24,9 @@ namespace traccc::device {
 template <typename algebra_t>
 TRACCC_HOST_DEVICE inline void fit_prelude(
     const global_index_t globalIndex,
-    vecmem::data::vector_view<const unsigned int> param_ids_view,
     typename edm::track_candidate_container<algebra_t>::const_view
         track_candidates_view,
-    typename edm::track_fit_container<algebra_t>::view tracks_view,
-    vecmem::data::vector_view<unsigned int> param_liveness_view) {
+    typename edm::track_fit_container<algebra_t>::view tracks_view) {
 
     typename edm::track_fit_collection<algebra_t>::device tracks(
         tracks_view.tracks);
@@ -40,16 +38,11 @@ TRACCC_HOST_DEVICE inline void fit_prelude(
     typename edm::track_state_collection<algebra_t>::device track_states(
         tracks_view.states);
 
-    vecmem::device_vector<const unsigned int> param_ids(param_ids_view);
-    vecmem::device_vector<unsigned int> param_liveness(param_liveness_view);
-
-    const unsigned int param_id = param_ids.at(globalIndex);
-
-    auto track = tracks.at(param_id);
+    auto track = tracks.at(globalIndex);
 
     const typename edm::track_candidate_collection<algebra_t>::const_device
         track_candidates{track_candidates_view.tracks};
-    const auto track_candidate = track_candidates.at(param_id);
+    const auto track_candidate = track_candidates.at(globalIndex);
     const auto track_candidate_measurement_indices =
         track_candidate.measurement_indices();
     const measurement_collection_types::const_device measurements{
@@ -62,7 +55,6 @@ TRACCC_HOST_DEVICE inline void fit_prelude(
 
     // TODO: Set other stuff in the header?
     track.params() = track_candidate.params();
-    param_liveness.at(param_id) = 1u;
 }
 
 }  // namespace traccc::device