Merge pull request #1050 from beomki-yeo/fast-measurement-lookup

Use a map to find the unique measurement ID

Author: beomki-yeo

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

@@ -56,6 +56,11 @@ struct count_removable_tracks_payload {
      */
     vecmem::data::vector_view<const measurement_id_type> unique_meas_view;
 
+    /**
+     * @brief View object to the meas id to unique id map
+     */
+    vecmem::data::vector_view<const unsigned int> meas_id_to_unique_id_view;
+
     /**
      * @brief View object to the number of accepted tracks per measurement
      */

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

@@ -31,9 +31,9 @@ struct count_shared_measurements_payload {
     vecmem::data::jagged_vector_view<const measurement_id_type> meas_ids_view;
 
     /**
-     * @brief View object to the unique measurement ids
+     * @brief View object to the meas id to unique id map
      */
-    vecmem::data::vector_view<const measurement_id_type> unique_meas_view;
+    vecmem::data::vector_view<const unsigned int> meas_id_to_unique_id_view;
 
     /**
      * @brief View object to the tracks per measurement

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

@@ -34,9 +34,9 @@ struct fill_tracks_per_measurement_payload {
     vecmem::data::jagged_vector_view<const measurement_id_type> meas_ids_view;
 
     /**
-     * @brief View object to the unique measurement ids
+     * @brief View object to the meas id to unique id map
      */
-    vecmem::data::vector_view<const measurement_id_type> unique_meas_view;
+    vecmem::data::vector_view<const unsigned int> meas_id_to_unique_id_view;
 
     /**
      * @brief View object to the tracks per measurement

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

@@ -0,0 +1,33 @@
+/** 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/vector_view.hpp>
+
+namespace traccc::device {
+
+/// (Event Data) Payload for the @c traccc::device::fill_unique_meas_id_map
+/// function
+struct fill_unique_meas_id_map_payload {
+
+    /**
+     * @brief View object to the unique measurement ids
+     */
+    vecmem::data::vector_view<const measurement_id_type> unique_meas_view;
+
+    /**
+     * @brief View object to the meas id to unique id map
+     */
+    vecmem::data::vector_view<unsigned int> meas_id_to_unique_id_view;
+};
+
+}  // namespace traccc::device

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

