Cut by ..._track_candidates_per_track as soon as possible

Author: fiedlerp

device/alpaka/CMakeLists.txt

@@ -52,7 +52,6 @@ traccc_add_alpaka_library( traccc_alpaka alpaka TYPE SHARED
   "src/finding/kernels/make_barcode_sequence.hpp"
   "src/finding/kernels/apply_interaction.hpp"
   "src/finding/kernels/fill_sort_keys.hpp"
-  "src/finding/kernels/prune_tracks.hpp"
   "src/finding/kernels/build_tracks.hpp"
   "src/finding/kernels/find_tracks.hpp"
   "src/finding/kernels/propagate_to_next_surface.hpp"

device/alpaka/src/finding/finding_algorithm.cpp

@@ -17,7 +17,6 @@
 #include "./kernels/find_tracks.hpp"
 #include "./kernels/make_barcode_sequence.hpp"
 #include "./kernels/propagate_to_next_surface.hpp"
-#include "./kernels/prune_tracks.hpp"
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/device/sort_key.hpp"
@@ -299,6 +298,10 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
             ::alpaka::wait(queue);
         }
 
+        if (step == m_cfg.max_track_candidates_per_track - 1) {
+            break;
+        }
+
         if (n_candidates > 0) {
             /*****************************************************************
              * Kernel4: Get key and value for parameter sorting
@@ -356,6 +359,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
                                 .params_view = in_params_buffer,
                                 .params_liveness_view = param_liveness_buffer,
                                 .param_ids_view = param_ids_buffer,
+                                .links_view = links_buffer,
                                 .prev_links_idx = step_to_link_idx_map[step],
                                 .step = step,
                                 .n_in_params = n_candidates,
@@ -403,65 +407,21 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
     m_copy.setup(track_candidates_buffer.headers)->ignore();
     m_copy.setup(track_candidates_buffer.items)->ignore();
 
-    // Create buffer for valid indices
-    vecmem::data::vector_buffer<unsigned int> valid_indices_buffer(n_tips_total,
-                                                                   m_mr.main);
-
-    // Count the number of valid tracks
-    auto bufHost_n_valid_tracks =
-        ::alpaka::allocBuf<unsigned int, Idx>(devHost, 1u);
-    unsigned int* n_valid_tracks =
-        ::alpaka::getPtrNative(bufHost_n_valid_tracks);
-    ::alpaka::memset(queue, bufHost_n_valid_tracks, 0);
-    ::alpaka::wait(queue);
-
     // @Note: nBlocks can be zero in case there is no tip. This happens when
     // chi2_max config is set tightly and no tips are found
     if (n_tips_total > 0) {
-        auto n_valid_tracks_device =
-            ::alpaka::allocBuf<unsigned int, Idx>(devAcc, 1u);
-        ::alpaka::memset(queue, n_valid_tracks_device, 0);
-
         Idx blocksPerGrid =
             (n_tips_total + threadsPerBlock - 1) / threadsPerBlock;
         auto workDiv = makeWorkDiv<Acc>(blocksPerGrid, threadsPerBlock);
 
         ::alpaka::exec<Acc>(
-            queue, workDiv, BuildTracksKernel{}, m_cfg,
-            device::build_tracks_payload{
-                measurements, seeds_view, links_buffer, tips_buffer,
-                track_candidates_buffer, valid_indices_buffer,
-                ::alpaka::getPtrNative(n_valid_tracks_device)});
-        ::alpaka::wait(queue);
-
-        // Global counter object: Device -> Host
-        ::alpaka::memcpy(queue, bufHost_n_valid_tracks, n_valid_tracks_device);
-        ::alpaka::wait(queue);
-    }
-
-    // Create pruned candidate buffer
-    track_candidate_container_types::buffer prune_candidates_buffer{
-        {*n_valid_tracks, m_mr.main},
-        {std::vector<std::size_t>(*n_valid_tracks,
-                                  m_cfg.max_track_candidates_per_track),
-         m_mr.main, m_mr.host, vecmem::data::buffer_type::resizable}};
-
-    m_copy.setup(prune_candidates_buffer.headers)->ignore();
-    m_copy.setup(prune_candidates_buffer.items)->ignore();
-
-    if (*n_valid_tracks > 0) {
-        Idx blocksPerGrid =
-            (*n_valid_tracks + threadsPerBlock - 1) / threadsPerBlock;
-        auto workDiv = makeWorkDiv<Acc>(blocksPerGrid, threadsPerBlock);
-
-        ::alpaka::exec<Acc>(queue, workDiv, PruneTracksKernel{},
-                            device::prune_tracks_payload{
-                                track_candidates_buffer, valid_indices_buffer,
-                                prune_candidates_buffer});
+            queue, workDiv, BuildTracksKernel{},
+            device::build_tracks_payload{measurements, seeds_view, links_buffer,
+                                         tips_buffer, track_candidates_buffer});
         ::alpaka::wait(queue);
     }
 
-    return prune_candidates_buffer;
+    return track_candidates_buffer;
 }
 
 // Explicit template instantiation

device/alpaka/src/finding/kernels/build_tracks.hpp

@@ -19,13 +19,13 @@ namespace traccc::alpaka {
 
 struct BuildTracksKernel {
     template <typename TAcc>
-    ALPAKA_FN_ACC void operator()(TAcc const& acc, const finding_config cfg,
+    ALPAKA_FN_ACC void operator()(TAcc const& acc,
                                   device::build_tracks_payload payload) const {
 
         device::global_index_t globalThreadIdx =
             ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0];
 
-        device::build_tracks(globalThreadIdx, cfg, payload);
+        device::build_tracks(globalThreadIdx, payload);
     }
 };
 

device/alpaka/src/finding/kernels/prune_tracks.hpp

@@ -64,14 +64,12 @@ traccc_add_library( traccc_device_common device_common TYPE INTERFACE
    "include/traccc/finding/device/fill_sort_keys.hpp"
    "include/traccc/finding/device/make_barcode_sequence.hpp"
    "include/traccc/finding/device/propagate_to_next_surface.hpp"
-   "include/traccc/finding/device/prune_tracks.hpp"
    "include/traccc/finding/device/impl/apply_interaction.ipp"
    "include/traccc/finding/device/impl/build_tracks.ipp"
    "include/traccc/finding/device/impl/find_tracks.ipp"
    "include/traccc/finding/device/impl/fill_sort_keys.ipp"
    "include/traccc/finding/device/impl/make_barcode_sequence.ipp"
    "include/traccc/finding/device/impl/propagate_to_next_surface.ipp"
-   "include/traccc/finding/device/impl/prune_tracks.ipp"
    # Track fitting funtions(s).
    "include/traccc/fitting/device/fit.hpp"
    "include/traccc/fitting/device/impl/fit.ipp"

device/common/CMakeLists.txt

@@ -52,17 +52,6 @@ struct build_tracks_payload {
      * @brief View object to the vector of track candidates
      */
     track_candidate_container_types::view track_candidates_view;
-
-    /**
-     * @brief View object to the vector of indices meeting the selection
-     * criteria
-     */
-    vecmem::data::vector_view<unsigned int> valid_indices_view;
-
-    /**
-     * @brief The number of valid tracks meeting criteria
-     */
-    unsigned int* n_valid_tracks;
 };
 
 /// Function for building full tracks from the link container:
