Skip to content

Commit

Permalink
fix style
Browse files Browse the repository at this point in the history
  • Loading branch information
@jnke2016
jnke2016 committed Nov 22, 2024
1 parent 061c8cc commit ed4c069
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 117 deletions.
78 changes: 40 additions & 38 deletions cpp/src/c_api/neighbor_sampling.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -778,7 +778,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
cugraph::c_api::cugraph_graph_t* graph_{nullptr};
cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
cugraph::c_api::cugraph_type_erased_device_array_view_t const* starting_vertex_label_offsets_{nullptr};
cugraph::c_api::cugraph_type_erased_device_array_view_t const* starting_vertex_label_offsets_{
nullptr};
cugraph::c_api::cugraph_type_erased_device_array_view_t const* vertex_type_offsets_{nullptr};
cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
int num_edge_types_{};
Expand All @@ -787,18 +788,19 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
bool do_expensive_check_{false};
cugraph::c_api::cugraph_sample_result_t* result_{nullptr};

neighbor_sampling_functor(cugraph_resource_handle_t const* handle,
cugraph_rng_state_t* rng_state,
cugraph_graph_t* graph,
cugraph_edge_property_view_t const* edge_biases,
cugraph_type_erased_device_array_view_t const* start_vertices,
cugraph_type_erased_device_array_view_t const* starting_vertex_label_offsets,
cugraph_type_erased_device_array_view_t const* vertex_type_offsets,
cugraph_type_erased_host_array_view_t const* fan_out,
int num_edge_types,
cugraph::c_api::cugraph_sampling_options_t options,
bool is_biased,
bool do_expensive_check)
neighbor_sampling_functor(
cugraph_resource_handle_t const* handle,
cugraph_rng_state_t* rng_state,
cugraph_graph_t* graph,
cugraph_edge_property_view_t const* edge_biases,
cugraph_type_erased_device_array_view_t const* start_vertices,
cugraph_type_erased_device_array_view_t const* starting_vertex_label_offsets,
cugraph_type_erased_device_array_view_t const* vertex_type_offsets,
cugraph_type_erased_host_array_view_t const* fan_out,
int num_edge_types,
cugraph::c_api::cugraph_sampling_options_t options,
bool is_biased,
bool do_expensive_check)
: abstract_functor(),
handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
Expand Down Expand Up @@ -878,7 +880,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
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
std::nullopt}; // global after allgatherv

std::optional<rmm::device_uvector<edge_t>> renumbered_and_sorted_edge_id_renumber_map(
std::nullopt);
Expand Down Expand Up @@ -1159,10 +1161,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
? std::make_optional(raft::device_span<vertex_t const>{
start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
: std::nullopt,
options_.retain_seeds_
? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(), starting_vertex_label_offsets_->size_})
: std::nullopt,
options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(),
starting_vertex_label_offsets_->size_})
: std::nullopt,
offsets ? std::make_optional(
raft::device_span<size_t const>{offsets->data(), offsets->size()})
: std::nullopt,
Expand Down Expand Up @@ -1204,10 +1206,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
? std::make_optional(raft::device_span<vertex_t const>{
start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
: std::nullopt,
options_.retain_seeds_
? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(), starting_vertex_label_offsets_->size_})
: std::nullopt,
options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(),
starting_vertex_label_offsets_->size_})
: std::nullopt,
offsets ? std::make_optional(
raft::device_span<size_t const>{offsets->data(), offsets->size()})
: std::nullopt,
Expand All @@ -1234,10 +1236,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
// If no 'vertex_type_offsets' is provided, all vertices are assumed to have
// a vertex type of value 1.
cugraph::detail::stride_fill(handle_.get_stream(),
vertex_type_offsets.begin(),
vertex_type_offsets.size(),
vertex_t{0},
vertex_t{graph_view.local_vertex_partition_range_size()}
vertex_type_offsets.begin(),
vertex_type_offsets.size(),
vertex_t{0},
vertex_t{graph_view.local_vertex_partition_range_size()}

);
}
Expand Down Expand Up @@ -1266,27 +1268,27 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
? std::make_optional(raft::device_span<vertex_t const>{
start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
: std::nullopt,
options_.retain_seeds_
? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(), starting_vertex_label_offsets_->size_})
: std::nullopt,
options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
starting_vertex_label_offsets_->as_type<size_t>(),
starting_vertex_label_offsets_->size_})
: std::nullopt,
offsets ? std::make_optional(
raft::device_span<size_t const>{offsets->data(), offsets->size()})
: std::nullopt,

