src/fitting/kalman_fitting_algorithm.cu

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 "../utils/utils.hpp"
0010 #include "./kernels/fill_fitting_sort_keys.hpp"
0011 #include "traccc/cuda/fitting/kalman_fitting_algorithm.hpp"
0012
0013 // Thrust include(s).
0014 #include <thrust/execution_policy.h>
0015 #include <thrust/sort.h>
0016
0017 // System include(s).
0018 #include <cassert>
0019 #include <memory_resource>
0020
0021 namespace traccc::cuda {
0022
0023 void kalman_fitting_algorithm::prepare_track_fit_order(
0024     const edm::track_collection<default_algebra>::const_view& tracks,
0025     vecmem::data::vector_view<device::sort_key>& track_sort_keys,
0026     vecmem::data::vector_view<unsigned int>& track_indices) const {
0027   // Get the number of tracks.
0028   const unsigned int n_tracks = tracks.capacity();
0029   assert(n_tracks == copy().get_size(tracks));
0030   assert(n_tracks == track_indices.capacity());
0031   assert(track_indices.size_ptr() == nullptr);
0032
0033   // Launch parameters for the kernel.
0034   const unsigned int nThreads = warp_size() * 4;
0035   const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads;
0036
0037   // Fill the keys and indices buffers.
0038   fill_fitting_sort_keys(nBlocks, nThreads, details::get_stream(stream()),
0039                          tracks, track_sort_keys, track_indices);
0040
0041   // Sort the key to get the sorted parameter ids
0042   vecmem::device_vector<device::sort_key> keys_device(track_sort_keys);
0043   vecmem::device_vector<unsigned int> track_indices_device(track_indices);
0044   thrust::sort_by_key(
0045       thrust::cuda::par_nosync(std::pmr::polymorphic_allocator(&mr().main))
0046           .on(details::get_stream(stream())),
0047       keys_device.begin(), keys_device.end(), track_indices_device.begin());
0048 }
0049
0050 }  // namespace traccc::cuda