Merge pull request #1023 from krasznaa/DeviceClusterReco-main-20250617

Device Cluster Reconstruction, main branch (2025.06.17.)

Author: krasznaa

device/alpaka/src/clusterization/clusterization_algorithm.cpp

@@ -35,6 +35,8 @@ struct CCLKernel {
         vecmem::data::vector_view<unsigned char> adjc_backup_view,
         vecmem::data::vector_view<device::details::index_t> adjv_backup_view,
         uint32_t* backup_mutex_ptr,
+        vecmem::data::vector_view<unsigned int> disjoint_set_view,
+        vecmem::data::vector_view<unsigned int> cluster_size_view,
         measurement_collection_types::view measurements_view,
         vecmem::data::vector_view<unsigned int> cell_links) const {
 
@@ -57,11 +59,11 @@ struct CCLKernel {
 
         alpaka::barrier<TAcc> barry_r(&acc);
 
-        device::ccl_kernel(cfg, thread_id, cells_view, det_descr_view,
-                           partition_start, partition_end, outi, f_view,
-                           gf_view, f_backup_view, gf_backup_view,
-                           adjc_backup_view, adjv_backup_view, backup_mutex,
-                           barry_r, measurements_view, cell_links);
+        device::ccl_kernel(
+            cfg, thread_id, cells_view, det_descr_view, partition_start,
+            partition_end, outi, f_view, gf_view, f_backup_view, gf_backup_view,
+            adjc_backup_view, adjv_backup_view, backup_mutex, disjoint_set_view,
+            cluster_size_view, barry_r, measurements_view, cell_links);
     }
 };
 
@@ -136,12 +138,13 @@ clusterization_algorithm::output_type clusterization_algorithm::operator()(
                   "with support for multi-thread blocks.");
     auto workDiv = makeWorkDiv<Acc>(num_blocks, m_config.threads_per_partition);
 
+    vecmem::data::vector_view<unsigned int> dummy_view{0u, nullptr};
     ::alpaka::exec<Acc>(
         queue, workDiv, CCLKernel{}, m_config, cells, det_descr,
         vecmem::get_data(m_f_backup), vecmem::get_data(m_gf_backup),
         vecmem::get_data(m_adjc_backup), vecmem::get_data(m_adjv_backup),
-        m_backup_mutex.get(), vecmem::get_data(measurements),
-        vecmem::get_data(cell_links));
+        m_backup_mutex.get(), dummy_view, dummy_view,
+        vecmem::get_data(measurements), vecmem::get_data(cell_links));
     ::alpaka::wait(queue);
 
     return measurements;

device/common/include/traccc/clusterization/device/aggregate_cluster.hpp

