Merge find_max_shared and gather_tracks kernel

Author: beomki-yeo

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

@@ -13,8 +13,8 @@
 namespace traccc::device {
 
 /// (Event Data) Payload for the @c
-/// traccc::device::gather_tracks function
-struct gather_tracks_payload {
+/// traccc::device::update_status function
+struct update_status_payload {
 
     /**
      * @brief Whether to terminate the calculation
@@ -50,6 +50,16 @@ struct gather_tracks_payload {
      * @brief View object to the whether track id is updated
      */
     vecmem::data::vector_view<int> is_updated_view;
+
+    /**
+     * @brief View object to the vector of number of shared measurements
+     */
+    vecmem::data::vector_view<const unsigned int> n_shared_view;
+
+    /**
+     * @brief The number of max shared
+     */
+    unsigned int* max_shared;
 };
 
 }  // namespace traccc::device

device/common/include/traccc/ambiguity_resolution/device/gather_tracks.hpp renamed to device/common/include/traccc/ambiguity_resolution/device/update_status.hpp

@@ -91,8 +91,6 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/ambiguity_resolution/kernels/block_inclusive_scan.cuh"
   "src/ambiguity_resolution/kernels/count_shared_measurements.cu"
   "src/ambiguity_resolution/kernels/count_shared_measurements.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"
   "src/ambiguity_resolution/kernels/fill_inverted_ids.cuh"
   "src/ambiguity_resolution/kernels/fill_tracks_per_measurement.cu"
@@ -103,8 +101,8 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/ambiguity_resolution/kernels/fill_track_candidates.cuh"
   "src/ambiguity_resolution/kernels/fill_vectors.cu"
   "src/ambiguity_resolution/kernels/fill_vectors.cuh"
-  "src/ambiguity_resolution/kernels/gather_tracks.cu"
-  "src/ambiguity_resolution/kernels/gather_tracks.cuh"
+  "src/ambiguity_resolution/kernels/update_status.cu"
+  "src/ambiguity_resolution/kernels/update_status.cuh"
   "src/ambiguity_resolution/kernels/rearrange_tracks.cu"
   "src/ambiguity_resolution/kernels/rearrange_tracks.cuh"
   "src/ambiguity_resolution/kernels/scan_block_offsets.cu"

device/cuda/CMakeLists.txt

@@ -16,13 +16,12 @@
 #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"
 #include "./kernels/rearrange_tracks.cuh"
 #include "./kernels/remove_tracks.cuh"
 #include "./kernels/scan_block_offsets.cuh"
 #include "./kernels/sort_tracks_per_measurement.cuh"
 #include "./kernels/sort_updated_tracks.cuh"
+#include "./kernels/update_status.cuh"
 #include "traccc/cuda/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.hpp"
 #include "traccc/definitions/math.hpp"
 
@@ -406,20 +405,21 @@ greedy_ambiguity_resolution_algorithm::operator()(
     int terminate = 0;
     vecmem::unique_alloc_ptr<int> terminate_device =
         vecmem::make_unique_alloc<int>(m_mr.main);
+    auto max_shared = thrust::max_element(thrust::device, n_shared_buffer.ptr(),
+                                          n_shared_buffer.ptr() + n_tracks);
     vecmem::unique_alloc_ptr<unsigned int> max_shared_device =
         vecmem::make_unique_alloc<unsigned int>(m_mr.main);
+    cudaMemcpyAsync(max_shared_device.get(), max_shared, sizeof(unsigned int),
+                    cudaMemcpyHostToDevice, stream);
+
     vecmem::unique_alloc_ptr<unsigned int> n_updated_tracks_device =
         vecmem::make_unique_alloc<unsigned int>(m_mr.main);
 
     // Thread block size
-    unsigned int nThreads_adaptive = m_warp_size * 4;
+    unsigned int nThreads_adaptive = m_warp_size;
     unsigned int nBlocks_adaptive =
         (n_accepted + nThreads_adaptive - 1) / nThreads_adaptive;
 
-    unsigned int nThreads_warp = m_warp_size;
-    unsigned int nBlocks_warp =
-        (n_accepted + nThreads_warp - 1) / nThreads_warp;
-
     unsigned int nThreads_rearrange = 1024;
     unsigned int nBlocks_rearrange =
         (n_accepted + (nThreads_rearrange / kernels::nThreads_per_track) - 1) /
@@ -462,7 +462,6 @@ greedy_ambiguity_resolution_algorithm::operator()(
     while (!terminate && n_accepted > 0) {
         nBlocks_adaptive =
             (n_accepted + nThreads_adaptive - 1) / nThreads_adaptive;
-        nBlocks_warp = (n_accepted + nThreads_warp - 1) / nThreads_warp;
 
         scan_dim = compute_scan_config(n_accepted);
         nThreads_scan = scan_dim.first;
@@ -478,14 +477,6 @@ greedy_ambiguity_resolution_algorithm::operator()(
 
         cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
 
-        kernels::find_max_shared<<<nBlocks_warp, nThreads_warp, 0, stream>>>(
-            device::find_max_shared_payload{
-                .sorted_ids_view = sorted_ids_buffer,
-                .n_accepted = n_accepted_device.get(),
-                .n_shared_view = n_shared_buffer,
-                .terminate = terminate_device.get(),
-                .max_shared = max_shared_device.get()});
-
         kernels::remove_tracks<<<1, 512, 0, stream>>>(
             device::remove_tracks_payload{
                 .sorted_ids_view = sorted_ids_buffer,
@@ -589,15 +580,17 @@ greedy_ambiguity_resolution_algorithm::operator()(
         });
 
         kernels::
-            gather_tracks<<<nBlocks_adaptive, nThreads_adaptive, 0, stream>>>(
-                device::gather_tracks_payload{
+            update_status<<<nBlocks_adaptive, nThreads_adaptive, 0, stream>>>(
+                device::update_status_payload{
                     .terminate = terminate_device.get(),
                     .n_accepted = n_accepted_device.get(),
                     .n_updated_tracks = n_updated_tracks_device.get(),
                     .temp_sorted_ids_view = temp_sorted_ids_buffer,
                     .sorted_ids_view = sorted_ids_buffer,
                     .updated_tracks_view = updated_tracks_buffer,
-                    .is_updated_view = is_updated_buffer});
+                    .is_updated_view = is_updated_buffer,
+                    .n_shared_view = n_shared_buffer,
+                    .max_shared = max_shared_device.get()});
 
         cudaStreamEndCapture(stream, &graph);
         cudaGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0);

device/cuda/src/ambiguity_resolution/greedy_ambiguity_resolution_algorithm.cu

@@ -0,0 +1,78 @@
+/** 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 "update_status.cuh"
+
+// VecMem include(s).
+#include <vecmem/containers/device_vector.hpp>
+
+namespace traccc::cuda::kernels {
+
+__global__ void update_status(device::update_status_payload payload) {
+
+    if (*(payload.terminate) == 1) {
+        return;
+    }
+
+    vecmem::device_vector<const unsigned int> temp_sorted_ids(
+        payload.temp_sorted_ids_view);
+    vecmem::device_vector<unsigned int> sorted_ids(payload.sorted_ids_view);
+    vecmem::device_vector<unsigned int> updated_tracks(
+        payload.updated_tracks_view);
+    vecmem::device_vector<int> is_updated(payload.is_updated_view);
+    vecmem::device_vector<const unsigned int> n_shared(payload.n_shared_view);
+
+    auto globalIndex = threadIdx.x + blockIdx.x * blockDim.x;
+    const unsigned int n_accepted = *(payload.n_accepted);
+    const unsigned int n_updated = *(payload.n_updated_tracks);
+
+    /***********************
+     * Update Max Shared
+     ***********************/
+
+    unsigned int shared = 0;
+
+    if (globalIndex < n_accepted) {
+        auto tid = sorted_ids[globalIndex];
+        shared = n_shared[tid];
+    }
+
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        unsigned int other_shared =
+            __shfl_down_sync(0xffffffff, shared, offset);
+
+        if (other_shared > shared) {
+            shared = other_shared;
+        }
+    }
+
+    if (threadIdx.x == 0) {
+        atomicMax(payload.max_shared, shared);
+    }
+
+    /***********************
+     * Update Sorted Ids
+     ***********************/
+
+    if (n_updated == 0) {
+        return;
+    }
+
+    // Reset is_updated vector
+    if (globalIndex < n_updated) {
+        is_updated[updated_tracks[globalIndex]] = 0;
+    }
+
+    if (globalIndex < n_accepted) {
+        auto tid = temp_sorted_ids[globalIndex];
+        sorted_ids[globalIndex] = tid;
+    }
+}
+
+}  // namespace traccc::cuda::kernels

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

@@ -8,9 +8,9 @@
 #pragma once
 
 // Project include(s).
-#include "traccc/ambiguity_resolution/device/gather_tracks.hpp"
+#include "traccc/ambiguity_resolution/device/update_status.hpp"
 
 namespace traccc::cuda::kernels {
 
-__global__ void gather_tracks(device::gather_tracks_payload payload);
+__global__ void update_status(device::update_status_payload payload);
 }