(vertex_type_offsets_ != nullptr)
? raft::device_span<vertex_t const>{vertex_type_offsets_->as_type<vertex_t>(), vertex_type_offsets_->size_}
? raft::device_span<vertex_t const>{vertex_type_offsets_->as_type<vertex_t>(),
vertex_type_offsets_->size_}
: raft::device_span<vertex_t const>{vertex_type_offsets.data(),
vertex_type_offsets.size()},
vertex_type_offsets.size()},

edge_label ? (*offsets).size() - 1 : size_t{1},
hop ? fan_out_->size_ : size_t{1},

(vertex_type_offsets_ != nullptr)
? vertex_type_offsets_->size_ - 1
: vertex_type_offsets.size() - 1,
(vertex_type_offsets_ != nullptr) ? vertex_type_offsets_->size_ - 1
: vertex_type_offsets.size() - 1,

// num_vertex_type is by default 1 if 'vertex_type_offsets' is not provided
// num_vertex_type is by default 1 if 'vertex_type_offsets' is not provided
num_edge_types_,
src_is_major,
do_expensive_check_);
Expand Down
138 changes: 64 additions & 74 deletions cpp/src/sampling/neighbor_sampling_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
#include <raft/core/handle.hpp>

#include <rmm/device_uvector.hpp>

#include <thrust/unique.h>