@@ -1,6 +1,6 @@
 /** TRACCC library, part of the ACTS project (R&D line)
  *
- * (c) 2022-2024 CERN for the benefit of the ACTS project
+ * (c) 2022-2025 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
@@ -30,6 +30,12 @@ namespace traccc::device {
 /// @param[in] end       partition end point this cell belongs to
 /// @param[in] cid       current cell id
 /// @param[out] out      cluster to fill
+/// @param[out] disjoint_set Array of unsigned integers of
+///                      length $|cells|$ to which an integer is written
+///                      identifying the measurement index to which each cell
+///                      belongs.
+/// @param[out] cluster_size Optional integer which is filled with the size of
+///                      the measurement that is created.
 ///
 TRACCC_HOST_DEVICE
 inline void aggregate_cluster(
@@ -38,7 +44,9 @@ inline void aggregate_cluster(
     const silicon_detector_description::const_device& det_descr,
     const vecmem::device_vector<details::index_t>& f, unsigned int start,
     unsigned int end, unsigned short cid, measurement& out,
-    vecmem::data::vector_view<unsigned int> cell_links, unsigned int link);
+    vecmem::data::vector_view<unsigned int> cell_links, unsigned int link,
+    vecmem::device_vector<unsigned int>& disjoint_set,
+    std::optional<std::reference_wrapper<unsigned int>> cluster_size);
 
 }  // namespace traccc::device
 

device/common/include/traccc/clusterization/device/ccl_kernel.hpp

@@ -49,6 +49,13 @@ namespace traccc::device {
 ///     matrix fragment storage
 /// @param backup_mutex mutex lock to mediate control over the backup global
 ///     memory data structures.
+/// @param[out] disjoint_set_view Array of unsigned integers of
+///     length $|cells|$ to which an integer is written identifying the
+///     measurement index to which each cell belongs.
+/// @param[out] cluster_size_view Array of unsigned integers of
+///     size $|cells|$; the first $N$ elements - where $N$ is the number of
+///     output measurements - will have their value set to the size of the
+///     corresponding measurement.
 /// @param barrier  A generic object for block-wide synchronisation
 /// @param[out] measurements_view collection of measurements
 /// @param[out] cell_links    collection of links to measurements each cell is
@@ -66,7 +73,10 @@ TRACCC_DEVICE inline void ccl_kernel(
     vecmem::data::vector_view<details::index_t> gf_backup_view,
     vecmem::data::vector_view<unsigned char> adjc_backup_view,
     vecmem::data::vector_view<details::index_t> adjv_backup_view,
-    vecmem::device_atomic_ref<uint32_t> backup_mutex, const barrier_t& barrier,
+    vecmem::device_atomic_ref<uint32_t> backup_mutex,
+    vecmem::data::vector_view<unsigned int> disjoint_set_view,
+    vecmem::data::vector_view<unsigned int> cluster_size_view,
+    const barrier_t& barrier,
     measurement_collection_types::view measurements_view,
     vecmem::data::vector_view<unsigned int> cell_links);
 

device/common/include/traccc/clusterization/device/impl/aggregate_cluster.ipp

@@ -1,6 +1,6 @@
 /** TRACCC library, part of the ACTS project (R&D line)
  *
- * (c) 2022-2024 CERN for the benefit of the ACTS project
+ * (c) 2022-2025 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
@@ -19,8 +19,9 @@ inline void aggregate_cluster(
     const silicon_detector_description::const_device& det_descr,
     const vecmem::device_vector<details::index_t>& f, const unsigned int start,
     const unsigned int end, const unsigned short cid, measurement& out,
-    vecmem::data::vector_view<unsigned int> cell_links,
-    const unsigned int link) {
+    vecmem::data::vector_view<unsigned int> cell_links, const unsigned int link,
+    vecmem::device_vector<unsigned int>& disjoint_set,
+    std::optional<std::reference_wrapper<unsigned int>> cluster_size) {
     vecmem::device_vector<unsigned int> cell_links_device(cell_links);
 
     /*
@@ -63,6 +64,7 @@ inline void aggregate_cluster(
     const unsigned int module_idx = cells.module_index().at(cid + start);
     const auto module_descr = det_descr.at(module_idx);
     const auto partition_size = static_cast<unsigned short>(end - start);
+    unsigned int tmp_cluster_size = 0;
 
     bool first_processed = false;
 
@@ -125,6 +127,12 @@ inline void aggregate_cluster(
             }
 
             cell_links_device.at(pos) = link;
+
+            tmp_cluster_size++;
+
+            if (disjoint_set.capacity()) {
+                disjoint_set.at(pos) = link;
+            }
         }
 
         /*
@@ -134,6 +142,10 @@ inline void aggregate_cluster(
         if (cell.channel1() > maxChannel1 + 1) {
             break;
         }
+
+        if (cluster_size.has_value()) {
+            (*cluster_size).get() = tmp_cluster_size;
+        }
     }
 
     var = var + point2{module_descr.pitch_x() * module_descr.pitch_x() /

device/common/include/traccc/clusterization/device/impl/ccl_kernel.ipp

@@ -145,7 +145,8 @@ TRACCC_DEVICE inline void ccl_core(
     const edm::silicon_cell_collection::const_device& cells_device,
     const silicon_detector_description::const_device& det_descr,
     measurement_collection_types::device measurements_device,
-    const barrier_t& barrier) {
+    const barrier_t& barrier, vecmem::device_vector<unsigned int>& disjoint_set,
+    vecmem::device_vector<unsigned int>& cluster_size) {
     const auto size =
         static_cast<details::index_t>(partition_end - partition_start);
 
@@ -197,17 +198,23 @@ TRACCC_DEVICE inline void ccl_core(
     for (details::index_t tst = 0; tst < thread_cell_count; ++tst) {
         const auto cid = static_cast<details::index_t>(
             tst * thread_id.getBlockDimX() + thread_id.getLocalThreadIdX());
+
         if (f.at(cid) == cid) {
             // Add a new measurement to the output buffer. Remembering its
             // position inside of the container.
             const measurement_collection_types::device::size_type meas_pos =
                 measurements_device.push_back({});
             // Set up the measurement under the appropriate index.
-            aggregate_cluster(cfg, cells_device, det_descr, f,
-                              static_cast<unsigned int>(partition_start),
-                              static_cast<unsigned int>(partition_end), cid,
-                              measurements_device.at(meas_pos), cell_links,
-                              meas_pos);
+            aggregate_cluster(
+                cfg, cells_device, det_descr, f,
+                static_cast<unsigned int>(partition_start),
+                static_cast<unsigned int>(partition_end), cid,
+                measurements_device.at(meas_pos), cell_links, meas_pos,
+                disjoint_set,
+                (cluster_size.capacity()
+                     ? std::optional<std::reference_wrapper<
+                           unsigned int>>{cluster_size.at(meas_pos)}
+                     : std::nullopt));
         }
     }
 }
@@ -225,7 +232,10 @@ TRACCC_DEVICE inline void ccl_kernel(
     vecmem::data::vector_view<details::index_t> gf_backup_view,
     vecmem::data::vector_view<unsigned char> adjc_backup_view,
     vecmem::data::vector_view<details::index_t> adjv_backup_view,
-    vecmem::device_atomic_ref<uint32_t> backup_mutex, const barrier_t& barrier,
+    vecmem::device_atomic_ref<uint32_t> backup_mutex,
+    vecmem::data::vector_view<unsigned int> disjoint_set_view,
+    vecmem::data::vector_view<unsigned int> cluster_size_view,
+    const barrier_t& barrier,
     measurement_collection_types::view measurements_view,
     vecmem::data::vector_view<unsigned int> cell_links) {
 
@@ -239,6 +249,8 @@ TRACCC_DEVICE inline void ccl_kernel(
     vecmem::device_vector<details::index_t> gf_backup(gf_backup_view);
     vecmem::device_vector<unsigned char> adjc_backup(adjc_backup_view);
     vecmem::device_vector<details::index_t> adjv_backup(adjv_backup_view);
+    vecmem::device_vector<unsigned int> disjoint_set(disjoint_set_view);
+    vecmem::device_vector<unsigned int> cluster_size(cluster_size_view);
 
     mutex<uint32_t> mutex(backup_mutex);
     unique_lock lock(mutex, std::defer_lock);
@@ -353,7 +365,8 @@ TRACCC_DEVICE inline void ccl_kernel(
     ccl_core(cfg, thread_id, partition_start, partition_end,
              use_scratch ? f_backup : f_primary,
              use_scratch ? gf_backup : gf_primary, cell_links, adjv, adjc,
-             cells_device, det_descr, measurements_device, barrier);
+             cells_device, det_descr, measurements_device, barrier,
+             disjoint_set, cluster_size);
 
     barrier.blockBarrier();
 }

device/common/include/traccc/clusterization/device/impl/reify_cluster_data.ipp

@@ -0,0 +1,30 @@
+/** 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
+ */
+
+#pragma once
+
+namespace traccc::device {
+
+TRACCC_HOST_DEVICE inline void reify_cluster_data(
+    global_index_t thread_id,
+    vecmem::data::vector_view<const unsigned int> disjoint_set_view,
+    traccc::edm::silicon_cluster_collection::view cluster_view) {
+
+    // Create the device objects.
+    const vecmem::device_vector<const unsigned int> disjoint_set(
+        disjoint_set_view);
+    traccc::edm::silicon_cluster_collection::device clusters(cluster_view);
+
+    // Fill the output container.
+    if (thread_id < disjoint_set.size()) {
+        clusters.cell_indices()
+            .at(disjoint_set.at(thread_id))
+            .push_back(thread_id);
+    }
+}
+
+}  // namespace traccc::device