@@ -75,8 +64,7 @@ struct build_tracks_payload {
 /// @param[inout] payload      The function call payload
 ///
 TRACCC_HOST_DEVICE inline void build_tracks(
-    global_index_t globalIndex, const finding_config& cfg,
-    const build_tracks_payload& payload);
+    global_index_t globalIndex, const build_tracks_payload& payload);
 
 }  // namespace traccc::device
 

device/common/include/traccc/finding/device/build_tracks.hpp

@@ -13,8 +13,7 @@
 namespace traccc::device {
 
 TRACCC_HOST_DEVICE inline void build_tracks(
-    const global_index_t globalIndex, const finding_config& cfg,
-    const build_tracks_payload& payload) {
+    const global_index_t globalIndex, const build_tracks_payload& payload) {
 
     const measurement_collection_types::const_device measurements(
         payload.measurements_view);
@@ -29,9 +28,6 @@ TRACCC_HOST_DEVICE inline void build_tracks(
     track_candidate_container_types::device track_candidates(
         payload.track_candidates_view);
 
-    vecmem::device_vector<unsigned int> valid_indices(
-        payload.valid_indices_view);
-
     if (globalIndex >= tips.size()) {
         return;
     }
@@ -50,8 +46,6 @@ TRACCC_HOST_DEVICE inline void build_tracks(
     // Resize the candidates with the exact size
     cands_per_track.resize(n_cands);
 
-    bool success = true;
-
     // Track summary variables
     scalar ndf_sum = 0.f;
     scalar chi2_sum = 0.f;
@@ -67,11 +61,7 @@ TRACCC_HOST_DEVICE inline void build_tracks(
             L = links.at(L.previous_candidate_idx);
         }
 
-        // Break if the measurement is still invalid
-        if (L.meas_idx >= measurements.size()) {
-            success = false;
-            break;
-        }
+        assert(L.meas_idx < n_meas);
 
         *it = {measurements.at(L.meas_idx)};
         num_inserted++;
@@ -97,36 +87,21 @@ TRACCC_HOST_DEVICE inline void build_tracks(
     }
 
 #ifndef NDEBUG
-    if (success) {
-        // Assert that we inserted exactly as many elements as we reserved
-        // space for.
-        assert(num_inserted == cands_per_track.size());
-
-        // Assert that we did not make any duplicate track states.
-        for (unsigned int i = 0; i < cands_per_track.size(); ++i) {
-            for (unsigned int j = 0; j < cands_per_track.size(); ++j) {
-                if (i != j) {
-                    // TODO: Re-enable me!
-                    // assert(cands_per_track.at(i).measurement_id !=
-                    //       cands_per_track.at(j).measurement_id);
-                }
+    // Assert that we inserted exactly as many elements as we reserved
+    // space for.
+    assert(num_inserted == cands_per_track.size());
+
+    // Assert that we did not make any duplicate track states.
+    for (unsigned int i = 0; i < cands_per_track.size(); ++i) {
+        for (unsigned int j = 0; j < cands_per_track.size(); ++j) {
+            if (i != j) {
+                // TODO: Re-enable me!
+                // assert(cands_per_track.at(i).measurement_id !=
+                //       cands_per_track.at(j).measurement_id);
             }
         }
     }
 #endif
-
-    // NOTE: We may at some point want to assert that `success` is true
-
-    // Criteria for valid tracks
-    if (n_cands >= cfg.min_track_candidates_per_track &&
-        n_cands <= cfg.max_track_candidates_per_track && success) {
-
-        vecmem::device_atomic_ref<unsigned int> num_valid_tracks(
-            *payload.n_valid_tracks);
-
-        const unsigned int pos = num_valid_tracks.fetch_add(1);
-        valid_indices[pos] = globalIndex;
-    }
 }
 
 }  // namespace traccc::device

device/common/include/traccc/finding/device/impl/build_tracks.ipp

@@ -77,6 +77,9 @@ TRACCC_HOST_DEVICE inline void find_tracks(
 
     const unsigned int in_param_id = thread_id.getGlobalThreadIdX();
 
+    const bool last_step =
+        payload.step == cfg.max_track_candidates_per_track - 1;
+
     /*
      * Step 1 of this kernel is to determine which indices belong to which
      * parameter. Because the measurements are guaranteed to be grouped, we can
@@ -262,10 +265,19 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                                    .seed_idx = seed_idx,
                                    .n_skipped = n_skipped,
                                    .chi2 = chi2};
-
                 out_params.at(l_pos - payload.curr_links_idx) =
                     trk_state.filtered();
-                out_params_liveness.at(l_pos - payload.curr_links_idx) = 1u;
+                out_params_liveness.at(l_pos - payload.curr_links_idx) =
+                    static_cast<unsigned int>(!last_step);
+
+                const unsigned int n_cands = payload.step + 1 - n_skipped;
+
+                // If no more CKF step is expected, current candidate is kept as
+                // a tip
+                if (last_step &&
+                    n_cands >= cfg.min_track_candidates_per_track) {
+                    tips.push_back(l_pos);
+                }
             }
         }
 
@@ -312,11 +324,15 @@ TRACCC_HOST_DEVICE inline void find_tracks(
             const unsigned int n_skipped =
                 payload.step == 0 ? 0 : links.at(prev_link_idx).n_skipped;
 
-            if (n_skipped >= cfg.max_num_skipping_per_cand) {
-                // In case of max skipping being 0 and first step being skipped,
-                // the links are empty, and the tip has nowhere to point
-                assert(payload.step > 0);
-                tips.push_back(prev_link_idx);
+            if (n_skipped >= cfg.max_num_skipping_per_cand || last_step) {
+                const unsigned int n_cands = payload.step - n_skipped;
+                if (n_cands >= cfg.min_track_candidates_per_track) {
+                    // In case of max skipping and min length being 0, and first
+                    // step being skipped, the links are empty, and the tip has
+                    // nowhere to point
+                    assert(payload.step > 0);
+                    tips.push_back(prev_link_idx);
+                }
             } else {
                 // Add measurement candidates to link
                 const unsigned int l_pos = links.bulk_append_implicit(1);

device/common/include/traccc/finding/device/impl/find_tracks.ipp

@@ -32,6 +32,14 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(
 
     const unsigned int param_id = param_ids.at(globalIndex);
 
+    // Links
+    vecmem::device_vector<const candidate_link> links(payload.links_view);
+
+    const unsigned int link_idx = payload.prev_links_idx + param_id;
+    const auto& link = links.at(link_idx);
+    assert(link.step == payload.step);
+    const unsigned int n_cands = link.step + 1 - link.n_skipped;
+
     // Parameter liveness
     vecmem::device_vector<unsigned int> params_liveness(
         payload.params_liveness_view);
@@ -93,18 +101,12 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(
         assert(propagation._navigation.is_on_sensitive());
 
         params[param_id] = propagation._stepping.bound_params();
-
-        if (payload.step == cfg.max_track_candidates_per_track - 1) {
-            tips.push_back(payload.prev_links_idx + param_id);
-            params_liveness[param_id] = 0u;
-        } else {
-            params_liveness[param_id] = 1u;
-        }
+        params_liveness[param_id] = 1u;
     } else {
         params_liveness[param_id] = 0u;
 
-        if (payload.step >= cfg.min_track_candidates_per_track - 1) {
-            tips.push_back(payload.prev_links_idx + param_id);
+        if (n_cands >= cfg.min_track_candidates_per_track) {
+            tips.push_back(link_idx);
         }
     }
 }