Skip to content

Improve memory usage in track finding postamble #908

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 22, 2025
Merged

Improve memory usage in track finding postamble #908

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions device/alpaka/src/finding/finding_algorithm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,9 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
m_cfg.max_num_branches_per_seed * n_seeds, m_mr.main,
vecmem::data::buffer_type::resizable};
m_copy.setup(tips_buffer)->wait();
vecmem::data::vector_buffer<unsigned int> tip_length_buffer{
m_cfg.max_num_branches_per_seed * n_seeds, m_mr.main};
m_copy.setup(tip_length_buffer)->wait();

std::map<unsigned int, unsigned int> step_to_link_idx_map;
step_to_link_idx_map[0] = 0;
Expand Down Expand Up @@ -276,6 +279,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
.out_params_view = updated_params_buffer,
.out_params_liveness_view = updated_liveness_buffer,
.tips_view = tips_buffer,
.tip_lengths_view = tip_length_buffer,
.n_tracks_per_seed_view = n_tracks_per_seed_buffer};

auto bufAcc_payload =
Expand Down Expand Up @@ -363,7 +367,8 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
.prev_links_idx = step_to_link_idx_map[step],
.step = step,
.n_in_params = n_candidates,
.tips_view = tips_buffer};
.tips_view = tips_buffer,
.tip_lengths_view = tip_length_buffer};

auto bufAcc_payload =
::alpaka::allocBuf<PayloadType, Idx>(devAcc, 1u);
Expand Down Expand Up @@ -397,12 +402,16 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
// Get the number of tips
auto n_tips_total = m_copy.get_size(tips_buffer);

std::vector<unsigned int> tips_length_host;

if (n_tips_total > 0) {
m_copy(tip_length_buffer, tips_length_host)->wait();
tips_length_host.resize(n_tips_total);
}

// Create track candidate buffer
track_candidate_container_types::buffer track_candidates_buffer{
{n_tips_total, m_mr.main},
{std::vector<std::size_t>(n_tips_total,
m_cfg.max_track_candidates_per_track),
m_mr.main, m_mr.host, vecmem::data::buffer_type::resizable}};
{n_tips_total, m_mr.main}, {tips_length_host, m_mr.main, m_mr.host}};

