Skip to content

Commit

Permalink
Fix GPU CKF to skip surfaces in case of no branches
Browse files Browse the repository at this point in the history
  • Loading branch information
beomki-yeo committed Apr 26, 2024
1 parent 2150b3f commit a90e721
Show file tree
Hide file tree
Showing 15 changed files with 459 additions and 179 deletions.
3 changes: 0 additions & 3 deletions core/include/traccc/finding/candidate_link.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,6 @@ struct candidate_link {

// Index to the initial seed
unsigned int seed_idx;

// How many times it skipped a surface
unsigned int n_skipped;
};

} // namespace traccc
242 changes: 93 additions & 149 deletions core/include/traccc/finding/finding_algorithm.ipp
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

// System include
#include <algorithm>
#include <iostream>
#include <limits>

namespace traccc {
Expand All @@ -28,8 +29,6 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
const measurement_collection_types::host& measurements,
const bound_track_parameters_collection_types::host& seeds) const {

track_candidate_container_types::host output_candidates;

/*****************************************************************
* Measurement Operations
*****************************************************************/
Expand All @@ -50,6 +49,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
uniques[i], measurement_sort_comp());
upper_bounds.push_back(std::distance(measurements.begin(), up));
}
const auto n_meas = measurements.size();

// Get the number of measurements of each module
std::vector<unsigned int> sizes(n_modules);
Expand All @@ -63,10 +63,6 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
barcodes.push_back(uniques[i].surface_link);
}

/**********************
* Find tracks
**********************/

std::vector<std::vector<candidate_link>> links;
links.resize(m_cfg.max_track_candidates_per_track);

