Skip to content
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

Fix bug in MG Neighborhood sampling #4827

Open
wants to merge 3 commits into
base: branch-25.02
Choose a base branch
from
Open

Fix bug in MG Neighborhood sampling #4827

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
42a2d4a
fix style
jnke2016 Dec 10, 2024
8676f09
Merge branch 'branch-25.02' into branch-25.02_fix-bug-mg-nbr-sampling
nv-rliu Dec 12, 2024
5a1679b
Merge branch 'branch-25.02' into branch-25.02_fix-bug-mg-nbr-sampling
alexbarghi-nv Dec 17, 2024
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
16 changes: 9 additions & 7 deletions cpp/src/c_api/neighbor_sampling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -880,7 +880,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
handle_.get_stream());

std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
std::optional<rmm::device_uvector<label_t>> local_label_to_comm_rank{std::nullopt};
std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{
std::nullopt}; // global after allgatherv

Expand Down Expand Up @@ -932,12 +931,13 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
handle_.get_stream(),
raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});

(*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
rmm::device_uvector<label_t> local_label_to_comm_rank(num_unique_labels,
handle_.get_stream());

cugraph::detail::scalar_fill(
handle_.get_stream(),
(*local_label_to_comm_rank).begin(), // This should be rename to rank
(*local_label_to_comm_rank).size(),
local_label_to_comm_rank.begin(), // This should be rename to rank
local_label_to_comm_rank.size(),
label_t{handle_.get_comms().get_rank()});

// Perform allgather to get global_label_to_comm_rank_d_vector
Expand All @@ -948,11 +948,13 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
std::exclusive_scan(
recvcounts.begin(), recvcounts.end(), displacements.begin(), size_t{0});

(*label_to_comm_rank)
.resize(displacements.back() + recvcounts.back(), handle_.get_stream());
rmm::device_uvector<label_t> tmp_label_to_comm_rank(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just use label_to_comm_rank.resize(...) here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is an optional which is initialized to std::nullopt

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

label_to_comm_rank =
  rmm::device_uvector<label_t>(displacements.back() + recvcounts.back(), handle_.get_stream());

is sufficient (no need to create a temporary and move).

displacements.back() + recvcounts.back(), handle_.get_stream());

label_to_comm_rank = std::move(tmp_label_to_comm_rank);

cugraph::device_allgatherv(handle_.get_comms(),
(*local_label_to_comm_rank).begin(),
local_label_to_comm_rank.begin(),
(*label_to_comm_rank).begin(),
recvcounts,
displacements,
Expand Down
Loading