Merge pull request #1093 from beomki-yeo/optimize-remove-tracks

Optimize `remove_tracks` by skipping unnecessary iteration

Author: stephenswat

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

@@ -86,6 +86,11 @@ struct count_removable_tracks_payload {
      * @brief View object to thread id of measurements to remove
      */
     vecmem::data::vector_view<unsigned int> threads_view;
+
+    /**
+     * @brief The number of threads that can remove its corresponding track
+     */
+    unsigned int* n_valid_threads;
 };
 
 }  // namespace traccc::device

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

@@ -79,11 +79,6 @@ struct remove_tracks_payload {
      */
     unsigned int* n_removable_tracks;
 
-    /**
-     * @brief The number of measurements to remove
-     */
-    unsigned int* n_meas_to_remove;
-
     /**
      * @brief View object to measurements to remove
      */
@@ -113,6 +108,11 @@ struct remove_tracks_payload {
      * @brief View object to the whether track id is updated
      */
     vecmem::data::vector_view<int> is_updated_view;
+
+    /**
+     * @brief The number of threads that can remove its corresponding track
+     */
+    unsigned int* n_valid_threads;
 };
 
 }  // namespace traccc::device

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -375,6 +375,8 @@ greedy_ambiguity_resolution_algorithm::operator()(
         vecmem::make_unique_alloc<unsigned int>(m_mr.main);
     vecmem::unique_alloc_ptr<unsigned int> n_meas_to_remove_device =
         vecmem::make_unique_alloc<unsigned int>(m_mr.main);
+    vecmem::unique_alloc_ptr<unsigned int> n_valid_threads_device =
+        vecmem::make_unique_alloc<unsigned int>(m_mr.main);
 
     // Device objects
     int is_first_iteration = 1;
@@ -459,9 +461,10 @@ greedy_ambiguity_resolution_algorithm::operator()(
                 .n_removable_tracks = n_removable_tracks_device.get(),
                 .n_meas_to_remove = n_meas_to_remove_device.get(),
                 .meas_to_remove_view = meas_to_remove_buffer,
-                .threads_view = threads_buffer});
+                .threads_view = threads_buffer,
+                .n_valid_threads = n_valid_threads_device.get()});
 
-        kernels::remove_tracks<<<1, 1024, 0, stream>>>(
+        kernels::remove_tracks<<<1, 512, 0, stream>>>(
             device::remove_tracks_payload{
                 .sorted_ids_view = sorted_ids_buffer,
                 .n_accepted = n_accepted_device.get(),
@@ -476,13 +479,13 @@ greedy_ambiguity_resolution_algorithm::operator()(
                 .n_shared_view = n_shared_buffer,
                 .rel_shared_view = rel_shared_buffer,
                 .n_removable_tracks = n_removable_tracks_device.get(),
-                .n_meas_to_remove = n_meas_to_remove_device.get(),
                 .meas_to_remove_view = meas_to_remove_buffer,
                 .threads_view = threads_buffer,
                 .terminate = terminate_device.get(),
                 .n_updated_tracks = n_updated_tracks_device.get(),
                 .updated_tracks_view = updated_tracks_buffer,
-                .is_updated_view = is_updated_buffer});
+                .is_updated_view = is_updated_buffer,
+                .n_valid_threads = n_valid_threads_device.get()});
 
         // 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/count_removable_tracks.cu

@@ -102,6 +102,7 @@ __launch_bounds__(512) __global__ void count_removable_tracks(
     if (threadIndex == 0) {
         *(payload.n_removable_tracks) = 0;
         *(payload.n_meas_to_remove) = 0;
+        *(payload.n_valid_threads) = 0;
         n_meas_total = 0;
         bound = 512;
         N = 1;
@@ -242,20 +243,12 @@ __launch_bounds__(512) __global__ void count_removable_tracks(
 
     __syncthreads();
 
-    auto n_meas_to_remove_temp = *(payload.n_meas_to_remove);
-
-    if (threadIndex == 0) {
-        *(payload.n_meas_to_remove) = 0;
-    }
-
-    __syncthreads();
-
     int is_valid =
         (threads[threadIndex] < *(payload.n_removable_tracks)) ? 1 : 0;
 
     // TODO: Use better reduction algorithm
     if (is_valid) {
-        atomicAdd(payload.n_meas_to_remove, 1);
+        atomicAdd(payload.n_valid_threads, 1);
     }
 
     __syncthreads();
@@ -264,7 +257,7 @@ __launch_bounds__(512) __global__ void count_removable_tracks(
     prefix[threadIndex] = is_valid;  // copy input
     __syncthreads();
 
-    for (int offset = 1; offset < n_meas_to_remove_temp; offset <<= 1) {
+    for (int offset = 1; offset < *(payload.n_meas_to_remove); offset <<= 1) {
         int val = 0;
         if (threadIndex >= offset) {
             val = prefix[threadIndex - offset];

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

@@ -26,15 +26,16 @@
 
 namespace traccc::cuda::kernels {
 
-__global__ void remove_tracks(device::remove_tracks_payload payload) {
+__launch_bounds__(512) __global__
+    void remove_tracks(device::remove_tracks_payload payload) {
 
     if (*(payload.terminate) == 1) {
         return;
     }
 
-    __shared__ unsigned int shared_tids[1024];
-    __shared__ measurement_id_type sh_meas_ids[1024];
-    __shared__ unsigned int sh_threads[1024];
+    __shared__ unsigned int shared_tids[512];
+    __shared__ measurement_id_type sh_meas_ids[512];
+    __shared__ unsigned int sh_threads[512];
 
     auto threadIndex = threadIdx.x;
 
@@ -73,7 +74,7 @@ __global__ void remove_tracks(device::remove_tracks_payload payload) {
         (*payload.n_accepted) -= *(payload.n_removable_tracks);
     }
 
-    if (threadIndex < *(payload.n_meas_to_remove)) {
+    if (threadIndex < *(payload.n_valid_threads)) {
         sh_meas_ids[threadIndex] = meas_to_remove[threadIndex];
         sh_threads[threadIndex] = threads[threadIndex];
         is_valid_thread = true;
@@ -82,16 +83,8 @@ __global__ void remove_tracks(device::remove_tracks_payload payload) {
     __syncthreads();
 
     if (is_valid_thread) {
-
         const auto id = sh_meas_ids[threadIndex];
-        is_duplicate = false;
-
-        for (unsigned int i = 0; i < threadIndex; ++i) {
-            if (sh_meas_ids[i] == id) {
-                is_duplicate = true;
-                break;
-            }
-        }
+        is_duplicate = (threadIndex > 0 && sh_meas_ids[threadIndex - 1] == id);
     }
 
     bool active = false;
@@ -116,7 +109,7 @@ __global__ void remove_tracks(device::remove_tracks_payload payload) {
         track_status[worst_idx] = 0;
 
         int n_sharing_tracks = 1;
-        for (unsigned int i = threadIndex + 1; i < *(payload.n_meas_to_remove);
+        for (unsigned int i = threadIndex + 1; i < *(payload.n_valid_threads);
              ++i) {
 
             if (sh_meas_ids[i] == id && sh_threads[i] != sh_threads[i - 1]) {