Debug MG egonet issues (#2926)

@jnke2016 identified some issues while creating the python bindings for egonet.  This PR addresses issues found while testing MG egonet in different MG configurations.

Authors:
  - Chuck Hastings (https://github.com/ChuckHastings)

Approvers:
  - Seunghwa Kang (https://github.com/seunghwak)
  - Joseph Nke (https://github.com/jnke2016)

URL: #2926

Author: ChuckHastings

cpp/src/community/egonet_impl.cuh

@@ -21,12 +21,7 @@
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/graph_view.hpp>
 #include <cugraph/utilities/error.hpp>
-
-#include <cstddef>
-#include <cugraph/algorithms.hpp>
-#include <memory>
-#include <tuple>
-#include <utility>
+#include <cugraph/utilities/host_scalar_comm.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_vector.hpp>
@@ -42,6 +37,13 @@
 
 #include <utilities/high_res_timer.hpp>
 
+#include <cstddef>
+#include <cugraph/algorithms.hpp>
+#include <memory>
+#include <numeric>
+#include <tuple>
+#include <utility>
+
 namespace {
 
 /*
@@ -69,20 +71,25 @@ extract(raft::handle_t const& handle,
         bool do_expensive_check)
 {
   auto user_stream_view = handle.get_stream();
-  rmm::device_uvector<size_t> neighbors_offsets(source_vertex.size() + 1, user_stream_view);
+
+  size_t num_sources = source_vertex.size();
+  [[maybe_unused]] std::vector<size_t> source_start{{0}};
+
+  if constexpr (multi_gpu) {
+    source_start =
+      cugraph::host_scalar_allgather(handle.get_comms(), num_sources, handle.get_stream());
+    num_sources = std::reduce(source_start.begin(), source_start.end());
+    std::exclusive_scan(source_start.begin(), source_start.end(), source_start.begin(), size_t{0});
+  }
+
+  rmm::device_uvector<size_t> neighbors_offsets(num_sources + 1, user_stream_view);
   rmm::device_uvector<vertex_t> neighbors(0, user_stream_view);
 
-  std::vector<size_t> h_neighbors_offsets(source_vertex.size() + 1);
+  std::vector<size_t> h_neighbors_offsets(num_sources + 1);
 
   // Streams will allocate concurrently later
   std::vector<rmm::device_uvector<vertex_t>> reached{};
-  reached.reserve(source_vertex.size());
-  for (size_t i = 0; i < source_vertex.size(); i++) {
-    // Allocations and operations are attached to the worker stream
-    rmm::device_uvector<vertex_t> local_reach(graph_view.local_vertex_partition_range_size(),
-                                              handle.get_next_usable_stream(i));
-    reached.push_back(std::move(local_reach));
-  }
+  reached.reserve(num_sources);
 
   user_stream_view.synchronize();
 #ifdef TIMING
@@ -96,11 +103,15 @@ extract(raft::handle_t const& handle,
   // the vertices and search for matches until the frontiers
   // are large enough to use this approach?
 
-  for (size_t i = 0; i < source_vertex.size(); i++) {
+  for (size_t i = 0; i < num_sources; i++) {
     // get light handle from worker pool
     raft::handle_t light_handle(handle.get_next_usable_stream(i));
     auto worker_stream_view = multi_gpu ? handle.get_stream() : light_handle.get_stream();
 
+    rmm::device_uvector<vertex_t> local_reach(graph_view.local_vertex_partition_range_size(),
+                                              worker_stream_view);
+    reached.push_back(std::move(local_reach));
+
     // BFS with cutoff
     // consider adding a device API to BFS (ie. accept source on the device)
     bool direction_optimizing = false;
@@ -109,12 +120,24 @@ extract(raft::handle_t const& handle,
                  reached[i].end(),
                  std::numeric_limits<vertex_t>::max());
 
+    raft::device_span<vertex_t const> source{source_vertex.data() + i, 1};
+
+    if constexpr (multi_gpu) {
+      auto it  = std::upper_bound(source_start.begin(), source_start.end(), i);
+      auto gpu = thrust::distance(source_start.begin(), it) - 1;
+
+      if (gpu == handle.get_comms().get_rank())
+        source = raft::device_span<vertex_t const>{source_vertex.data() + i - source_start[gpu], 1};
+      else
+        source = raft::device_span<vertex_t const>{source_vertex.data(), size_t{0}};
+    }
+
     cugraph::bfs<vertex_t, edge_t, weight_t, multi_gpu>(multi_gpu ? handle : light_handle,
                                                         graph_view,
                                                         reached[i].data(),
                                                         nullptr,
-                                                        source_vertex.data() + i,
-                                                        1,
+                                                        source.data(),
+                                                        source.size(),
                                                         direction_optimizing,
                                                         radius,
                                                         do_expensive_check);
@@ -146,18 +169,16 @@ extract(raft::handle_t const& handle,
 
   // Construct neighbors offsets (just a scan on neighborhod vector sizes)
   h_neighbors_offsets[0] = 0;
-  for (size_t i = 0; i < source_vertex.size(); i++) {
+  for (size_t i = 0; i < num_sources; i++) {
     h_neighbors_offsets[i + 1] = h_neighbors_offsets[i] + reached[i].size();
   }
-  raft::update_device(neighbors_offsets.data(),
-                      &h_neighbors_offsets[0],
-                      source_vertex.size() + 1,
-                      user_stream_view.value());
-  neighbors.resize(h_neighbors_offsets[source_vertex.size()], user_stream_view.value());
+  raft::update_device(
+    neighbors_offsets.data(), &h_neighbors_offsets[0], num_sources + 1, user_stream_view.value());
+  neighbors.resize(h_neighbors_offsets[num_sources], user_stream_view.value());
   user_stream_view.synchronize();
 
   // Construct the neighbors list concurrently
-  for (size_t i = 0; i < source_vertex.size(); i++) {
+  for (size_t i = 0; i < num_sources; i++) {
     auto worker_stream_view = handle.get_next_usable_stream(i);
     thrust::copy(rmm::exec_policy(worker_stream_view),
                  reached[i].begin(),
@@ -226,9 +247,6 @@ extract_ego(raft::handle_t const& handle,
             vertex_t radius,
             bool do_expensive_check)
 {
-  CUGRAPH_EXPECTS(source_vertex.size() > 0, "Need at least one source to extract the egonet from");
-  CUGRAPH_EXPECTS(source_vertex.size() < static_cast<size_t>(graph_view.number_of_vertices()),
-                  "Can't have more sources to extract from than vertices in the graph");
   CUGRAPH_EXPECTS(radius > 0, "Radius should be at least 1");
   CUGRAPH_EXPECTS(radius < graph_view.number_of_vertices(), "radius is too large");
 

cpp/src/structure/decompress_to_edgelist_impl.cuh

@@ -59,12 +59,11 @@ decompress_to_edgelist_impl(
   std::optional<raft::device_span<vertex_t const>> renumber_map,
   bool do_expensive_check)
 {
-  CUGRAPH_EXPECTS(renumber_map.has_value(),
-                  "Invalid input arguments: renumber_map.has_value() should be true if multi-GPU.");
-  CUGRAPH_EXPECTS(
-    (*renumber_map).size() == static_cast<size_t>(graph_view.local_vertex_partition_range_size()),
-    "Invalid input arguments: (*renumber_map).size() should match with the local "
-    "vertex partition range size.");
+  CUGRAPH_EXPECTS(!renumber_map.has_value() ||
+                    ((*renumber_map).size() ==
+                     static_cast<size_t>(graph_view.local_vertex_partition_range_size())),
+                  "Invalid input arguments: (*renumber_map).size() should match with the local "
+                  "vertex partition range size.");
 
   if (do_expensive_check) { /* currently, nothing to do */
   }

cpp/tests/CMakeLists.txt

@@ -476,7 +476,7 @@ if(BUILD_CUGRAPH_MG_TESTS)
 
     ###############################################################################################
     # - MG LOUVAIN tests --------------------------------------------------------------------------
-    ConfigureTestMG(MG_EGO_TEST community/mg_egonet_test.cpp)
+    ConfigureTestMG(MG_EGO_TEST community/mg_egonet_test.cu)
 
     ###############################################################################################
     # - MG WEAKLY CONNECTED COMPONENTS tests ------------------------------------------------------

cpp/tests/community/mg_egonet_test.cpp cpp/tests/community/mg_egonet_test.cucpp/tests/community/mg_egonet_test.cpp renamed to cpp/tests/community/mg_egonet_test.cu

@@ -16,6 +16,8 @@
 
 #include "egonet_validate.hpp"
 
+#include <structure/detail/structure_utils.cuh>
+
 #include <utilities/base_fixture.hpp>
 #include <utilities/device_comm_wrapper.hpp>
 #include <utilities/high_res_clock.h>
@@ -34,6 +36,7 @@
 #include <thrust/execution_policy.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/sequence.h>
+#include <thrust/sort.h>
 
 #include <gtest/gtest.h>
 
@@ -126,7 +129,7 @@ class Tests_MGEgonet
       cugraph::extract_ego(
         *handle_,
         mg_graph_view,
-        raft::device_span<vertex_t const>{d_ego_sources.data(), egonet_usecase.ego_sources_.size()},
+        raft::device_span<vertex_t const>{d_ego_sources.data(), d_ego_sources.size()},
         static_cast<vertex_t>(egonet_usecase.radius_));
 
     if (cugraph::test::g_perf) {
@@ -138,11 +141,6 @@ class Tests_MGEgonet
     }
 
     if (egonet_usecase.check_correctness_) {
-      *d_renumber_map_labels = cugraph::test::device_gatherv(
-        *handle_,
-        raft::device_span<vertex_t const>(d_renumber_map_labels->data(),
-                                          d_renumber_map_labels->size()));
-
       d_ego_edgelist_src = cugraph::test::device_gatherv(
         *handle_,
         raft::device_span<vertex_t const>(d_ego_edgelist_src.data(), d_ego_edgelist_src.size()));
@@ -157,54 +155,48 @@ class Tests_MGEgonet
                                           d_ego_edgelist_wgt->data(), d_ego_edgelist_wgt->size()));
       }
 
-      d_ego_edgelist_offsets = cugraph::test::device_gatherv(
-        *handle_,
+      size_t offsets_size = d_ego_edgelist_offsets.size();
+
+      auto graph_ids_v = cugraph::detail::expand_sparse_offsets(
         raft::device_span<size_t const>(d_ego_edgelist_offsets.data(),
-                                        d_ego_edgelist_offsets.size()));
+                                        d_ego_edgelist_offsets.size()),
+        vertex_t{0},
+        handle_->get_stream());
 
-      auto [sg_graph, sg_number_map] =
-        cugraph::test::mg_graph_to_sg_graph(*handle_, mg_graph_view, d_renumber_map_labels, false);
+      graph_ids_v = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>(graph_ids_v.data(), graph_ids_v.size()));
 
