MG C-API test failure fixes

cpp/CMakeLists.txt

@@ -423,7 +423,7 @@ add_library(cugraph_c
         src/c_api/core_result.cpp
         src/c_api/extract_ego.cpp
         src/c_api/k_core.cpp
-	src/c_api/hierarchical_clustering_result.cpp
+        src/c_api/hierarchical_clustering_result.cpp
         src/c_api/induced_subgraph.cpp
         src/c_api/capi_helper.cu
         src/c_api/legacy_spectral.cpp

cpp/src/c_api/capi_helper.cu

@@ -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

cpp/src/c_api/capi_helper.hpp

@@ -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

cpp/src/c_api/extract_ego.cpp

@@ -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>
@@ -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 {
@@ -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>(
@@ -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_),

cpp/src/community/leiden_impl.cuh

@@ -568,17 +568,17 @@ void flatten_leiden_dendrogram(raft::handle_t const& handle,
   leiden_partition_at_level<vertex_t, multi_gpu>(
     handle, dendrogram, clustering, dendrogram.num_levels());
 
-  rmm::device_uvector<vertex_t> unique_cluster_ids(graph_view.number_of_vertices(),
+  rmm::device_uvector<vertex_t> unique_cluster_ids(graph_view.local_vertex_partition_range_size(),
                                                    handle.get_stream());
   thrust::copy(handle.get_thrust_policy(),
                clustering,
-               clustering + graph_view.number_of_vertices(),
+               clustering + graph_view.local_vertex_partition_range_size(),
                unique_cluster_ids.begin());
 
   remove_duplicates<vertex_t, multi_gpu>(handle, unique_cluster_ids);
 
   relabel_cluster_ids<vertex_t, multi_gpu>(
-    handle, unique_cluster_ids, clustering, graph_view.number_of_vertices());
+    handle, unique_cluster_ids, clustering, graph_view.local_vertex_partition_range_size());
 }
 
 }  // namespace detail

cpp/src/detail/shuffle_vertices.cu

@@ -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,
@@ -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,

cpp/tests/CMakeLists.txt

@@ -257,7 +257,7 @@ ConfigureTest(BALANCED_TEST community/balanced_edge_test.cpp)
 
 ###################################################################################################
 # - EGO tests -------------------------------------------------------------------------------------
-ConfigureTest(EGO_TEST community/egonet_test.cpp GPUS 1 PERCENT 75)
+ConfigureTest(EGONET_TEST community/egonet_test.cpp GPUS 1 PERCENT 75)
 
 ###################################################################################################
 # - FORCE ATLAS 2  tests --------------------------------------------------------------------------
@@ -531,7 +531,7 @@ if(BUILD_CUGRAPH_MG_TESTS)
 
     ###############################################################################################
     # - MG LOUVAIN tests --------------------------------------------------------------------------
-    ConfigureTestMG(MG_EGO_TEST community/mg_egonet_test.cu)
+    ConfigureTestMG(MG_EGONET_TEST community/mg_egonet_test.cu)
 
     ###############################################################################################
     # - MG WEAKLY CONNECTED COMPONENTS tests ------------------------------------------------------

cpp/tests/c_api/mg_create_graph_test.c

@@ -175,18 +175,18 @@ int test_create_mg_graph_multiple_edge_lists(const cugraph_resource_handle_t* ha
   int my_rank = cugraph_resource_handle_get_rank(handle);
   int comm_size = cugraph_resource_handle_get_comm_size(handle);
 
-  size_t local_num_vertices = (num_vertices + comm_size - 1) / comm_size;
+  size_t local_num_vertices = num_vertices / comm_size;
   size_t local_start_vertex = my_rank * local_num_vertices;
-  size_t local_num_edges = (num_edges + comm_size - 1) / comm_size;
+  size_t local_num_edges = num_edges / comm_size;
   size_t local_start_edge = my_rank * local_num_edges;
 
-  local_num_edges = (local_num_edges < (num_edges - local_start_edge)) ? local_num_edges : (num_edges - local_start_edge);
-  local_num_vertices = (local_num_vertices < (num_vertices - local_start_vertex)) ? local_num_vertices : (num_vertices - local_start_vertex);
+  local_num_edges = (my_rank != (comm_size - 1)) ? local_num_edges : (num_edges - local_start_edge);
+  local_num_vertices = (my_rank != (comm_size - 1)) ? local_num_vertices : (num_vertices - local_start_vertex);
 
   for (size_t i = 0 ; i < num_local_arrays ; ++i) {
-    size_t vertex_count = (local_num_vertices + num_local_arrays - 1) / num_local_arrays;
+    size_t vertex_count = local_num_vertices / num_local_arrays;
     size_t vertex_start = i * vertex_count;
-    vertex_count = (vertex_count < (local_num_vertices - vertex_start)) ? vertex_count : (local_num_vertices - vertex_start);
+    vertex_count = (i != (num_local_arrays - 1)) ? vertex_count : (local_num_vertices - vertex_start);
 
     ret_code =
       cugraph_type_erased_device_array_create(handle, vertex_count, vertex_tid, vertices + i, &ret_error);
@@ -363,18 +363,18 @@ int test_create_mg_graph_multiple_edge_lists_multi_edge(const cugraph_resource_h
   int my_rank = cugraph_resource_handle_get_rank(handle);
   int comm_size = cugraph_resource_handle_get_comm_size(handle);
 
-  size_t local_num_vertices = (num_vertices + comm_size - 1) / comm_size;
+  size_t local_num_vertices = num_vertices / comm_size;
   size_t local_start_vertex = my_rank * local_num_vertices;
-  size_t local_num_edges = (num_edges + comm_size - 1) / comm_size;
+  size_t local_num_edges = num_edges / comm_size;
   size_t local_start_edge = my_rank * local_num_edges;
 
-  local_num_edges = (local_num_edges < (num_edges - local_start_edge)) ? local_num_edges : (num_edges - local_start_edge);
-  local_num_vertices = (local_num_vertices < (num_vertices - local_start_vertex)) ? local_num_vertices : (num_vertices - local_start_vertex);
+  local_num_edges = (my_rank != (comm_size - 1)) ? local_num_edges : (num_edges - local_start_edge);
+  local_num_vertices = (my_rank != (comm_size - 1)) ? local_num_vertices : (num_vertices - local_start_vertex);
 
   for (size_t i = 0 ; i < num_local_arrays ; ++i) {
     size_t vertex_count = (local_num_vertices + num_local_arrays - 1) / num_local_arrays;
     size_t vertex_start = i * vertex_count;
-    vertex_count = (vertex_count < (local_num_vertices - vertex_start)) ? vertex_count : (local_num_vertices - vertex_start);
+    vertex_count = (i != (num_local_arrays - 1)) ? vertex_count : (local_num_vertices - vertex_start);
 
     ret_code =
       cugraph_type_erased_device_array_create(handle, vertex_count, vertex_tid, vertices + i, &ret_error);