Implement a best-for-n-spacepoints cut on seeds

This PR implements a cut that requires a seed to be the best seed for at
least $n$ of its constituent spacepoints, which mirrors the cut that is
used in ACTS. This cut is a specialisation of the cut in #1082,
implemented in a way that should allow it to run faster.

Author: stephenswat

core/include/traccc/seeding/detail/seeding_config.hpp

@@ -203,6 +203,10 @@ struct seedfilter_config {
     // how often do you want to increase the weight of a seed for finding a
     // compatible seed?
     size_t compatSeedLimit = 2;
+    // Number of consituents spacepoints for which a seed must be the best
+    // seed, where "best" is defined by the weight. Valid range is
+    // [0, 1, 2, 3] where 0 disables the cut.
+    unsigned int minNumTimesHighestWeight = 1;
     // Tool to apply experiment specific cuts on collected middle space points
 
     size_t max_triplets_per_spM = 5;

device/common/include/traccc/seeding/device/impl/select_seeds.ipp

@@ -66,7 +66,8 @@ inline void select_seeds(
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view,
     const device_triplet_collection_types::const_view& triplet_view,
-    triplet* data, edm::seed_collection::view seed_view) {
+    vecmem::data::vector_view<const scalar> highest_weight_view, triplet* data,
+    edm::seed_collection::view seed_view) {
 
     // Check if anything needs to be done.
     const triplet_counter_spM_collection_types::const_device triplet_counts_spM(
@@ -84,13 +85,18 @@ inline void select_seeds(
         sp_view);
 
     const device_triplet_collection_types::const_device triplets(triplet_view);
+
+    const vecmem::device_vector<const scalar> highest_weights(
+        highest_weight_view);
+
     edm::seed_collection::device seeds_device(seed_view);
 
     // Current work item = middle spacepoint
     const triplet_counter_spM spM_counter = triplet_counts_spM.at(globalIndex);
     const sp_location spM_loc = spM_counter.spM;
+    const unsigned int spM_idx = sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx];
     const edm::spacepoint_collection::const_device::const_proxy_type spM =
-        spacepoints.at(sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
+        spacepoints.at(spM_idx);
 
     // Number of triplets added for this spM
     unsigned int n_triplets_per_spM = 0;
@@ -106,10 +112,33 @@ inline void select_seeds(
             triplet_counts.at(static_cast<unsigned int>(aTriplet.counter_link))
                 .spB;
         const sp_location spT_loc = aTriplet.spT;
+        const unsigned int spB_idx =
+            sp_device.bin(spB_loc.bin_idx)[spB_loc.sp_idx];
         const edm::spacepoint_collection::const_device::const_proxy_type spB =
-            spacepoints.at(sp_device.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
+            spacepoints.at(spB_idx);
+        const unsigned int spT_idx =
+            sp_device.bin(spT_loc.bin_idx)[spT_loc.sp_idx];
         const edm::spacepoint_collection::const_device::const_proxy_type spT =
-            spacepoints.at(sp_device.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
+            spacepoints.at(spT_idx);
+
+        if (filter_config.minNumTimesHighestWeight > 0) {
+            assert(highest_weights.capacity() > 0);
+            unsigned int highest_for_n_spacepoints = 0;
+            if (highest_weights.at(spB_idx) <= aTriplet.weight) {
+                highest_for_n_spacepoints++;
+            }
+            if (highest_weights.at(spM_idx) <= aTriplet.weight) {
+                highest_for_n_spacepoints++;
+            }
+            if (highest_weights.at(spT_idx) <= aTriplet.weight) {
+                highest_for_n_spacepoints++;
+            }
+
+            if (highest_for_n_spacepoints <
+                filter_config.minNumTimesHighestWeight) {
+                continue;
+            }
+        }
 
         // update weight of triplet
         seed_selecting_helper::seed_weight(filter_config, spM, spB, spT,

device/common/include/traccc/seeding/device/select_seeds.hpp

@@ -41,7 +41,8 @@ inline void select_seeds(
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view,
     const device_triplet_collection_types::const_view& triplet_view,
-    triplet* data, edm::seed_collection::view seed_view);
+    vecmem::data::vector_view<const scalar> highest_weight_view, triplet* data,
+    edm::seed_collection::view seed_view);
 
 }  // namespace traccc::device
 

device/cuda/src/seeding/seed_finding.cu

@@ -135,6 +135,7 @@ __global__ void select_seeds(
     device::triplet_counter_spM_collection_types::const_view spM_tc,
     device::triplet_counter_collection_types::const_view midBot_tc,
     device::device_triplet_collection_types::view triplet_view,
+    vecmem::data::vector_view<const scalar> highest_weight_view,
     edm::seed_collection::view seed_view) {
 
     // Array for temporary storage of triplets for comparing within seed
@@ -144,10 +145,72 @@ __global__ void select_seeds(
     triplet* dataPos = &data2[threadIdx.x * filter_config.max_triplets_per_spM];
 
     device::select_seeds(details::global_index1(), filter_config, spacepoints,
-                         sp_view, spM_tc, midBot_tc, triplet_view, dataPos,
-                         seed_view);
+                         sp_view, spM_tc, midBot_tc, triplet_view,
+                         highest_weight_view, dataPos, seed_view);
 }
 
+__global__ void collect_highest_weight_per_spacepoint(
+    vecmem::data::vector_view<scalar> highest_weight_view,
+    traccc::details::spacepoint_grid_types::const_view sp_grid_view,
+    device::device_triplet_collection_types::const_view triplet_view) {
+    vecmem::device_vector<scalar> highest_weights(highest_weight_view);
+    const traccc::details::spacepoint_grid_types::const_device sp_grid(
+        sp_grid_view);
+    const device::device_triplet_collection_types::const_device triplets(
+        triplet_view);
+
+    const unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    const unsigned int seed_idx = tid / 3;
+    const unsigned int local_spacepoint_idx = tid % 3;
+
+    if (seed_idx >= triplets.size()) {
+        return;
+    }
+
+    sp_location global_spacepoint_loc;
+
+    if (local_spacepoint_idx == 0) {
+        global_spacepoint_loc = triplets.at(seed_idx).spB;
+    } else if (local_spacepoint_idx == 1) {
+        global_spacepoint_loc = triplets.at(seed_idx).spM;
+    } else if (local_spacepoint_idx == 2) {
+        global_spacepoint_loc = triplets.at(seed_idx).spT;
+    } else {
+        __builtin_unreachable();
+    }
+
+    const unsigned int global_spacepoint_idx = sp_grid.bin(
+        global_spacepoint_loc.bin_idx)[global_spacepoint_loc.sp_idx];
+
+    static_assert(sizeof(scalar) == 4 || sizeof(scalar) == 8);
+    using cas_type = std::conditional_t<sizeof(scalar) == 4, unsigned int,
+                                        unsigned long long int>;
+
+    /*
+     * The following is simply an implementation of atomic max.
+     */
+    scalar& weight_loc = highest_weights.at(global_spacepoint_idx);
+    cas_type* weight_raw = reinterpret_cast<cas_type*>(&weight_loc);
+
+    vecmem::device_atomic_ref<cas_type> atomic(*weight_raw);
+
+    cas_type current_weight_raw = atomic.load();
+    scalar current_weight = std::bit_cast<scalar>(current_weight_raw);
+    const scalar own_weight = triplets.at(seed_idx).weight;
+    const cas_type own_weight_raw = std::bit_cast<cas_type>(own_weight);
+
+    while (own_weight > current_weight) {
+        const bool res =
+            atomic.compare_exchange_strong(current_weight_raw, own_weight_raw);
+
+        if (res) {
+            return;
+        } else {
+            current_weight = std::bit_cast<scalar>(current_weight_raw);
+        }
+    }
+}
 }  // namespace kernels
 
 namespace details {
@@ -184,6 +247,9 @@ edm::seed_collection::buffer seed_finding::operator()(
         return {0, m_mr.main};
     }
 
+    // Total number of spacepoints, including those not in the grid.
+    const auto total_num_spacepoints = m_copy.get_size(spacepoints_view);
+
     // Set up the doublet counter buffer.
     device::doublet_counter_collection_types::buffer doublet_counter_buffer = {
         num_spacepoints, m_mr.main, vecmem::data::buffer_type::resizable};
@@ -340,6 +406,32 @@ edm::seed_collection::buffer seed_finding::operator()(
                   triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
                   triplet_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+    TRACCC_CUDA_ERROR_CHECK(cudaStreamSynchronize(stream));
+
+    vecmem::data::vector_buffer<scalar> highest_weight_buffer(0, m_mr.main);
+
+    if (m_seedfilter_config.minNumTimesHighestWeight > 0) {
+        highest_weight_buffer = {total_num_spacepoints, m_mr.main};
+
+        /*
+         * NOTE: The value 0b11100000 is chosen because it, when repeated four
+         * times to create a 32-bit floating point value, evaluates to the
+         * number 0b11100000111000001110000011100000 which is very small, much
+         * smaller than any seed weight should ever be. When broadcast eight
+         * times to a 64-bit number it also produces an incredibly small value.
+         */
+        m_copy.memset(highest_weight_buffer, 0b11100000)->wait();
+
+        const unsigned int num_votes = 3 * globalCounter_host->m_nTriplets;
+        const unsigned int num_threads = 256;
+        const unsigned int num_blocks =
+            (num_votes + num_threads - 1) / num_threads;
+
+        kernels::collect_highest_weight_per_spacepoint<<<
+            num_blocks, num_threads, 0, stream>>>(highest_weight_buffer,
+                                                  g2_view, triplet_buffer);
+        TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+    }
 
     // Create result object: collection of seeds
     edm::seed_collection::buffer seed_buffer(
@@ -359,10 +451,10 @@ edm::seed_collection::buffer seed_finding::operator()(
                             sizeof(triplet) *
                                 m_seedfilter_config.max_triplets_per_spM *
                                 nSeedSelectingThreads,
-                            stream>>>(m_seedfilter_config, spacepoints_view,
-                                      g2_view, triplet_counter_spM_buffer,
-                                      triplet_counter_midBot_buffer,
-                                      triplet_buffer, seed_buffer);
+                            stream>>>(
+        m_seedfilter_config, spacepoints_view, g2_view,
+        triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
+        triplet_buffer, highest_weight_buffer, seed_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 
     return seed_buffer;