m_copy.setup(track_candidates_buffer.headers)->ignore();
m_copy.setup(track_candidates_buffer.items)->ignore();
Expand Down
5 changes: 5 additions & 0 deletions device/common/include/traccc/finding/device/find_tracks.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,11 @@ struct find_tracks_payload {
*/
vecmem::data::vector_view<unsigned int> tips_view;

/**
* @brief Vector to hold the number of track states per tip
*/
vecmem::data::vector_view<unsigned int> tip_lengths_view;

/**
* @brief View object to the vector of the number of tracks per initial
* input seed
Expand Down
12 changes: 0 additions & 12 deletions device/common/include/traccc/finding/device/impl/build_tracks.ipp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,17 +41,10 @@ TRACCC_HOST_DEVICE inline void build_tracks(
auto L = links.at(tip);
const unsigned int n_meas = measurements.size();

const unsigned int n_cands = L.step + 1 - L.n_skipped;

// Resize the candidates with the exact size
cands_per_track.resize(n_cands);

// Track summary variables
scalar ndf_sum = 0.f;
scalar chi2_sum = 0.f;

[[maybe_unused]] std::size_t num_inserted = 0;

// Reversely iterate to fill the track candidates
for (auto it = cands_per_track.rbegin(); it != cands_per_track.rend();
it++) {
Expand All @@ -64,7 +57,6 @@ TRACCC_HOST_DEVICE inline void build_tracks(
assert(L.meas_idx < n_meas);

*it = {measurements.at(L.meas_idx)};
num_inserted++;

// Sanity check on chi2
assert(L.chi2 < std::numeric_limits<traccc::scalar>::max());
Expand All @@ -87,10 +79,6 @@ TRACCC_HOST_DEVICE inline void build_tracks(
}

#ifndef NDEBUG
// 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) {
Expand Down
7 changes: 5 additions & 2 deletions device/common/include/traccc/finding/device/impl/find_tracks.ipp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ TRACCC_HOST_DEVICE inline void find_tracks(
vecmem::device_vector<const unsigned int> upper_bounds(
payload.upper_bounds_view);
vecmem::device_vector<unsigned int> tips(payload.tips_view);
vecmem::device_vector<unsigned int> tip_lengths(payload.tip_lengths_view);
vecmem::device_vector<unsigned int> n_tracks_per_seed(
payload.n_tracks_per_seed_view);

Expand Down Expand Up @@ -276,7 +277,8 @@ TRACCC_HOST_DEVICE inline void find_tracks(
// a tip
if (last_step &&
n_cands >= cfg.min_track_candidates_per_track) {
tips.push_back(l_pos);
auto tip_pos = tips.push_back(l_pos);
tip_lengths.at(tip_pos) = n_cands;
}
}
}
Expand Down Expand Up @@ -331,7 +333,8 @@ TRACCC_HOST_DEVICE inline void find_tracks(
// step being skipped, the links are empty, and the tip has
// nowhere to point
assert(payload.step > 0);
tips.push_back(prev_link_idx);
auto tip_pos = tips.push_back(prev_link_idx);
tip_lengths.at(tip_pos) = n_cands;
}
} else {
// Add measurement candidates to link
Expand Down
4 changes: 3 additions & 1 deletion device/common/include/traccc/finding/device/impl/propagate_to_next_surface.ipp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(

// tips
vecmem::device_vector<unsigned int> tips(payload.tips_view);
vecmem::device_vector<unsigned int> tip_lengths(payload.tip_lengths_view);

// Detector
typename propagator_t::detector_type det(payload.det_data);
Expand Down Expand Up @@ -106,7 +107,8 @@ TRACCC_HOST_DEVICE inline void propagate_to_next_surface(
params_liveness[param_id] = 0u;

if (n_cands >= cfg.min_track_candidates_per_track) {
tips.push_back(link_idx);
auto tip_pos = tips.push_back(link_idx);
tip_lengths.at(tip_pos) = n_cands;
}
}
}
Expand Down
5 changes: 5 additions & 0 deletions device/common/include/traccc/finding/device/propagate_to_next_surface.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,11 @@ struct propagate_to_next_surface_payload {
* @brief View object to the vector of tips
*/
vecmem::data::vector_view<unsigned int> tips_view;

/**
* @brief Vector to hold the number of track states per tip
*/
vecmem::data::vector_view<unsigned int> tip_lengths_view;
};

/// Function for propagating the kalman-updated tracks to the next surface
Expand Down
19 changes: 14 additions & 5 deletions device/cuda/src/finding/finding_algorithm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,9 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
m_cfg.max_num_branches_per_seed * n_seeds, m_mr.main,
vecmem::data::buffer_type::resizable};
m_copy.setup(tips_buffer)->wait();
vecmem::data::vector_buffer<unsigned int> tip_length_buffer{
m_cfg.max_num_branches_per_seed * n_seeds, m_mr.main};
m_copy.setup(tip_length_buffer)->wait();

std::map<unsigned int, unsigned int> step_to_link_idx_map;
step_to_link_idx_map[0] = 0;
Expand Down Expand Up @@ -279,6 +282,7 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
.out_params_view = updated_params_buffer,
.out_params_liveness_view = updated_liveness_buffer,
.tips_view = tips_buffer,
.tip_lengths_view = tip_length_buffer,
.n_tracks_per_seed_view = n_tracks_per_seed_buffer});
TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());

Expand Down Expand Up @@ -359,7 +363,8 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
.prev_links_idx = step_to_link_idx_map[step],
.step = step,
.n_in_params = n_candidates,
.tips_view = tips_buffer});
.tips_view = tips_buffer,
.tip_lengths_view = tip_length_buffer});
TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());