device/common/include/traccc/clusterization/device/reify_cluster_data.hpp

@@ -0,0 +1,36 @@
+/** 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
+ */
+
+#pragma once
+
+// Local include(s).
+#include "traccc/device/global_index.hpp"
+
+// Project include(s).
+#include "traccc/definitions/qualifiers.hpp"
+#include "traccc/edm/silicon_cluster_collection.hpp"
+
+// VecMem include(s).
+#include <vecmem/containers/data/vector_view.hpp>
+
+namespace traccc::device {
+
+/// Fill a cluster collection with the cell indices
+///
+/// @param thread_id The thread identifier
+/// @param disjoint_set_view The cluster/measurement index of each cell
+/// @param cluster_view The collection to fill
+///
+TRACCC_HOST_DEVICE inline void reify_cluster_data(
+    global_index_t thread_id,
+    vecmem::data::vector_view<const unsigned int> disjoint_set_view,
+    traccc::edm::silicon_cluster_collection::view cluster_view);
+
+}  // namespace traccc::device
+
+// Include the implementation.
+#include "traccc/clusterization/device/impl/reify_cluster_data.ipp"

device/common/include/traccc/clusterization/device/tags.hpp

@@ -0,0 +1,29 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2022-2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+namespace traccc::device {
+/**
+ * @defgroup Clustering algorithm cluster retention parameters
+ *
+ * Optional parameters to clustering algorithms which convert directly from
+ * cells to measurements, determining whether to reconstruct the intermediate
+ * cluster data or not.
+ *
+ * @{
+ * @brief Explicitly discard cluster information.
+ */
+struct clustering_discard_disjoint_set {};
+/*
+ * @brief Explicitly reconstruct and return cluster information.
+ */
+struct clustering_keep_disjoint_set {};
+/*
+ * @}
+ */
+}  // namespace traccc::device

