Merge pull request #843 from krasznaa/AsyncThrust-main-20250205

Asynchronous Thrust, main branch (2025.02.05.)

Author: stephenswat

device/cuda/src/clusterization/measurement_sorting_algorithm.cu

@@ -1,6 +1,6 @@
 /** TRACCC library, part of the ACTS project (R&D line)
  *
- * (c) 2024 CERN for the benefit of the ACTS project
+ * (c) 2024-2025 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
@@ -12,6 +12,9 @@
 // Thrust include(s).
 #include <thrust/sort.h>
 
+// System include(s).
+#include <memory_resource>
+
 namespace traccc::cuda {
 
 measurement_sorting_algorithm::measurement_sorting_algorithm(
@@ -33,7 +36,7 @@ measurement_sorting_algorithm::operator()(
 
     // Sort the measurements in place
     thrust::sort(
-        thrust::cuda::par(std::pmr::polymorphic_allocator(&(m_mr.main)))
+        thrust::cuda::par_nosync(std::pmr::polymorphic_allocator(&(m_mr.main)))
             .on(stream),
         measurements_view.ptr(), measurements_view.ptr() + n_measurements,
         measurement_sort_comp());

device/cuda/src/finding/finding_algorithm.cu

@@ -49,6 +49,7 @@
 
 // System include(s).
 #include <cassert>
+#include <memory_resource>
 #include <vector>
 
 namespace traccc::cuda {
@@ -77,6 +78,11 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
     // Copy setup
     m_copy.setup(seeds_buffer)->ignore();
 
+    // The Thrust policy to use.
+    auto thrust_policy =
+        thrust::cuda::par_nosync(std::pmr::polymorphic_allocator(&(m_mr.main)))
+            .on(stream);
+
     /*****************************************************************
      * Measurement Operations
      *****************************************************************/
@@ -97,11 +103,11 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
         measurement_collection_types::device uniques(uniques_buffer);
 
-        measurement* uniques_end = thrust::unique_copy(
-            thrust::cuda::par(std::pmr::polymorphic_allocator(&(m_mr.main)))
-                .on(stream),
-            measurements.ptr(), measurements.ptr() + n_measurements,
-            uniques.begin(), measurement_equal_comp());
+        measurement* uniques_end =
+            thrust::unique_copy(thrust_policy, measurements.ptr(),
+                                measurements.ptr() + n_measurements,
+                                uniques.begin(), measurement_equal_comp());
+        m_stream.synchronize();
         n_modules = static_cast<unsigned int>(uniques_end - uniques.begin());
     }
 
@@ -115,12 +121,10 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
         measurement_collection_types::device uniques(uniques_buffer);
 
-        thrust::upper_bound(
-            thrust::cuda::par(std::pmr::polymorphic_allocator(&(m_mr.main)))
-                .on(stream),
-            measurements.ptr(), measurements.ptr() + n_measurements,
-            uniques.begin(), uniques.begin() + n_modules, upper_bounds.begin(),
-            measurement_sort_comp());
+        thrust::upper_bound(thrust_policy, measurements.ptr(),
+                            measurements.ptr() + n_measurements,
+                            uniques.begin(), uniques.begin() + n_modules,
+                            upper_bounds.begin(), measurement_sort_comp());
     }
 
     /*****************************************************************
@@ -285,12 +289,9 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
                     keys_buffer);
                 vecmem::device_vector<unsigned int> param_ids_device(
                     param_ids_buffer);
-                thrust::sort_by_key(
-                    thrust::cuda::par(
-                        std::pmr::polymorphic_allocator(&(m_mr.main)))
-                        .on(stream),
-                    keys_device.begin(), keys_device.end(),
-                    param_ids_device.begin());
+                thrust::sort_by_key(thrust_policy, keys_device.begin(),
+                                    keys_device.end(),
+                                    param_ids_device.begin());
 
                 m_stream.synchronize();
             }
@@ -338,10 +339,8 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
         vecmem::device_vector<candidate_link> out(
             *(links_buffer.host_ptr() + it));
 
-        thrust::copy(
-            thrust::cuda::par(std::pmr::polymorphic_allocator(&(m_mr.main)))
-                .on(stream),
-            in.begin(), in.begin() + n_candidates_per_step[it], out.begin());
+        thrust::copy(thrust_policy, in.begin(),
+                     in.begin() + n_candidates_per_step[it], out.begin());
     }
 
     /*****************************************************************

device/cuda/src/fitting/fitting_algorithm.cu

@@ -111,10 +111,11 @@ track_state_container_types::buffer fitting_algorithm<fitter_t>::operator()(
         vecmem::device_vector<device::sort_key> keys_device(keys_buffer);
         vecmem::device_vector<unsigned int> param_ids_device(param_ids_buffer);
 
-        thrust::sort_by_key(
-            thrust::cuda::par(std::pmr::polymorphic_allocator(&m_mr.main))
-                .on(stream),
-            keys_device.begin(), keys_device.end(), param_ids_device.begin());
+        thrust::sort_by_key(thrust::cuda::par_nosync(
+                                std::pmr::polymorphic_allocator(&m_mr.main))
+                                .on(stream),
+                            keys_device.begin(), keys_device.end(),
+                            param_ids_device.begin());
 
         // Run the track fitting
         kernels::fit<fitter_t><<<nBlocks, nThreads, 0, stream>>>(