Skip to content

Commit

Permalink
fix C egonet test failure
Browse files Browse the repository at this point in the history
  • Loading branch information
@seunghwak
seunghwak committed Dec 8, 2023
1 parent 9f8b06a commit ab2f561
Show file tree
Hide file tree
Showing 4 changed files with 166 additions and 2 deletions.
98 changes: 98 additions & 0 deletions cpp/src/c_api/capi_helper.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,104 @@ template void sort_by_key(raft::handle_t const& handle,
raft::device_span<int64_t> keys,
raft::device_span<int64_t> values);

template <typename vertex_t, typename weight_t>
std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<vertex_t>,
rmm::device_uvector<vertex_t>,
std::optional<rmm::device_uvector<weight_t>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<vertex_t>&& edge_srcs,
rmm::device_uvector<vertex_t>&& edge_dsts,
std::optional<rmm::device_uvector<weight_t>>&& edge_weights)
{
rmm::device_uvector<size_t> sort_indices(edge_srcs.size(), handle.get_stream());
thrust::tabulate(
handle.get_thrust_policy(),
sort_indices.begin(),
sort_indices.end(),
[offset_lasts = raft::device_span<size_t const>(offsets.begin() + 1, offsets.end()),
source_indices = raft::device_span<size_t const>(source_indices.data(),
source_indices.size())] __device__(size_t i) {
auto idx = static_cast<size_t>(thrust::distance(
offset_lasts.begin(),
thrust::upper_bound(thrust::seq, offset_lasts.begin(), offset_lasts.end(), i)));
return source_indices[idx];
});
source_indices.resize(0, handle.get_stream());
source_indices.shrink_to_fit(handle.get_stream());

auto triplet_first =
thrust::make_zip_iterator(sort_indices.begin(), edge_srcs.begin(), edge_dsts.begin());
if (edge_weights) {
thrust::sort_by_key(handle.get_thrust_policy(),
triplet_first,
triplet_first + sort_indices.size(),
(*edge_weights).begin());
} else {
thrust::sort(handle.get_thrust_policy(), triplet_first, triplet_first + sort_indices.size());
}

thrust::tabulate(
handle.get_thrust_policy(),
offsets.begin() + 1,
offsets.end(),
[sort_indices = raft::device_span<size_t const>(sort_indices.data(),
sort_indices.size())] __device__(size_t i) {
return static_cast<size_t>(thrust::distance(
sort_indices.begin(),
thrust::upper_bound(thrust::seq, sort_indices.begin(), sort_indices.end(), i)));
});

return std::make_tuple(
std::move(offsets), std::move(edge_srcs), std::move(edge_dsts), std::move(edge_weights));
}

template std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<int32_t>,
rmm::device_uvector<int32_t>,
std::optional<rmm::device_uvector<float>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<int32_t>&& edge_srcs,
rmm::device_uvector<int32_t>&& edge_dsts,
std::optional<rmm::device_uvector<float>>&& edge_weights);

template std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<int32_t>,
rmm::device_uvector<int32_t>,
std::optional<rmm::device_uvector<double>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<int32_t>&& edge_srcs,
rmm::device_uvector<int32_t>&& edge_dsts,
std::optional<rmm::device_uvector<double>>&& edge_weights);

template std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<int64_t>,
rmm::device_uvector<int64_t>,
std::optional<rmm::device_uvector<float>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<int64_t>&& edge_srcs,
rmm::device_uvector<int64_t>&& edge_dsts,
std::optional<rmm::device_uvector<float>>&& edge_weights);

template std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<int64_t>,
rmm::device_uvector<int64_t>,
std::optional<rmm::device_uvector<double>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<int64_t>&& edge_srcs,
rmm::device_uvector<int64_t>&& edge_dsts,
std::optional<rmm::device_uvector<double>>&& edge_weights);

} // namespace detail
} // namespace c_api
} // namespace cugraph
12 changes: 12 additions & 0 deletions cpp/src/c_api/capi_helper.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,18 @@ void sort_by_key(raft::handle_t const& handle,
raft::device_span<key_t> keys,
raft::device_span<value_t> values);