@@ -42,9 +42,9 @@ struct remove_tracks_payload {
     vecmem::data::vector_view<const unsigned int> n_meas_view;
 
     /**
-     * @brief View object to the unique measurement ids
+     * @brief View object to the meas id to unique id map
      */
-    vecmem::data::vector_view<const measurement_id_type> unique_meas_view;
+    vecmem::data::vector_view<const unsigned int> meas_id_to_unique_id_view;
 
     /**
      * @brief View object to the tracks per measurement

device/cuda/CMakeLists.txt

@@ -97,6 +97,8 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/ambiguity_resolution/kernels/fill_inverted_ids.cuh"
   "src/ambiguity_resolution/kernels/fill_tracks_per_measurement.cu"
   "src/ambiguity_resolution/kernels/fill_tracks_per_measurement.cuh"
+  "src/ambiguity_resolution/kernels/fill_unique_meas_id_map.cu"
+  "src/ambiguity_resolution/kernels/fill_unique_meas_id_map.cuh"
   "src/ambiguity_resolution/kernels/fill_track_candidates.cu"
   "src/ambiguity_resolution/kernels/fill_track_candidates.cuh"
   "src/ambiguity_resolution/kernels/fill_vectors.cu"

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -16,6 +16,7 @@
 #include "./kernels/fill_inverted_ids.cuh"
 #include "./kernels/fill_track_candidates.cuh"
 #include "./kernels/fill_tracks_per_measurement.cuh"
+#include "./kernels/fill_unique_meas_id_map.cuh"
 #include "./kernels/fill_vectors.cuh"
 #include "./kernels/find_max_shared.cuh"
 #include "./kernels/gather_tracks.cuh"
@@ -28,6 +29,7 @@
 
 // Thrust include(s).
 #include <thrust/execution_policy.h>
+#include <thrust/extrema.h>
 #include <thrust/fill.h>
 #include <thrust/functional.h>
 #include <thrust/iterator/constant_iterator.h>
@@ -62,6 +64,13 @@ struct track_comparator {
     }
 };
 
+struct measurement_id_comparator {
+    TRACCC_HOST_DEVICE bool operator()(const measurement& a,
+                                       const measurement& b) const {
+        return a.measurement_id < b.measurement_id;
+    }
+};
+
 greedy_ambiguity_resolution_algorithm::greedy_ambiguity_resolution_algorithm(
     const config_type& cfg, const traccc::memory_resource& mr,
     vecmem::copy& copy, stream& str, std::unique_ptr<const Logger> logger)
@@ -77,6 +86,30 @@ greedy_ambiguity_resolution_algorithm::operator()(
     const edm::track_candidate_container<default_algebra>::const_view&
         track_candidates_view) const {
 
+    measurement_collection_types::const_device measurements(
+        track_candidates_view.measurements);
+
+    auto n_meas_total =
+        m_copy.get().get_size(track_candidates_view.measurements);
+
+    // Make sure that max_measurement_id = number_of_measurement -1
+    // @TODO: More robust way is to assert that measurement id ranges from 0, 1,
+    // ..., number_of_measurement - 1
+    [[maybe_unused]] auto max_meas_it = thrust::max_element(
+        thrust::device, track_candidates_view.measurements.ptr(),
+        track_candidates_view.measurements.ptr() + n_meas_total,
+        measurement_id_comparator{});
+
+    measurement max_meas;
+    cudaMemcpy(&max_meas, thrust::raw_pointer_cast(&(*max_meas_it)),
+               sizeof(measurement), cudaMemcpyDeviceToHost);
+
+    if (max_meas.measurement_id != n_meas_total - 1) {
+        throw std::runtime_error(
+            "max measurement id should be equal to (the number of measurements "
+            "- 1)");
+    }
+
     // Get a convenience variable for the stream that we'll be using.
     cudaStream_t stream = details::get_stream(m_stream);
 
@@ -199,6 +232,24 @@ greedy_ambiguity_resolution_algorithm::operator()(
                         unique_meas_buffer.ptr() + meas_count,
                         unique_meas_counts_buffer.ptr());
 
+    // Unique measurement ids
+    vecmem::data::vector_buffer<measurement_id_type>
+        meas_id_to_unique_id_buffer{max_meas.measurement_id, m_mr.main};
+
+    // Make meas_id to meas vector
+    {
+        const unsigned int nThreads = m_warp_size * 2;
+        const unsigned int nBlocks = (meas_count + nThreads - 1) / nThreads;
+
+        kernels::fill_unique_meas_id_map<<<nBlocks, nThreads, 0, stream>>>(
+            device::fill_unique_meas_id_map_payload{
+                .unique_meas_view = unique_meas_buffer,
+                .meas_id_to_unique_id_view = meas_id_to_unique_id_buffer});
+        TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+
+        m_stream.get().synchronize();
+    }
+
     // Retreive the counting vector to host
     std::vector<std::size_t> unique_meas_counts;
     m_copy
@@ -235,7 +286,7 @@ greedy_ambiguity_resolution_algorithm::operator()(
             device::fill_tracks_per_measurement_payload{
                 .accepted_ids_view = pre_accepted_ids_buffer,
                 .meas_ids_view = meas_ids_buffer,
-                .unique_meas_view = unique_meas_buffer,
+                .meas_id_to_unique_id_view = meas_id_to_unique_id_buffer,
                 .tracks_per_measurement_view = tracks_per_measurement_buffer,
                 .track_status_per_measurement_view =
                     track_status_per_measurement_buffer,
@@ -262,7 +313,7 @@ greedy_ambiguity_resolution_algorithm::operator()(
             device::count_shared_measurements_payload{
                 .accepted_ids_view = pre_accepted_ids_buffer,
                 .meas_ids_view = meas_ids_buffer,
-                .unique_meas_view = unique_meas_buffer,
+                .meas_id_to_unique_id_view = meas_id_to_unique_id_buffer,
                 .n_accepted_tracks_per_measurement_view =
                     n_accepted_tracks_per_measurement_buffer,
                 .n_shared_view = n_shared_buffer});
@@ -405,7 +456,7 @@ greedy_ambiguity_resolution_algorithm::operator()(
             .n_accepted = n_accepted_device.get(),
             .meas_ids_view = meas_ids_buffer,
             .n_meas_view = n_meas_buffer,
-            .unique_meas_view = unique_meas_buffer,
+            .meas_id_to_unique_id_view = meas_id_to_unique_id_buffer,
             .n_accepted_tracks_per_measurement_view =
                 n_accepted_tracks_per_measurement_buffer,
             .n_removable_tracks = n_removable_tracks_device.get(),
@@ -432,7 +483,7 @@ greedy_ambiguity_resolution_algorithm::operator()(
             .n_accepted = n_accepted_device.get(),
             .meas_ids_view = meas_ids_buffer,
             .n_meas_view = n_meas_buffer,
-            .unique_meas_view = unique_meas_buffer,
+            .meas_id_to_unique_id_view = meas_id_to_unique_id_buffer,
             .tracks_per_measurement_view = tracks_per_measurement_buffer,
             .track_status_per_measurement_view =
                 track_status_per_measurement_buffer,

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

@@ -85,10 +85,10 @@ __global__ void count_removable_tracks(
     vecmem::device_vector<measurement_id_type> meas_to_remove(
         payload.meas_to_remove_view);
     vecmem::device_vector<unsigned int> threads(payload.threads_view);
-    vecmem::device_vector<const measurement_id_type> unique_meas(
-        payload.unique_meas_view);
     vecmem::device_vector<const unsigned int> n_accepted_tracks_per_measurement(
         payload.n_accepted_tracks_per_measurement_view);
+    vecmem::device_vector<const unsigned int> meas_id_to_unique_id(
+        payload.meas_id_to_unique_id_view);
 
     auto threadIndex = threadIdx.x;
 
@@ -193,10 +193,7 @@ __global__ void count_removable_tracks(
         auto mid = sh_meas_ids[threadIndex];
         bool is_start =
             (threadIndex == 0) || (sh_meas_ids[threadIndex - 1] != mid);
-        const std::size_t unique_meas_idx =
-            thrust::lower_bound(thrust::seq, unique_meas.begin(),
-                                unique_meas.end(), mid) -
-            unique_meas.begin();
+        const auto unique_meas_idx = meas_id_to_unique_id.at(mid);
 
         if (is_start) {
 

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

@@ -33,8 +33,8 @@ __global__ void count_shared_measurements(
 
     vecmem::jagged_device_vector<const measurement_id_type> meas_ids(
         payload.meas_ids_view);
-    vecmem::device_vector<const measurement_id_type> unique_meas(
-        payload.unique_meas_view);
+    vecmem::device_vector<const unsigned int> meas_id_to_unique_id(
+        payload.meas_id_to_unique_id_view);
     vecmem::device_vector<const unsigned int> n_accepted_tracks_per_measurement(
         payload.n_accepted_tracks_per_measurement_view);
     vecmem::device_vector<unsigned int> n_shared(payload.n_shared_view);
@@ -43,10 +43,7 @@ __global__ void count_shared_measurements(
 
     for (const auto& meas_id : meas_ids[id]) {
 
-        const auto it = thrust::lower_bound(thrust::seq, unique_meas.begin(),
-                                            unique_meas.end(), meas_id);
-        const auto unique_meas_idx = static_cast<unsigned int>(
-            thrust::distance(unique_meas.begin(), it));
+        const auto unique_meas_idx = meas_id_to_unique_id.at(meas_id);
 
         if (n_accepted_tracks_per_measurement.at(unique_meas_idx) > 1) {
             vecmem::device_atomic_ref<unsigned int>(n_shared.at(id))

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

@@ -33,8 +33,8 @@ __global__ void fill_tracks_per_measurement(
 
     vecmem::jagged_device_vector<const measurement_id_type> meas_ids(
         payload.meas_ids_view);
-    vecmem::device_vector<const measurement_id_type> unique_meas(
-        payload.unique_meas_view);
+    vecmem::device_vector<const unsigned int> meas_id_to_unique_id(
+        payload.meas_id_to_unique_id_view);
     vecmem::jagged_device_vector<unsigned int> tracks_per_measurement(
         payload.tracks_per_measurement_view);
     vecmem::jagged_device_vector<int> track_status_per_measurement(
@@ -53,10 +53,7 @@ __global__ void fill_tracks_per_measurement(
             continue;
         }
 
-        const auto it = thrust::lower_bound(thrust::seq, unique_meas.begin(),
-                                            unique_meas.end(), meas_id);
-        const std::size_t unique_meas_idx =
-            static_cast<std::size_t>(thrust::distance(unique_meas.begin(), it));
+        const auto unique_meas_idx = meas_id_to_unique_id.at(meas_id);
 
         auto tracks = tracks_per_measurement.at(unique_meas_idx);