src/seeding/triplet_seeding_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 // Library include(s).
0009 #include "traccc/alpaka/seeding/triplet_seeding_algorithm.hpp"
0010
0011 #include "../utils/get_queue.hpp"
0012 #include "../utils/utils.hpp"
0013
0014 // Project include(s).
0015 #include "traccc/seeding/device/count_doublets.hpp"
0016 #include "traccc/seeding/device/count_grid_capacities.hpp"
0017 #include "traccc/seeding/device/count_triplets.hpp"
0018 #include "traccc/seeding/device/find_doublets.hpp"
0019 #include "traccc/seeding/device/find_triplets.hpp"
0020 #include "traccc/seeding/device/populate_grid.hpp"
0021 #include "traccc/seeding/device/reduce_triplet_counts.hpp"
0022 #include "traccc/seeding/device/select_seeds.hpp"
0023 #include "traccc/seeding/device/update_triplet_weights.hpp"
0024
0025 namespace traccc::alpaka {
0026 namespace kernels {
0027
0028 /// Kernel for running @c traccc::device::count_grid_capacities
0029 struct count_grid_capacity {
0030   template <typename TAcc>
0031   ALPAKA_FN_ACC void operator()(
0032       TAcc const& acc, const seedfinder_config& config,
0033       const traccc::details::spacepoint_grid_types::host::axis_p0_type phi_axis,
0034       const traccc::details::spacepoint_grid_types::host::axis_p1_type z_axis,
0035       const edm::spacepoint_collection::const_view spacepoints_view,
0036       vecmem::data::vector_view<unsigned int> grid_capacities_view) const {
0037     auto const globalThreadIdx =
0038         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0039     device::count_grid_capacities(globalThreadIdx, config, phi_axis, z_axis,
0040                                   spacepoints_view, grid_capacities_view);
0041   }
0042 };
0043
0044 /// Kernel for running @c traccc::device::populate_grid
0045 struct populate_grid {
0046   template <typename TAcc>
0047   ALPAKA_FN_ACC void operator()(
0048       TAcc const& acc, seedfinder_config config,
0049       edm::spacepoint_collection::const_view spacepoints_view,
0050       traccc::details::spacepoint_grid_types::view grid_view,
0051       vecmem::data::vector_view<device::prefix_sum_element_t> grid_prefix_sum)
0052       const {
0053     auto const globalThreadIdx =
0054         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0055     device::populate_grid(globalThreadIdx, config, spacepoints_view, grid_view,
0056                           grid_prefix_sum);
0057   }
0058 };
0059
0060 /// Kernel for running @c traccc::device::count_doublets
0061 struct count_doublets {
0062   template <typename TAcc>
0063   ALPAKA_FN_ACC void operator()(
0064       TAcc const& acc, seedfinder_config config,
0065       const edm::spacepoint_collection::const_view spacepoints,
0066       const traccc::details::spacepoint_grid_types::const_view sp_grid,
0067       const vecmem::data::vector_view<const device::prefix_sum_element_t>
0068           sp_prefix_sum,
0069       device::doublet_counter_collection_types::view doublet_counter,
0070       unsigned int* nMidBot, unsigned int* nMidTop) const {
0071     auto const globalThreadIdx =
0072         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0073     device::count_doublets(globalThreadIdx, config, spacepoints, sp_grid,
0074                            sp_prefix_sum, doublet_counter, *nMidBot, *nMidTop);
0075   }
0076 };
0077
0078 /// Kernel for running @c traccc::device::find_doublets
0079 struct find_doublets {
0080   template <typename TAcc>
0081   ALPAKA_FN_ACC void operator()(
0082       TAcc const& acc, seedfinder_config config,
0083       edm::spacepoint_collection::const_view spacepoints,
0084       traccc::details::spacepoint_grid_types::const_view sp_grid,
0085       device::doublet_counter_collection_types::const_view doublet_counter,
0086       device::device_doublet_collection_types::view mb_doublets,
0087       device::device_doublet_collection_types::view mt_doublets) const {
0088     auto const globalThreadIdx =
0089         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0090     device::find_doublets(globalThreadIdx, config, spacepoints, sp_grid,
0091                           doublet_counter, mb_doublets, mt_doublets);
0092   }
0093 };
0094
0095 /// Kernel for running @c traccc::device::count_triplets
0096 struct count_triplets {
0097   template <typename TAcc>
0098   ALPAKA_FN_ACC void operator()(
0099       TAcc const& acc, seedfinder_config config,
0100       edm::spacepoint_collection::const_view spacepoints,
0101       traccc::details::spacepoint_grid_types::const_view sp_grid,
0102       device::doublet_counter_collection_types::const_view doublet_counter,
0103       device::device_doublet_collection_types::const_view mb_doublets,
0104       device::device_doublet_collection_types::const_view mt_doublets,
0105       device::triplet_counter_spM_collection_types::view spM_counter,
0106       device::triplet_counter_collection_types::view midBot_counter) const {
0107     auto const globalThreadIdx =
0108         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0109     device::count_triplets(globalThreadIdx, config, spacepoints, sp_grid,
0110                            doublet_counter, mb_doublets, mt_doublets,
0111                            spM_counter, midBot_counter);
0112   }
0113 };
0114
0115 /// Kernel for running @c traccc::device::reduce_triplet_counts
0116 struct reduce_triplet_counts {
0117   template <typename TAcc>
0118   ALPAKA_FN_ACC void operator()(
0119       TAcc const& acc,
0120       device::doublet_counter_collection_types::const_view doublet_counter,
0121       device::triplet_counter_spM_collection_types::view spM_counter,
0122       unsigned int* nTriplets) const {
0123     auto const globalThreadIdx =
0124         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0125     device::reduce_triplet_counts(globalThreadIdx, doublet_counter, spM_counter,
0126                                   *nTriplets);
0127   }
0128 };
0129
0130 /// Kernel for running @c traccc::device::find_triplets
0131 struct find_triplets {
0132   template <typename TAcc>
0133   ALPAKA_FN_ACC void operator()(
0134       TAcc const& acc, seedfinder_config config,
0135       seedfilter_config filter_config,
0136       edm::spacepoint_collection::const_view spacepoints,
0137       traccc::details::spacepoint_grid_types::const_view sp_grid,
0138       device::doublet_counter_collection_types::const_view doublet_counter,
0139       device::device_doublet_collection_types::const_view mt_doublets,
0140       device::triplet_counter_spM_collection_types::const_view spM_tc,
0141       device::triplet_counter_collection_types::const_view midBot_tc,
0142       device::device_triplet_collection_types::view triplet_view) const {
0143     auto const globalThreadIdx =
0144         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0145     device::find_triplets(globalThreadIdx, config, filter_config, spacepoints,
0146                           sp_grid, doublet_counter, mt_doublets, spM_tc,
0147                           midBot_tc, triplet_view);
0148   }
0149 };
0150
0151 /// Kernel for running @c traccc::device::update_triplet_weights
0152 struct update_triplet_weights {
0153   template <typename TAcc>
0154   ALPAKA_FN_ACC void operator()(
0155       TAcc const& acc, seedfilter_config filter_config,
0156       edm::spacepoint_collection::const_view spacepoints,
0157       device::triplet_counter_spM_collection_types::const_view spM_tc,
0158       device::triplet_counter_collection_types::const_view midBot_tc,
0159       device::device_triplet_collection_types::view triplet_view) const {
0160     auto const globalThreadIdx =
0161         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0162     auto const localThreadIdx =
0163         ::alpaka::getIdx<::alpaka::Block, ::alpaka::Threads>(acc)[0u];
0164
0165     // Array for temporary storage of quality parameters for comparing
0166     // triplets within weight updating kernel
0167     scalar* const data = ::alpaka::getDynSharedMem<scalar>(acc);
0168
0169     // Each thread uses compatSeedLimit elements of the array
0170     scalar* dataPos = &data[localThreadIdx * filter_config.compatSeedLimit];
0171
0172     device::update_triplet_weights(globalThreadIdx, filter_config, spacepoints,
0173                                    spM_tc, midBot_tc, dataPos, triplet_view);
0174   }
0175 };
0176
0177 /// Kernel for running @c traccc::device::select_seeds
0178 struct select_seeds {
0179   template <typename TAcc>
0180   ALPAKA_FN_ACC void operator()(
0181       TAcc const& acc, seedfinder_config finder_config,
0182       seedfilter_config filter_config,
0183       edm::spacepoint_collection::const_view spacepoints,
0184       traccc::details::spacepoint_grid_types::const_view sp_view,
0185       device::triplet_counter_spM_collection_types::const_view spM_tc,
0186       device::triplet_counter_collection_types::const_view midBot_tc,
0187       device::device_triplet_collection_types::const_view triplet_view,
0188       edm::seed_collection::view seed_view) const {
0189     auto const globalThreadIdx =
0190         ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0u];
0191     auto const localThreadIdx =
0192         ::alpaka::getIdx<::alpaka::Block, ::alpaka::Threads>(acc)[0u];
0193
0194     // Array for temporary storage of quality parameters for comparing
0195     // triplets within weight updating kernel
0196     device::device_triplet* const data =
0197         ::alpaka::getDynSharedMem<device::device_triplet>(acc);
0198
0199     // Each thread uses max_triplets_per_spM elements of the array
0200     device::device_triplet* dataPos =
0201         &data[localThreadIdx * finder_config.maxSeedsPerSpM];
0202
0203     device::select_seeds(globalThreadIdx, finder_config, filter_config,
0204                          spacepoints, sp_view, spM_tc, midBot_tc, triplet_view,
0205                          dataPos, seed_view);
0206   }
0207 };
0208
0209 }  // namespace kernels
0210
0211 triplet_seeding_algorithm::triplet_seeding_algorithm(
0212     const seedfinder_config& finder_config,
0213     const spacepoint_grid_config& grid_config,
0214     const seedfilter_config& filter_config, const traccc::memory_resource& mr,
0215     const vecmem::copy& copy, alpaka::queue& q,
0216     std::unique_ptr<const Logger> logger)
0217     : device::triplet_seeding_algorithm(finder_config, grid_config,
0218                                         filter_config, mr, copy,
0219                                         std::move(logger)),
0220       alpaka::algorithm_base{q} {}
0221
0222 void triplet_seeding_algorithm::count_grid_capacities_kernel(
0223     const count_grid_capacities_kernel_payload& payload) const {
0224   const unsigned int n_threads = warp_size() * 8;
0225   const unsigned int n_blocks =
0226       (payload.n_spacepoints + n_threads - 1) / n_threads;
0227   ::alpaka::exec<Acc>(
0228       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0229       kernels::count_grid_capacity{}, payload.config, payload.phi_axis,
0230       payload.z_axis, payload.spacepoints, payload.grid_capacities);
0231 }
0232
0233 void triplet_seeding_algorithm::populate_grid_kernel(
0234     const populate_grid_kernel_payload& payload) const {
0235   const unsigned int n_threads = warp_size() * 8;
0236   const unsigned int n_blocks =
0237       (payload.n_spacepoints + n_threads - 1) / n_threads;
0238   ::alpaka::exec<Acc>(
0239       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0240       kernels::populate_grid{}, payload.config, payload.spacepoints,
0241       payload.grid, payload.grid_prefix_sum);
0242 }
0243
0244 void triplet_seeding_algorithm::count_doublets_kernel(
0245     const count_doublets_kernel_payload& payload) const {
0246   const unsigned int n_threads = warp_size() * 2;
0247   const unsigned int n_blocks =
0248       (payload.n_spacepoints + n_threads - 1) / n_threads;
0249   ::alpaka::exec<Acc>(
0250       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0251       kernels::count_doublets{}, payload.config, payload.spacepoints,
0252       payload.grid, payload.grid_prefix_sum, payload.doublet_counter,
0253       &(payload.nMidBot), &(payload.nMidTop));
0254 }
0255
0256 void triplet_seeding_algorithm::find_doublets_kernel(
0257     const find_doublets_kernel_payload& payload) const {
0258   const unsigned int n_threads = warp_size() * 2;
0259   const unsigned int n_blocks =
0260       (payload.n_doublets + n_threads - 1) / n_threads;
0261   ::alpaka::exec<Acc>(
0262       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0263       kernels::find_doublets{}, payload.config, payload.spacepoints,
0264       payload.grid, payload.doublet_counter, payload.mb_doublets,
0265       payload.mt_doublets);
0266 }
0267
0268 void triplet_seeding_algorithm::count_triplets_kernel(
0269     const count_triplets_kernel_payload& payload) const {
0270   const unsigned int n_threads = warp_size() * 2;
0271   const unsigned int n_blocks = (payload.nMidBot + n_threads - 1) / n_threads;
0272   ::alpaka::exec<Acc>(
0273       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0274       kernels::count_triplets{}, payload.config, payload.spacepoints,
0275       payload.grid, payload.doublet_counter, payload.mb_doublets,
0276       payload.mt_doublets, payload.spM_counter, payload.midBot_counter);
0277 }
0278
0279 void triplet_seeding_algorithm::triplet_counts_reduction_kernel(
0280     const triplet_counts_reduction_kernel_payload& payload) const {
0281   const unsigned int n_threads = warp_size() * 2;
0282   const unsigned int n_blocks =
0283       (payload.n_doublets + n_threads - 1) / n_threads;
0284   ::alpaka::exec<Acc>(details::get_queue(queue()),
0285                       makeWorkDiv<Acc>(n_blocks, n_threads),
0286                       kernels::reduce_triplet_counts{}, payload.doublet_counter,
0287                       payload.spM_counter, &(payload.nTriplets));
0288 }
0289
0290 void triplet_seeding_algorithm::find_triplets_kernel(
0291     const find_triplets_kernel_payload& payload) const {
0292   const unsigned int n_threads = warp_size() * 2;
0293   const unsigned int n_blocks = (payload.nMidBot + n_threads - 1) / n_threads;
0294   ::alpaka::exec<Acc>(
0295       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0296       kernels::find_triplets{}, payload.finding_config, payload.filter_config,
0297       payload.spacepoints, payload.grid, payload.doublet_counter,
0298       payload.mt_doublets, payload.spM_tc, payload.midBot_tc, payload.triplets);
0299 }
0300
0301 void triplet_seeding_algorithm::update_triplet_weights_kernel(
0302     const update_triplet_weights_kernel_payload& payload) const {
0303   const unsigned int n_threads = warp_size() * 2;
0304   const unsigned int n_blocks =
0305       (payload.n_triplets + n_threads - 1) / n_threads;
0306   ::alpaka::exec<Acc>(
0307       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0308       kernels::update_triplet_weights{}, payload.config, payload.spacepoints,
0309       payload.spM_tc, payload.midBot_tc, payload.triplets);
0310 }
0311
0312 void triplet_seeding_algorithm::select_seeds_kernel(
0313     const select_seeds_kernel_payload& payload) const {
0314   const unsigned int n_threads = warp_size() * 2;
0315   const unsigned int n_blocks =
0316       (payload.n_doublets + n_threads - 1) / n_threads;
0317   ::alpaka::exec<Acc>(
0318       details::get_queue(queue()), makeWorkDiv<Acc>(n_blocks, n_threads),
0319       kernels::select_seeds{}, payload.finder_config, payload.filter_config,
0320       payload.spacepoints, payload.grid, payload.spM_tc, payload.midBot_tc,
0321       payload.triplets, payload.seeds);
0322 }
0323
0324 }  // namespace traccc::alpaka
0325
0326 // Define the required trait needed for Dynamic shared memory allocation.
0327 namespace alpaka::trait {
0328
0329 template <typename TAcc>
0330 struct BlockSharedMemDynSizeBytes<
0331     traccc::alpaka::kernels::update_triplet_weights, TAcc> {
0332   template <typename TVec, typename... TArgs>
0333   ALPAKA_FN_HOST_ACC static auto getBlockSharedMemDynSizeBytes(
0334       traccc::alpaka::kernels::update_triplet_weights const& /* kernel */,
0335       TVec const& blockThreadExtent, TVec const& /* threadElemExtent */,
0336       traccc::seedfilter_config config, TArgs const&... /* args */
0337       ) -> std::size_t {
0338     return static_cast<std::size_t>(config.compatSeedLimit *
0339                                     blockThreadExtent.prod()) *
0340            sizeof(traccc::scalar);
0341   }
0342 };
0343
0344 template <typename TAcc>
0345 struct BlockSharedMemDynSizeBytes<traccc::alpaka::kernels::select_seeds, TAcc> {
0346   template <typename TVec, typename... TArgs>
0347   ALPAKA_FN_HOST_ACC static auto getBlockSharedMemDynSizeBytes(
0348       traccc::alpaka::kernels::select_seeds const& /* kernel */,
0349       TVec const& blockThreadExtent, TVec const& /* threadElemExtent */,
0350       traccc::seedfinder_config config, TArgs const&... /* args */
0351       ) -> std::size_t {
0352     return static_cast<std::size_t>(config.maxSeedsPerSpM *
0353                                     blockThreadExtent.prod()) *
0354            sizeof(traccc::device::device_triplet);
0355   }
0356 };
0357
0358 }  // namespace alpaka::trait