template <typename vertex_t, typename weight_t>
std::tuple<rmm::device_uvector<size_t>,
rmm::device_uvector<vertex_t>,
rmm::device_uvector<vertex_t>,
std::optional<rmm::device_uvector<weight_t>>>
reorder_extracted_egonets(raft::handle_t const& handle,
rmm::device_uvector<size_t>&& source_indices,
rmm::device_uvector<size_t>&& offsets,
rmm::device_uvector<vertex_t>&& edge_srcs,
rmm::device_uvector<vertex_t>&& edge_dsts,
std::optional<rmm::device_uvector<weight_t>>&& edge_weights);

} // namespace detail
} // namespace c_api
} // namespace cugraph
46 changes: 44 additions & 2 deletions cpp/src/c_api/extract_ego.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <cugraph_c/algorithms.h>

#include <c_api/abstract_functor.hpp>
#include <c_api/capi_helper.hpp>
#include <c_api/graph.hpp>
#include <c_api/induced_subgraph_result.hpp>
#include <c_api/resource_handle.hpp>
Expand All @@ -26,7 +27,10 @@
#include <cugraph/detail/shuffle_wrappers.hpp>
#include <cugraph/detail/utility_wrappers.hpp>
#include <cugraph/graph_functions.hpp>
#include <cugraph/utilities/device_comm.hpp>
#include <cugraph/utilities/host_scalar_comm.hpp>

#include <numeric>
#include <optional>

namespace {
Expand Down Expand Up @@ -91,9 +95,22 @@ struct extract_ego_functor : public cugraph::c_api::abstract_functor {
source_vertices.size(),
handle_.get_stream());

std::optional<rmm::device_uvector<size_t>> source_indices{std::nullopt};

if constexpr (multi_gpu) {
source_vertices = cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
handle_, std::move(source_vertices));
auto displacements = cugraph::host_scalar_allgather(
handle_.get_comms(), source_vertices.size(), handle_.get_stream());
std::exclusive_scan(
displacements.begin(), displacements.end(), displacements.begin(), size_t{0});
source_indices = rmm::device_uvector<size_t>(source_vertices.size(), handle_.get_stream());
cugraph::detail::sequence_fill(handle_.get_stream(),
(*source_indices).data(),
(*source_indices).size(),
displacements[handle_.get_comms().get_rank()]);

std::tie(source_vertices, source_indices) =
cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
handle_, std::move(source_vertices), std::move(*source_indices));
}

cugraph::renumber_ext_vertices<vertex_t, multi_gpu>(
Expand Down Expand Up @@ -130,6 +147,31 @@ struct extract_ego_functor : public cugraph::c_api::abstract_functor {
graph_view.vertex_partition_range_lasts(),
do_expensive_check_);

if constexpr (multi_gpu) {
auto recvcounts = cugraph::host_scalar_allgather(
handle_.get_comms(), (*source_indices).size(), handle_.get_stream());
std::vector<size_t> displacements(recvcounts.size());
std::exclusive_scan(recvcounts.begin(), recvcounts.end(), displacements.begin(), size_t{0});
rmm::device_uvector<size_t> allgathered_indices(displacements.back() + recvcounts.back(),
handle_.get_stream());
cugraph::device_allgatherv(handle_.get_comms(),
(*source_indices).begin(),
allgathered_indices.begin(),
recvcounts,
displacements,
handle_.get_stream());
source_indices = std::move(allgathered_indices);

std::tie(edge_offsets, src, dst, wgt) =
cugraph::c_api::detail::reorder_extracted_egonets<vertex_t, weight_t>(
handle_,
std::move(*source_indices),
std::move(edge_offsets),
std::move(src),
std::move(dst),
std::move(wgt));
}

result_ = new cugraph::c_api::cugraph_induced_subgraph_result_t{
new cugraph::c_api::cugraph_type_erased_device_array_t(src, graph_->vertex_type_),
new cugraph::c_api::cugraph_type_erased_device_array_t(dst, graph_->vertex_type_),
Expand Down
12 changes: 12 additions & 0 deletions cpp/src/detail/shuffle_vertices.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,12 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
rmm::device_uvector<int32_t>&& vertices,
rmm::device_uvector<int32_t>&& values);

template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>>
shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
raft::handle_t const& handle,
rmm::device_uvector<int32_t>&& vertices,
rmm::device_uvector<size_t>&& values);

template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
raft::handle_t const& handle,
Expand All @@ -224,6 +230,12 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
rmm::device_uvector<int64_t>&& vertices,
rmm::device_uvector<int64_t>&& values);

template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>>
shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
raft::handle_t const& handle,
rmm::device_uvector<int64_t>&& vertices,
rmm::device_uvector<size_t>&& values);

template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>>
shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
raft::handle_t const& handle,
Expand Down

0 comments on commit ab2f561

Please sign in to comment.