-      if (my_rank == 0) {
-        cugraph::unrenumber_int_vertices<vertex_t, false>(
-          *handle_,
-          d_ego_edgelist_src.data(),
-          d_ego_edgelist_src.size(),
-          d_renumber_map_labels->data(),
-          std::vector<vertex_t>{mg_graph_view.number_of_vertices()});
-
-        cugraph::unrenumber_int_vertices<vertex_t, false>(
-          *handle_,
-          d_ego_edgelist_dst.data(),
-          d_ego_edgelist_dst.size(),
-          d_renumber_map_labels->data(),
-          std::vector<vertex_t>{mg_graph_view.number_of_vertices()});
-
-        rmm::device_uvector<vertex_t> d_sg_ego_sources(egonet_usecase.ego_sources_.size(),
-                                                       handle_->get_stream());
-
-        if constexpr (std::is_same<int32_t, vertex_t>::value) {
-          raft::update_device(d_sg_ego_sources.data(),
-                              egonet_usecase.ego_sources_.data(),
-                              egonet_usecase.ego_sources_.size(),
-                              handle_->get_stream());
-        } else {
-          std::vector<vertex_t> h_ego_sources(d_sg_ego_sources.size());
-          std::transform(egonet_usecase.ego_sources_.begin(),
-                         egonet_usecase.ego_sources_.end(),
-                         h_ego_sources.begin(),
-                         [](auto v) { return static_cast<vertex_t>(v); });
-          raft::update_device(d_sg_ego_sources.data(),
-                              h_ego_sources.data(),
-                              h_ego_sources.size(),
-                              handle_->get_stream());
-        }
+      if (d_ego_edgelist_wgt) {
+        thrust::sort_by_key(
+          handle_->get_thrust_policy(),
+          thrust::make_zip_iterator(
+            graph_ids_v.begin(), d_ego_edgelist_src.begin(), d_ego_edgelist_dst.begin()),
+          thrust::make_zip_iterator(
+            graph_ids_v.end(), d_ego_edgelist_src.end(), d_ego_edgelist_dst.end()),
+          d_ego_edgelist_wgt->begin());
+      } else {
+        thrust::sort(handle_->get_thrust_policy(),
+                     thrust::make_zip_iterator(
+                       graph_ids_v.begin(), d_ego_edgelist_src.begin(), d_ego_edgelist_dst.begin()),
+                     thrust::make_zip_iterator(
+                       graph_ids_v.end(), d_ego_edgelist_src.end(), d_ego_edgelist_dst.end()));
+      }
+
+      d_ego_edgelist_offsets = cugraph::detail::compute_sparse_offsets<size_t>(
+        graph_ids_v.begin(), graph_ids_v.end(), size_t{0}, offsets_size - 1, handle_->get_stream());
 
+      auto [sg_graph, sg_number_map] = cugraph::test::mg_graph_to_sg_graph(
+        *handle_, mg_graph_view, std::optional<rmm::device_uvector<vertex_t>>{std::nullopt}, false);
+
+      d_ego_sources = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>(d_ego_sources.data(), d_ego_sources.size()));
+
+      if (my_rank == 0) {
         auto [d_reference_src, d_reference_dst, d_reference_wgt, d_reference_offsets] =
           cugraph::extract_ego(
             *handle_,
             sg_graph.view(),
-            raft::device_span<vertex_t const>{d_sg_ego_sources.data(), d_sg_ego_sources.size()},
+            raft::device_span<vertex_t const>{d_ego_sources.data(), d_ego_sources.size()},
             static_cast<vertex_t>(egonet_usecase.radius_));
 
         cugraph::test::egonet_validate(*handle_,
@@ -289,7 +281,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGEgonet_File,
   ::testing::Combine(
     // disable correctness checks for large graphs
-    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, true}),
+    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, false}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
 INSTANTIATE_TEST_SUITE_P(
@@ -301,7 +293,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGEgonet_File64,
   ::testing::Combine(
     // disable correctness checks for large graphs
-    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, true}),
+    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, false}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
 INSTANTIATE_TEST_SUITE_P(
@@ -313,12 +305,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGEgonet_Rmat,
   ::testing::Combine(
     // disable correctness checks for large graphs
-    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{0}, 1, false, true},
-                      Egonet_Usecase{std::vector<int32_t>{0}, 2, false, true},
-                      Egonet_Usecase{std::vector<int32_t>{0}, 3, false, true},
-                      Egonet_Usecase{std::vector<int32_t>{10, 0, 5}, 2, false, true},
-                      Egonet_Usecase{std::vector<int32_t>{9, 3, 10}, 2, false, true},
-                      Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, true}),
+    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, false}),
     ::testing::Values(
       cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, true, false, 0, true))));
 
@@ -331,7 +318,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGEgonet_Rmat64,
   ::testing::Combine(
     // disable correctness checks for large graphs
-    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, true}),
+    ::testing::Values(Egonet_Usecase{std::vector<int32_t>{5, 9, 3, 10, 12, 13}, 2, true, false}),
     ::testing::Values(
       cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, true, false, 0, true))));