m_stream.synchronize();
Expand All @@ -384,12 +389,16 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
// Get the number of tips
auto n_tips_total = m_copy.get_size(tips_buffer);

std::vector<unsigned int> tips_length_host;

if (n_tips_total > 0) {
m_copy(tip_length_buffer, tips_length_host)->wait();
tips_length_host.resize(n_tips_total);
}

// Create track candidate buffer
track_candidate_container_types::buffer track_candidates_buffer{
{n_tips_total, m_mr.main},
{std::vector<std::size_t>(n_tips_total,
m_cfg.max_track_candidates_per_track),
m_mr.main, m_mr.host, vecmem::data::buffer_type::resizable}};
{n_tips_total, m_mr.main}, {tips_length_host, m_mr.main, m_mr.host}};

m_copy.setup(track_candidates_buffer.headers)->ignore();
m_copy.setup(track_candidates_buffer.items)->ignore();
Expand Down
23 changes: 16 additions & 7 deletions device/sycl/src/finding/find_tracks.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,9 @@ track_candidate_container_types::buffer find_tracks(
config.max_num_branches_per_seed * n_seeds, mr.main,
vecmem::data::buffer_type::resizable};
copy.setup(tips_buffer)->wait();
vecmem::data::vector_buffer<unsigned int> tip_length_buffer{
config.max_num_branches_per_seed * n_seeds, mr.main};
copy.setup(tip_length_buffer)->wait();

std::map<unsigned int, unsigned int> step_to_link_idx_map;
step_to_link_idx_map[0] = 0;
Expand Down Expand Up @@ -289,6 +292,7 @@ track_candidate_container_types::buffer find_tracks(
updated_liveness =
vecmem::get_data(updated_liveness_buffer),
tips = vecmem::get_data(tips_buffer),
tip_lengths = vecmem::get_data(tip_length_buffer),
n_tracks_per_seed =
vecmem::get_data(n_tracks_per_seed_buffer),
shared_candidates_size, shared_num_candidates,
Expand All @@ -305,7 +309,7 @@ track_candidate_container_types::buffer find_tracks(
n_in_params, barcodes, upper_bounds, links_view,
prev_links_idx, curr_links_idx, step,
updated_params, updated_liveness, tips,
n_tracks_per_seed},
tip_lengths, n_tracks_per_seed},
{&(shared_num_candidates[0]),
&(shared_candidates[0]),
shared_candidates_size[0]});
Expand Down Expand Up @@ -397,15 +401,16 @@ track_candidate_container_types::buffer find_tracks(
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)](
n_candidates, tips = vecmem::get_data(tips_buffer),
tip_lengths = vecmem::get_data(tip_length_buffer)](
::sycl::nd_item<1> item) {
device::propagate_to_next_surface<
propagator_type,
typename stepper_t::magnetic_field_type>(
details::global_index(item), config,
{det, field, in_params, param_liveness,
param_ids, links_view, prev_links_idx, step,
n_candidates, tips});
n_candidates, tips, tip_lengths});
});
})
.wait_and_throw();
Expand All @@ -421,12 +426,16 @@ track_candidate_container_types::buffer find_tracks(
// Get the number of tips
auto n_tips_total = copy.get_size(tips_buffer);

std::vector<unsigned int> tips_length_host;

if (n_tips_total > 0) {
copy(tip_length_buffer, tips_length_host)->wait();
tips_length_host.resize(n_tips_total);
}

// Create track candidate buffer
track_candidate_container_types::buffer track_candidates_buffer{
{n_tips_total, mr.main},
{std::vector<std::size_t>(n_tips_total,
config.max_track_candidates_per_track),
mr.main, mr.host, vecmem::data::buffer_type::resizable}};
{n_tips_total, mr.main}, {tips_length_host, mr.main, mr.host}};
copy.setup(track_candidates_buffer.headers)->wait();
copy.setup(track_candidates_buffer.items)->wait();
track_candidate_container_types::view track_candidates =
Expand Down
Loading