src/fitting/kalman_fitting_algorithm.cpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2022-2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // Local include(s).
0009 #include "traccc/cuda/fitting/kalman_fitting_algorithm.hpp"
0010
0011 #include "../utils/magnetic_field_types.hpp"
0012 #include "../utils/utils.hpp"
0013 #include "./kernels/fit_backward.hpp"
0014 #include "./kernels/fit_forward.hpp"
0015 #include "./kernels/fit_prelude.hpp"
0016
0017 // Project include(s).
0018 #include "traccc/fitting/details/kalman_fitting_types.hpp"
0019 #include "traccc/geometry/detector.hpp"
0020 #include "traccc/utils/detector_buffer_bfield_visitor.hpp"
0021
0022 namespace traccc::cuda {
0023
0024 kalman_fitting_algorithm::kalman_fitting_algorithm(
0025     const config_type& config, const traccc::memory_resource& mr,
0026     const vecmem::copy& copy, const stream_wrapper& str,
0027     std::unique_ptr<const Logger> logger)
0028     : device::kalman_fitting_algorithm{config, mr, copy, std::move(logger)},
0029       cuda::algorithm_base{str} {}
0030
0031 void kalman_fitting_algorithm::fit_prelude_kernel(
0032     const device::fit_prelude_payload& payload) const {
0033   // Get the number of tracks.
0034   const unsigned int n_tracks = payload.input_tracks.tracks.capacity();
0035   assert(n_tracks == copy().get_size(payload.input_tracks.tracks));
0036   assert(n_tracks == payload.track_indices.capacity());
0037   assert(payload.track_indices.size_ptr() == nullptr);
0038   assert(n_tracks == copy().get_size(payload.output_tracks.tracks));
0039
0040   // Launch parameters for the kernel.
0041   const unsigned int nThreads = warp_size() * 4;
0042   const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads;
0043
0044   // Run the fitting, using the sorted parameter IDs.
0045   fit_prelude(nBlocks, nThreads, 0, details::get_stream(stream()), payload);
0046 }
0047
0048 auto kalman_fitting_algorithm::prepare_fit_payload(
0049     const detector_buffer& det, const magnetic_field& field,
0050     const std::vector<unsigned int>& n_surfaces,
0051     const device::fit_payload& payload) const -> fit_payload {
0052   return prepare_fit_payload_helper<detector_type_list,
0053                                     cuda::bfield_type_list<scalar>>(
0054       det, field, n_surfaces, payload);
0055 }
0056
0057 void kalman_fitting_algorithm::fit_forward_kernel(
0058     const fitting_config& config, const fit_payload& payload) const {
0059   return detector_buffer_magnetic_field_visitor<detector_type_list,
0060                                                 cuda::bfield_type_list<scalar>>(
0061       payload.detector, payload.field,
0062       [&]<typename detector_traits_t, typename bfield_view_t>(
0063           const typename detector_traits_t::view&, const bfield_view_t&) {
0064         // Get the number of tracks.
0065         const unsigned int n_tracks = payload.payload.tracks.tracks.capacity();
0066         assert(n_tracks == copy().get_size(payload.payload.tracks.tracks));
0067
0068         // Launch parameters for the kernel.
0069         const unsigned int nThreads = warp_size() * 4;
0070         const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads;
0071
0072         // Fitter type to use.
0073         using fitter_t =
0074             traccc::details::kalman_fitter_t<typename detector_traits_t::device,
0075                                              bfield_view_t>;
0076
0077         // Run the track fitting
0078         fit_forward<fitter_t>(
0079             nBlocks, nThreads, 0, details::get_stream(stream()), config,
0080             payload.payload, payload.get_tpayload<fitter_t>());
0081       });
0082 }
0083
0084 void kalman_fitting_algorithm::fit_backward_kernel(
0085     const fitting_config& config, const fit_payload& payload) const {
0086   return detector_buffer_magnetic_field_visitor<detector_type_list,
0087                                                 cuda::bfield_type_list<scalar>>(
0088       payload.detector, payload.field,
0089       [&]<typename detector_traits_t, typename bfield_view_t>(
0090           const typename detector_traits_t::view&, const bfield_view_t&) {
0091         // Get the number of tracks.
0092         const unsigned int n_tracks = payload.payload.tracks.tracks.capacity();
0093         assert(n_tracks == copy().get_size(payload.payload.tracks.tracks));
0094
0095         // Launch parameters for the kernel.
0096         const unsigned int nThreads = warp_size() * 4;
0097         const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads;
0098
0099         // Fitter type to use.
0100         using fitter_t =
0101             traccc::details::kalman_fitter_t<typename detector_traits_t::device,
0102                                              bfield_view_t>;
0103
0104         // Run the track fitting
0105         fit_backward<fitter_t>(
0106             nBlocks, nThreads, 0, details::get_stream(stream()), config,
0107             payload.payload, payload.get_tpayload<fitter_t>());
0108       });
0109 }
0110
0111 }  // namespace traccc::cuda