diff --git a/.vscode/settings.json b/.vscode/settings.json index 3fc135b9d6..8ab96c424d 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,7 +2,8 @@ "files.associations": { "*.hip": "cpp", "*.sycl": "cpp", - "*.ipp": "cpp" + "*.ipp": "cpp", + "*.cu.in": "cpp" }, "sonarlint.connectedMode.project": { "connectionId": "acts-project", diff --git a/core/include/traccc/edm/track_container.hpp b/core/include/traccc/edm/track_container.hpp index f7061e1ed5..7478084789 100644 --- a/core/include/traccc/edm/track_container.hpp +++ b/core/include/traccc/edm/track_container.hpp @@ -74,6 +74,8 @@ struct track_container { }; struct view { + /// Default constructor + view() = default; /// Constructor from a buffer TRACCC_HOST_DEVICE view(const buffer& b) @@ -95,6 +97,8 @@ struct track_container { }; struct const_view { + /// Default constructor + const_view() = default; /// Constructor from a buffer TRACCC_HOST_DEVICE const_view(const buffer& b) diff --git a/core/include/traccc/geometry/move_only_any.hpp b/core/include/traccc/geometry/move_only_any.hpp index ca91e07377..ac8df3866f 100644 --- a/core/include/traccc/geometry/move_only_any.hpp +++ b/core/include/traccc/geometry/move_only_any.hpp @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2024 CERN for the benefit of the ACTS project + * (c) 2024-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -113,6 +113,9 @@ class move_only_any { "Value for `traccc::move_only_any` requested, but no value " "exists."); } + if (!is()) { + throw std::bad_any_cast(); + } return *static_cast(m_obj); } diff --git a/device/alpaka/CMakeLists.txt b/device/alpaka/CMakeLists.txt index 40811c9ff0..c8b9b55ef3 100644 --- a/device/alpaka/CMakeLists.txt +++ b/device/alpaka/CMakeLists.txt @@ -60,7 +60,6 @@ traccc_add_alpaka_library( traccc_alpaka alpaka TYPE SHARED # Track fitting algorithm(s). "include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp" "src/fitting/kalman_fitting_algorithm.cpp" - "src/fitting/kalman_fitting.hpp" ) # Set up Thrust specifically for the traccc::alpaka library. diff --git a/device/alpaka/include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp b/device/alpaka/include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp index 2f1c6d50fb..e644e8df50 100644 --- a/device/alpaka/include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp +++ b/device/alpaka/include/traccc/alpaka/fitting/kalman_fitting_algorithm.hpp @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -8,37 +8,19 @@ #pragma once // Library include(s). +#include "traccc/alpaka/utils/algorithm_base.hpp" #include "traccc/alpaka/utils/queue.hpp" // Project include(s). -#include "traccc/bfield/magnetic_field.hpp" -#include "traccc/edm/track_container.hpp" -#include "traccc/fitting/fitting_config.hpp" -#include "traccc/geometry/detector.hpp" -#include "traccc/geometry/detector_buffer.hpp" -#include "traccc/utils/algorithm.hpp" -#include "traccc/utils/memory_resource.hpp" -#include "traccc/utils/messaging.hpp" - -// VecMem include(s). -#include - -// System include(s). -#include +#include "traccc/fitting/device/kalman_fitting_algorithm.hpp" namespace traccc::alpaka { -/// Kalman filter based track fitting algorithm -class kalman_fitting_algorithm - : public algorithm::buffer( - const detector_buffer&, const magnetic_field&, - const edm::track_container::const_view&)>, - public messaging { +/// Kalman filter based track fitting algorithm using Alpaka +class kalman_fitting_algorithm : public device::kalman_fitting_algorithm, + public alpaka::algorithm_base { public: - /// Configuration type - using config_type = fitting_config; - /// Constructor with the algorithm's configuration /// /// @param config The configuration object @@ -49,31 +31,63 @@ class kalman_fitting_algorithm /// kalman_fitting_algorithm( const config_type& config, const traccc::memory_resource& mr, - const vecmem::copy& copy, queue& q, + const vecmem::copy& copy, alpaka::queue& q, std::unique_ptr logger = getDummyLogger().clone()); - /// Execute the algorithm + private: + /// @name Function(s) implemented from @c device::kalman_fitting_algorithm + /// @{ + + /// Prepare a buffer with the index order with which to fit the tracks /// - /// @param det The detector object - /// @param bfield The magnetic field object - /// @param track_candidates All track candidates to fit + /// @param[in] tracks The tracks to be fitted + /// @param[out] track_sort_keys Buffer storing temporary sorting keys + /// @param[out] track_indices The buffer to write the fitting order into /// - /// @return A container of the fitted track states + void prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const override; + + /// Kernel to prepare the fitting payloads + /// + /// @param payload The payload for the kernel(s) /// - output_type operator()( - const detector_buffer& det, const magnetic_field& bfield, - const edm::track_container::const_view& - track_candidates) const override; + void fit_prelude_kernel( + const device::fit_prelude_payload& payload) const override; - private: - /// Algorithm configuration - config_type m_config; - /// Memory resource used by the algorithm - traccc::memory_resource m_mr; - /// Copy object used by the algorithm - std::reference_wrapper m_copy; - /// Queue wrapper - std::reference_wrapper m_queue; + /// Function preparing the fitting payload + /// + /// @param det The detector buffer to prepare the payload for + /// @param field The magnetic field to prepare the payload for + /// @param n_surfaces The number of surfaces for each track to be + /// fitted + /// @param payload The (non-templated) payload for the kernel(s) + /// + /// @return The prepared payload for the fitting kernel(s) + /// + fit_payload prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const override; + + /// Function launching the "forward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_forward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// Function launching the "backward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_backward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// @} }; // class kalman_fitting_algorithm diff --git a/device/alpaka/src/fitting/kalman_fitting.hpp b/device/alpaka/src/fitting/kalman_fitting.hpp deleted file mode 100644 index c9e3915237..0000000000 --- a/device/alpaka/src/fitting/kalman_fitting.hpp +++ /dev/null @@ -1,242 +0,0 @@ -/** TRACCC library, part of the ACTS project (R&D line) - * - * (c) 2025-2026 CERN for the benefit of the ACTS project - * - * Mozilla Public License Version 2.0 - */ - -#pragma once - -// Local include(s). -#include "../utils/parallel_algorithms.hpp" -#include "../utils/utils.hpp" - -// Project include(s). -#include "traccc/edm/device/sort_key.hpp" -#include "traccc/edm/track_container.hpp" -#include "traccc/fitting/details/kalman_fitting_types.hpp" -#include "traccc/fitting/device/fill_fitting_sort_keys.hpp" -#include "traccc/fitting/device/fit.hpp" -#include "traccc/fitting/device/fit_backward.hpp" -#include "traccc/fitting/device/fit_forward.hpp" -#include "traccc/fitting/device/fit_prelude.hpp" -#include "traccc/fitting/fitting_config.hpp" -#include "traccc/utils/memory_resource.hpp" - -// VecMem include(s). -#include - -namespace traccc::alpaka::details { -namespace kernels { - -/// Alpaka kernel functor for @c traccc::device::fill_fitting_sort_keys -struct fill_fitting_sort_keys { - template - ALPAKA_FN_ACC void operator()( - TAcc const& acc, - edm::track_collection::const_view - track_candidates_view, - vecmem::data::vector_view keys_view, - vecmem::data::vector_view 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, - keys_view, ids_view); - } -}; - -/// Alpaka kernel functor for @c traccc::device::fit_prelude -struct fit_prelude { - template - ALPAKA_FN_ACC void operator()( - TAcc const& acc, - vecmem::data::vector_view param_ids_view, - edm::track_container::const_view track_candidates_view, - edm::track_container::view track_states_view, - vecmem::data::vector_view param_liveness_view) const { - - const device::global_index_t globalThreadIdx = - ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; - device::fit_prelude( - globalThreadIdx, param_ids_view, track_candidates_view, - track_states_view, param_liveness_view); - } -}; - -/// Alpaka kernel functor for @c traccc::device::fit_forward -template -struct fit_forward { - template - ALPAKA_FN_ACC void operator()( - TAcc const& acc, const typename fitter_t::config_type cfg, - const device::fit_payload* payload) const { - - const device::global_index_t globalThreadIdx = - ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; - device::fit_forward(globalThreadIdx, cfg, *payload); - } -}; - -/// Alpaka kernel functor for @c traccc::device::fit_backward -template -struct fit_backward { - template - ALPAKA_FN_ACC void operator()( - TAcc const& acc, const typename fitter_t::config_type cfg, - const device::fit_payload* payload) const { - - const device::global_index_t globalThreadIdx = - ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; - device::fit_backward(globalThreadIdx, cfg, *payload); - } -}; - -} // namespace kernels - -/// Templated implementation of the Alpaka track fitting algorithm. -/// -/// @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 edm::track_container::buffer -kalman_fitting( - const typename detector_t::const_view_type& det_view, - const bfield_t& field_view, - const typename edm::track_container< - typename detector_t::algebra_type>::const_view& track_candidates_view, - const fitting_config& config, const memory_resource& mr, - const vecmem::copy& copy, Queue& queue) { - - // Number of threads per block to use. - const Idx threadsPerBlock = getWarpSize() * 2; - - // Get the number of tracks. - const unsigned int n_tracks = copy.get_size(track_candidates_view.tracks); - - // Get the sizes of the track candidates in each track. - const std::vector 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_container::buffer - track_states_buffer{ - {candidate_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}, - {n_states, mr.main, vecmem::data::buffer_type::resizable}, - track_candidates_view.measurements}; - 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 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.surface_sequence_size_factor, - config.min_surface_sequence_capacity); - }); - vecmem::data::jagged_vector_buffer - seqs_buffer{seqs_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}; - copy.setup(seqs_buffer)->wait(); - - // Create the buffers for sorting the parameter IDs. - vecmem::data::vector_buffer keys_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_ids_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_liveness_buffer(n_tracks, - mr.main); - vecmem::copy::event_type keys_setup_event = copy.setup(keys_buffer); - vecmem::copy::event_type param_ids_setup_event = - copy.setup(param_ids_buffer); - vecmem::copy::event_type param_liveness_setup_event = - copy.setup(param_liveness_buffer); - keys_setup_event->wait(); - param_ids_setup_event->wait(); - param_liveness_setup_event->wait(); - - // The execution range for the two kernels of the function. - const Idx blocksPerGrid = - (n_tracks + threadsPerBlock - 1) / threadsPerBlock; - const auto workDiv = makeWorkDiv(blocksPerGrid, threadsPerBlock); - - // Fill the keys and param_ids buffers. - ::alpaka::exec(queue, workDiv, kernels::fill_fitting_sort_keys{}, - track_candidates_view.tracks, - vecmem::get_data(keys_buffer), - vecmem::get_data(param_ids_buffer)); - ::alpaka::wait(queue); - - // Sort the key to get the sorted parameter ids - vecmem::device_vector keys_device(keys_buffer); - vecmem::device_vector param_ids_device(param_ids_buffer); - 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_container::view - track_states_view{track_states_buffer}; - tracks_setup_event->wait(); - track_states_setup_event->wait(); - - ::alpaka::exec(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; - device::fit_payload 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, - .surfaces_view = seqs_buffer}; - // Now copy it to device memory. - vecmem::data::vector_buffer> device_payload( - 1u, mr.main); - copy.setup(device_payload)->wait(); - copy(vecmem::data::vector_view>( - 1u, &host_payload), - device_payload) - ->wait(); - - for (std::size_t i = 0; i < config.n_iterations; ++i) { - // Run the track fitting - ::alpaka::exec(queue, workDiv, kernels::fit_forward{}, - config, device_payload.ptr()); - ::alpaka::wait(queue); - ::alpaka::exec(queue, workDiv, kernels::fit_backward{}, - config, device_payload.ptr()); - ::alpaka::wait(queue); - } - - // Return the fitted tracks. - return track_states_buffer; -} - -} // namespace traccc::alpaka::details diff --git a/device/alpaka/src/fitting/kalman_fitting_algorithm.cpp b/device/alpaka/src/fitting/kalman_fitting_algorithm.cpp index 3a8316e359..b61cbb900e 100644 --- a/device/alpaka/src/fitting/kalman_fitting_algorithm.cpp +++ b/device/alpaka/src/fitting/kalman_fitting_algorithm.cpp @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -10,37 +10,210 @@ #include "../utils/get_queue.hpp" #include "../utils/magnetic_field_types.hpp" -#include "kalman_fitting.hpp" -#include "traccc/alpaka/fitting/kalman_fitting_algorithm.hpp" -#include "traccc/bfield/magnetic_field_types.hpp" +#include "../utils/parallel_algorithms.hpp" +#include "../utils/utils.hpp" + +// Project include(s). +#include "traccc/fitting/details/kalman_fitting_types.hpp" +#include "traccc/fitting/device/fill_fitting_sort_keys.hpp" +#include "traccc/fitting/device/fit_backward.hpp" +#include "traccc/fitting/device/fit_forward.hpp" +#include "traccc/fitting/device/fit_prelude.hpp" #include "traccc/utils/detector_buffer_bfield_visitor.hpp" namespace traccc::alpaka { +namespace kernels { + +/// Alpaka kernel functor for @c traccc::device::fill_fitting_sort_keys +struct fill_fitting_sort_keys { + template + ALPAKA_FN_ACC void operator()( + TAcc const& acc, + edm::track_collection::const_view + track_candidates_view, + vecmem::data::vector_view keys_view, + vecmem::data::vector_view 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, + keys_view, ids_view); + } +}; + +/// Alpaka kernel functor for @c traccc::device::fit_prelude +struct fit_prelude { + template + ALPAKA_FN_ACC void operator()(TAcc const& acc, + device::fit_prelude_payload payload) const { + + const device::global_index_t globalThreadIdx = + ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; + device::fit_prelude(globalThreadIdx, payload); + } +}; + +/// Alpaka kernel functor for @c traccc::device::fit_forward +template +struct fit_forward { + template + ALPAKA_FN_ACC void operator()( + TAcc const& acc, const typename fitter_t::config_type cfg, + const device::fit_payload payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, typename fitter_t::surface_type>* + tpayload) const { + + const device::global_index_t globalThreadIdx = + ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; + device::fit_forward(globalThreadIdx, cfg, payload, *tpayload); + } +}; + +/// Alpaka kernel functor for @c traccc::device::fit_backward +template +struct fit_backward { + template + ALPAKA_FN_ACC void operator()( + TAcc const& acc, const typename fitter_t::config_type cfg, + const device::fit_payload payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, typename fitter_t::surface_type>* + tpayload) const { + + const device::global_index_t globalThreadIdx = + ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0]; + device::fit_backward(globalThreadIdx, cfg, payload, + *tpayload); + } +}; + +} // namespace kernels kalman_fitting_algorithm::kalman_fitting_algorithm( const config_type& config, const traccc::memory_resource& mr, - const vecmem::copy& copy, queue& q, std::unique_ptr logger) - : messaging(std::move(logger)), - m_config{config}, - m_mr{mr}, - m_copy{copy}, - m_queue{q} {} - -kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()( - const detector_buffer& det, const magnetic_field& bfield, - const edm::track_container::const_view& track_candidates) - const { - - // Run the track fitting. + const vecmem::copy& copy, alpaka::queue& q, + std::unique_ptr logger) + : device::kalman_fitting_algorithm{config, mr, copy, std::move(logger)}, + alpaka::algorithm_base{q} {} + +void kalman_fitting_algorithm::prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const { + + // Get the number of tracks. + const unsigned int n_tracks = tracks.capacity(); + assert(n_tracks == copy().get_size(tracks)); + assert(n_tracks == track_indices.capacity()); + assert(track_indices.size_ptr() == nullptr); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + auto workDiv = makeWorkDiv(nBlocks, nThreads); + + // Fill the keys and indices buffers. + ::alpaka::exec(details::get_queue(queue()), workDiv, + kernels::fill_fitting_sort_keys{}, tracks, + track_sort_keys, track_indices); + + // Sort the key to get the sorted parameter ids + vecmem::device_vector keys_device(track_sort_keys); + vecmem::device_vector track_indices_device(track_indices); + details::sort_by_key(details::get_queue(queue()), mr(), keys_device.begin(), + keys_device.end(), track_indices_device.begin()); +} + +void kalman_fitting_algorithm::fit_prelude_kernel( + const device::fit_prelude_payload& payload) const { + + // Get the number of tracks. + const unsigned int n_tracks = payload.input_tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.input_tracks.tracks)); + assert(n_tracks == payload.track_indices.capacity()); + assert(payload.track_indices.size_ptr() == nullptr); + assert(n_tracks == copy().get_size(payload.output_tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + auto workDiv = makeWorkDiv(nBlocks, nThreads); + + // Run the fitting, using the sorted parameter IDs. + ::alpaka::exec(details::get_queue(queue()), workDiv, + kernels::fit_prelude{}, payload); +} + +auto kalman_fitting_algorithm::prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const -> fit_payload { + + return prepare_fit_payload_helper>( + det, field, n_surfaces, payload); +} + +void kalman_fitting_algorithm::fit_forward_kernel( + const fitting_config& config, const fit_payload& payload) const { + return detector_buffer_magnetic_field_visitor< detector_type_list, alpaka::bfield_type_list>( - det, bfield, - [&]( - const typename detector_t::view& detector, - const bfield_view_t& field) { - return details::kalman_fitting( - detector, field, track_candidates, m_config, m_mr, m_copy.get(), - details::get_queue(m_queue.get())); + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + auto workDiv = makeWorkDiv(nBlocks, nThreads); + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + ::alpaka::exec(details::get_queue(queue()), workDiv, + kernels::fit_forward{}, config, + payload.payload, + payload.get_tpayload()); + }); +} + +void kalman_fitting_algorithm::fit_backward_kernel( + const fitting_config& config, const fit_payload& payload) const { + + return detector_buffer_magnetic_field_visitor< + detector_type_list, alpaka::bfield_type_list>( + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + auto workDiv = makeWorkDiv(nBlocks, nThreads); + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + ::alpaka::exec(details::get_queue(queue()), workDiv, + kernels::fit_backward{}, config, + payload.payload, + payload.get_tpayload()); }); } diff --git a/device/common/CMakeLists.txt b/device/common/CMakeLists.txt index 0ce7ea5311..4d5a849c01 100644 --- a/device/common/CMakeLists.txt +++ b/device/common/CMakeLists.txt @@ -132,10 +132,18 @@ traccc_add_library( traccc_device_common device_common "include/traccc/finding/device/impl/update_tip_length_buffer.ipp" "include/traccc/finding/device/combinatorial_kalman_filter_algorithm.hpp" "src/finding/combinatorial_kalman_filter_algorithm.cpp" - # Track fitting funtions(s). - "include/traccc/fitting/device/fit.hpp" + # Track fitting code. + "include/traccc/fitting/device/fit_payload.hpp" + "include/traccc/fitting/device/fit_prelude.hpp" + "include/traccc/fitting/device/impl/fit_prelude.ipp" + "include/traccc/fitting/device/fit_forward.hpp" + "include/traccc/fitting/device/impl/fit_forward.ipp" + "include/traccc/fitting/device/fit_backward.hpp" + "include/traccc/fitting/device/impl/fit_backward.ipp" "include/traccc/fitting/device/fill_fitting_sort_keys.hpp" "include/traccc/fitting/device/impl/fill_fitting_sort_keys.ipp" + "include/traccc/fitting/device/kalman_fitting_algorithm.hpp" + "src/fitting/kalman_fitting_algorithm.cpp" ) target_link_libraries( traccc_device_common PUBLIC traccc::core vecmem::core ) diff --git a/device/common/include/traccc/fitting/device/fit_backward.hpp b/device/common/include/traccc/fitting/device/fit_backward.hpp index 1f86664d09..f289145c55 100644 --- a/device/common/include/traccc/fitting/device/fit_backward.hpp +++ b/device/common/include/traccc/fitting/device/fit_backward.hpp @@ -1,64 +1,28 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -#include "traccc/fitting/device/fit.hpp" -#include "traccc/fitting/status_codes.hpp" +// Local include(s). +#include "traccc/device/global_index.hpp" +#include "traccc/fitting/device/fit_payload.hpp" namespace traccc::device { +/// Function performing a backward fit iteration template TRACCC_HOST_DEVICE inline void fit_backward( - const global_index_t globalIndex, const typename fitter_t::config_type cfg, - const fit_payload& payload) { - - typename fitter_t::detector_type det(payload.det_data); - - vecmem::device_vector param_ids(payload.param_ids_view); - vecmem::device_vector param_liveness( - payload.param_liveness_view); - typename edm::track_container< - typename fitter_t::detector_type::algebra_type>::device - tracks(payload.tracks_view); - - if (globalIndex >= tracks.tracks.size()) { - return; - } - - const unsigned int param_id = param_ids.at(globalIndex); - edm::track track = tracks.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, tracks.states, tracks.measurements, - *(payload.surfaces_view.ptr() + param_id), - fitter.config().propagation, fitter.config().meas_calibration); - - kalman_fitter_status fit_status = fitter.smooth(fitter_state); - - fitter.update_statistics(fitter_state); - - // Assume that this branch is only called if the forward fit was - // successfull (track param are alive) - fitter.check_fitting_result(fitter_state, kalman_fitter_status::SUCCESS, - fit_status); - - if (fit_status == kalman_fitter_status::SUCCESS) { - track = fitter_state.m_fit_res; - } else { - param_liveness.at(param_id) = 0u; - } - - // TODO: Grab the smoothed state for next it - } -} + const global_index_t globalIndex, const typename fitter_t::config_type& cfg, + const fit_payload& payload, + const fit_tpayload& tpayload); } // namespace traccc::device + +// Include the implementation. +#include "traccc/fitting/device/impl/fit_backward.ipp" diff --git a/device/common/include/traccc/fitting/device/fit_forward.hpp b/device/common/include/traccc/fitting/device/fit_forward.hpp index cc741e2b62..c81a6e21ab 100644 --- a/device/common/include/traccc/fitting/device/fit_forward.hpp +++ b/device/common/include/traccc/fitting/device/fit_forward.hpp @@ -1,58 +1,28 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -#include "traccc/fitting/device/fit.hpp" -#include "traccc/fitting/status_codes.hpp" +// Local include(s). +#include "traccc/device/global_index.hpp" +#include "traccc/fitting/device/fit_payload.hpp" namespace traccc::device { +/// Function performing a forward fit iteration template TRACCC_HOST_DEVICE inline void fit_forward( - const global_index_t globalIndex, const typename fitter_t::config_type cfg, - const fit_payload& payload) { - - typename fitter_t::detector_type det(payload.det_data); - - vecmem::device_vector param_ids(payload.param_ids_view); - vecmem::device_vector param_liveness( - payload.param_liveness_view); - typename edm::track_container< - typename fitter_t::detector_type::algebra_type>::device - tracks(payload.tracks_view); - - if (globalIndex >= tracks.tracks.size()) { - return; - } - - const unsigned int param_id = param_ids.at(globalIndex); - - fitter_t fitter(det, payload.field_data, cfg); - - edm::track track = tracks.tracks.at(param_id); - bound_track_parameters<> params = track.params(); - - // TODO: Merge into filter? - inflate_covariance(params, fitter.config().covariance_inflation_factor); - - typename fitter_t::state fitter_state( - track, tracks.states, tracks.measurements, - *(payload.surfaces_view.ptr() + param_id), fitter.config().propagation, - fitter.config().meas_calibration); - - kalman_fitter_status fit_status = fitter.filter(params, fitter_state); - - fitter.check_fitting_result(fitter_state, fit_status, - kalman_fitter_status::SUCCESS); - - if (fit_status != kalman_fitter_status::SUCCESS) { - param_liveness.at(param_id) = 0u; - } -} + const global_index_t globalIndex, const typename fitter_t::config_type& cfg, + const fit_payload& payload, + const fit_tpayload& tpayload); } // namespace traccc::device + +// Include the implementation. +#include "traccc/fitting/device/impl/fit_forward.ipp" diff --git a/device/common/include/traccc/fitting/device/fit.hpp b/device/common/include/traccc/fitting/device/fit_payload.hpp similarity index 55% rename from device/common/include/traccc/fitting/device/fit.hpp rename to device/common/include/traccc/fitting/device/fit_payload.hpp index bdd25b2fb0..0d883c1dde 100644 --- a/device/common/include/traccc/fitting/device/fit.hpp +++ b/device/common/include/traccc/fitting/device/fit_payload.hpp @@ -1,17 +1,13 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -// Local include(s). -#include "traccc/device/global_index.hpp" - // Project include(s). -#include "traccc/definitions/qualifiers.hpp" #include "traccc/edm/track_container.hpp" // VecMem include(s). @@ -20,40 +16,43 @@ namespace traccc::device { -// Payload for the fitting algorithm -template +/// (Non-Templated) Payload for the fitting function(s) + struct fit_payload { /** - * @brief View object to the detector description + * @brief View object to the input track parameters */ - typename fitter_t::detector_type::const_view_type det_data; + vecmem::data::vector_view track_indices; /** - * @brief View object to the magnetic field description + * @brief View object to the vector of parameter liveness */ - typename fitter_t::bfield_type field_data; + vecmem::data::vector_view track_liveness; /** - * @brief View object to the input track parameters + * @brief View object to the output tracks */ - vecmem::data::vector_view param_ids_view; + edm::track_container::view tracks; +}; + +/// (Templated) Payload for the fitting function(s) +template +struct fit_tpayload { /** - * @brief View object to the vector of parameter liveness + * @brief View object to the detector description */ - vecmem::data::vector_view param_liveness_view; + detector_t det; /** - * @brief View object to the output tracks + * @brief View object to the magnetic field description */ - typename edm::track_container< - typename fitter_t::detector_type::algebra_type>::view tracks_view; + bfield_t field; /** * @brief View object to the output geometry identifer sequence */ - vecmem::data::jagged_vector_view - surfaces_view; + vecmem::data::jagged_vector_view surfaces; }; } // namespace traccc::device diff --git a/device/common/include/traccc/fitting/device/fit_prelude.hpp b/device/common/include/traccc/fitting/device/fit_prelude.hpp index a27d813920..0d58c4ddd5 100644 --- a/device/common/include/traccc/fitting/device/fit_prelude.hpp +++ b/device/common/include/traccc/fitting/device/fit_prelude.hpp @@ -12,58 +12,31 @@ // Project include(s). #include "traccc/edm/track_container.hpp" -#include "traccc/edm/track_state_helpers.hpp" -#include "traccc/fitting/status_codes.hpp" // VecMem include(s). #include namespace traccc::device { -template -TRACCC_HOST_DEVICE inline void fit_prelude( - const global_index_t globalIndex, - vecmem::data::vector_view param_ids_view, - typename edm::track_container::const_view track_candidates_view, - typename edm::track_container::view tracks_view, - vecmem::data::vector_view param_liveness_view) { +/// Payload for the @c traccc::device::fit_prelude function +struct fit_prelude_payload { - typename edm::track_collection::device tracks( - tracks_view.tracks); + /// Input track parameter IDs + vecmem::data::vector_view track_indices; + /// Input tracks + edm::track_container::const_view input_tracks; + /// Output tracks + edm::track_container::view output_tracks; + /// Output track liveness + vecmem::data::vector_view track_liveness; - if (globalIndex >= tracks.size()) { - return; - } +}; // struct fit_prelude_payload - typename edm::track_state_collection::device track_states( - tracks_view.states); - - vecmem::device_vector param_ids(param_ids_view); - vecmem::device_vector param_liveness(param_liveness_view); - - const unsigned int param_id = param_ids.at(globalIndex); - - edm::track track = tracks.at(param_id); - - const typename edm::track_collection::const_device - track_candidates{track_candidates_view.tracks}; - const edm::track track_candidate = track_candidates.at(param_id); - const auto track_candidate_constituent_links = - track_candidate.constituent_links(); - const edm::measurement_collection::const_device measurements{ - track_candidates_view.measurements}; - for (const edm::track_constituent_link& link : - track_candidate_constituent_links) { - assert(link.type == edm::track_constituent_link::measurement); - const unsigned int track_state_index = track_states.push_back( - edm::make_track_state(measurements, link.index)); - track.constituent_links().push_back( - {edm::track_constituent_link::track_state, track_state_index}); - } - - // TODO: Set other stuff in the header? - track.params() = track_candidate.params(); - param_liveness.at(param_id) = 1u; -} +/// Function to prepare the fitting payloads for the fitting algorithm +TRACCC_HOST_DEVICE inline void fit_prelude(const global_index_t globalIndex, + const fit_prelude_payload& payload); } // namespace traccc::device + +// Include the implementation. +#include "traccc/fitting/device/impl/fit_prelude.ipp" diff --git a/device/common/include/traccc/fitting/device/impl/fit_backward.ipp b/device/common/include/traccc/fitting/device/impl/fit_backward.ipp new file mode 100644 index 0000000000..9a1d8e564f --- /dev/null +++ b/device/common/include/traccc/fitting/device/impl/fit_backward.ipp @@ -0,0 +1,66 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +#pragma once + +// Project include(s). +#include "traccc/fitting/status_codes.hpp" + +namespace traccc::device { + +template +TRACCC_HOST_DEVICE inline void fit_backward( + const global_index_t globalIndex, const typename fitter_t::config_type& cfg, + const fit_payload& payload, + const fit_tpayload& tpayload) { + + typename fitter_t::detector_type det(tpayload.det); + + vecmem::device_vector param_ids(payload.track_indices); + vecmem::device_vector param_liveness(payload.track_liveness); + typename edm::track_container< + typename fitter_t::detector_type::algebra_type>::device + tracks(payload.tracks); + + if (globalIndex >= tracks.tracks.size()) { + return; + } + + const unsigned int param_id = param_ids.at(globalIndex); + edm::track track = tracks.tracks.at(param_id); + + // Run fitting + fitter_t fitter(det, tpayload.field, cfg); + + if (param_liveness.at(param_id) > 0u) { + typename fitter_t::state fitter_state( + track, tracks.states, tracks.measurements, + *(tpayload.surfaces.ptr() + param_id), fitter.config().propagation, + fitter.config().meas_calibration); + + kalman_fitter_status fit_status = fitter.smooth(fitter_state); + + fitter.update_statistics(fitter_state); + + // Assume that this branch is only called if the forward fit was + // successfull (track param are alive) + fitter.check_fitting_result(fitter_state, kalman_fitter_status::SUCCESS, + fit_status); + + if (fit_status == kalman_fitter_status::SUCCESS) { + track = fitter_state.m_fit_res; + } else { + param_liveness.at(param_id) = 0u; + } + + // TODO: Grab the smoothed state for next it + } +} + +} // namespace traccc::device diff --git a/device/common/include/traccc/fitting/device/impl/fit_forward.ipp b/device/common/include/traccc/fitting/device/impl/fit_forward.ipp new file mode 100644 index 0000000000..d38cc8e145 --- /dev/null +++ b/device/common/include/traccc/fitting/device/impl/fit_forward.ipp @@ -0,0 +1,60 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +#pragma once + +// Project include(s). +#include "traccc/fitting/status_codes.hpp" + +namespace traccc::device { + +template +TRACCC_HOST_DEVICE inline void fit_forward( + const global_index_t globalIndex, const typename fitter_t::config_type& cfg, + const fit_payload& payload, + const fit_tpayload& tpayload) { + + typename fitter_t::detector_type det(tpayload.det); + + vecmem::device_vector param_ids(payload.track_indices); + vecmem::device_vector param_liveness(payload.track_liveness); + typename edm::track_container< + typename fitter_t::detector_type::algebra_type>::device + tracks(payload.tracks); + + if (globalIndex >= tracks.tracks.size()) { + return; + } + + const unsigned int param_id = param_ids.at(globalIndex); + + fitter_t fitter(det, tpayload.field, cfg); + + edm::track track = tracks.tracks.at(param_id); + bound_track_parameters<> params = track.params(); + + // TODO: Merge into filter? + inflate_covariance(params, fitter.config().covariance_inflation_factor); + + typename fitter_t::state fitter_state( + track, tracks.states, tracks.measurements, + *(tpayload.surfaces.ptr() + param_id), fitter.config().propagation, + fitter.config().meas_calibration); + + kalman_fitter_status fit_status = fitter.filter(params, fitter_state); + + fitter.check_fitting_result(fitter_state, fit_status, + kalman_fitter_status::SUCCESS); + + if (fit_status != kalman_fitter_status::SUCCESS) { + param_liveness.at(param_id) = 0u; + } +} + +} // namespace traccc::device diff --git a/device/common/include/traccc/fitting/device/impl/fit_prelude.ipp b/device/common/include/traccc/fitting/device/impl/fit_prelude.ipp new file mode 100644 index 0000000000..7e0f16e6c7 --- /dev/null +++ b/device/common/include/traccc/fitting/device/impl/fit_prelude.ipp @@ -0,0 +1,58 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +#pragma once + +// Project include(s). +#include "traccc/edm/track_state_helpers.hpp" +#include "traccc/fitting/status_codes.hpp" + +namespace traccc::device { + +TRACCC_HOST_DEVICE inline void fit_prelude(const global_index_t globalIndex, + const fit_prelude_payload& payload) { + + edm::track_collection::device tracks( + payload.output_tracks.tracks); + + if (globalIndex >= tracks.size()) { + return; + } + + edm::track_state_collection::device track_states( + payload.output_tracks.states); + + vecmem::device_vector track_indices( + payload.track_indices); + vecmem::device_vector track_liveness(payload.track_liveness); + + const unsigned int track_id = track_indices.at(globalIndex); + + edm::track track = tracks.at(track_id); + + const edm::track_collection::const_device track_candidates{ + payload.input_tracks.tracks}; + const edm::track track_candidate = track_candidates.at(track_id); + const auto track_candidate_constituent_links = + track_candidate.constituent_links(); + const edm::measurement_collection::const_device measurements{ + payload.input_tracks.measurements}; + for (const edm::track_constituent_link& link : + track_candidate_constituent_links) { + assert(link.type == edm::track_constituent_link::measurement); + const unsigned int track_state_index = track_states.push_back( + edm::make_track_state(measurements, link.index)); + track.constituent_links().push_back( + {edm::track_constituent_link::track_state, track_state_index}); + } + + // TODO: Set other stuff in the header? + track.params() = track_candidate.params(); + track_liveness.at(track_id) = 1u; +} + +} // namespace traccc::device diff --git a/device/common/include/traccc/fitting/device/impl/kalman_fitting_algorithm.ipp b/device/common/include/traccc/fitting/device/impl/kalman_fitting_algorithm.ipp new file mode 100644 index 0000000000..31d4e7c4ed --- /dev/null +++ b/device/common/include/traccc/fitting/device/impl/kalman_fitting_algorithm.ipp @@ -0,0 +1,86 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +#pragma once + +// System include(s). +#include + +namespace traccc::device { + +template +const device::fit_tpayload* +kalman_fitting_algorithm::fit_payload::get_tpayload() const { + + return device_tpayload + .as>>() + .ptr(); +} + +template +auto kalman_fitting_algorithm::prepare_fit_payload_helper( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const -> fit_payload { + + return detector_buffer_magnetic_field_visitor( + det, field, + [&]( + const typename detector_traits_t::view& detector, + const bfield_view_t& bfield) -> fit_payload { + // Create the surface buffer used during the fitting. + vecmem::data::jagged_vector_buffer< + typename detector_traits_t::device::surface_type> + surfaces{n_surfaces, mr().main, mr().host, + vecmem::data::buffer_type::resizable}; + copy().setup(surfaces)->ignore(); + + // Create the (templated) host payload. + device::fit_tpayload< + typename detector_traits_t::device::const_view_type, + bfield_view_t, typename detector_traits_t::device::surface_type> + host_tpayload{ + .det = detector, .field = bfield, .surfaces = surfaces}; + + // Create the (templated) device payload buffer, and copy the host + // payload into it. + vecmem::data::vector_buffer> + device_tpayload{1u, mr().main}; + copy().setup(device_tpayload)->ignore(); + copy()(vecmem::data::vector_view>( + 1u, &host_tpayload), + device_tpayload) + ->ignore(); + + // Create the result payload object. + fit_payload result{det, field}; + + // Save the (non-templated) host payload. + result.payload = payload; + + // Save all the type erased payloads into it. + result.surfaces.set(std::move(surfaces)); + result.host_tpayload = host_tpayload; + result.device_tpayload.set(std::move(device_tpayload)); + + // All done, we can return the created payload. + return result; + }); +} + +} // namespace traccc::device diff --git a/device/common/include/traccc/fitting/device/kalman_fitting_algorithm.hpp b/device/common/include/traccc/fitting/device/kalman_fitting_algorithm.hpp new file mode 100644 index 0000000000..ae7cfaf87c --- /dev/null +++ b/device/common/include/traccc/fitting/device/kalman_fitting_algorithm.hpp @@ -0,0 +1,197 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +#pragma once + +// Local include(s). +#include "traccc/device/algorithm_base.hpp" +#include "traccc/fitting/device/fit_payload.hpp" + +// Project include(s). +#include "traccc/bfield/magnetic_field.hpp" +#include "traccc/edm/device/sort_key.hpp" +#include "traccc/edm/track_container.hpp" +#include "traccc/fitting/details/kalman_fitting_types.hpp" +#include "traccc/fitting/device/fit_prelude.hpp" +#include "traccc/fitting/fitting_config.hpp" +#include "traccc/geometry/detector_buffer.hpp" +#include "traccc/geometry/move_only_any.hpp" +#include "traccc/utils/algorithm.hpp" +#include "traccc/utils/memory_resource.hpp" +#include "traccc/utils/messaging.hpp" + +// System include(s). +#include + +namespace traccc::device { + +/// Kalman filter based track fitting algorithm +class kalman_fitting_algorithm + : public algorithm::buffer( + const detector_buffer&, const magnetic_field&, + const edm::track_container::const_view&)>, + public messaging, + public algorithm_base { + + public: + /// Configuration type + using config_type = fitting_config; + + /// Constructor with the algorithm's configuration + /// + /// @param config The configuration object + /// @param mr The memory resource(s) used by the algorithm + /// @param copy The copy object used by the algorithm + /// @param logger The logger used by the algorithm + /// + kalman_fitting_algorithm( + const config_type& config, const traccc::memory_resource& mr, + const vecmem::copy& copy, + std::unique_ptr logger = getDummyLogger().clone()); + /// Destructor + virtual ~kalman_fitting_algorithm(); + + /// Operator executing the algorithm. + /// + /// @param det The detector object + /// @param field 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& field, + const edm::track_container::const_view& + track_candidates) const override; + + protected: + /// @name Type(s)/function(s) used internally by the algorithm + /// @{ + + /// Payload of the fitting kernels + struct fit_payload { + + /// Constructor + fit_payload(const detector_buffer& det, const magnetic_field& field); + + /// Helper function to get the device resident, templated payload + template + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, typename fitter_t::surface_type>* + get_tpayload() const; + + /// Polymorphic tracking geometry buffer + const detector_buffer& detector; + /// Polymorphic magnetic field object + const magnetic_field& field; + + /// The host-resident, non-templated payload + device::fit_payload payload; + + /// Surface buffer used during the fitting + move_only_any surfaces; + + /// The host-resident, templated payload + std::any host_tpayload; + /// The device-resident, templated payload (buffer) + move_only_any device_tpayload; + }; + + /// Prepare a detector+bfield specific payload for the fitting kernel(s) + /// + /// Function to be used by the specific @c prepare_fit_payload functions + /// for preparing the payload. Since apart from different template types, + /// they all work the same way. + /// + /// @tparam detector_list_t The list of supported detector types to use for + /// the visitor + /// @tparam bfield_list_t The list of supported b-field types to use for + /// the visitor + /// + /// @param det The detector buffer to prepare the payload for + /// @param field The magnetic field to prepare the payload for + /// @param n_surfaces The number of surfaces for each track to be + /// fitted + /// @param payload The (non-templated) payload for the kernel(s) + /// + /// @return The prepared payload for the fitting kernel(s) + /// + template + fit_payload prepare_fit_payload_helper( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const; + + /// @} + + /// @name Function(s) to be implemented by derived classes + /// @{ + + /// Prepare a buffer with the index order with which to fit the tracks + /// + /// @param[in] tracks The tracks to be fitted + /// @param[out] track_sort_keys Buffer storing temporary sorting keys + /// @param[out] track_indices The buffer to write the fitting order into + /// + virtual void prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const = 0; + + /// Kernel to prepare the fitting payloads + /// + /// @param payload The payload for the kernel(s) + /// + virtual void fit_prelude_kernel( + const fit_prelude_payload& payload) const = 0; + + /// Function preparing the fitting payload + /// + /// @param det The detector buffer to prepare the payload for + /// @param field The magnetic field to prepare the payload for + /// @param n_surfaces The number of surfaces for each track to be + /// fitted + /// @param payload The (non-templated) payload for the kernel(s) + /// + /// @return The prepared payload for the fitting kernel(s) + /// + virtual fit_payload prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const = 0; + + /// Function launching the "forward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + virtual void fit_forward_kernel(const fitting_config& config, + const fit_payload& payload) const = 0; + + /// Function launching the "backward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + virtual void fit_backward_kernel(const fitting_config& config, + const fit_payload& payload) const = 0; + + /// @} + + private: + /// Internal data type + struct data; + /// Pointer to internal data + std::unique_ptr m_data; + +}; // class kalman_fitting_algorithm + +} // namespace traccc::device + +// Include the implementation. +#include "traccc/fitting/device/impl/kalman_fitting_algorithm.ipp" diff --git a/device/common/src/fitting/kalman_fitting_algorithm.cpp b/device/common/src/fitting/kalman_fitting_algorithm.cpp new file mode 100644 index 0000000000..31a74b0692 --- /dev/null +++ b/device/common/src/fitting/kalman_fitting_algorithm.cpp @@ -0,0 +1,128 @@ +/** TRACCC library, part of the ACTS project (R&D line) + * + * (c) 2022-2026 CERN for the benefit of the ACTS project + * + * Mozilla Public License Version 2.0 + */ + +// Local include(s). +#include "traccc/fitting/device/kalman_fitting_algorithm.hpp" + +// VecMem include(s). +#include + +// System include(s). +#include +#include + +namespace traccc::device { + +struct kalman_fitting_algorithm::data { + + /// @name Configuration object(s) + /// @{ + + /// Configuration for the fitting algorithm + fitting_config m_config; + + /// @} + +}; // struct kalman_fitting_algorithm::data + +kalman_fitting_algorithm::fit_payload::fit_payload(const detector_buffer& det, + const magnetic_field& f) + : detector(det), field(f) {} + +kalman_fitting_algorithm::kalman_fitting_algorithm( + const config_type& config, const traccc::memory_resource& mr, + const vecmem::copy& copy, std::unique_ptr logger) + : messaging(std::move(logger)), + algorithm_base{mr, copy}, + m_data{std::make_unique(config)} {} + +kalman_fitting_algorithm::~kalman_fitting_algorithm() = default; + +kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()( + const detector_buffer& det, const magnetic_field& field, + const edm::track_container::const_view& input_tracks) + const { + + // Get the number of tracks and the number of constituens (measurements) + // in each track. Note that the number of tracks is not "resizable". That + // we can just get directly from the view object. But the number of + // constituents need to be taken from the buffer itself. + assert(copy().get_size(input_tracks.tracks) == + input_tracks.tracks.capacity()); + const edm::track_collection::const_view::size_type + n_tracks = input_tracks.tracks.capacity(); + if (n_tracks == 0) { + // Return early, if there are no tracks. + return {}; + } + std::vector candidate_sizes; + if (mr().host) { + vecmem::async_sizes sizes = + copy().get_sizes(input_tracks.tracks, *(mr().host)); + // Here we could give control back to the caller, once our code allows + // for it. (coroutines...) + auto& temp = sizes.get(); + candidate_sizes = {temp.begin(), temp.end()}; + } else { + candidate_sizes = copy().get_sizes(input_tracks.tracks); + } + + // Get the total number of states (measurements) to fit. + const unsigned int n_states = + std::accumulate(candidate_sizes.begin(), candidate_sizes.end(), 0u); + + // Create the result buffer. + edm::track_container::buffer output_tracks{ + {candidate_sizes, mr().main, mr().host, + vecmem::data::buffer_type::resizable}, + {n_states, mr().main, vecmem::data::buffer_type::resizable}, + input_tracks.measurements}; + copy().setup(output_tracks.tracks)->ignore(); + copy().setup(output_tracks.states)->ignore(); + + // Create the order to fit the tracks in. + vecmem::data::vector_buffer track_sort_keys(n_tracks, + mr().main); + vecmem::data::vector_buffer track_indices{n_tracks, + mr().main}; + prepare_track_fit_order(input_tracks.tracks, track_sort_keys, + track_indices); + + // Create the buffer(s) used during the fitting. + vecmem::data::vector_buffer track_liveness(n_tracks, + mr().main); + + // Run "fitting prelude" kernel. + fit_prelude_kernel( + {track_indices, input_tracks, output_tracks, track_liveness}); + + // Calculate the number of surfaces to use during the fit for each track. + // Then create the concrete buffer such that it could be passed to the + // fitting functions through a polymorphic pointer. + std::vector n_surfaces(candidate_sizes.size()); + std::transform(candidate_sizes.begin(), candidate_sizes.end(), + n_surfaces.begin(), [&](const unsigned int sz) { + return std::max( + sz * m_data->m_config.surface_sequence_size_factor, + m_data->m_config.min_surface_sequence_capacity); + }); + + // Prepare the payload for the fitting kernel(s). + const fit_payload payload = prepare_fit_payload( + det, field, n_surfaces, {track_indices, track_liveness, output_tracks}); + + // Run the iterative track fitting. + for (std::size_t i = 0; i < m_data->m_config.n_iterations; ++i) { + fit_forward_kernel(m_data->m_config, payload); + fit_backward_kernel(m_data->m_config, payload); + } + + // Return the fitted tracks. + return output_tracks; +} + +} // namespace traccc::device diff --git a/device/cuda/CMakeLists.txt b/device/cuda/CMakeLists.txt index 5c72fe09f5..32757dc2aa 100644 --- a/device/cuda/CMakeLists.txt +++ b/device/cuda/CMakeLists.txt @@ -103,7 +103,6 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED "include/traccc/cuda/fitting/kalman_fitting_algorithm.hpp" "src/fitting/kalman_fitting_algorithm.cpp" "src/fitting/kalman_fitting_algorithm.cu" - "src/fitting/kalman_fitting.cuh" "src/fitting/kernels/fill_fitting_sort_keys.cu" "src/fitting/kernels/fit_prelude.cu" ) diff --git a/device/cuda/include/traccc/cuda/fitting/kalman_fitting_algorithm.hpp b/device/cuda/include/traccc/cuda/fitting/kalman_fitting_algorithm.hpp index ac2008e426..7e55583b3a 100644 --- a/device/cuda/include/traccc/cuda/fitting/kalman_fitting_algorithm.hpp +++ b/device/cuda/include/traccc/cuda/fitting/kalman_fitting_algorithm.hpp @@ -8,37 +8,19 @@ #pragma once // Library include(s). +#include "traccc/cuda/utils/algorithm_base.hpp" #include "traccc/cuda/utils/stream_wrapper.hpp" // Project include(s). -#include "traccc/bfield/magnetic_field.hpp" -#include "traccc/edm/track_container.hpp" -#include "traccc/fitting/fitting_config.hpp" -#include "traccc/geometry/detector.hpp" -#include "traccc/geometry/detector_buffer.hpp" -#include "traccc/utils/algorithm.hpp" -#include "traccc/utils/memory_resource.hpp" -#include "traccc/utils/messaging.hpp" - -// VecMem include(s). -#include - -// System include(s). -#include +#include "traccc/fitting/device/kalman_fitting_algorithm.hpp" namespace traccc::cuda { -/// Kalman filter based track fitting algorithm -class kalman_fitting_algorithm - : public algorithm::buffer( - const detector_buffer&, const magnetic_field&, - const edm::track_container::const_view&)>, - public messaging { +/// Kalman filter based track fitting algorithm using CUDA +class kalman_fitting_algorithm : public device::kalman_fitting_algorithm, + public cuda::algorithm_base { public: - /// Configuration type - using config_type = fitting_config; - /// Constructor with the algorithm's configuration /// /// @param config The configuration object @@ -52,30 +34,60 @@ class kalman_fitting_algorithm const vecmem::copy& copy, const stream_wrapper& str, std::unique_ptr logger = getDummyLogger().clone()); - /// Execute the algorithm + private: + /// @name Function(s) implemented from @c device::kalman_fitting_algorithm + /// @{ + + /// Prepare a buffer with the index order with which to fit the tracks + /// + /// @param[in] tracks The tracks to be fitted + /// @param[out] track_sort_keys Buffer storing temporary sorting keys + /// @param[out] track_indices The buffer to write the fitting order into + /// + void prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const override; + + /// Kernel to prepare the fitting payloads + /// + /// @param payload The payload for the kernel(s) + /// + void fit_prelude_kernel( + const device::fit_prelude_payload& payload) const override; + + /// Function preparing the fitting payload /// - /// @param det The detector object - /// @param field The magnetic field object - /// @param track_candidates All track candidates to fit + /// @param det The detector buffer to prepare the payload for + /// @param field The magnetic field to prepare the payload for + /// @param n_surfaces The number of surfaces for each track to be + /// fitted + /// @param payload The (non-templated) payload for the kernel(s) /// - /// @return A container of the fitted track states + /// @return The prepared payload for the fitting kernel(s) /// - output_type operator()( + fit_payload prepare_fit_payload( const detector_buffer& det, const magnetic_field& field, - const edm::track_container::const_view& - track_candidates) const override; + const std::vector& n_surfaces, + const device::fit_payload& payload) const override; - private: - /// Algorithm configuration - config_type m_config; - /// Memory resource used by the algorithm - traccc::memory_resource m_mr; - /// Copy object used by the algorithm - std::reference_wrapper m_copy; - /// The CUDA stream to use - stream_wrapper m_stream; - /// Warp size of the GPU being used - unsigned int m_warp_size; + /// Function launching the "forward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_forward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// Function launching the "backward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_backward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// @} }; // class kalman_fitting_algorithm diff --git a/device/cuda/src/fitting/kalman_fitting.cuh b/device/cuda/src/fitting/kalman_fitting.cuh deleted file mode 100644 index 06b97305cd..0000000000 --- a/device/cuda/src/fitting/kalman_fitting.cuh +++ /dev/null @@ -1,167 +0,0 @@ -/** TRACCC library, part of the ACTS project (R&D line) - * - * (c) 2022-2025 CERN for the benefit of the ACTS project - * - * Mozilla Public License Version 2.0 - */ - -#pragma once - -// Local include(s). -#include "../utils/cuda_error_handling.hpp" -#include "../utils/global_index.hpp" -#include "../utils/utils.hpp" -#include "./kernels/fill_fitting_sort_keys.hpp" -#include "./kernels/fit_backward.hpp" -#include "./kernels/fit_forward.hpp" -#include "./kernels/fit_prelude.hpp" - -// Project include(s). -#include "traccc/edm/device/sort_key.hpp" -#include "traccc/edm/track_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" -#include "traccc/utils/memory_resource.hpp" - -// VecMem include(s). -#include - -// Thrust include(s). -#include -#include - -// System include(s). -#include - -namespace traccc::cuda::details { - -/// Templated implementation of the CUDA track fitting algorithm. -/// -/// @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 edm::track_container::buffer -kalman_fitting( - const typename detector_t::const_view_type& det_view, - const bfield_t& field_view, - const typename edm::track_container< - typename detector_t::algebra_type>::const_view& track_candidates_view, - const fitting_config& config, const memory_resource& mr, - const vecmem::copy& copy, const stream_wrapper& str, - unsigned int warp_size) { - - // Get a convenience variable for the stream that we'll be using. - cudaStream_t stream = details::get_stream(str); - - // Get the number of tracks. - const unsigned int n_tracks = copy.get_size(track_candidates_view.tracks); - - // Get the sizes of the track candidates in each track. - const std::vector 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_container::buffer - track_states_buffer{ - {candidate_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}, - {n_states, mr.main, vecmem::data::buffer_type::resizable}, - track_candidates_view.measurements}; - 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) { - return track_states_buffer; - } - - std::vector 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.surface_sequence_size_factor, - config.min_surface_sequence_capacity); - }); - vecmem::data::jagged_vector_buffer - seqs_buffer{seqs_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}; - copy.setup(seqs_buffer)->ignore(); - - // Create the buffers for sorting the parameter IDs. - vecmem::data::vector_buffer keys_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_ids_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_liveness_buffer(n_tracks, - mr.main); - vecmem::copy::event_type keys_setup_event = copy.setup(keys_buffer); - vecmem::copy::event_type param_ids_setup_event = - copy.setup(param_ids_buffer); - vecmem::copy::event_type param_liveness_setup_event = - copy.setup(param_liveness_buffer); - keys_setup_event->ignore(); - param_ids_setup_event->ignore(); - param_liveness_setup_event->ignore(); - - // Launch parameters for all the kernels. - const unsigned int nThreads = warp_size * 4; - const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; - - // Fill the keys and param_ids buffers. - fill_fitting_sort_keys(nBlocks, nThreads, stream, - track_candidates_view.tracks, keys_buffer, - param_ids_buffer); - - // Sort the key to get the sorted parameter ids - vecmem::device_vector keys_device(keys_buffer); - vecmem::device_vector param_ids_device(param_ids_buffer); - thrust::sort_by_key( - thrust::cuda::par_nosync(std::pmr::polymorphic_allocator(&mr.main)) - .on(stream), - keys_device.begin(), keys_device.end(), param_ids_device.begin()); - - // Run the fitting, using the sorted parameter IDs. - fit_prelude(nBlocks, nThreads, 0, stream, param_ids_buffer, - track_candidates_view, track_states_buffer, - param_liveness_buffer); - TRACCC_CUDA_ERROR_CHECK(cudaGetLastError()); - str.synchronize(); - - // Allocate the fitting kernels's payload in host memory. - using fitter_t = traccc::details::kalman_fitter_t; - device::fit_payload 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_buffer, - .surfaces_view = seqs_buffer}; - - for (std::size_t i = 0; i < config.n_iterations; ++i) { - // Run the track fitting - fit_forward(nBlocks, nThreads, 0, stream, config, - host_payload); - TRACCC_CUDA_ERROR_CHECK(cudaGetLastError()); - fit_backward(nBlocks, nThreads, 0, stream, config, - host_payload); - TRACCC_CUDA_ERROR_CHECK(cudaGetLastError()); - } - - // Return the fitted tracks. - return track_states_buffer; -} - -} // namespace traccc::cuda::details diff --git a/device/cuda/src/fitting/kalman_fitting_algorithm.cpp b/device/cuda/src/fitting/kalman_fitting_algorithm.cpp index db018eb764..49bb85db9b 100644 --- a/device/cuda/src/fitting/kalman_fitting_algorithm.cpp +++ b/device/cuda/src/fitting/kalman_fitting_algorithm.cpp @@ -8,7 +8,16 @@ // Local include(s). #include "traccc/cuda/fitting/kalman_fitting_algorithm.hpp" +#include "../utils/magnetic_field_types.hpp" #include "../utils/utils.hpp" +#include "./kernels/fit_backward.hpp" +#include "./kernels/fit_forward.hpp" +#include "./kernels/fit_prelude.hpp" + +// Project include(s). +#include "traccc/fitting/details/kalman_fitting_types.hpp" +#include "traccc/geometry/detector.hpp" +#include "traccc/utils/detector_buffer_bfield_visitor.hpp" namespace traccc::cuda { @@ -16,11 +25,91 @@ kalman_fitting_algorithm::kalman_fitting_algorithm( const config_type& config, const traccc::memory_resource& mr, const vecmem::copy& copy, const stream_wrapper& str, std::unique_ptr logger) - : messaging(std::move(logger)), - m_config{config}, - m_mr{mr}, - m_copy{copy}, - m_stream{str}, - m_warp_size(details::get_warp_size(str.device())) {} + : device::kalman_fitting_algorithm{config, mr, copy, std::move(logger)}, + cuda::algorithm_base{str} {} + +void kalman_fitting_algorithm::fit_prelude_kernel( + const device::fit_prelude_payload& payload) const { + + // Get the number of tracks. + const unsigned int n_tracks = payload.input_tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.input_tracks.tracks)); + assert(n_tracks == payload.track_indices.capacity()); + assert(payload.track_indices.size_ptr() == nullptr); + assert(n_tracks == copy().get_size(payload.output_tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + + // Run the fitting, using the sorted parameter IDs. + fit_prelude(nBlocks, nThreads, 0, details::get_stream(stream()), payload); +} + +auto kalman_fitting_algorithm::prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const -> fit_payload { + + return prepare_fit_payload_helper>( + det, field, n_surfaces, payload); +} + +void kalman_fitting_algorithm::fit_forward_kernel( + const fitting_config& config, const fit_payload& payload) const { + + return detector_buffer_magnetic_field_visitor< + detector_type_list, cuda::bfield_type_list>( + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + fit_forward( + nBlocks, nThreads, 0, details::get_stream(stream()), config, + payload.payload, payload.get_tpayload()); + }); +} + +void kalman_fitting_algorithm::fit_backward_kernel( + const fitting_config& config, const fit_payload& payload) const { + + return detector_buffer_magnetic_field_visitor< + detector_type_list, cuda::bfield_type_list>( + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + fit_backward( + nBlocks, nThreads, 0, details::get_stream(stream()), config, + payload.payload, payload.get_tpayload()); + }); +} } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kalman_fitting_algorithm.cu b/device/cuda/src/fitting/kalman_fitting_algorithm.cu index a0ecb2f7aa..da6514d1b7 100644 --- a/device/cuda/src/fitting/kalman_fitting_algorithm.cu +++ b/device/cuda/src/fitting/kalman_fitting_algorithm.cu @@ -1,35 +1,51 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ // Local include(s). -#include "../utils/magnetic_field_types.hpp" -#include "kalman_fitting.cuh" +#include "../utils/utils.hpp" +#include "./kernels/fill_fitting_sort_keys.hpp" #include "traccc/cuda/fitting/kalman_fitting_algorithm.hpp" -#include "traccc/geometry/detector.hpp" -#include "traccc/utils/detector_buffer_bfield_visitor.hpp" + +// Thrust include(s). +#include +#include + +// System include(s). +#include +#include namespace traccc::cuda { -kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()( - const detector_buffer& det, const magnetic_field& field, - const edm::track_container::const_view& track_candidates) - const { - - // Run the track fitting. - return detector_buffer_magnetic_field_visitor< - detector_type_list, cuda::bfield_type_list>( - det, field, - [&]( - const typename detector_t::view& detector, - const bfield_view_t& bfield) { - return details::kalman_fitting( - detector, bfield, track_candidates, m_config, m_mr, - m_copy.get(), m_stream, m_warp_size); - }); +void kalman_fitting_algorithm::prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const { + + // Get the number of tracks. + const unsigned int n_tracks = tracks.capacity(); + assert(n_tracks == copy().get_size(tracks)); + assert(n_tracks == track_indices.capacity()); + assert(track_indices.size_ptr() == nullptr); + + // Launch parameters for the kernel. + const unsigned int nThreads = warp_size() * 4; + const unsigned int nBlocks = (n_tracks + nThreads - 1) / nThreads; + + // Fill the keys and indices buffers. + fill_fitting_sort_keys(nBlocks, nThreads, details::get_stream(stream()), + tracks, track_sort_keys, track_indices); + + // Sort the key to get the sorted parameter ids + vecmem::device_vector keys_device(track_sort_keys); + vecmem::device_vector track_indices_device(track_indices); + thrust::sort_by_key( + thrust::cuda::par_nosync(std::pmr::polymorphic_allocator(&mr().main)) + .on(details::get_stream(stream())), + keys_device.begin(), keys_device.end(), track_indices_device.begin()); } } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/fit_backward.hpp b/device/cuda/src/fitting/kernels/fit_backward.hpp index 205d6e2490..77265ed6fc 100644 --- a/device/cuda/src/fitting/kernels/fit_backward.hpp +++ b/device/cuda/src/fitting/kernels/fit_backward.hpp @@ -1,23 +1,28 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -#include - -#include "traccc/fitting/device/fit.hpp" +// Project include(s). +#include "traccc/fitting/device/fit_payload.hpp" #include "traccc/fitting/fitting_config.hpp" +// CUDA include(s). +#include + namespace traccc::cuda { template void fit_backward(const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload); + const fitting_config& cfg, const device::fit_payload& payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, + typename fitter_t::surface_type>* tpayload); } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/fit_forward.hpp b/device/cuda/src/fitting/kernels/fit_forward.hpp index 694ebb175b..2384e5c55e 100644 --- a/device/cuda/src/fitting/kernels/fit_forward.hpp +++ b/device/cuda/src/fitting/kernels/fit_forward.hpp @@ -1,23 +1,28 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -#include - -#include "traccc/fitting/device/fit.hpp" +// Project include(s). +#include "traccc/fitting/device/fit_payload.hpp" #include "traccc/fitting/fitting_config.hpp" +// CUDA include(s). +#include + namespace traccc::cuda { template void fit_forward(const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload); + const fitting_config& cfg, const device::fit_payload& payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, + typename fitter_t::surface_type>* tpayload); } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/fit_prelude.cu b/device/cuda/src/fitting/kernels/fit_prelude.cu index 84a3e2ac43..ad3fb3b476 100644 --- a/device/cuda/src/fitting/kernels/fit_prelude.cu +++ b/device/cuda/src/fitting/kernels/fit_prelude.cu @@ -1,36 +1,35 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ +// Local include(s). +#include "../../utils/cuda_error_handling.hpp" #include "../../utils/global_index.hpp" #include "./fit_prelude.hpp" + +// Project include(s). #include "traccc/fitting/device/fit_prelude.hpp" namespace traccc::cuda { namespace kernels { -__global__ void fit_prelude( - vecmem::data::vector_view param_ids_view, - edm::track_container::const_view track_candidates_view, - edm::track_container::view tracks_view, - vecmem::data::vector_view param_liveness_view) { - device::fit_prelude(details::global_index1(), - param_ids_view, track_candidates_view, - tracks_view, param_liveness_view); + +__global__ void fit_prelude(const device::fit_prelude_payload payload) { + + device::fit_prelude(details::global_index1(), payload); } + } // namespace kernels -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 param_ids_view, - edm::track_container::const_view track_candidates_view, - edm::track_container::view tracks_view, - vecmem::data::vector_view param_liveness_view) { +void fit_prelude(const dim3& grid_size, const dim3& block_size, + std::size_t shared_mem_size, const cudaStream_t& stream, + const device::fit_prelude_payload& payload) { + kernels::fit_prelude<<>>( - param_ids_view, track_candidates_view, tracks_view, - param_liveness_view); + payload); + TRACCC_CUDA_ERROR_CHECK(cudaGetLastError()); } + } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/fit_prelude.hpp b/device/cuda/src/fitting/kernels/fit_prelude.hpp index 40111d7c11..3475c1563c 100644 --- a/device/cuda/src/fitting/kernels/fit_prelude.hpp +++ b/device/cuda/src/fitting/kernels/fit_prelude.hpp @@ -1,22 +1,20 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once -#include +// Project include(s). +#include "traccc/fitting/device/fit_prelude.hpp" -#include "traccc/edm/track_container.hpp" +// CUDA include(s). +#include 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 param_ids_view, - edm::track_container::const_view track_candidates_view, - edm::track_container::view tracks_view, - vecmem::data::vector_view param_liveness_view); +void fit_prelude(const dim3& grid_size, const dim3& block_size, + std::size_t shared_mem_size, const cudaStream_t& stream, + const device::fit_prelude_payload& payload); } diff --git a/device/cuda/src/fitting/kernels/specializations/fit_backward.cu.in b/device/cuda/src/fitting/kernels/specializations/fit_backward.cu.in index e079f17941..27a1de6f0b 100644 --- a/device/cuda/src/fitting/kernels/specializations/fit_backward.cu.in +++ b/device/cuda/src/fitting/kernels/specializations/fit_backward.cu.in @@ -1,13 +1,13 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2023-2025 CERN for the benefit of the ACTS project + * (c) 2023-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ // Local include(s). -#include "src/utils/magnetic_field_types.hpp" #include "src/fitting/kernels/specializations/fit_backward_src.cuh" +#include "src/utils/magnetic_field_types.hpp" // Project include(s). #include "traccc/bfield/magnetic_field_types.hpp" @@ -18,16 +18,17 @@ #include namespace traccc::cuda { -using scalar = default_detector::device::scalar_type; -using fitter = traccc::details::kalman_fitter_t< - @DETECTOR_NAME@::device, - covfie::field<@BFIELD_NAME@>::view_t>; -template void fit_backward(const dim3& grid_size, - const dim3& block_size, - std::size_t shared_mem_size, - const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload); +using fitter = + traccc::details::kalman_fitter_t<@DETECTOR_NAME@::device, + covfie::field<@BFIELD_NAME@>::view_t>; + +template void fit_backward( + const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, + const cudaStream_t& stream, const fitting_config& cfg, + const device::fit_payload& payload, + const device::fit_tpayload* + tpayload); } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/specializations/fit_backward_src.cuh b/device/cuda/src/fitting/kernels/specializations/fit_backward_src.cuh index 28286cfa30..34dff0737b 100644 --- a/device/cuda/src/fitting/kernels/specializations/fit_backward_src.cuh +++ b/device/cuda/src/fitting/kernels/specializations/fit_backward_src.cuh @@ -1,32 +1,48 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once +// Local include(s). #include "../../../utils/global_index.hpp" #include "../fit_backward.hpp" + +// Project include(s). #include "traccc/fitting/device/fit_backward.hpp" namespace traccc::cuda { namespace kernels { + template __global__ __launch_bounds__(128) void fit_backward( - const fitting_config cfg, const device::fit_payload payload) { - device::fit_backward(details::global_index1(), cfg, payload); + const fitting_config cfg, const device::fit_payload payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, typename fitter_t::surface_type>* + tpayload) { + + device::fit_backward(details::global_index1(), cfg, payload, + *tpayload); } + } // namespace kernels template void fit_backward(const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload) { + const fitting_config& cfg, const device::fit_payload& payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, + typename fitter_t::surface_type>* tpayload) { + kernels::fit_backward - <<>>(cfg, payload); + <<>>(cfg, payload, + tpayload); } } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/specializations/fit_forward.cu.in b/device/cuda/src/fitting/kernels/specializations/fit_forward.cu.in index 9b5ecd169b..813f8031f2 100644 --- a/device/cuda/src/fitting/kernels/specializations/fit_forward.cu.in +++ b/device/cuda/src/fitting/kernels/specializations/fit_forward.cu.in @@ -1,13 +1,13 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2023-2025 CERN for the benefit of the ACTS project + * (c) 2023-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ // Local include(s). -#include "src/utils/magnetic_field_types.hpp" #include "src/fitting/kernels/specializations/fit_forward_src.cuh" +#include "src/utils/magnetic_field_types.hpp" // Project include(s). #include "traccc/bfield/magnetic_field_types.hpp" @@ -18,15 +18,17 @@ #include namespace traccc::cuda { -using scalar = default_detector::device::scalar_type; -using fitter = traccc::details::kalman_fitter_t< - @DETECTOR_NAME@::device, - covfie::field<@BFIELD_NAME@>::view_t>; -template void fit_forward(const dim3& grid_size, const dim3& block_size, - std::size_t shared_mem_size, - const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload); +using fitter = + traccc::details::kalman_fitter_t<@DETECTOR_NAME@::device, + covfie::field<@BFIELD_NAME@>::view_t>; + +template void fit_forward( + const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, + const cudaStream_t& stream, const fitting_config& cfg, + const device::fit_payload& payload, + const device::fit_tpayload* + tpayload); } // namespace traccc::cuda diff --git a/device/cuda/src/fitting/kernels/specializations/fit_forward_src.cuh b/device/cuda/src/fitting/kernels/specializations/fit_forward_src.cuh index 2727c331b6..f0b264e721 100644 --- a/device/cuda/src/fitting/kernels/specializations/fit_forward_src.cuh +++ b/device/cuda/src/fitting/kernels/specializations/fit_forward_src.cuh @@ -1,32 +1,48 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2025 CERN for the benefit of the ACTS project + * (c) 2025-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ #pragma once +// Local include(s). #include "../../../utils/global_index.hpp" #include "../fit_forward.hpp" + +// Project include(s). #include "traccc/fitting/device/fit_forward.hpp" namespace traccc::cuda { namespace kernels { + template __global__ __launch_bounds__(128) void fit_forward( - const fitting_config cfg, const device::fit_payload payload) { - device::fit_forward(details::global_index1(), cfg, payload); + const fitting_config cfg, const device::fit_payload payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, typename fitter_t::surface_type>* + tpayload) { + + device::fit_forward(details::global_index1(), cfg, payload, + *tpayload); } + } // namespace kernels template void fit_forward(const dim3& grid_size, const dim3& block_size, std::size_t shared_mem_size, const cudaStream_t& stream, - const fitting_config& cfg, - const device::fit_payload& payload) { - kernels::fit_forward<<>>( - cfg, payload); + const fitting_config& cfg, const device::fit_payload& payload, + const device::fit_tpayload< + typename fitter_t::detector_type::const_view_type, + typename fitter_t::bfield_type, + typename fitter_t::surface_type>* tpayload) { + + kernels::fit_forward + <<>>(cfg, payload, + tpayload); } } // namespace traccc::cuda diff --git a/device/sycl/CMakeLists.txt b/device/sycl/CMakeLists.txt index a62074edd0..9da0c3c67c 100644 --- a/device/sycl/CMakeLists.txt +++ b/device/sycl/CMakeLists.txt @@ -38,7 +38,6 @@ traccc_add_library( traccc_sycl sycl TYPE SHARED "include/traccc/sycl/fitting/kalman_fitting_algorithm.hpp" "src/fitting/kalman_fitting_algorithm.cpp" "src/fitting/kalman_fitting_algorithm.sycl" - "src/fitting/kalman_fitting.hpp" # Utilities. "include/traccc/sycl/utils/algorithm_base.hpp" "src/utils/algorithm_base.sycl" diff --git a/device/sycl/include/traccc/sycl/fitting/kalman_fitting_algorithm.hpp b/device/sycl/include/traccc/sycl/fitting/kalman_fitting_algorithm.hpp index eaac0f823c..760f63e3bc 100644 --- a/device/sycl/include/traccc/sycl/fitting/kalman_fitting_algorithm.hpp +++ b/device/sycl/include/traccc/sycl/fitting/kalman_fitting_algorithm.hpp @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -8,68 +8,86 @@ #pragma once // SYCL library include(s). +#include "traccc/sycl/utils/algorithm_base.hpp" #include "traccc/sycl/utils/queue_wrapper.hpp" // Project include(s). -#include "traccc/bfield/magnetic_field.hpp" -#include "traccc/edm/track_container.hpp" -#include "traccc/fitting/fitting_config.hpp" -#include "traccc/geometry/detector.hpp" -#include "traccc/geometry/detector_buffer.hpp" -#include "traccc/utils/algorithm.hpp" -#include "traccc/utils/memory_resource.hpp" -#include "traccc/utils/messaging.hpp" - -// VecMem include(s). -#include - -// System include(s). -#include +#include "traccc/fitting/device/kalman_fitting_algorithm.hpp" namespace traccc::sycl { -/// Kalman filter based track fitting algorithm -class kalman_fitting_algorithm - : public algorithm::buffer( - const detector_buffer&, const magnetic_field&, - const edm::track_container::const_view&)>, - public messaging { +/// Kalman filter based track fitting algorithm using SYCL +class kalman_fitting_algorithm : public device::kalman_fitting_algorithm, + public sycl::algorithm_base { public: - /// Configuration type - using config_type = fitting_config; - /// Constructor with the algorithm's configuration /// /// @param config The configuration object + /// @param mr The memory resource(s) used by the algorithm + /// @param copy The copy object used by the algorithm + /// @param queue The SYCL queue used by the algorithm + /// @param logger The logger used by the algorithm /// kalman_fitting_algorithm( const config_type& config, const traccc::memory_resource& mr, - const vecmem::copy& copy, queue_wrapper queue, + const vecmem::copy& copy, queue_wrapper& queue, std::unique_ptr logger = getDummyLogger().clone()); - /// Execute the algorithm + private: + /// @name Function(s) implemented from @c device::kalman_fitting_algorithm + /// @{ + + /// Prepare a buffer with the index order with which to fit the tracks /// - /// @param det The detector object - /// @param bfield The magnetic field object - /// @param track_candidates All track candidates to fit + /// @param[in] tracks The tracks to be fitted + /// @param[out] track_sort_keys Buffer storing temporary sorting keys + /// @param[out] track_indices The buffer to write the fitting order into /// - /// @return A container of the fitted track states + void prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const override; + + /// Kernel to prepare the fitting payloads + /// + /// @param payload The payload for the kernel(s) /// - output_type operator()( - const detector_buffer& det, const magnetic_field& bfield, - const edm::track_container::const_view& - track_candidates) const override; + void fit_prelude_kernel( + const device::fit_prelude_payload& payload) const override; - private: - /// Algorithm configuration - config_type m_config; - /// Memory resource used by the algorithm - traccc::memory_resource m_mr; - /// Copy object used by the algorithm - std::reference_wrapper m_copy; - /// Queue wrapper - mutable queue_wrapper m_queue; + /// Function preparing the fitting payload + /// + /// @param det The detector buffer to prepare the payload for + /// @param field The magnetic field to prepare the payload for + /// @param n_surfaces The number of surfaces for each track to be + /// fitted + /// @param payload The (non-templated) payload for the kernel(s) + /// + /// @return The prepared payload for the fitting kernel(s) + /// + fit_payload prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const override; + + /// Function launching the "forward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_forward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// Function launching the "backward fitting" kernel(s) + /// + /// @param config The fitting configuration + /// @param payload The payload for the fitting kernel(s) + /// + void fit_backward_kernel(const fitting_config& config, + const fit_payload& payload) const override; + + /// @} }; // class kalman_fitting_algorithm diff --git a/device/sycl/src/fitting/kalman_fitting.hpp b/device/sycl/src/fitting/kalman_fitting.hpp deleted file mode 100644 index 234c197019..0000000000 --- a/device/sycl/src/fitting/kalman_fitting.hpp +++ /dev/null @@ -1,226 +0,0 @@ -/** TRACCC library, part of the ACTS project (R&D line) - * - * (c) 2022-2025 CERN for the benefit of the ACTS project - * - * Mozilla Public License Version 2.0 - */ - -#pragma once - -// Local include(s). -#include "../utils/calculate1DimNdRange.hpp" -#include "../utils/global_index.hpp" -#include "../utils/oneDPL.hpp" - -// Project include(s). -#include "traccc/edm/device/sort_key.hpp" -#include "traccc/edm/track_container.hpp" -#include "traccc/fitting/details/kalman_fitting_types.hpp" -#include "traccc/fitting/device/fill_fitting_sort_keys.hpp" -#include "traccc/fitting/device/fit.hpp" -#include "traccc/fitting/device/fit_backward.hpp" -#include "traccc/fitting/device/fit_forward.hpp" -#include "traccc/fitting/device/fit_prelude.hpp" -#include "traccc/fitting/fitting_config.hpp" -#include "traccc/utils/memory_resource.hpp" - -// VecMem include(s). -#include - -// SYCL include(s). -#include - -// System include(s). -#include - -namespace traccc::sycl { -namespace kernels { - -template -struct fill_fitting_sort_keys; -template -struct fit_prelude; -template -struct fit_forward; -template -struct fit_backward; - -} // namespace kernels - -namespace details { - -/// Templated implementation of the SYCL track fitting algorithm. -/// -/// @tparam kernel_t Structure to generate unique kernel names with -/// @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 SYCL queue to use for execution -/// -/// @return A container of the fitted track states -/// -template -typename edm::track_container::buffer -kalman_fitting( - const typename detector_t::const_view_type& det_view, - const bfield_t& field_view, - const typename edm::track_container< - typename detector_t::algebra_type>::const_view& track_candidates_view, - const fitting_config& config, const memory_resource& mr, - const vecmem::copy& copy, ::sycl::queue& queue) { - - // Get the number of tracks. - const unsigned int n_tracks = copy.get_size(track_candidates_view.tracks); - - // Get the sizes of the track candidates in each track. - const std::vector 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_container::buffer - track_states_buffer{ - {candidate_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}, - {n_states, mr.main, vecmem::data::buffer_type::resizable}, - track_candidates_view.measurements}; - 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 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.surface_sequence_size_factor, - config.min_surface_sequence_capacity); - }); - vecmem::data::jagged_vector_buffer - seqs_buffer{seqs_sizes, mr.main, mr.host, - vecmem::data::buffer_type::resizable}; - copy.setup(seqs_buffer)->wait(); - - // Create the buffers for sorting the parameter IDs. - vecmem::data::vector_buffer keys_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_ids_buffer(n_tracks, - mr.main); - vecmem::data::vector_buffer param_liveness_buffer(n_tracks, - mr.main); - vecmem::copy::event_type keys_setup_event = copy.setup(keys_buffer); - vecmem::copy::event_type param_ids_setup_event = - copy.setup(param_ids_buffer); - vecmem::copy::event_type param_liveness_setup_event = - copy.setup(param_liveness_buffer); - keys_setup_event->wait(); - param_ids_setup_event->wait(); - param_liveness_setup_event->wait(); - - // The execution range for the two kernels of the function. - static constexpr unsigned int localSize = 64; - ::sycl::nd_range<1> range = calculate1DimNdRange(n_tracks, localSize); - - // Fill the keys and param_ids buffers. - ::sycl::event fill_keys_event = queue.submit([&](::sycl::handler& h) { - h.parallel_for>( - range, - [track_candidates_view, keys_view = vecmem::get_data(keys_buffer), - param_ids_view = - vecmem::get_data(param_ids_buffer)](::sycl::nd_item<1> item) { - device::fill_fitting_sort_keys(details::global_index(item), - track_candidates_view.tracks, - keys_view, param_ids_view); - }); - }); - - // Sort the key to get the sorted parameter ids - vecmem::device_vector keys_device(keys_buffer); - vecmem::device_vector param_ids_device(param_ids_buffer); - fill_keys_event.wait_and_throw(); - oneapi::dpl::sort_by_key(oneapi::dpl::execution::dpcpp_default, - keys_device.begin(), keys_device.end(), - param_ids_device.begin()); - - // Run the fitting, using the sorted parameter IDs. - typename edm::track_container::view - track_states_view{track_states_buffer}; - tracks_setup_event->wait(); - track_states_setup_event->wait(); - - queue - .submit([&](::sycl::handler& h) { - h.parallel_for>( - range, [param_ids_view = vecmem::get_data(param_ids_buffer), - track_candidates_view, track_states_view, - param_liveness_view = vecmem::get_data( - param_liveness_buffer)](::sycl::nd_item<1> item) { - device::fit_prelude( - details::global_index(item), param_ids_view, - track_candidates_view, track_states_view, - param_liveness_view); - }); - }) - .wait_and_throw(); - - // Allocate the fitting kernels's payload in host memory. - using fitter_t = traccc::details::kalman_fitter_t; - device::fit_payload 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, - .surfaces_view = seqs_buffer}; - // Now copy it to device memory. - vecmem::data::vector_buffer> device_payload( - 1u, mr.main); - copy.setup(device_payload)->wait(); - copy(vecmem::data::vector_view>( - 1u, &host_payload), - device_payload) - ->wait(); - - for (std::size_t i = 0; i < config.n_iterations; ++i) { - queue - .submit([&](::sycl::handler& h) { - h.parallel_for>( - range, [config, payload = device_payload.ptr()]( - ::sycl::nd_item<1> item) { - device::fit_forward( - details::global_index(item), config, *payload); - }); - }) - .wait_and_throw(); - - queue - .submit([&](::sycl::handler& h) { - h.parallel_for>( - range, [config, payload = device_payload.ptr()]( - ::sycl::nd_item<1> item) { - device::fit_backward( - details::global_index(item), config, *payload); - }); - }) - .wait_and_throw(); - } - - // Return the fitted tracks. - return track_states_buffer; -} - -} // namespace details -} // namespace traccc::sycl diff --git a/device/sycl/src/fitting/kalman_fitting_algorithm.cpp b/device/sycl/src/fitting/kalman_fitting_algorithm.cpp index 3aa5ca2b7a..9bf8919eeb 100644 --- a/device/sycl/src/fitting/kalman_fitting_algorithm.cpp +++ b/device/sycl/src/fitting/kalman_fitting_algorithm.cpp @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2024 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -12,12 +12,9 @@ namespace traccc::sycl { kalman_fitting_algorithm::kalman_fitting_algorithm( const config_type& config, const traccc::memory_resource& mr, - const vecmem::copy& copy, queue_wrapper queue, + const vecmem::copy& copy, queue_wrapper& q, std::unique_ptr logger) - : messaging(std::move(logger)), - m_config{config}, - m_mr{mr}, - m_copy{copy}, - m_queue{queue} {} + : device::kalman_fitting_algorithm{config, mr, copy, std::move(logger)}, + sycl::algorithm_base{q} {} } // namespace traccc::sycl diff --git a/device/sycl/src/fitting/kalman_fitting_algorithm.sycl b/device/sycl/src/fitting/kalman_fitting_algorithm.sycl index 218746742f..df90f1a7a3 100644 --- a/device/sycl/src/fitting/kalman_fitting_algorithm.sycl +++ b/device/sycl/src/fitting/kalman_fitting_algorithm.sycl @@ -1,6 +1,6 @@ /** TRACCC library, part of the ACTS project (R&D line) * - * (c) 2022-2025 CERN for the benefit of the ACTS project + * (c) 2022-2026 CERN for the benefit of the ACTS project * * Mozilla Public License Version 2.0 */ @@ -9,40 +9,176 @@ #include // Local include(s). +#include "../utils/calculate1DimNdRange.hpp" #include "../utils/detector_types.hpp" #include "../utils/get_queue.hpp" +#include "../utils/global_index.hpp" #include "../utils/magnetic_field_types.hpp" -#include "kalman_fitting.hpp" +#include "../utils/oneDPL.hpp" #include "traccc/sycl/fitting/kalman_fitting_algorithm.hpp" // Project include(s). #include "traccc/bfield/magnetic_field_types.hpp" +#include "traccc/fitting/device/fill_fitting_sort_keys.hpp" +#include "traccc/fitting/device/fit_backward.hpp" +#include "traccc/fitting/device/fit_forward.hpp" +#include "traccc/fitting/device/fit_prelude.hpp" #include "traccc/utils/detector_buffer_bfield_visitor.hpp" namespace traccc::sycl { namespace kernels { -template -struct fit_tracks; + +struct fill_fitting_sort_keys; +struct fit_prelude; +template +struct fit_forward; +template +struct fit_backward; + } // namespace kernels -kalman_fitting_algorithm::output_type kalman_fitting_algorithm::operator()( - const detector_buffer& detector, const magnetic_field& bfield, - const edm::track_container::const_view& track_candidates) - const { +void kalman_fitting_algorithm::prepare_track_fit_order( + const edm::track_collection::const_view& tracks, + vecmem::data::vector_view& track_sort_keys, + vecmem::data::vector_view& track_indices) const { + + // Get the number of tracks. + const unsigned int n_tracks = tracks.capacity(); + assert(n_tracks == copy().get_size(tracks)); + assert(n_tracks == track_indices.capacity()); + assert(track_indices.size_ptr() == nullptr); + + // Launch parameters for the kernel. + static constexpr unsigned int localSize = 64; + ::sycl::nd_range<1> range = + details::calculate1DimNdRange(n_tracks, localSize); + + // Fill the keys and indices buffers. + details::get_queue(queue()).submit([&](::sycl::handler& h) { + h.parallel_for( + range, + [tracks, track_sort_keys, track_indices](::sycl::nd_item<1> item) { + device::fill_fitting_sort_keys(details::global_index(item), + tracks, track_sort_keys, + track_indices); + }); + }); + + // Sort the key to get the sorted parameter ids + vecmem::device_vector keys_device(track_sort_keys); + vecmem::device_vector track_indices_device(track_indices); + oneapi::dpl::sort_by_key( + oneapi::dpl::execution::device_policy{details::get_queue(queue())}, + keys_device.begin(), keys_device.end(), track_indices_device.begin()); +} + +void kalman_fitting_algorithm::fit_prelude_kernel( + const device::fit_prelude_payload& payload) const { + + // Get the number of tracks. + const unsigned int n_tracks = payload.input_tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.input_tracks.tracks)); + assert(n_tracks == payload.track_indices.capacity()); + assert(payload.track_indices.size_ptr() == nullptr); + assert(n_tracks == copy().get_size(payload.output_tracks.tracks)); + + // Launch parameters for the kernel. + static constexpr unsigned int localSize = 64; + ::sycl::nd_range<1> range = + details::calculate1DimNdRange(n_tracks, localSize); + + // Run the fitting, using the sorted parameter IDs. + details::get_queue(queue()).submit([&](::sycl::handler& h) { + h.parallel_for( + range, [payload](::sycl::nd_item<1> item) { + device::fit_prelude(details::global_index(item), payload); + }); + }); +} + +auto kalman_fitting_algorithm::prepare_fit_payload( + const detector_buffer& det, const magnetic_field& field, + const std::vector& n_surfaces, + const device::fit_payload& payload) const -> fit_payload { + + return prepare_fit_payload_helper>( + det, field, n_surfaces, payload); +} + +void kalman_fitting_algorithm::fit_forward_kernel( + const fitting_config& config, const fit_payload& payload) const { + + return detector_buffer_magnetic_field_visitor< + detector_type_list, sycl::bfield_type_list>( + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + static constexpr unsigned int localSize = 64; + ::sycl::nd_range<1> range = + details::calculate1DimNdRange(n_tracks, localSize); + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + details::get_queue(queue()).submit([&](::sycl::handler& h) { + h.parallel_for, + bfield_tag_selector_t>>( + range, [config, payload = payload.payload, + tpayload = payload.get_tpayload()]( + ::sycl::nd_item<1> item) { + device::fit_forward( + details::global_index(item), config, payload, + *tpayload); + }); + }); + }); +} + +void kalman_fitting_algorithm::fit_backward_kernel( + const fitting_config& config, const fit_payload& payload) const { - // Perform the track fitting using the templated implementation. return detector_buffer_magnetic_field_visitor< detector_type_list, sycl::bfield_type_list>( - detector, bfield, - [&]( - const typename detector_t::view& det, const bfield_view_t& field) { - return details::kalman_fitting< - kernels::fit_tracks< - detector_tag_selector_t, - bfield_tag_selector_t>, - typename detector_t::device>(det, field, track_candidates, - m_config, m_mr, m_copy.get(), - details::get_queue(m_queue)); + payload.detector, payload.field, + [&]( + const typename detector_traits_t::view&, const bfield_view_t&) { + // Get the number of tracks. + const unsigned int n_tracks = + payload.payload.tracks.tracks.capacity(); + assert(n_tracks == copy().get_size(payload.payload.tracks.tracks)); + + // Launch parameters for the kernel. + static constexpr unsigned int localSize = 64; + ::sycl::nd_range<1> range = + details::calculate1DimNdRange(n_tracks, localSize); + + // Fitter type to use. + using fitter_t = traccc::details::kalman_fitter_t< + typename detector_traits_t::device, bfield_view_t>; + + // Run the track fitting + details::get_queue(queue()).submit([&](::sycl::handler& h) { + h.parallel_for, + bfield_tag_selector_t>>( + range, [config, payload = payload.payload, + tpayload = payload.get_tpayload()]( + ::sycl::nd_item<1> item) { + device::fit_backward( + details::global_index(item), config, payload, + *tpayload); + }); + }); }); }