Cut by ..._track_candidates_per_track as soon as possible

Author: fiedlerp

device/alpaka/src/finding/finding_algorithm.cpp

@@ -309,6 +309,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

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

@@ -47,11 +47,13 @@ TRACCC_HOST_DEVICE inline void build_tracks(
 
     const unsigned int n_cands = L.step + 1 - L.n_skipped;
 
+    // Criteria for valid tracks
+    assert(n_cands >= cfg.min_track_candidates_per_track &&
+           n_cands <= cfg.max_track_candidates_per_track);
+
     // 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 +69,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 +95,27 @@ 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
+    vecmem::device_atomic_ref<unsigned int> num_valid_tracks(
+        *payload.n_valid_tracks);
 
-    // 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;
-    }
+    const unsigned int pos = num_valid_tracks.fetch_add(1);
+    valid_indices[pos] = globalIndex;
 }
 
 }  // namespace traccc::device

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

@@ -256,16 +256,28 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                 const unsigned int n_skipped =
                     payload.step == 0 ? 0 : links.at(prev_link_idx).n_skipped;
 
+                const bool last_step =
+                    payload.step == cfg.max_track_candidates_per_track - 1;
+
                 links.at(l_pos) = {.step = payload.step,
                                    .previous_candidate_idx = prev_link_idx,
                                    .meas_idx = meas_idx,
                                    .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);
+                }
             }
         }
 
@@ -318,10 +330,14 @@ TRACCC_HOST_DEVICE inline void find_tracks(
         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);
+        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/propagate_to_next_surface.ipp

@@ -30,6 +30,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);
@@ -82,18 +90,12 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(
     // If a surface found, add the parameter for the next step
     if (s4.success) {
         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);
         }
     }
 }

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

@@ -50,6 +50,11 @@ struct propagate_to_next_surface_payload {
      */
     vecmem::data::vector_view<const unsigned int> param_ids_view;
 
+    /**
+     * @brief View object to the vector of candidate links
+     */
+    vecmem::data::vector_view<const candidate_link> links_view;
+
     /**
      * @brief Index in the link vector at which the current step starts
      */

device/cuda/src/finding/finding_algorithm.cu

@@ -189,9 +189,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
     unsigned int n_in_params = n_seeds;
 
-    for (unsigned int step = 0;
-         step < m_cfg.max_track_candidates_per_track && n_in_params > 0;
-         step++) {
+    for (unsigned int step = 0; n_in_params > 0; step++) {
 
         /*****************************************************************
          * Kernel2: Apply material interaction
@@ -268,13 +266,12 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
             const unsigned int nThreads = m_warp_size * 2;
             const unsigned int nBlocks =
                 (n_in_params + nThreads - 1) / nThreads;
+            const std::size_t shared_size =
+                nThreads * sizeof(unsigned int) +
+                2 * nThreads * sizeof(std::pair<unsigned int, unsigned int>);
 
             kernels::find_tracks<std::decay_t<detector_type>>
-                <<<nBlocks, nThreads,
-                   nThreads * sizeof(unsigned int) +
-                       2 * nThreads *
-                           sizeof(std::pair<unsigned int, unsigned int>),
-                   stream>>>(
+                <<<nBlocks, nThreads, shared_size, stream>>>(
                     m_cfg,
                     device::find_tracks_payload<std::decay_t<detector_type>>{
                         .det_data = det_view,
@@ -306,6 +303,12 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
             m_stream.synchronize();
         }
 
+        // If no more CKF step is expected, the tips and links are populated,
+        // and any further time-consuming action is avoided
+        if (step == m_cfg.max_track_candidates_per_track - 1) {
+            break;
+        }
+
         if (n_candidates > 0) {
             /*****************************************************************
              * Kernel4: Get key and value for parameter sorting
@@ -362,6 +365,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,

device/sycl/src/finding/find_tracks.hpp

@@ -325,6 +325,10 @@ track_candidate_container_types::buffer find_tracks(
         n_candidates =
             step_to_link_idx_map[step + 1] - step_to_link_idx_map[step];
 
+        if (step == config.max_track_candidates_per_track - 1) {
+            break;
+        }
+
         if (n_candidates > 0) {
             /*****************************************************************
              * Kernel4: Get key and value for parameter sorting
@@ -396,6 +400,7 @@ track_candidate_container_types::buffer find_tracks(
                          param_liveness =
                              vecmem::get_data(param_liveness_buffer),
                          param_ids = vecmem::get_data(param_ids_buffer),
+                         links_view = vecmem::get_data(links_buffer),
                          prev_links_idx = step_to_link_idx_map[step], step,
                          n_candidates, tips = vecmem::get_data(tips_buffer)](
                             ::sycl::nd_item<1> item) {
@@ -404,8 +409,8 @@ track_candidate_container_types::buffer find_tracks(
                                 typename stepper_t::magnetic_field_type>(
                                 details::global_index(item), config,
                                 {det, field, in_params, param_liveness,
-                                 param_ids, prev_links_idx, step, n_candidates,
-                                 tips});
+                                 param_ids, links_view, prev_links_idx, step,
+                                 n_candidates, tips});
                         });
                 })
                 .wait_and_throw();