Sort tracks per measurement for fast search (acts-project#1110)

Author: beomki-yeo

device/common/include/traccc/ambiguity_resolution/device/sort_tracks_per_measurement.hpp

@@ -0,0 +1,32 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/definitions/primitives.hpp"
+
+// VecMem include(s).
+#include <vecmem/containers/data/jagged_vector_view.hpp>
+#include <vecmem/containers/data/vector_view.hpp>
+
+// System include(s).
+#include <cstddef>
+
+namespace traccc::device {
+
+/// (Event Data) Payload for the @c traccc::device::sort_tracks_per_measurement
+/// function
+struct sort_tracks_per_measurement_payload {
+
+    /**
+     * @brief View object to the tracks per measurement
+     */
+    vecmem::data::jagged_vector_view<unsigned int> tracks_per_measurement_view;
+};
+
+}  // namespace traccc::device

device/cuda/CMakeLists.txt

@@ -113,6 +113,8 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/ambiguity_resolution/kernels/reset_status.cuh"
   "src/ambiguity_resolution/kernels/scan_block_offsets.cu"
   "src/ambiguity_resolution/kernels/scan_block_offsets.cuh"
+  "src/ambiguity_resolution/kernels/sort_tracks_per_measurement.cu"
+  "src/ambiguity_resolution/kernels/sort_tracks_per_measurement.cuh"
   "src/ambiguity_resolution/kernels/sort_updated_tracks.cu"
   "src/ambiguity_resolution/kernels/sort_updated_tracks.cuh"
   "src/ambiguity_resolution/kernels/remove_tracks.cu"

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -23,6 +23,7 @@
 #include "./kernels/remove_tracks.cuh"
 #include "./kernels/reset_status.cuh"
 #include "./kernels/scan_block_offsets.cuh"
+#include "./kernels/sort_tracks_per_measurement.cuh"
 #include "./kernels/sort_updated_tracks.cuh"
 #include "traccc/cuda/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.hpp"
 
@@ -296,6 +297,22 @@ greedy_ambiguity_resolution_algorithm::operator()(
         m_stream.get().synchronize();
     }
 
+    // Sort tracks per measurement vector
+    // @TODO: For the case where the measurement is shared by more than 1024
+    // tracks, the tracks need to be sorted again using thrust::sort
+    {
+        const unsigned int nThreads = 1024;
+        const unsigned int nBlocks = meas_count;
+
+        kernels::sort_tracks_per_measurement<<<nBlocks, nThreads, 0, stream>>>(
+            device::sort_tracks_per_measurement_payload{
+                .tracks_per_measurement_view = tracks_per_measurement_buffer,
+            });
+        TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+
+        m_stream.get().synchronize();
+    }
+
     // Make shared number of measurements vector
     vecmem::data::vector_buffer<unsigned int> n_shared_buffer{n_tracks,
                                                               m_mr.main};

device/cuda/src/ambiguity_resolution/kernels/remove_tracks.cu

@@ -106,7 +106,8 @@ __launch_bounds__(512) __global__
             sorted_ids.at(n_accepted_prev - 1 - sh_threads[threadIndex]);
 
         unsigned int worst_idx =
-            thrust::find(thrust::seq, tracks.begin(), tracks.end(), trk_id) -
+            thrust::lower_bound(thrust::seq, tracks.begin(), tracks.end(),
+                                trk_id) -
             tracks.begin();
 
         track_status[worst_idx] = 0;
@@ -120,8 +121,8 @@ __launch_bounds__(512) __global__
 
                 trk_id = sorted_ids[n_accepted_prev - 1 - sh_threads[i]];
 
-                worst_idx = thrust::find(thrust::seq, tracks.begin(),
-                                         tracks.end(), trk_id) -
+                worst_idx = thrust::lower_bound(thrust::seq, tracks.begin(),
+                                                tracks.end(), trk_id) -
                             tracks.begin();
 
                 track_status[worst_idx] = 0;

device/cuda/src/ambiguity_resolution/kernels/sort_tracks_per_measurement.cu

@@ -0,0 +1,63 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+// Local include(s).
+#include "../../utils/global_index.hpp"
+#include "sort_tracks_per_measurement.cuh"
+
+// VecMem include(s).
+#include <vecmem/containers/jagged_device_vector.hpp>
+
+namespace traccc::cuda::kernels {
+
+__global__ void sort_tracks_per_measurement(
+    device::sort_tracks_per_measurement_payload payload) {
+
+    __shared__ unsigned int sh_trk_ids[1024];
+
+    vecmem::jagged_device_vector<unsigned int> tracks_per_measurement(
+        payload.tracks_per_measurement_view);
+
+    auto tracks = tracks_per_measurement.at(blockIdx.x);
+    const unsigned int tid = threadIdx.x;
+    const unsigned int n_tracks = tracks.size();
+
+    sh_trk_ids[tid] = std::numeric_limits<unsigned int>::max();
+
+    if (tid < n_tracks) {
+        sh_trk_ids[tid] = tracks[tid];
+    }
+
+    // Bitonic sort
+    const unsigned int N = 1 << (32 - __clz(n_tracks - 1));
+    for (int k = 2; k <= N; k <<= 1) {
+
+        bool ascending = ((tid & k) == 0);
+
+        for (int j = k >> 1; j > 0; j >>= 1) {
+            int ixj = tid ^ j;
+
+            if (ixj > tid && ixj < N && tid < N) {
+                auto trk_i = sh_trk_ids[tid];
+                auto trk_j = sh_trk_ids[ixj];
+
+                bool should_swap = (trk_i > trk_j) == ascending;
+
+                if (should_swap) {
+                    sh_trk_ids[tid] = trk_j;
+                    sh_trk_ids[ixj] = trk_i;
+                }
+            }
+            __syncthreads();
+        }
+    }
+
+    if (tid < n_tracks) {
+        tracks[tid] = sh_trk_ids[tid];
+    }
+}
+}  // namespace traccc::cuda::kernels

device/cuda/src/ambiguity_resolution/kernels/sort_tracks_per_measurement.cuh

@@ -0,0 +1,17 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/ambiguity_resolution/device/sort_tracks_per_measurement.hpp"
+
+namespace traccc::cuda::kernels {
+
+__global__ void sort_tracks_per_measurement(
+    device::sort_tracks_per_measurement_payload payload);
+}