Adapted traccc::device_common and traccc::cuda to the SoA track fit EDM.

Author: krasznaa

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

@@ -178,8 +178,10 @@ TRACCC_HOST_DEVICE inline void find_tracks(
 
         barrier.blockBarrier();
 
-        std::optional<std::tuple<track_state<typename detector_t::algebra_type>,
-                                 unsigned int, unsigned int>>
+        std::optional<std::tuple<
+            typename edm::track_state_collection<
+                typename detector_t::algebra_type>::device::object_type,
+            unsigned int, unsigned int>>
             result = std::nullopt;
 
         /*
@@ -216,17 +218,21 @@ TRACCC_HOST_DEVICE inline void find_tracks(
             }
 
             if (use_measurement) {
-                const auto& meas = measurements.at(meas_idx);
 
-                track_state<typename detector_t::algebra_type> trk_state(meas);
+                typename edm::track_state_collection<
+                    typename detector_t::algebra_type>::device::object_type
+                    trk_state(0u, 0.f, 0.f, 0.f, {}, {}, meas_idx);
+                trk_state.filtered_params().set_surface_link(
+                    measurements.at(meas_idx).surface_link);
+
                 const detray::tracking_surface sf{det, in_par.surface_link()};
 
                 const bool is_line = sf.template visit_mask<is_line_visitor>();
 
                 // Run the Kalman update
                 const kalman_fitter_status res =
                     gain_matrix_updater<typename detector_t::algebra_type>{}(
-                        trk_state, in_par, is_line);
+                        trk_state, measurements, in_par, is_line);
 
                 /*
                  * The $\chi^2$ value from the Kalman update should be less than
@@ -246,8 +252,7 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                 }
 
                 if (use_measurement) {
-                    result.emplace(std::move(trk_state), meas_idx,
-                                   owner_local_thread_id);
+                    result.emplace(trk_state, meas_idx, owner_local_thread_id);
                 }
             }
         }
@@ -441,10 +446,12 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                         .chi2_sum = prev_chi2_sum + chi2,
                         .ndf_sum =
                             prev_ndf_sum +
-                            std::get<0>(*result).get_measurement().meas_dim};
+                            measurements
+                                .at(std::get<0>(*result).measurement_index())
+                                .meas_dim};
 
                     tmp_params.at(p_offset + l_pos) =
-                        std::get<0>(*result).filtered();
+                        std::get<0>(*result).filtered_params();
 
                     /*
                      * Reset the temporary state storage, as this is no longer

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

@@ -12,7 +12,7 @@
 
 // Project include(s).
 #include "traccc/definitions/qualifiers.hpp"
-#include "traccc/edm/track_state.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 
 // VecMem include(s).
 #include <vecmem/containers/data/jagged_vector_view.hpp>
@@ -44,9 +44,10 @@ struct fit_payload {
     vecmem::data::vector_view<unsigned int> param_liveness_view;
 
     /**
-     * @brief View object to the output track states
+     * @brief View object to the output tracks
      */
-    track_state_container_types::view track_states_view;
+    typename edm::track_fit_container<
+        typename fitter_t::detector_type::algebra_type>::view tracks_view;
 
     /**
      * @brief View object to the output barcode sequence

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

@@ -22,20 +22,28 @@ TRACCC_HOST_DEVICE inline void fit_backward(
     vecmem::device_vector<const unsigned int> param_ids(payload.param_ids_view);
     vecmem::device_vector<unsigned int> param_liveness(
         payload.param_liveness_view);
-    track_state_container_types::device track_states(payload.track_states_view);
-
-    if (globalIndex >= track_states.size()) {
+    typename edm::track_fit_collection<
+        typename fitter_t::detector_type::algebra_type>::device
+        tracks(payload.tracks_view.tracks);
+    typename edm::track_state_collection<
+        typename fitter_t::detector_type::algebra_type>::device
+        track_states(payload.tracks_view.states);
+    measurement_collection_types::const_device measurements{
+        payload.tracks_view.measurements};
+
+    if (globalIndex >= tracks.size()) {
         return;
     }
 
     const unsigned int param_id = param_ids.at(globalIndex);
+    auto track = tracks.at(param_id);
 
     // Run fitting
     fitter_t fitter(det, payload.field_data, cfg);
 
     if (param_liveness.at(param_id) > 0u) {
         typename fitter_t::state fitter_state(
-            track_states.at(param_id).items,
+            track, track_states, measurements,
             *(payload.barcodes_view.ptr() + param_id));
 
         kalman_fitter_status fit_status = fitter.smooth(fitter_state);
@@ -48,7 +56,7 @@ TRACCC_HOST_DEVICE inline void fit_backward(
 
             assert(fit_status == kalman_fitter_status::SUCCESS);
 
-            track_states.at(param_id).header = fitter_state.m_fit_res;
+            track = fitter_state.m_fit_res;
         } else {
             param_liveness.at(param_id) = 0u;
         }

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

@@ -22,23 +22,31 @@ TRACCC_HOST_DEVICE inline void fit_forward(
     vecmem::device_vector<const unsigned int> param_ids(payload.param_ids_view);
     vecmem::device_vector<unsigned int> param_liveness(
         payload.param_liveness_view);
-    track_state_container_types::device track_states(payload.track_states_view);
-
-    if (globalIndex >= track_states.size()) {
+    typename edm::track_fit_collection<
+        typename fitter_t::detector_type::algebra_type>::device
+        tracks(payload.tracks_view.tracks);
+    typename edm::track_state_collection<
+        typename fitter_t::detector_type::algebra_type>::device
+        track_states(payload.tracks_view.states);
+    measurement_collection_types::const_device measurements{
+        payload.tracks_view.measurements};
+
+    if (globalIndex >= tracks.size()) {
         return;
     }
 
     const unsigned int param_id = param_ids.at(globalIndex);
 
     fitter_t fitter(det, payload.field_data, cfg);
 
-    auto params = track_states.at(param_id).header.fit_params;
+    auto track = tracks.at(param_id);
+    auto params = track.params();
 
     // TODO: Merge into filter?
     inflate_covariance(params, fitter.config().covariance_inflation_factor);
 
     typename fitter_t::state fitter_state(
-        track_states.at(param_id).items,
+        track, track_states, measurements,
         *(payload.barcodes_view.ptr() + param_id));
 
     kalman_fitter_status fit_status = fitter.filter(params, fitter_state);

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

@@ -7,45 +7,66 @@
 
 #pragma once
 
+// Local include(s).
+#include "traccc/device/global_index.hpp"
+
+// Project include(s).
 #include "traccc/edm/track_candidate_container.hpp"
-#include "traccc/fitting/device/fit.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 #include "traccc/fitting/status_codes.hpp"
 
+// VecMem include(s).
+#include <vecmem/containers/data/vector_view.hpp>
+
 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,
-    edm::track_candidate_container<default_algebra>::const_view
+    typename edm::track_candidate_container<algebra_t>::const_view
         track_candidates_view,
-    track_state_container_types::view track_states_view,
+    typename edm::track_fit_container<algebra_t>::view tracks_view,
     vecmem::data::vector_view<unsigned int> param_liveness_view) {
 
-    track_state_container_types::device track_states(track_states_view);
+    typename edm::track_fit_collection<algebra_t>::device tracks(
+        tracks_view.tracks);
 
-    if (globalIndex >= track_states.size()) {
+    if (globalIndex >= tracks.size()) {
         return;
     }
 
+    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_states_per_track = track_states.at(param_id).items;
+    auto track = tracks.at(param_id);
 
-    const edm::track_candidate_collection<default_algebra>::const_device
+    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_measurement_indices =
+        track_candidate.measurement_indices();
     const measurement_collection_types::const_device measurements{
         track_candidates_view.measurements};
-    for (unsigned int meas_idx :
-         track_candidates.measurement_indices().at(param_id)) {
-        track_states_per_track.emplace_back(measurements.at(meas_idx));
+    for (unsigned int meas_idx : track_candidate_measurement_indices) {
+        const unsigned int track_state_index =
+            track_states.push_back({0, 0.f, 0.f, 0.f, {}, {}, meas_idx});
+        auto state = track_states.at(track_state_index);
+        state.set_hole(true);
+        state.set_smoothed(false);
+        const auto surface_link = measurements.at(meas_idx).surface_link;
+        state.filtered_params().set_surface_link(surface_link);
+        state.smoothed_params().set_surface_link(surface_link);
+        track.state_indices().push_back(track_state_index);
     }
 
     // TODO: Set other stuff in the header?
-    track_states.at(param_id).header.fit_params =
-        track_candidates.at(param_id).params();
+    track.params() = track_candidate.params();
     param_liveness.at(param_id) = 1u;
 }
 

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

@@ -13,7 +13,7 @@
 // Project include(s).
 #include "traccc/bfield/magnetic_field.hpp"
 #include "traccc/edm/track_candidate_container.hpp"
-#include "traccc/edm/track_state.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 #include "traccc/fitting/fitting_config.hpp"
 #include "traccc/geometry/detector.hpp"
 #include "traccc/utils/algorithm.hpp"
@@ -30,10 +30,10 @@ namespace traccc::cuda {
 
 /// Kalman filter based track fitting algorithm
 class kalman_fitting_algorithm
-    : public algorithm<track_state_container_types::buffer(
+    : public algorithm<edm::track_fit_container<default_algebra>::buffer(
           const default_detector::view&, const magnetic_field&,
           const edm::track_candidate_container<default_algebra>::const_view&)>,
-      public algorithm<track_state_container_types::buffer(
+      public algorithm<edm::track_fit_container<default_algebra>::buffer(
           const telescope_detector::view&, const magnetic_field&,
           const edm::track_candidate_container<default_algebra>::const_view&)>,
       public messaging {
@@ -42,7 +42,7 @@ class kalman_fitting_algorithm
     /// Configuration type
     using config_type = fitting_config;
     /// Output type
-    using output_type = track_state_container_types::buffer;
+    using output_type = edm::track_fit_container<default_algebra>::buffer;
 
     /// Constructor with the algorithm's configuration
     ///

device/cuda/src/fitting/kalman_fitting.cuh

@@ -19,7 +19,7 @@
 // Project include(s).
 #include "traccc/edm/device/sort_key.hpp"
 #include "traccc/edm/track_candidate_container.hpp"
-#include "traccc/edm/track_state.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 #include "traccc/fitting/details/kalman_fitting_types.hpp"
 #include "traccc/fitting/device/fill_fitting_sort_keys.hpp"
 #include "traccc/fitting/fitting_config.hpp"
@@ -32,6 +32,9 @@
 #include <thrust/execution_policy.h>
 #include <thrust/sort.h>
 
+// System include(s).
+#include <numeric>
+
 namespace traccc::cuda::details {
 
 /// Templated implementation of the CUDA track fitting algorithm.
@@ -50,7 +53,8 @@ namespace traccc::cuda::details {
 /// @return A container of the fitted track states
 ///
 template <typename detector_t, typename bfield_t>
-track_state_container_types::buffer kalman_fitting(
+typename edm::track_fit_container<typename detector_t::algebra_type>::buffer
+kalman_fitting(
     const typename detector_t::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,
@@ -68,14 +72,17 @@ track_state_container_types::buffer kalman_fitting(
     // 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.
-    track_state_container_types::buffer track_states_buffer{
-        {n_tracks, mr.main},
-        {candidate_sizes, mr.main, mr.host,
-         vecmem::data::buffer_type::resizable}};
-    copy.setup(track_states_buffer.headers)->ignore();
-    copy.setup(track_states_buffer.items)->ignore();
+    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}};
+    copy.setup(track_states_buffer.tracks)->ignore();
+    copy.setup(track_states_buffer.states)->ignore();
 
     // Return early, if there are no tracks.
     if (n_tracks == 0) {
@@ -127,7 +134,9 @@ track_state_container_types::buffer kalman_fitting(
 
     // Run the fitting, using the sorted parameter IDs.
     fit_prelude(nBlocks, nThreads, 0, stream, param_ids_buffer,
-                track_candidates_view, track_states_buffer,
+                track_candidates_view,
+                {track_states_buffer.tracks, track_states_buffer.states,
+                 track_candidates_view.measurements},
                 param_liveness_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
     str.synchronize();
@@ -139,7 +148,8 @@ track_state_container_types::buffer kalman_fitting(
         .field_data = field_view,
         .param_ids_view = param_ids_buffer,
         .param_liveness_view = param_liveness_buffer,
-        .track_states_view = track_states_buffer,
+        .tracks_view = {track_states_buffer.tracks, track_states_buffer.states,
+                        track_candidates_view.measurements},
         .barcodes_view = seqs_buffer};
 
     for (std::size_t i = 0; i < config.n_iterations; ++i) {

device/cuda/src/fitting/kernels/fit_prelude.cu

@@ -15,11 +15,11 @@ __global__ void fit_prelude(
     vecmem::data::vector_view<const unsigned int> param_ids_view,
     edm::track_candidate_container<default_algebra>::const_view
         track_candidates_view,
-    track_state_container_types::view track_states_view,
+    edm::track_fit_container<default_algebra>::view tracks_view,
     vecmem::data::vector_view<unsigned int> param_liveness_view) {
-    device::fit_prelude(details::global_index1(), param_ids_view,
-                        track_candidates_view, track_states_view,
-                        param_liveness_view);
+    device::fit_prelude<default_algebra>(details::global_index1(),
+                                         param_ids_view, track_candidates_view,
+                                         tracks_view, param_liveness_view);
 }
 }  // namespace kernels
 
@@ -28,10 +28,10 @@ void fit_prelude(const dim3& grid_size, const dim3& block_size,
                  vecmem::data::vector_view<const unsigned int> param_ids_view,
                  edm::track_candidate_container<default_algebra>::const_view
                      track_candidates_view,
-                 track_state_container_types::view track_states_view,
+                 edm::track_fit_container<default_algebra>::view tracks_view,
                  vecmem::data::vector_view<unsigned int> param_liveness_view) {
     kernels::fit_prelude<<<grid_size, block_size, shared_mem_size, stream>>>(
-        param_ids_view, track_candidates_view, track_states_view,
+        param_ids_view, track_candidates_view, tracks_view,
         param_liveness_view);
 }
 }  // namespace traccc::cuda

device/cuda/src/fitting/kernels/fit_prelude.hpp

@@ -10,14 +10,14 @@
 #include <cuda_runtime.h>
 
 #include "traccc/edm/track_candidate_container.hpp"
-#include "traccc/edm/track_state.hpp"
+#include "traccc/edm/track_fit_container.hpp"
 
 namespace traccc::cuda {
 void fit_prelude(const dim3& grid_size, const dim3& block_size,
                  std::size_t shared_mem_size, const cudaStream_t& stream,
                  vecmem::data::vector_view<const unsigned int> param_ids_view,
                  edm::track_candidate_container<default_algebra>::const_view
                      track_candidates_view,
-                 track_state_container_types::view track_states_view,
+                 edm::track_fit_container<default_algebra>::view tracks_view,
                  vecmem::data::vector_view<unsigned int> param_liveness_view);
 }