Merge pull request #1076 from beomki-yeo/remove-exclusive-scan

Remove exclusive_scan kernel in the greedy ambiguity resolver

Author: stephenswat

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

@@ -91,8 +91,6 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/ambiguity_resolution/kernels/count_removable_tracks.cuh"
   "src/ambiguity_resolution/kernels/count_shared_measurements.cu"
   "src/ambiguity_resolution/kernels/count_shared_measurements.cuh"
-  "src/ambiguity_resolution/kernels/exclusive_scan.cu"
-  "src/ambiguity_resolution/kernels/exclusive_scan.cuh"
   "src/ambiguity_resolution/kernels/find_max_shared.cu"
   "src/ambiguity_resolution/kernels/find_max_shared.cuh"
   "src/ambiguity_resolution/kernels/fill_inverted_ids.cu"

device/cuda/CMakeLists.txt

@@ -12,7 +12,6 @@
 #include "./kernels/block_inclusive_scan.cuh"
 #include "./kernels/count_removable_tracks.cuh"
 #include "./kernels/count_shared_measurements.cuh"
-#include "./kernels/exclusive_scan.cuh"
 #include "./kernels/fill_inverted_ids.cuh"
 #include "./kernels/fill_track_candidates.cuh"
 #include "./kernels/fill_tracks_per_measurement.cuh"
@@ -462,14 +461,6 @@ greedy_ambiguity_resolution_algorithm::operator()(
                 .meas_to_remove_view = meas_to_remove_buffer,
                 .threads_view = threads_buffer});
 
-        kernels::exclusive_scan<<<1, 1024, 0, stream>>>(
-            device::exclusive_scan_payload{
-                .terminate = terminate_device.get(),
-                .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});
-
         kernels::remove_tracks<<<1, 1024, 0, stream>>>(
             device::remove_tracks_payload{
                 .sorted_ids_view = sorted_ids_buffer,

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -70,6 +70,7 @@ __launch_bounds__(512) __global__ void count_removable_tracks(
     __shared__ int shared_n_meas[512];
     __shared__ measurement_id_type sh_meas_ids[512];
     __shared__ unsigned int sh_threads[512];
+    __shared__ int prefix[512];
     __shared__ unsigned int n_meas_total;
     __shared__ unsigned int bound;
     __shared__ unsigned int n_tracks_to_iterate;
@@ -237,6 +238,51 @@ __launch_bounds__(512) __global__ void count_removable_tracks(
     if (threadIndex == 0) {
         *(payload.n_meas_to_remove) = n_meas_total;
     }
+
+    __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);
+    }
+
+    __syncthreads();
+
+    // Exclusive scan (Hillis-Steele)
+    prefix[threadIndex] = is_valid;  // copy input
+    __syncthreads();
+
+    for (int offset = 1; offset < n_meas_to_remove_temp; offset <<= 1) {
+        int val = 0;
+        if (threadIndex >= offset) {
+            val = prefix[threadIndex - offset];
+        }
+        __syncthreads();
+        prefix[threadIndex] += val;
+        __syncthreads();
+    }
+
+    if (is_valid) {
+        prefix[threadIndex] -= 1;
+        sh_meas_ids[prefix[threadIndex]] = meas_to_remove[threadIndex];
+        sh_threads[prefix[threadIndex]] = threads[threadIndex];
+    }
+
+    __syncthreads();
+
+    meas_to_remove[threadIndex] = sh_meas_ids[threadIndex];
+    threads[threadIndex] = sh_threads[threadIndex];
 }
 
 }  // namespace traccc::cuda::kernels