ambiguity_resolution/kernels/block_inclusive_scan.cu

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2025 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 // Local include(s).
0009 #include "../../utils/barrier.hpp"
0010 #include "../../utils/global_index.hpp"
0011 #include "block_inclusive_scan.cuh"
0012
0013 // VecMem include(s).
0014 #include <vecmem/containers/device_vector.hpp>
0015
0016 namespace traccc::cuda::kernels {
0017
0018 __global__ void block_inclusive_scan(
0019     device::block_inclusive_scan_payload payload) {
0020   if (*(payload.terminate) == 1 || *(payload.n_updated_tracks) == 0) {
0021     return;
0022   }
0023
0024   // temporary buffer where the block-wise prefix sum will be calculated
0025   extern __shared__ int shared_temp[];
0026
0027   vecmem::device_vector<const unsigned int> sorted_ids(payload.sorted_ids_view);
0028   vecmem::device_vector<const int> is_updated(payload.is_updated_view);
0029   vecmem::device_vector<int> block_offsets(payload.block_offsets_view);
0030   vecmem::device_vector<int> prefix_sums(payload.prefix_sums_view);
0031
0032   auto globalIndex = threadIdx.x + blockIdx.x * blockDim.x;
0033   auto threadIndex = threadIdx.x;
0034
0035   const unsigned int n_accepted = *(payload.n_accepted);
0036
0037   if (globalIndex >= n_accepted) {
0038     shared_temp[threadIndex] = 0;
0039   }
0040   // Start with boolean number depending on whether track id corresponding to
0041   // the current thread is updated during the iteration
0042   else {
0043     shared_temp[threadIndex] = is_updated[sorted_ids[globalIndex]];
0044   }
0045
0046   __syncthreads();
0047
0048   // Inclusive scan the boolean numbers to calculate the block-wise prefix sum
0049   for (int stride = 1; stride < blockDim.x; stride *= 2) {
0050     int val = 0;
0051     if (threadIndex >= stride) {
0052       val = shared_temp[threadIndex - stride];
0053     }
0054     __syncthreads();
0055     if (threadIndex >= stride) {
0056       shared_temp[threadIndex] += val;
0057     }
0058     __syncthreads();
0059   }
0060
0061   // Move the block-wise prefix_sums to global memory
0062   if (globalIndex < n_accepted) {
0063     prefix_sums[globalIndex] = shared_temp[threadIndex];
0064   }
0065
0066   __syncthreads();
0067
0068   // Block offset, the last element of block-wise prefix sums, is also
0069   // recorded to calculate full prefix sums later
0070   if (threadIndex == blockDim.x - 1) {
0071     block_offsets[blockIdx.x] = shared_temp[threadIndex];
0072   }
0073 }
0074
0075 }  // namespace traccc::cuda::kernels