alpaka/src/full_chain_algorithm.cpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2023-2026 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // Local include(s).
0009 #include "traccc/examples/alpaka/full_chain_algorithm.hpp"
0010
0011 // Project include(s).
0012 #include "traccc/alpaka/utils/make_magnetic_field.hpp"
0013
0014 // System include(s).
0015 #include <iostream>
0016 #include <stdexcept>
0017
0018 namespace traccc::alpaka {
0019
0020 full_chain_algorithm::full_chain_algorithm(
0021     vecmem::memory_resource& host_mr,
0022     const clustering_config& clustering_config,
0023     const seedfinder_config& finder_config,
0024     const spacepoint_grid_config& grid_config,
0025     const seedfilter_config& filter_config,
0026     const gbts_seedfinder_config& gbts_config,
0027     const track_params_estimation_config& track_params_estimation_config,
0028     const finding_algorithm::config_type& finding_config,
0029     const fitting_algorithm::config_type& fitting_config,
0030     const detector_design_description::host& det_descr,
0031     const detector_conditions_description::host& det_cond,
0032     const magnetic_field& field, host_detector* detector,
0033     std::unique_ptr<const traccc::Logger> logger, bool useGBTS)
0034     : messaging(logger->clone()),
0035       m_queue(),
0036       m_vecmem_objects(m_queue),
0037       m_host_mr(host_mr),
0038       m_cached_pinned_host_mr(m_vecmem_objects.host_mr()),
0039       m_cached_device_mr(m_vecmem_objects.device_mr()),
0040       m_field_vec{0.f, 0.f, finder_config.bFieldInZ},
0041       m_field(make_magnetic_field(field, m_queue)),
0042       m_det_descr(det_descr),
0043       m_device_det_descr(
0044           [&]() {
0045             // number of elements in the detector design description
0046             std::vector<unsigned int> sizes(det_descr.size());
0047             for (std::size_t i = 0; i < det_descr.size(); ++i) {
0048               auto this_design = det_descr.at(i);
0049               // now for each element, set the size to the largest size of
0050               // that element across all modules
0051               sizes[i] = std::max(static_cast<unsigned int>(
0052                                       ((this_design.bin_edges_x()).size())),
0053                                   static_cast<unsigned int>(
0054                                       ((this_design.bin_edges_y()).size())));
0055             }
0056             return sizes;
0057           }(),
0058           m_cached_device_mr, &m_cached_pinned_host_mr,
0059           vecmem::data::buffer_type::resizable),
0060       m_det_cond(det_cond),
0061       m_device_det_cond(
0062           static_cast<detector_conditions_description::buffer::size_type>(
0063               m_det_cond.get().size()),
0064           m_cached_device_mr),
0065       m_detector(detector),
0066       m_clusterization({m_cached_device_mr, &m_cached_pinned_host_mr},
0067                        m_vecmem_objects.async_copy(), m_queue,
0068                        clustering_config),
0069       m_measurement_sorting({m_cached_device_mr, &m_cached_pinned_host_mr},
0070                             m_vecmem_objects.async_copy(), m_queue,
0071                             logger->cloneWithSuffix("MeasSortingAlg")),
0072       m_spacepoint_formation({m_cached_device_mr, &m_cached_pinned_host_mr},
0073                              m_vecmem_objects.async_copy(), m_queue,
0074                              logger->cloneWithSuffix("SpFormationAlg")),
0075       m_seeding(finder_config, grid_config, filter_config,
0076                 {m_cached_device_mr, &m_cached_pinned_host_mr},
0077                 m_vecmem_objects.async_copy(), m_queue,
0078                 logger->cloneWithSuffix("SeedingAlg")),
0079       m_track_parameter_estimation(
0080           track_params_estimation_config,
0081           {m_cached_device_mr, &m_cached_pinned_host_mr},
0082           m_vecmem_objects.async_copy(), m_queue,
0083           logger->cloneWithSuffix("TrackParamEstAlg")),
0084       m_finding(finding_config, {m_cached_device_mr, &m_cached_pinned_host_mr},
0085                 m_vecmem_objects.async_copy(), m_queue,
0086                 logger->cloneWithSuffix("TrackFindingAlg")),
0087       m_fitting(fitting_config, {m_cached_device_mr, &m_cached_pinned_host_mr},
0088                 m_vecmem_objects.async_copy(), m_queue,
0089                 logger->cloneWithSuffix("TrackFittingAlg")),
0090       m_clustering_config(clustering_config),
0091       m_finder_config(finder_config),
0092       m_grid_config(grid_config),
0093       m_filter_config(filter_config),
0094       m_gbts_config(gbts_config),
0095       m_track_params_estimation_config(track_params_estimation_config),
0096       m_finding_config(finding_config),
0097       m_fitting_config(fitting_config),
0098       usingGBTS(useGBTS) {
0099   if (usingGBTS) {
0100     std::cout << "GBTS not implemented for alpaka, this will run with "
0101                  "triplet seeding"
0102               << std::endl;
0103   }
0104   std::cout << traccc::alpaka::get_device_info() << std::endl;
0105
0106   // Copy the detector (description) to the device.
0107   m_vecmem_objects.async_copy().setup(m_device_det_descr)->wait();
0108   m_vecmem_objects
0109       .async_copy()(::vecmem::get_data(m_det_descr.get()), m_device_det_descr)
0110       ->wait();
0111   m_vecmem_objects
0112       .async_copy()(::vecmem::get_data(m_det_cond.get()), m_device_det_cond)
0113       ->wait();
0114   if (m_detector != nullptr) {
0115     m_device_detector = traccc::buffer_from_host_detector(
0116         *m_detector, m_vecmem_objects.device_mr(),
0117         m_vecmem_objects.async_copy());
0118   }
0119 }
0120
0121 full_chain_algorithm::full_chain_algorithm(const full_chain_algorithm& parent)
0122     : messaging(parent.logger().clone()),
0123       m_queue(),
0124       m_vecmem_objects(m_queue),
0125       m_host_mr(parent.m_host_mr),
0126       m_cached_pinned_host_mr(m_vecmem_objects.host_mr()),
0127       m_cached_device_mr(m_vecmem_objects.device_mr()),
0128       m_field_vec(parent.m_field_vec),
0129       m_field(parent.m_field),
0130       m_det_descr(parent.m_det_descr),
0131       m_device_det_descr(
0132           [&]() {
0133             // number of elements in the detector design description
0134             std::vector<unsigned int> sizes(parent.m_det_descr.get().size());
0135             for (std::size_t i = 0; i < parent.m_det_descr.get().size(); ++i) {
0136               auto this_design = parent.m_det_descr.get().at(i);
0137               // now for each element, set the size to the largest size of
0138               // that element across all modules
0139               sizes[i] = std::max(static_cast<unsigned int>(
0140                                       ((this_design.bin_edges_x()).size())),
0141                                   static_cast<unsigned int>(
0142                                       ((this_design.bin_edges_y()).size())));
0143             }
0144             return sizes;
0145           }(),
0146           m_cached_device_mr, &m_cached_pinned_host_mr,
0147           vecmem::data::buffer_type::resizable),
0148       m_det_cond(parent.m_det_cond),
0149       m_device_det_cond(
0150           static_cast<detector_conditions_description::buffer::size_type>(
0151               m_det_cond.get().size()),
0152           m_cached_device_mr),
0153       m_detector(parent.m_detector),
0154       m_clusterization({m_cached_device_mr, &m_cached_pinned_host_mr},
0155                        m_vecmem_objects.async_copy(), m_queue,
0156                        parent.m_clustering_config),
0157       m_measurement_sorting({m_cached_device_mr, &m_cached_pinned_host_mr},
0158                             m_vecmem_objects.async_copy(), m_queue,
0159                             parent.logger().cloneWithSuffix("MeasSortingAlg")),
0160       m_spacepoint_formation({m_cached_device_mr, &m_cached_pinned_host_mr},
0161                              m_vecmem_objects.async_copy(), m_queue,
0162                              parent.logger().cloneWithSuffix("SpFormationAlg")),
0163       m_seeding(parent.m_finder_config, parent.m_grid_config,
0164                 parent.m_filter_config,
0165                 {m_cached_device_mr, &m_cached_pinned_host_mr},
0166                 m_vecmem_objects.async_copy(), m_queue,
0167                 parent.logger().cloneWithSuffix("SeedingAlg")),
0168       m_track_parameter_estimation(
0169           parent.m_track_params_estimation_config,
0170           {m_cached_device_mr, &m_cached_pinned_host_mr},
0171           m_vecmem_objects.async_copy(), m_queue,
0172           parent.logger().cloneWithSuffix("TrackParamEstAlg")),
0173       m_finding(parent.m_finding_config,
0174                 {m_cached_device_mr, &m_cached_pinned_host_mr},
0175                 m_vecmem_objects.async_copy(), m_queue,
0176                 parent.logger().cloneWithSuffix("TrackFindingAlg")),
0177       m_fitting(parent.m_fitting_config,
0178                 {m_cached_device_mr, &m_cached_pinned_host_mr},
0179                 m_vecmem_objects.async_copy(), m_queue,
0180                 parent.logger().cloneWithSuffix("TrackFittingAlg")),
0181       m_clustering_config(parent.m_clustering_config),
0182       m_finder_config(parent.m_finder_config),
0183       m_grid_config(parent.m_grid_config),
0184       m_filter_config(parent.m_filter_config),
0185       m_gbts_config(parent.m_gbts_config),
0186       m_track_params_estimation_config(parent.m_track_params_estimation_config),
0187       m_finding_config(parent.m_finding_config),
0188       m_fitting_config(parent.m_fitting_config),
0189       usingGBTS(parent.usingGBTS) {
0190   // Copy the detector (description) to the device.
0191   m_vecmem_objects.async_copy().setup(m_device_det_descr)->wait();
0192   m_vecmem_objects
0193       .async_copy()(::vecmem::get_data(m_det_descr.get()), m_device_det_descr)
0194       ->wait();
0195   m_vecmem_objects
0196       .async_copy()(::vecmem::get_data(m_det_cond.get()), m_device_det_cond)
0197       ->wait();
0198   if (m_detector != nullptr) {
0199     m_device_detector = traccc::buffer_from_host_detector(
0200         *m_detector, m_vecmem_objects.device_mr(),
0201         m_vecmem_objects.async_copy());
0202   }
0203 }
0204
0205 full_chain_algorithm::~full_chain_algorithm() = default;
0206
0207 full_chain_algorithm::output_type full_chain_algorithm::operator()(
0208     const edm::silicon_cell_collection::host& cells) const {
0209   // Create device copy of input collections
0210   edm::silicon_cell_collection::buffer cells_buffer(
0211       static_cast<unsigned int>(cells.size()), m_cached_device_mr);
0212   m_vecmem_objects.async_copy()(::vecmem::get_data(cells), cells_buffer)
0213       ->ignore();
0214
0215   // Run the clusterization (asynchronously).
0216   const auto unsorted_measurements =
0217       m_clusterization(cells_buffer, m_device_det_descr, m_device_det_cond);
0218   const measurement_sorting_algorithm::output_type measurements =
0219       m_measurement_sorting(unsorted_measurements);
0220
0221   // If we have a Detray detector, run the track finding and fitting.
0222   if (m_detector != nullptr) {
0223     // Run the seed-finding (asynchronously).
0224     const spacepoint_formation_algorithm::output_type spacepoints =
0225         m_spacepoint_formation(m_device_detector, measurements);
0226     const seed_parameter_estimation_algorithm::output_type track_params =
0227         m_track_parameter_estimation(m_field, measurements, spacepoints,
0228                                      m_seeding(spacepoints));
0229
0230     // Run the track finding (asynchronously).
0231     const finding_algorithm::output_type track_candidates =
0232         m_finding(m_device_detector, m_field, measurements, track_params);
0233
0234     // Run the track fitting (asynchronously).
0235     const fitting_algorithm::output_type track_states =
0236         m_fitting(m_device_detector, m_field, track_candidates);
0237
0238     // Copy a limited amount of result data back to the host.
0239     const auto host_tracks = m_vecmem_objects.async_copy().to(
0240         track_states.tracks, m_cached_pinned_host_mr, nullptr,
0241         ::vecmem::copy::type::device_to_host);
0242     output_type result{m_host_mr};
0243     ::vecmem::copy host_copy;
0244     host_copy(host_tracks, result)->wait();
0245     return result;
0246
0247   }
0248   // If not, copy the track parameters back to the host, and return a dummy
0249   // object.
0250   else {
0251     // Copy the measurements back to the host.
0252     edm::measurement_collection::host measurements_host(m_host_mr);
0253     m_vecmem_objects.async_copy()(measurements, measurements_host)->wait();
0254
0255     // Return an empty object.
0256     return output_type{m_host_mr};
0257   }
0258 }
0259
0260 bound_track_parameters_collection_types::host full_chain_algorithm::seeding(
0261     const edm::silicon_cell_collection::host& cells) const {
0262   // Create device copy of input collections
0263   edm::silicon_cell_collection::buffer cells_buffer(
0264       static_cast<unsigned int>(cells.size()), m_cached_device_mr);
0265   m_vecmem_objects.async_copy()(::vecmem::get_data(cells), cells_buffer)
0266       ->ignore();
0267
0268   // Run the clusterization (asynchronously).
0269   const auto unsorted_measurements =
0270       m_clusterization(cells_buffer, m_device_det_descr, m_device_det_cond);
0271   const measurement_sorting_algorithm::output_type measurements =
0272       m_measurement_sorting(unsorted_measurements);
0273
0274   // If we have a Detray detector, run the track finding and fitting.
0275   if (m_detector != nullptr) {
0276     // Run the seed-finding (asynchronously).
0277     const spacepoint_formation_algorithm::output_type spacepoints =
0278         m_spacepoint_formation(m_device_detector, measurements);
0279     const seed_parameter_estimation_algorithm::output_type track_params =
0280         m_track_parameter_estimation(m_field, measurements, spacepoints,
0281                                      m_seeding(spacepoints));
0282
0283     // Copy a limited amount of result data back to the host.
0284     const auto host_seeds =
0285         m_vecmem_objects.async_copy().to(track_params, m_cached_pinned_host_mr,
0286                                          ::vecmem::copy::type::device_to_host);
0287     bound_track_parameters_collection_types::host result{&m_host_mr};
0288     ::vecmem::copy host_copy;
0289     host_copy(host_seeds, result)->wait();
0290     return result;
0291
0292   }
0293   // If not, copy the track parameters back to the host, and return a dummy
0294   // object.
0295   else {
0296     // Copy the measurements back to the host.
0297     edm::measurement_collection::host measurements_host(m_host_mr);
0298     m_vecmem_objects.async_copy()(measurements, measurements_host)->wait();
0299
0300     // Return an empty object.
0301     return {};
0302   }
0303 }
0304
0305 }  // namespace traccc::alpaka