diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp index a87fd0cfd8..600c2fe1a7 100644 --- a/cpp/src/c_api/neighbor_sampling.cpp +++ b/cpp/src/c_api/neighbor_sampling.cpp @@ -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_{}; @@ -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(handle)->handle_), rng_state_(reinterpret_cast(rng_state)), @@ -878,7 +880,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor { std::optional> start_vertex_labels{std::nullopt}; std::optional> local_label_to_comm_rank{std::nullopt}; std::optional> label_to_comm_rank{ - std::nullopt}; // global after allgatherv + std::nullopt}; // global after allgatherv std::optional> renumbered_and_sorted_edge_id_renumber_map( std::nullopt); @@ -1159,10 +1161,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor { ? std::make_optional(raft::device_span{ start_vertices_->as_type(), start_vertices_->size_}) : std::nullopt, - options_.retain_seeds_ - ? std::make_optional(raft::device_span{ - starting_vertex_label_offsets_->as_type(), starting_vertex_label_offsets_->size_}) - : std::nullopt, + options_.retain_seeds_ ? std::make_optional(raft::device_span{ + starting_vertex_label_offsets_->as_type(), + starting_vertex_label_offsets_->size_}) + : std::nullopt, offsets ? std::make_optional( raft::device_span{offsets->data(), offsets->size()}) : std::nullopt, @@ -1204,10 +1206,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor { ? std::make_optional(raft::device_span{ start_vertices_->as_type(), start_vertices_->size_}) : std::nullopt, - options_.retain_seeds_ - ? std::make_optional(raft::device_span{ - starting_vertex_label_offsets_->as_type(), starting_vertex_label_offsets_->size_}) - : std::nullopt, + options_.retain_seeds_ ? std::make_optional(raft::device_span{ + starting_vertex_label_offsets_->as_type(), + starting_vertex_label_offsets_->size_}) + : std::nullopt, offsets ? std::make_optional( raft::device_span{offsets->data(), offsets->size()}) : std::nullopt, @@ -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()} ); } @@ -1266,27 +1268,27 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor { ? std::make_optional(raft::device_span{ start_vertices_->as_type(), start_vertices_->size_}) : std::nullopt, - options_.retain_seeds_ - ? std::make_optional(raft::device_span{ - starting_vertex_label_offsets_->as_type(), starting_vertex_label_offsets_->size_}) - : std::nullopt, + options_.retain_seeds_ ? std::make_optional(raft::device_span{ + starting_vertex_label_offsets_->as_type(), + starting_vertex_label_offsets_->size_}) + : std::nullopt, offsets ? std::make_optional( raft::device_span{offsets->data(), offsets->size()}) : std::nullopt, - + (vertex_type_offsets_ != nullptr) - ? raft::device_span{vertex_type_offsets_->as_type(), vertex_type_offsets_->size_} + ? raft::device_span{vertex_type_offsets_->as_type(), + vertex_type_offsets_->size_} : raft::device_span{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_); diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp index 142d8054ab..a09cf757a5 100644 --- a/cpp/src/sampling/neighbor_sampling_impl.hpp +++ b/cpp/src/sampling/neighbor_sampling_impl.hpp @@ -30,6 +30,7 @@ #include #include + #include namespace cugraph { @@ -105,38 +106,30 @@ neighbor_sample_impl(raft::handle_t const& handle, edge_masks_vector{}; graph_view_t modified_graph_view = graph_view; edge_masks_vector.reserve(num_edge_types); - + label_t num_unique_labels = 0; std::optional> cp_starting_vertex_labels{std::nullopt}; - if (starting_vertex_labels) { // Find the number of unique lables - cp_starting_vertex_labels = rmm::device_uvector(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(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, bool> @@ -399,60 +392,57 @@ neighbor_sample_impl(raft::handle_t const& handle, if (result_labels) { cp_result_labels = rmm::device_uvector(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(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(0), + thrust::make_counting_iterator(result_offsets->size() - 1), + result_offsets->begin() + 1, + [result_labels = raft::device_span( + 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(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(0), - thrust::make_counting_iterator(result_offsets->size() - 1), - result_offsets->begin() + 1, - [ - result_labels = raft::device_span( - 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), diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx index e506a413ab..ee0e85fa3b 100644 --- a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx +++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx @@ -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. @@ -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 = \ @@ -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 = NULL if vertex_type_offsets is not None: vertex_type_offsets_ptr = \ diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx index 6e8ab93de6..dbee65323d 100644 --- a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx +++ b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx @@ -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. @@ -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 = \