Merge pull request #742 from stephenswat/refactor/ckf_tus

Split the CUDA CKF into different TUs

core/include/traccc/finding/actors/ckf_aborter.hpp

@@ -11,6 +11,7 @@
 #include "detray/definitions/detail/qualifiers.hpp"
 #include "detray/propagator/base_actor.hpp"
 #include "detray/propagator/base_stepper.hpp"
+#include "traccc/definitions/primitives.hpp"
 
 // System include(s)
 #include <limits>
@@ -51,4 +52,4 @@ struct ckf_aborter : detray::actor {
     }
 };
 
-}  // namespace traccc
+}  // namespace traccc

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

@@ -8,29 +8,47 @@
 #pragma once
 
 // Project include(s).
+#include "detray/navigation/navigator.hpp"
+#include "detray/propagator/actors/pointwise_material_interactor.hpp"
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/finding/finding_config.hpp"
+#include "traccc/utils/particle.hpp"
 
 namespace traccc::device {
+template <typename detector_t>
+struct apply_interaction_payload {
+    /**
+     * @brief View object describing the tracking detector
+     */
+    typename detector_t::view_type det_data;
+
+    /**
+     * @brief Total number of input parameters (including non-live ones)
+     */
+    const int n_params;
+
+    /**
+     * @brief View object to the vector of bound track parameters
+     */
+    bound_track_parameters_collection_types::view params_view;
+
+    /**
+     * @brief View object to the vector of boolean-like integers describing
+     * whether each parameter is live. Has the same size as \ref params_view
+     */
+    vecmem::data::vector_view<const unsigned int> params_liveness_view;
+};
 
 /// Function applying the Pre material interaction to tracks spawned by bound
 /// track parameters
 ///
 /// @param[in] globalIndex     The index of the current thread
 /// @param[in] cfg             Track finding config object
-/// @param[in] det_data        Detector view object
-/// @param[in] n_params        The number of parameters (or tracks)
-/// @param[out] params_view    Collection of output bound track_parameters
-/// @param[in] params_liveness_view Vector of parameter liveness indicators
-///
+/// @param[inout] payload      The function call payload
 template <typename detector_t>
 TRACCC_DEVICE inline void apply_interaction(
     std::size_t globalIndex, const finding_config& cfg,
-    typename detector_t::view_type det_data, const int n_params,
-    bound_track_parameters_collection_types::view params_view,
-    vecmem::data::vector_view<const unsigned int> params_liveness_view);
-
+    const apply_interaction_payload<detector_t>& payload);
 }  // namespace traccc::device
 
-// Include the implementation.
-#include "traccc/finding/device/impl/apply_interaction.ipp"
+#include "./impl/apply_interaction.ipp"

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

@@ -9,8 +9,52 @@
 
 // Project include(s).
 #include "traccc/definitions/qualifiers.hpp"
