src/seeding/triplet_seeding_algorithm.cpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2021-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/seeding/device/triplet_seeding_algorithm.hpp"
0010
0011 #include "traccc/edm/device/doublet_counter.hpp"
0012 #include "traccc/edm/device/seeding_global_counter.hpp"
0013
0014 // Project include(s).
0015 #include "traccc/seeding/spacepoint_binning_helper.hpp"
0016
0017 // System include(s).
0018 #include <cassert>
0019
0020 namespace traccc::device {
0021
0022 struct triplet_seeding_algorithm::data {
0023   /// @name Configuration objects
0024   /// @{
0025
0026   /// Configuration for the spacepoint grid forming step
0027   spacepoint_grid_config m_grid_config;
0028   /// Configuration for the seed finding step
0029   seedfinder_config m_finder_config;
0030   /// Configuration for the seed filtering step
0031   seedfilter_config m_filter_config;
0032
0033   /// @}
0034
0035   /// Axes for the internal spacepoint grid
0036   std::pair<details::spacepoint_grid_types::host::axis_p0_type,
0037             details::spacepoint_grid_types::host::axis_p1_type>
0038       m_axes;
0039
0040 };  // struct triplet_seeding_algorithm::data
0041
0042 triplet_seeding_algorithm::triplet_seeding_algorithm(
0043     const seedfinder_config& finder_config,
0044     const spacepoint_grid_config& grid_config,
0045     const seedfilter_config& filter_config, const memory_resource& mr,
0046     const vecmem::copy& copy, std::unique_ptr<const Logger> logger)
0047     : messaging(std::move(logger)),
0048       algorithm_base{mr, copy},
0049       m_data{std::make_unique<data>(
0050           grid_config, finder_config, filter_config,
0051           get_axes(grid_config, (mr.host ? *(mr.host) : mr.main)))} {}
0052
0053 triplet_seeding_algorithm::~triplet_seeding_algorithm() = default;
0054
0055 auto triplet_seeding_algorithm::operator()(
0056     const edm::spacepoint_collection::const_view& spacepoints) const
0057     -> output_type {
0058   // A small sanity check.
0059   assert(m_data);
0060
0061   // Get the number of spacepoints. In an asynchronous way if possible.
0062   edm::spacepoint_collection::const_view::size_type n_spacepoints = 0u;
0063   if (mr().host) {
0064     vecmem::async_size size = copy().get_size(spacepoints, *(mr().host));
0065     // Here we could give control back to the caller, once our code allows
0066     // for it. (coroutines...)
0067     n_spacepoints = size.get();
0068   } else {
0069     n_spacepoints = copy().get_size(spacepoints);
0070   }
0071
0072   // If there are no spacepoints, return right away.
0073   if (n_spacepoints == 0) {
0074     return {};
0075   }
0076
0077   // Set up the container that will be filled with the required capacities for
0078   // the spacepoint grid.
0079   const unsigned int grid_bins =
0080       m_data->m_axes.first.n_bins * m_data->m_axes.second.n_bins;
0081   vecmem::data::vector_buffer<unsigned int> grid_capacities_buffer(grid_bins,
0082                                                                    mr().main);
0083   copy().setup(grid_capacities_buffer)->ignore();
0084   copy().memset(grid_capacities_buffer, 0)->ignore();
0085
0086   // Launch the grid capacity counting kernel.
0087   count_grid_capacities_kernel({n_spacepoints, m_data->m_finder_config,
0088                                 m_data->m_axes.first, m_data->m_axes.second,
0089                                 spacepoints, grid_capacities_buffer});
0090
0091   // Copy grid capacities back to the host
0092   vecmem::vector<unsigned int> grid_capacities_host(mr().host ? mr().host
0093                                                               : &(mr().main));
0094   copy()(grid_capacities_buffer, grid_capacities_host)->wait();
0095
0096   // Create the spacepoint grid buffer and a prefix sum buffer that describes
0097   // it. (The latter is needed for the next few steps.)
0098   details::spacepoint_grid_types::buffer grid_buffer(
0099       m_data->m_axes.first, m_data->m_axes.second,
0100       std::vector<std::size_t>(grid_capacities_host.begin(),
0101                                grid_capacities_host.end()),
0102       mr().main, mr().host, vecmem::data::buffer_type::resizable);
0103   copy().setup(grid_buffer._buffer)->ignore();
0104   vecmem::data::vector_buffer<prefix_sum_element_t> grid_prefix_sum_buffer(
0105       n_spacepoints, mr().main, vecmem::data::buffer_type::resizable);
0106   copy().setup(grid_prefix_sum_buffer)->ignore();
0107
0108   // Launch the grid population kernel.
0109   populate_grid_kernel({n_spacepoints, m_data->m_finder_config, spacepoints,
0110                         grid_buffer, grid_prefix_sum_buffer});
0111
0112   // Update the spacepoint counter. So that it would only count the "good"
0113   // spacepoints that ended up in the spacepoint grid.
0114   if (mr().host) {
0115     vecmem::async_size size =
0116         copy().get_size(grid_prefix_sum_buffer, *(mr().host));
0117     // Here we could give control back to the caller, once our code allows
0118     // for it. (coroutines...)
0119     n_spacepoints = size.get();
0120   } else {
0121     n_spacepoints = copy().get_size(grid_prefix_sum_buffer);
0122   }
0123
0124   // Set up the doublet counter buffer.
0125   device::doublet_counter_collection_types::buffer doublet_counter_buffer{
0126       n_spacepoints, mr().main, vecmem::data::buffer_type::resizable};
0127   copy().setup(doublet_counter_buffer)->ignore();
0128
0129   // Set up a global counter used in the seeding kernels.
0130   vecmem::unique_alloc_ptr<device::seeding_global_counter>
0131       globalCounter_device =
0132           vecmem::make_unique_alloc<device::seeding_global_counter>(mr().main);
0133   copy()
0134       .memset(vecmem::data::vector_view<device::seeding_global_counter>(
0135                   1u, globalCounter_device.get()),
0136               0)
0137       ->ignore();
0138
0139   // Launch the doublet counting kernel.
0140   count_doublets_kernel(
0141       {n_spacepoints, m_data->m_finder_config, spacepoints, grid_buffer,
0142        grid_prefix_sum_buffer, doublet_counter_buffer,
0143        globalCounter_device->m_nMidBot, globalCounter_device->m_nMidTop});
0144
0145   // Get the number of doublets found.
0146   device::doublet_counter_collection_types::buffer::size_type n_doublets = 0u;
0147   if (mr().host) {
0148     vecmem::async_size size =
0149         copy().get_size(doublet_counter_buffer, *(mr().host));
0150     // Here we could give control back to the caller, once our code allows
0151     // for it. (coroutines...)
0152     n_doublets = size.get();
0153   } else {
0154     n_doublets = copy().get_size(doublet_counter_buffer);
0155   }
0156   vecmem::unique_alloc_ptr<device::seeding_global_counter> globalCounter_host =
0157       vecmem::make_unique_alloc<device::seeding_global_counter>(
0158           mr().host ? *(mr().host) : mr().main);
0159   copy()(vecmem::data::vector_view<device::seeding_global_counter>(
0160              1u, globalCounter_device.get()),
0161          vecmem::data::vector_view<device::seeding_global_counter>(
0162              1u, globalCounter_host.get()))
0163       ->wait();
0164
0165   // Exit already here if we won't find any triplets anyway.
0166   if ((globalCounter_host->m_nMidBot == 0) ||
0167       (globalCounter_host->m_nMidTop == 0)) {
0168     return {};
0169   }
0170
0171   // Set up the doublet buffers.
0172   device_doublet_collection_types::buffer doublet_buffer_mb{
0173       globalCounter_host->m_nMidBot, mr().main};
0174   copy().setup(doublet_buffer_mb)->ignore();
0175   device_doublet_collection_types::buffer doublet_buffer_mt{
0176       globalCounter_host->m_nMidTop, mr().main};
0177   copy().setup(doublet_buffer_mt)->ignore();
0178
0179   // Launch the doublet finding kernel.
0180   find_doublets_kernel({n_doublets, m_data->m_finder_config, spacepoints,
0181                         grid_buffer, doublet_counter_buffer, doublet_buffer_mb,
0182                         doublet_buffer_mt});
0183
0184   // Set up the triplet counter buffers
0185   triplet_counter_spM_collection_types::buffer triplet_counter_spM_buffer{
0186       n_doublets, mr().main};
0187   copy().setup(triplet_counter_spM_buffer)->ignore();
0188   copy().memset(triplet_counter_spM_buffer, 0)->ignore();
0189   triplet_counter_collection_types::buffer triplet_counter_midBot_buffer{
0190       globalCounter_host->m_nMidBot, mr().main,
0191       vecmem::data::buffer_type::resizable};
0192   copy().setup(triplet_counter_midBot_buffer)->ignore();
0193
0194   // Launch the triplet counting kernel.
0195   count_triplets_kernel({globalCounter_host->m_nMidBot, m_data->m_finder_config,
0196                          spacepoints, grid_buffer, doublet_counter_buffer,
0197                          doublet_buffer_mb, doublet_buffer_mt,
0198                          triplet_counter_spM_buffer,
0199                          triplet_counter_midBot_buffer});
0200   // Launch the triplet count reduction kernel.
0201   triplet_counts_reduction_kernel({n_doublets, doublet_counter_buffer,
0202                                    triplet_counter_spM_buffer,
0203                                    globalCounter_device->m_nTriplets});
0204
0205   // Get the number of triplets found.
0206   triplet_counter_collection_types::buffer::size_type n_midBotTriplets = 0u;
0207   if (mr().host) {
0208     vecmem::async_size size =
0209         copy().get_size(triplet_counter_midBot_buffer, *(mr().host));
0210     // Here we could give control back to the caller, once our code allows
0211     // for it. (coroutines...)
0212     n_midBotTriplets = size.get();
0213   } else {
0214     n_midBotTriplets = copy().get_size(triplet_counter_midBot_buffer);
0215   }
0216   copy()(vecmem::data::vector_view<device::seeding_global_counter>(
0217              1u, globalCounter_device.get()),
0218          vecmem::data::vector_view<device::seeding_global_counter>(
0219              1u, globalCounter_host.get()))
0220       ->wait();
0221
0222   // If no triplets could be found, exit already here.
0223   if (globalCounter_host->m_nTriplets == 0) {
0224     return {};
0225   }
0226
0227   // Set up the triplet buffer.
0228   device_triplet_collection_types::buffer triplet_buffer{
0229       globalCounter_host->m_nTriplets, mr().main};
0230   copy().setup(triplet_buffer)->ignore();
0231
0232   // Launch the triplet finding kernel.
0233   find_triplets_kernel({n_midBotTriplets, m_data->m_finder_config,
0234                         m_data->m_filter_config, spacepoints, grid_buffer,
0235                         doublet_counter_buffer, doublet_buffer_mt,
0236                         triplet_counter_spM_buffer,
0237                         triplet_counter_midBot_buffer, triplet_buffer});
0238   // Launch the triplet weight updating/filling kernel.
0239   update_triplet_weights_kernel(
0240       {globalCounter_host->m_nTriplets, m_data->m_filter_config, spacepoints,
0241        triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
0242        triplet_buffer});
0243
0244   // Create the result object.
0245   edm::seed_collection::buffer seed_buffer(
0246       globalCounter_host->m_nTriplets, mr().main,
0247       vecmem::data::buffer_type::resizable);
0248   copy().setup(seed_buffer)->ignore();
0249
0250   // Launch the seed selecting/filling kernel.
0251   select_seeds_kernel(
0252       {n_doublets, m_data->m_finder_config, m_data->m_filter_config,
0253        spacepoints, grid_buffer, triplet_counter_spM_buffer,
0254        triplet_counter_midBot_buffer, triplet_buffer, seed_buffer});
0255
0256   // Return the seed buffer.
0257   return seed_buffer;
0258 }
0259
0260 }  // namespace traccc::device