Remove the iterations in deduplication (#1126)

Author: beomki-yeo

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

@@ -113,6 +113,11 @@ struct remove_tracks_payload {
      * @brief The number of threads that can remove its corresponding track
      */
     unsigned int* n_valid_threads;
+
+    /**
+     * @brief View object to the vector of track count during removal process
+     */
+    vecmem::data::vector_view<int> track_count_view;
 };
 
 }  // namespace traccc::device

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -365,6 +365,11 @@ greedy_ambiguity_resolution_algorithm::operator()(
     vecmem::data::vector_buffer<int> is_updated_buffer{n_tracks, m_mr.main};
     m_copy.get().setup(inverted_ids_buffer)->ignore();
 
+    // Count track id apperance during removal process
+    vecmem::data::vector_buffer<int> track_count_buffer{n_tracks, m_mr.main};
+    m_copy.get().setup(track_count_buffer)->ignore();
+    m_copy.get().memset(track_count_buffer, 0)->ignore();
+
     // Prefix sum buffer
     vecmem::data::vector_buffer<int> prefix_sums_buffer{n_tracks, m_mr.main};
     m_copy.get().setup(prefix_sums_buffer)->ignore();
@@ -511,7 +516,8 @@ greedy_ambiguity_resolution_algorithm::operator()(
                 .n_updated_tracks = n_updated_tracks_device.get(),
                 .updated_tracks_view = updated_tracks_buffer,
                 .is_updated_view = is_updated_buffer,
-                .n_valid_threads = n_valid_threads_device.get()});
+                .n_valid_threads = n_valid_threads_device.get(),
+                .track_count_view = track_count_buffer});
 
         // The seven kernels below are to keep sorted_ids sorted based on
         // the relative shared measurements and pvalues. This can be reduced

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

@@ -111,6 +111,7 @@ __launch_bounds__(512) __global__
     vecmem::device_vector<unsigned int> updated_tracks(
         payload.updated_tracks_view);
     vecmem::device_vector<int> is_updated(payload.is_updated_view);
+    vecmem::device_vector<int> track_count(payload.track_count_view);
 
     if (threadIndex == 0) {
         *(payload.n_removable_tracks) = 0;
@@ -342,6 +343,7 @@ __launch_bounds__(512) __global__
 
     bool active = false;
     unsigned int pos1;
+    int alive_trk_id = 0;
 
     if (!is_duplicate && is_valid_thread) {
 
@@ -397,46 +399,40 @@ __launch_bounds__(512) __global__
                 track_status.begin();
 
             pos1 = atomicAdd(&n_updating_threads, 1);
+            alive_trk_id = static_cast<int>(tracks[alive_idx]);
 
-            sh_buffer[pos1] = static_cast<int>(tracks[alive_idx]);
+            sh_buffer[pos1] = alive_trk_id;
+            atomicAdd(&track_count[alive_trk_id], 1);
 
-            auto tid = sh_buffer[pos1];
+            const auto m_count = static_cast<unsigned int>(
+                thrust::count(thrust::seq, meas_ids[alive_trk_id].begin(),
+                              meas_ids[alive_trk_id].end(), id));
 
-            const auto m_count = static_cast<unsigned int>(thrust::count(
-                thrust::seq, meas_ids[tid].begin(), meas_ids[tid].end(), id));
-
-            const unsigned int N_S =
-                vecmem::device_atomic_ref<unsigned int>(n_shared.at(tid))
-                    .fetch_sub(m_count);
+            const unsigned int N_S = vecmem::device_atomic_ref<unsigned int>(
+                                         n_shared.at(alive_trk_id))
+                                         .fetch_sub(m_count);
         }
     }
 
     __syncthreads();
 
     if (active) {
-        auto tid = sh_buffer[pos1];
-        bool already_pushed = false;
-        for (unsigned int i = 0; i < pos1; ++i) {
-            if (sh_buffer[i] == tid) {
-                already_pushed = true;
-                break;
-            }
-        }
 
-        if (!already_pushed) {
+        auto count = atomicAdd(&track_count[alive_trk_id], -1);
+        if (count == 1) {
 
             // Write updated track IDs
             vecmem::device_atomic_ref<unsigned int> num_updated_tracks(
                 *(payload.n_updated_tracks));
 
             const unsigned int pos2 = num_updated_tracks.fetch_add(1);
 
-            updated_tracks[pos2] = tid;
-            is_updated[tid] = 1;
+            updated_tracks[pos2] = alive_trk_id;
+            is_updated[alive_trk_id] = 1;
 
-            rel_shared.at(tid) =
-                math::div_ieee754(static_cast<traccc::scalar>(n_shared.at(tid)),
-                                  static_cast<traccc::scalar>(n_meas.at(tid)));
+            rel_shared.at(alive_trk_id) = math::div_ieee754(
+                static_cast<traccc::scalar>(n_shared.at(alive_trk_id)),
+                static_cast<traccc::scalar>(n_meas.at(alive_trk_id)));
         }
     }
 }