device/cuda/CMakeLists.txt

@@ -51,6 +51,7 @@ traccc_add_library( traccc_cuda cuda TYPE SHARED
   "src/clusterization/measurement_sorting_algorithm.cu"
   "src/clusterization/kernels/ccl_kernel.cu"
   "src/clusterization/kernels/ccl_kernel.cuh"
+  "src/clusterization/kernels/reify_cluster_data.cu"
   # Finding
   "include/traccc/cuda/finding/finding_algorithm.hpp"
   "src/finding/finding_algorithm.cu"

device/cuda/include/traccc/cuda/clusterization/clusterization_algorithm.hpp

@@ -13,8 +13,10 @@
 // Project include(s).
 #include "traccc/clusterization/clustering_config.hpp"
 #include "traccc/clusterization/device/ccl_kernel_definitions.hpp"
+#include "traccc/clusterization/device/tags.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/silicon_cell_collection.hpp"
+#include "traccc/edm/silicon_cluster_collection.hpp"
 #include "traccc/geometry/silicon_detector_description.hpp"
 #include "traccc/utils/algorithm.hpp"
 #include "traccc/utils/memory_resource.hpp"
@@ -26,6 +28,7 @@
 
 // System include(s).
 #include <functional>
+#include <optional>
 
 namespace traccc::cuda {
 
@@ -41,6 +44,16 @@ class clusterization_algorithm
     : public algorithm<measurement_collection_types::buffer(
           const edm::silicon_cell_collection::const_view&,
           const silicon_detector_description::const_view&)>,
+      public algorithm<measurement_collection_types::buffer(
+          const edm::silicon_cell_collection::const_view&,
+          const silicon_detector_description::const_view&,
+          device::clustering_discard_disjoint_set&&)>,
+      public algorithm<
+          std::pair<measurement_collection_types::buffer,
+                    traccc::edm::silicon_cluster_collection::buffer>(
+              const edm::silicon_cell_collection::const_view&,
+              const silicon_detector_description::const_view&,
+              device::clustering_keep_disjoint_set&&)>,
       public messaging {
 
     public:
@@ -67,12 +80,31 @@ class clusterization_algorithm
     /// @param det_descr The detector description
     /// @return a measurement collection (buffer)
     ///
-    output_type operator()(
+    /// @{
+    measurement_collection_types::buffer operator()(
         const edm::silicon_cell_collection::const_view& cells,
         const silicon_detector_description::const_view& det_descr)
         const override;
 
+    measurement_collection_types::buffer operator()(
+        const edm::silicon_cell_collection::const_view& cells,
+        const silicon_detector_description::const_view& det_descr,
+        device::clustering_discard_disjoint_set&&) const override;
+
+    std::pair<measurement_collection_types::buffer,
+              traccc::edm::silicon_cluster_collection::buffer>
+    operator()(const edm::silicon_cell_collection::const_view& cells,
+               const silicon_detector_description::const_view& det_descr,
+               device::clustering_keep_disjoint_set&&) const override;
+    /// @}
+
     private:
+    std::pair<measurement_collection_types::buffer,
+              std::optional<traccc::edm::silicon_cluster_collection::buffer>>
+    execute_impl(const edm::silicon_cell_collection::const_view& cells,
+                 const silicon_detector_description::const_view& det_descr,
+                 bool keep_disjoint_set) const;
+
     /// The memory resource(s) to use
     traccc::memory_resource m_mr;
     /// The copy object to use