+#include "traccc/edm/measurement.hpp"
+#include "traccc/edm/track_candidate.hpp"
+#include "traccc/edm/track_parameters.hpp"
+#include "traccc/finding/candidate_link.hpp"
 
 namespace traccc::device {
+struct build_tracks_payload {
+    /**
+     * @brief View object to the vector of measurements
+     *
+     * @warning Measurements on the same surface must be adjacent
+     */
+    measurement_collection_types::const_view measurements_view;
+
+    /**
+     * @brief View object to the vector of measurements
+     */
+    bound_track_parameters_collection_types::const_view seeds_view;
+
+    /**
+     * @brief View object to the vector of candidate links
+     */
+    vecmem::data::jagged_vector_view<const candidate_link> links_view;
+
+    /**
+     * @brief View object to the vector of tips
+     */
+    vecmem::data::vector_view<const typename candidate_link::link_index_type>
+        tips_view;
+
+    /**
+     * @brief View object to the vector of 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;
+};
 
 /// Function for building full tracks from the link container:
 /// The full tracks are built using the link container and tip link container.
@@ -19,28 +63,12 @@ namespace traccc::device {
 ///
 /// @param[in] globalIndex         The index of the current thread
 /// @param[in] cfg                    Track finding config object
-/// @param[in] measurements_view   Measurements container view
-/// @param[in] seeds_view          Seed container view
-/// @param[in] link_view           Link container view
-/// @param[in] param_to_link_view  Container for param index -> link index
-/// @param[in] tips_view           Tip link container view
-/// @param[out] track_candidates_view  Track candidate container view
-/// @param[out] valid_indices_view   Valid indices meeting criteria
-/// @param[out] n_valid_tracks       The number of valid tracks meeting criteria
-
+/// @param[inout] payload      The function call payload
 template <typename config_t>
-TRACCC_DEVICE inline void build_tracks(
-    std::size_t globalIndex, const config_t cfg,
-    measurement_collection_types::const_view measurements_view,
-    bound_track_parameters_collection_types::const_view seeds_view,
-    vecmem::data::jagged_vector_view<const candidate_link> links_view,
-    vecmem::data::vector_view<const typename candidate_link::link_index_type>
-        tips_view,
-    track_candidate_container_types::view track_candidates_view,
-    vecmem::data::vector_view<unsigned int> valid_indices_view,
-    unsigned int& n_valid_tracks);
+TRACCC_DEVICE inline void build_tracks(std::size_t globalIndex,
+                                       const config_t cfg,
+                                       const build_tracks_payload& payload);
 
 }  // namespace traccc::device
 
-// Include the implementation.
-#include "traccc/finding/device/impl/build_tracks.ipp"
+#include "./impl/build_tracks.ipp"

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

@@ -12,21 +12,30 @@
 #include "traccc/edm/track_candidate.hpp"
 
 namespace traccc::device {
+struct fill_sort_keys_payload {
+    /**
+     * @brief View object to the vector of bound track parameters
+     */
+    bound_track_parameters_collection_types::const_view params_view;
+
+    /**
+     * @brief View object to the vector of sort keys
+     */
+    vecmem::data::vector_view<device::sort_key> keys_view;
+
+    /**
+     * @brief View object to the vector of parameter indices, which is the
+     * output to the algorithm
+     */
+    vecmem::data::vector_view<unsigned int> ids_view;
+};
 
 /// 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
-///
+/// @param[inout] payload      The function call payload
 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);
-
+    std::size_t globalIndex, const fill_sort_keys_payload& payload);
 }  // namespace traccc::device
 
-// Include the implementation.
-#include "traccc/finding/device/impl/fill_sort_keys.ipp"
+#include "./impl/fill_sort_keys.ipp"

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

@@ -14,11 +14,114 @@
 #include "traccc/device/concepts/thread_id.hpp"
 #include "traccc/edm/measurement.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"