namespace cugraph {
Expand Down Expand Up @@ -105,38 +106,30 @@ neighbor_sample_impl(raft::handle_t const& handle,
edge_masks_vector{};
graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
edge_masks_vector.reserve(num_edge_types);

label_t num_unique_labels = 0;

std::optional<rmm::device_uvector<label_t>> cp_starting_vertex_labels{std::nullopt};


if (starting_vertex_labels) {
// Find the number of unique lables
cp_starting_vertex_labels = rmm::device_uvector<label_t>(starting_vertex_labels->size(), handle.get_stream());

thrust::copy(
handle.get_thrust_policy(),
starting_vertex_labels->begin(),
starting_vertex_labels->end(),
cp_starting_vertex_labels->begin());

thrust::sort(
handle.get_thrust_policy(),
cp_starting_vertex_labels->begin(),
cp_starting_vertex_labels->end());

cp_starting_vertex_labels =
rmm::device_uvector<label_t>(starting_vertex_labels->size(), handle.get_stream());

thrust::copy(handle.get_thrust_policy(),
starting_vertex_labels->begin(),
starting_vertex_labels->end(),
cp_starting_vertex_labels->begin());

thrust::sort(handle.get_thrust_policy(),
cp_starting_vertex_labels->begin(),
cp_starting_vertex_labels->end());

num_unique_labels = thrust::unique_count(handle.get_thrust_policy(),
cp_starting_vertex_labels->begin(),
cp_starting_vertex_labels->end());


}





if (num_edge_types > 1) {
for (int i = 0; i < num_edge_types; i++) {
cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>
Expand Down Expand Up @@ -399,60 +392,57 @@ neighbor_sample_impl(raft::handle_t const& handle,
if (result_labels) {
cp_result_labels = rmm::device_uvector<label_t>(result_labels->size(), handle.get_stream());

thrust::copy(
handle.get_thrust_policy(),
result_labels->begin(),
result_labels->end(),
cp_result_labels->begin());
thrust::copy(handle.get_thrust_policy(),
result_labels->begin(),
result_labels->end(),
cp_result_labels->begin());
}
std::tie(result_srcs,
result_dsts,
result_weights,
result_edge_ids,
result_edge_types,
result_hops,
result_labels,
result_offsets) = detail::shuffle_and_organize_output(handle,
std::move(result_srcs),
std::move(result_dsts),
std::move(result_weights),
std::move(result_edge_ids),
std::move(result_edge_types),
std::move(result_hops),
std::move(result_labels),
label_to_output_comm_rank);

if (result_labels && (result_offsets->size() != num_unique_labels + 1)) {
// There are missing labels not sampled.
result_offsets = rmm::device_uvector<size_t>(num_unique_labels + 1, handle.get_stream());

// Sort labels
thrust::sort(handle.get_thrust_policy(), cp_result_labels->begin(), cp_result_labels->end());

thrust::transform(handle.get_thrust_policy(),
thrust::make_counting_iterator<edge_t>(0),
thrust::make_counting_iterator<edge_t>(result_offsets->size() - 1),
result_offsets->begin() + 1,
[result_labels = raft::device_span<label_t const>(
cp_result_labels->data(), cp_result_labels->size())] __device__(auto idx) {
auto itr_lower = thrust::lower_bound(
thrust::seq, result_labels.begin(), result_labels.end(), idx);

auto itr_upper = thrust::upper_bound(
thrust::seq, result_labels.begin(), result_labels.end(), idx);

auto sampled_label_size = thrust::distance(itr_lower, itr_upper);
return sampled_label_size;
});

// Run inclusive scan
thrust::inclusive_scan(handle.get_thrust_policy(),
result_offsets->begin() + 1,
result_offsets->end(),
result_offsets->begin() + 1);
}
std::tie(result_srcs, result_dsts, result_weights, result_edge_ids,
result_edge_types, result_hops, result_labels, result_offsets)
= detail::shuffle_and_organize_output(handle,
std::move(result_srcs),
std::move(result_dsts),
std::move(result_weights),
std::move(result_edge_ids),
std::move(result_edge_types),
std::move(result_hops),
std::move(result_labels),
label_to_output_comm_rank);

if (result_labels && (result_offsets->size() != num_unique_labels + 1)) {
// There are missing labels not sampled.
result_offsets = rmm::device_uvector<size_t>(num_unique_labels + 1, handle.get_stream());

// Sort labels
thrust::sort(
handle.get_thrust_policy(),
cp_result_labels->begin(),
cp_result_labels->end());

thrust::transform(
handle.get_thrust_policy(),
thrust::make_counting_iterator<edge_t>(0),
thrust::make_counting_iterator<edge_t>(result_offsets->size() - 1),
result_offsets->begin() + 1,
[
result_labels = raft::device_span<label_t const>(
cp_result_labels->data(),
cp_result_labels->size())
] __device__(auto idx) {
auto itr_lower = thrust::lower_bound(
thrust::seq, result_labels.begin(), result_labels.end(), idx);

auto itr_upper = thrust::upper_bound(
thrust::seq, result_labels.begin(), result_labels.end(), idx);

auto sampled_label_size = thrust::distance(itr_lower, itr_upper);
return sampled_label_size;
});

// Run inclusive scan
thrust::inclusive_scan(handle.get_thrust_policy(),
result_offsets->begin() + 1,
result_offsets->end(),
result_offsets->begin() + 1);
}
return std::make_tuple(std::move(result_srcs),
std::move(result_dsts),
std::move(result_weights),
Expand Down
6 changes: 3 additions & 3 deletions python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
Offsets of each label within the start vertex list. Expanding
'starting_vertex_label_offsets' must lead to an array of
len(start_vertex_list)
vertex_type_offsets: device array type (Optional)
Offsets for each vertex type in the graph.
Expand Down Expand Up @@ -281,7 +281,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
if starting_vertex_label_offsets is not None:
cai_starting_vertex_label_offsets_ptr = \
starting_vertex_label_offsets.__cuda_array_interface__['data'][0]

cdef uintptr_t cai_vertex_type_offsets_ptr
if vertex_type_offsets is not None:
cai_vertex_type_offsets_ptr = \
Expand All @@ -303,7 +303,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
len(starting_vertex_label_offsets),
SIZE_T
)

cdef cugraph_type_erased_device_array_view_t* vertex_type_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
if vertex_type_offsets is not None:
vertex_type_offsets_ptr = \
Expand Down
4 changes: 2 additions & 2 deletions python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
Offsets of each label within the start vertex list. Expanding
'starting_vertex_label_offsets' must lead to an array of
len(start_vertex_list)
vertex_type_offsets: device array type (Optional)
Offsets for each vertex type in the graph.
Expand Down Expand Up @@ -275,7 +275,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
if starting_vertex_label_offsets is not None:
cai_starting_vertex_label_offsets_ptr = \
starting_vertex_label_offsets.__cuda_array_interface__['data'][0]

cdef uintptr_t cai_vertex_type_offsets_ptr
if vertex_type_offsets is not None:
cai_vertex_type_offsets_ptr = \
Expand Down

0 comments on commit ed4c069

Please sign in to comment.