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

Biased sampling primitive bug fix #4607

Merged
merged 4 commits into from
Aug 9, 2024
Merged

Biased sampling primitive bug fix #4607

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cfda062
bug fix
seunghwak Aug 8, 2024
6091f71
udpate transform_v_frontier_e with edge mask to return the transforme…
seunghwak Aug 9, 2024
c574043
add tests that include zero biases
seunghwak Aug 9, 2024
9d566e8
Merge branch 'branch-24.10' of https://github.com/rapidsai/cugraph in…
seunghwak Aug 9, 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
2 changes: 1 addition & 1 deletion cpp/src/prims/detail/sample_and_compute_local_nbr_indices.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2041,7 +2041,7 @@ biased_sample_and_compute_local_nbr_indices(
zero_bias_frontier_indices.resize(zero_bias_count_inclusive_sums.back(),
handle.get_stream());
zero_bias_frontier_indices.shrink_to_fit(handle.get_stream());
zero_bias_local_nbr_indices.resize(frontier_indices.size(), handle.get_stream());
zero_bias_local_nbr_indices.resize(zero_bias_frontier_indices.size(), handle.get_stream());
zero_bias_local_nbr_indices.shrink_to_fit(handle.get_stream());
std::vector<size_t> zero_bias_counts(zero_bias_count_inclusive_sums.size());
std::adjacent_difference(zero_bias_count_inclusive_sums.begin(),
Expand Down
55 changes: 28 additions & 27 deletions cpp/src/prims/detail/transform_v_frontier_e.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,9 +209,6 @@ __global__ static void transform_v_frontier_e_mid_degree(
auto const lane_id = tid % raft::warp_size();
size_t idx = static_cast<size_t>(tid / raft::warp_size());

using WarpScan = cub::WarpScan<edge_t, raft::warp_size()>;
__shared__ typename WarpScan::TempStorage temp_storage;

while (idx < static_cast<size_t>(thrust::distance(edge_partition_frontier_key_index_first,
edge_partition_frontier_key_index_last))) {
auto key_idx = *(edge_partition_frontier_key_index_first + idx);
Expand All @@ -224,16 +221,15 @@ __global__ static void transform_v_frontier_e_mid_degree(
thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset);
auto this_key_value_first = value_first + edge_partition_frontier_local_degree_offsets[key_idx];
if (edge_partition_e_mask) {
// FIXME: it might be faster to update in warp-sync way
edge_t counter{0};
for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) {
if ((*edge_partition_e_mask).get(edge_offset + i)) { ++counter; }
}
edge_t offset_within_warp{};
WarpScan(temp_storage).ExclusiveSum(counter, offset_within_warp);
counter = 0;
for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) {
if ((*edge_partition_e_mask).get(edge_offset + i)) {
auto rounded_up_local_degree =
((static_cast<size_t>(local_degree) + (raft::warp_size() - 1)) / raft::warp_size()) *
raft::warp_size();
edge_t base_offset{0};
for (edge_t i = lane_id; i < rounded_up_local_degree; i += raft::warp_size()) {
auto valid = (i < local_degree) && (*edge_partition_e_mask).get(edge_offset + i);
auto ballot = __ballot_sync(raft::warp_full_mask(), valid ? uint32_t{1} : uint32_t{0});
if (valid) {
auto intra_warp_offset = __popc(ballot & ~(raft::warp_full_mask() << lane_id));
transform_v_frontier_e_update_buffer_element<key_t, GraphViewType>(
edge_partition,
key,
Expand All @@ -244,9 +240,9 @@ __global__ static void transform_v_frontier_e_mid_degree(
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
this_key_value_first + offset_within_warp + counter);
++counter;
this_key_value_first + base_offset + intra_warp_offset);
}
base_offset += __popc(ballot);
}
} else {
for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) {
Expand Down Expand Up @@ -300,6 +296,7 @@ __global__ static void transform_v_frontier_e_high_degree(

using BlockScan = cub::BlockScan<edge_t, transform_v_frontier_e_kernel_block_size>;
__shared__ typename BlockScan::TempStorage temp_storage;
__shared__ edge_t increment;

while (idx < static_cast<size_t>(thrust::distance(edge_partition_frontier_key_index_first,
edge_partition_frontier_key_index_last))) {
Expand All @@ -313,16 +310,16 @@ __global__ static void transform_v_frontier_e_high_degree(
thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset);
auto this_key_value_first = value_first + edge_partition_frontier_local_degree_offsets[key_idx];
if (edge_partition_e_mask) {
// FIXME: it might be faster to update in block-sync way
edge_t counter{0};
for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) {
if ((*edge_partition_e_mask).get(edge_offset + i)) { ++counter; }
}
edge_t offset_within_block{};
BlockScan(temp_storage).ExclusiveSum(counter, offset_within_block);
counter = 0;
for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) {
if ((*edge_partition_e_mask).get(edge_offset + i)) {
auto rounded_up_local_degree =
((static_cast<size_t>(local_degree) + (transform_v_frontier_e_kernel_block_size - 1)) /
transform_v_frontier_e_kernel_block_size) *
transform_v_frontier_e_kernel_block_size;
edge_t base_offset{0};
for (size_t i = threadIdx.x; i < rounded_up_local_degree; i += blockDim.x) {
auto valid = (i < local_degree) && (*edge_partition_e_mask).get(edge_offset + i);
edge_t intra_block_offset{};
BlockScan(temp_storage).ExclusiveSum(valid ? edge_t{1} : edge_t{0}, intra_block_offset);
if (valid) {
transform_v_frontier_e_update_buffer_element<key_t, GraphViewType>(
edge_partition,
key,
Expand All @@ -333,9 +330,13 @@ __global__ static void transform_v_frontier_e_high_degree(
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
this_key_value_first + offset_within_block + counter);
++counter;
this_key_value_first + base_offset + intra_block_offset);
}
if (threadIdx.x == transform_v_frontier_e_kernel_block_size - 1) {
increment = intra_block_offset + (valid ? edge_t{1} : edge_t{0});
}
__syncthreads();
base_offset += increment;
}
} else {
for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) {
Expand Down
Loading
Loading