+#include "traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
 
 // Thrust include(s)
 #include <thrust/binary_search.h>
 
 namespace traccc::device {
+template <typename detector_t>
+struct find_tracks_payload {
+    /**
+     * @brief View object to the tracking detector description
+     */
+    typename detector_t::view_type det_data;
+
+    /**
+     * @brief View object to the vector of bound track parameters
+     *
+     * @warning Measurements on the same surface must be adjacent
+     */
+    measurement_collection_types::const_view measurements_view;
+
+    /**
+     * @brief View object to the vector of track parameters
+     */
+    bound_track_parameters_collection_types::const_view in_params_view;
+
+    /**
+     * @brief View object to the vector of boolean-like integers describing the
+     * liveness of each parameter
+     */
+    vecmem::data::vector_view<const unsigned int> in_params_liveness_view;
+
+    /**
+     * @brief The total number of input parameters
+     */
+    const unsigned int n_in_params;
+
+    /**
+     * @brief View object to the vector of barcodes for each measurement
+     */
+    vecmem::data::vector_view<const detray::geometry::barcode> barcodes_view;
+
+    /**
+     * @brief View object to the vector of upper bounds of measurement indices
+     * per surface
+     */
+    vecmem::data::vector_view<const unsigned int> upper_bounds_view;
+
+    /**
+     * @brief View object to the link vector of the previous step
+     */
+    vecmem::data::vector_view<const candidate_link> prev_links_view;
+
+    /**
+     * @brief The current step identifier
+     */
+    const unsigned int step;
+
+    /**
+     * @brief The maximum number of new tracks to find
+     */
+    const unsigned int n_max_candidates;
+
+    /**
+     * @brief View object to the output track parameter vector
+     */
+    bound_track_parameters_collection_types::view out_params_view;
+
+    /**
+     * @brief View object to the output track parameter liveness vector
+     */
+    vecmem::data::vector_view<unsigned int> out_params_liveness_view;
+
+    /**
+     * @brief View object to the output candidate links
+     */
+    vecmem::data::vector_view<candidate_link> links_view;
+
+    /**
+     * @brief Pointer to the total of number of candidates; to be set to zero
+     * before launching the kernel
+     */
+    unsigned int* n_total_candidates;
+};
+
+struct find_tracks_shared_payload {
+    /**
+     * @brief Shared-memory vector with the number of measurements found per
+     * track
+     */
+    unsigned int* shared_num_candidates;
+
+    /**
+     * @brief Shared-memory vector of measurement candidats with ID and
+     * original track parameter identifier
+     *
+     * @note Length is always twice the block size
+     */
+    std::pair<unsigned int, unsigned int>* shared_candidates;
+
+    /**
+     * @brief Shared-memory atomic variable to track the size of
+     * \ref shared_candidates
+     */
+    unsigned int& shared_candidates_size;
+};
 
 /// Function for combinatorial finding.
 /// If the chi2 of the measurement < chi2_max, its measurement index and the
@@ -27,47 +130,14 @@ namespace traccc::device {
 /// @param[in] thread_id          A thread identifier object
 /// @param[in] barrier            A block-wide barrier
 /// @param[in] cfg                Track finding config object
-/// @param[in] det_data           Detector view object
-/// @param[in] measurements_view  Measurements container view
-/// @param[in] in_params_view     Input parameters
-/// @param[in] n_in_params        The number of input params
-/// @param[in] barcodes_view      View of a measurement -> barcode map
-/// @param[in] upper_bounds_view  Upper bounds of measurements unique w.r.t
-/// barcode
-/// @param[in] prev_links_view    link container from the previous step
-/// @param[in] prev_param_to_link_view  param_to_link container from the
-/// previous step
-/// @param[in] step               Step index
-/// @param[in] n_max_candidates   Number of maximum candidates
-/// @param[out] out_params_view   Output parameters
-/// @param[out] links_view        link container for the current step
-/// @param[out] n_total_candidates  The number of total candidates for the
-/// current step
-/// @param shared_num_candidates  Shared memory scratch space
-/// @param shared_candidates      Shared memory scratch space
-/// @param shared_candidates_size Shared memory scratch space
-///
+/// @param[inout] payload         The global memory payload
+/// @param[inout] shared_payload  The shared memory payload
 template <concepts::thread_id1 thread_id_t, concepts::barrier barrier_t,
           typename detector_t, typename config_t>
 TRACCC_DEVICE inline void find_tracks(
     thread_id_t& thread_id, barrier_t& barrier, const config_t cfg,
-    typename detector_t::view_type det_data,
-    measurement_collection_types::const_view measurements_view,
-    bound_track_parameters_collection_types::const_view in_params_view,
-    vecmem::data::vector_view<const unsigned int> in_params_liveness_view,
-    const unsigned int n_in_params,
-    vecmem::data::vector_view<const detray::geometry::barcode> barcodes_view,
-    vecmem::data::vector_view<const unsigned int> upper_bounds_view,
-    vecmem::data::vector_view<const candidate_link> prev_links_view,
-    const unsigned int step, const unsigned int& n_max_candidates,
-    bound_track_parameters_collection_types::view out_params_view,
-    vecmem::data::vector_view<unsigned int> out_params_liveness_view,
-    vecmem::data::vector_view<candidate_link> links_view,
-    unsigned int& n_total_candidates, unsigned int* shared_num_candidates,
-    std::pair<unsigned int, unsigned int>* shared_candidates,
-    unsigned int& shared_candidates_size);
-
+    const find_tracks_payload<detector_t>& payload,
+    const find_tracks_shared_payload& shared_payload);
 }  // namespace traccc::device
 
-// Include the implementation.
-#include "traccc/finding/device/impl/find_tracks.ipp"
+#include "./impl/find_tracks.ipp"