Expand Down Expand Up @@ -109,6 +105,9 @@ finding_algorithm<stepper_t, navigator_t>::operator()(

std::fill(n_trks_per_seed.begin(), n_trks_per_seed.end(), 0);

// Parameters updated by Kalman fitter
std::vector<bound_track_parameters> updated_params;

for (unsigned int in_param_id = 0; in_param_id < n_in_params;
in_param_id++) {

Expand All @@ -118,11 +117,6 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
? in_param_id
: links[step - 1][param_to_link[step - 1][in_param_id]]
.seed_idx);
unsigned int skip_counter =
(step == 0
? 0
: links[step - 1][param_to_link[step - 1][in_param_id]]
.n_skipped);

/*************************
* Material interaction
Expand Down Expand Up @@ -151,10 +145,6 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
static_cast<int>(detray::navigation::direction::e_forward), sf,
sfi.cos_incidence_angle);

/*************************
* CKF
*************************/

// Get barcode and measurements range on surface
const auto bcd = in_param.surface_link();
std::pair<unsigned int, unsigned int> range;
Expand All @@ -180,6 +170,10 @@ finding_algorithm<stepper_t, navigator_t>::operator()(

unsigned int n_branches = 0;

/*****************************************************************
* Find tracks (CKF)
*****************************************************************/

// Iterate over the measurements
for (unsigned int item_id = range.first; item_id < range.second;
item_id++) {
Expand Down Expand Up @@ -207,149 +201,85 @@ finding_algorithm<stepper_t, navigator_t>::operator()(

// Found a good measurement
if (chi2 < m_cfg.chi2_max) {

// Current link ID
unsigned int cur_link_id =
static_cast<unsigned int>(links[step].size());

n_branches++;
n_trks_per_seed[orig_param_id]++;

links[step].push_back({{previous_step, in_param_id},
item_id,
orig_param_id,
skip_counter});

/*********************************
* Propagate to the next surface
*********************************/

// Create propagator state
typename propagator_type::state propagation(
trk_state.filtered(), field, det);
propagation._stepping.template set_constraint<
detray::step::constraint::e_accuracy>(
m_cfg.propagation.stepping.step_constraint);

typename detray::pathlimit_aborter::state s0;
typename detray::parameter_transporter<
transform3_type>::state s1;
typename interactor::state s3;
typename interaction_register<interactor>::state s2{s3};
typename detray::next_surface_aborter::state s4{
m_cfg.min_step_length_for_surface_aborter};
// typename propagation::print_inspector::state s5{};

// @TODO: Should be removed once detray is fixed to set the
// volume in the constructor
propagation._navigation.set_volume(
trk_state.filtered().surface_link().volume());

// Propagate to the next surface
propagator.propagate_sync(propagation,
std::tie(s0, s1, s2, s3, s4));

/*
propagator.propagate_sync(propagation,
std::tie(s0, s1, s2, s3, s4, s5));
*/

// If a surface found, add the parameter for the next
// step
if (s4.success) {
out_params.push_back(
propagation._stepping._bound_params);
param_to_link[step].push_back(cur_link_id);
}
// Unless the track found a surface, it is considered a
// tip
else if (!s4.success &&
step >= m_cfg.min_track_candidates_per_track - 1) {
tips.push_back({step, cur_link_id});
}

// If no more CKF step is expected, current candidate is
// kept as a tip
if (s4.success &&
step == m_cfg.max_track_candidates_per_track - 1) {
tips.push_back({step, cur_link_id});
}
links[step].push_back(
{{previous_step, in_param_id}, item_id, orig_param_id});
updated_params.push_back(trk_state.filtered());
}
}
// After the loop over the measurements

/*****************************************************************
* Add a dummy links in case of no branches
*****************************************************************/

if (n_branches == 0) {
// let's skip this CKF step for the current track candidate
if (n_trks_per_seed[orig_param_id] >=
m_cfg.max_num_branches_per_initial_seed) {

continue;
}

// Put an invalid link with max item id
links[step].push_back({{previous_step, in_param_id},
std::numeric_limits<unsigned int>::max(),
orig_param_id});
bound_track_parameters bound_param(in_param.surface_link(),
in_param.vector(),
in_param.covariance());
updated_params.push_back(bound_param);
n_branches++;
}
}

measurement dummy_meas;

dummy_meas.local = in_param.bound_local();
dummy_meas.variance = dummy_meas.variance + point2{10., 10.};
dummy_meas.surface_link = in_param.surface_link();

track_state<transform3_type> trk_state(dummy_meas);

// Run the Kalman update
sf.template visit_mask<gain_matrix_updater<transform3_type>>(
trk_state, bound_param);

unsigned int cur_link_id =
static_cast<unsigned int>(links[step].size());

links[step].push_back({{previous_step, in_param_id},
std::numeric_limits<unsigned int>::max(),
orig_param_id,
skip_counter + 1});
/*********************************
* Propagate to the next surface
*********************************/

if (skip_counter + 1 > m_cfg.max_num_skipping_per_cand) {
tips.push_back({step, cur_link_id});
continue;
// exit from param_id loop
}
const unsigned int n_links = links[step].size();
for (unsigned int link_id = 0; link_id < n_links; link_id++) {

// Create propagator state
typename propagator_type::state propagation(
trk_state.filtered(), field, det);
propagation._stepping.template set_constraint<
detray::step::constraint::e_accuracy>(
const auto& param = updated_params[link_id];
// Create propagator state
typename propagator_type::state propagation(param, field, det);
propagation._stepping
.template set_constraint<detray::step::constraint::e_accuracy>(
m_cfg.propagation.stepping.step_constraint);

typename detray::pathlimit_aborter::state s0;
typename detray::parameter_transporter<transform3_type>::state
s1;
typename interactor::state s3;
typename interaction_register<interactor>::state s2{s3};
typename detray::next_surface_aborter::state s4{
m_cfg.min_step_length_for_surface_aborter};

propagation._navigation.set_volume(
trk_state.filtered().surface_link().volume());

// Propagate to the next surface
propagator.propagate_sync(propagation,
std::tie(s0, s1, s2, s3, s4));

// If a surface found, add the parameter for the next
// step
if (s4.success) {
out_params.push_back(propagation._stepping._bound_params);
param_to_link[step].push_back(cur_link_id);
}
// Unless the track found a surface, it is considered a
// tip
else if (!s4.success &&
step >= m_cfg.min_track_candidates_per_track - 1) {
tips.push_back({step, cur_link_id});
}
typename detray::pathlimit_aborter::state s0;
typename detray::parameter_transporter<transform3_type>::state s1;
typename interactor::state s3;
typename interaction_register<interactor>::state s2{s3};
typename detray::next_surface_aborter::state s4{
m_cfg.min_step_length_for_surface_aborter};

// @TODO: Should be removed once detray is fixed to set the
// volume in the constructor
propagation._navigation.set_volume(param.surface_link().volume());

// Propagate to the next surface
propagator.propagate_sync(propagation,
std::tie(s0, s1, s2, s3, s4));

// If a surface found, add the parameter for the next
// step
if (s4.success) {
out_params.push_back(propagation._stepping._bound_params);
param_to_link[step].push_back(link_id);
}
// Unless the track found a surface, it is considered a
// tip
else if (!s4.success &&
step >= m_cfg.min_track_candidates_per_track - 1) {
tips.push_back({step, link_id});
}

// If no more CKF step is expected, current candidate is
// kept as a tip
if (s4.success &&
step == m_cfg.max_track_candidates_per_track - 1) {
tips.push_back({step, link_id});
}
}

Expand All @@ -362,35 +292,49 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
**********************/

// Number of found tracks = number of tips
track_candidate_container_types::host output_candidates;
output_candidates.reserve(tips.size());

for (const auto& tip : tips) {
// Get the link corresponding to tip
auto L = links[tip.first][tip.second];

// Skip if the number of tracks candidates is too small
if (tip.first + 1 < m_cfg.min_track_candidates_per_track) {
continue;
// Count the number of skipped steps
unsigned int n_skipped{0u};
while (true) {

if (L.meas_idx > n_meas) {
n_skipped++;
}

if (L.previous.first == 0u) {
break;
}

const unsigned int link_pos =
param_to_link[L.previous.first][L.previous.second];
L = links[L.previous.first][link_pos];
}

// Get the link corresponding to tip
auto L = links[tip.first][tip.second];
const unsigned int n_cands = tip.first + 1 - n_skipped;

// Skip if the number of tracks candidates is too small
if (tip.first + 1 - L.n_skipped <
m_cfg.min_track_candidates_per_track) {
if (n_cands < m_cfg.min_track_candidates_per_track ||
n_cands > m_cfg.max_track_candidates_per_track) {
continue;
}

// Retrieve tip
L = links[tip.first][tip.second];

vecmem::vector<track_candidate> cands_per_track;
cands_per_track.resize(tip.first + 1 - L.n_skipped);
cands_per_track.resize(n_cands);

// Reversely iterate to fill the track candidates
for (auto it = cands_per_track.rbegin(); it != cands_per_track.rend();
it++) {

while (L.meas_idx > measurements.size()) {

if (L.previous.first > tip.first + 1)
break; // should not happen.
while (L.meas_idx > n_meas) {

const auto link_pos =
param_to_link[L.previous.first][L.previous.second];
Expand Down
3 changes: 0 additions & 3 deletions core/include/traccc/finding/finding_config.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,6 @@ struct finding_config {
unsigned int max_num_branches_per_initial_seed =
std::numeric_limits<unsigned int>::max();

/// Maximum allowed number of skipped steps per candidate
unsigned int max_num_skipping_per_cand = 3;

/// Minimum step length that track should make to reach the next surface. It
/// should be set higher than the overstep tolerance not to make it stay on
/// the same surface
Expand Down
4 changes: 4 additions & 0 deletions device/common/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -63,14 +63,18 @@ traccc_add_library( traccc_device_common device_common TYPE SHARED
"include/traccc/finding/device/build_tracks.hpp"
"include/traccc/finding/device/count_measurements.hpp"
"include/traccc/finding/device/find_tracks.hpp"
"include/traccc/finding/device/add_links_for_holes.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/count_measurements.ipp"
"include/traccc/finding/device/impl/find_tracks.ipp"
"include/traccc/finding/device/impl/add_links_for_holes.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"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ struct finding_global_counter {

// Number of parameters for the next step
unsigned int n_out_params;

// Number of valid tracks that meet criteria
unsigned int n_valid_tracks;
};

} // namespace traccc::device
Loading

0 comments on commit a90e721

Please sign in to comment.