[REVIEW] multi seeds BFS with one seed per component

cpp/include/cugraph/algorithms.hpp

@@ -592,7 +592,8 @@ void bfs(raft::handle_t const& handle,
  * 32-bit)
  * @tparam weight_t                  Type of edge weights. Supported values : float or double.
  *
- * @param[in]  handle                Library handle (RAFT). If a communicator is set in the handle,
+ * @param[in]  handle                Library handle (RAFT). If a communicator is set in the
+ * handle,
  * @param[in]  graph                 cuGRAPH COO graph
  * @param[in]  num_workers           number of vertices in the worker set
  * @param[in]  workers               device pointer to an array of worker vertex ids
@@ -1140,8 +1141,10 @@ weight_t hungarian(raft::handle_t const& handle,
  * @param graph_view Graph view object.
  * @param distances Pointer to the output distance array.
  * @param predecessors Pointer to the output predecessor array or `nullptr`.
- * @param source_vertex Source vertex to start breadth-first search (root vertex of the breath-first
- * search tree).
+ * @param sources Source vertices to start breadth-first search (root vertex of the breath-first
+ * search tree). If more than one source is passed, there must be a single source per component.
+ * Device memory and host memory are accepted.
+ * @param n_sources number of sources (one source per component at most).
  * @param direction_optimizing If set to true, this algorithm switches between the push based
  * breadth-first search and pull based breadth-first search depending on the size of the
  * breadth-first search frontier (currently unsupported). This option is valid only for symmetric
@@ -1155,7 +1158,8 @@ void bfs(raft::handle_t const& handle,
          graph_view_t<vertex_t, edge_t, weight_t, false, multi_gpu> const& graph_view,
          vertex_t* distances,
          vertex_t* predecessors,
-         vertex_t source_vertex,
+         vertex_t* sources,
+         size_t n_sources          = 1,
          bool direction_optimizing = false,
          vertex_t depth_limit      = std::numeric_limits<vertex_t>::max(),
          bool do_expensive_check   = false);

cpp/include/cugraph/utilities/cython.hpp

@@ -497,7 +497,8 @@ void call_bfs(raft::handle_t const& handle,
               vertex_t* distances,
               vertex_t* predecessors,
               vertex_t depth_limit,
-              const vertex_t start_vertex,
+              vertex_t* sources,
+              size_t n_sources,
               bool direction_optimizing);
 
 // Wrapper for calling SSSP through a graph container

cpp/src/community/legacy/egonet.cu

@@ -70,11 +70,8 @@ extract(raft::handle_t const& handle,
   rmm::device_vector<size_t> neighbors_offsets(n_subgraphs + 1);
   rmm::device_vector<vertex_t> neighbors;
 
-  std::vector<vertex_t> h_source_vertex(n_subgraphs);
   std::vector<size_t> h_neighbors_offsets(n_subgraphs + 1);
 
-  raft::update_host(&h_source_vertex[0], source_vertex, n_subgraphs, user_stream_view.value());
-
   // Streams will allocate concurrently later
   std::vector<rmm::device_uvector<vertex_t>> reached{};
   reached.reserve(n_subgraphs);
@@ -84,7 +81,6 @@ extract(raft::handle_t const& handle,
     reached.push_back(std::move(local_reach));
   }
 
-  // h_source_vertex[i] is used by other streams in the for loop
   user_stream_view.synchronize();
 #ifdef TIMING
   HighResTimer hr_timer;
@@ -111,7 +107,8 @@ extract(raft::handle_t const& handle,
                                                     csr_view,
                                                     reached[i].data(),
                                                     predecessors.data(),
-                                                    h_source_vertex[i],
+                                                    source_vertex + i,
+                                                    1,
                                                     direction_optimizing,
                                                     radius);
 

cpp/src/traversal/bfs.cu

@@ -16,10 +16,12 @@
 
 #include <cugraph/algorithms.hpp>
 #include <cugraph/graph_view.hpp>
+#include <cugraph/prims/count_if_v.cuh>
 #include <cugraph/prims/reduce_op.cuh>
 #include <cugraph/prims/update_frontier_v_push_if_out_nbr.cuh>
 #include <cugraph/prims/vertex_frontier.cuh>
 #include <cugraph/utilities/error.hpp>
+#include <cugraph/vertex_partition_device.cuh>
 #include <cugraph/vertex_partition_device_view.cuh>
 
 #include <raft/handle.hpp>
@@ -45,7 +47,8 @@ void bfs(raft::handle_t const& handle,
          GraphViewType const& push_graph_view,
          typename GraphViewType::vertex_type* distances,
          PredecessorIterator predecessor_first,
-         typename GraphViewType::vertex_type source_vertex,
+         typename GraphViewType::vertex_type* sources,
+         size_t n_sources,
          bool direction_optimizing,
          typename GraphViewType::vertex_type depth_limit,
          bool do_expensive_check)
@@ -59,50 +62,97 @@ void bfs(raft::handle_t const& handle,
 
   auto const num_vertices = push_graph_view.get_number_of_vertices();
   if (num_vertices == 0) { return; }
+  // CUGRAPH_EXPECTS(sources != nullptr, "Invalid input argument: sources cannot be null");
 
-  // 1. check input arguments
+  auto aggregate_n_sources =
+    GraphViewType::is_multi_gpu
+      ? host_scalar_allreduce(handle.get_comms(), n_sources, handle.get_stream())
+      : n_sources;
+  CUGRAPH_EXPECTS(aggregate_n_sources > 0,
+                  "Invalid input argument: input should have at least one source");
 
+  // 1. check input arguments
   CUGRAPH_EXPECTS(
     push_graph_view.is_symmetric() || !direction_optimizing,
     "Invalid input argument: input graph should be symmetric for direction optimizing BFS.");
-  CUGRAPH_EXPECTS(push_graph_view.is_valid_vertex(source_vertex),
-                  "Invalid input argument: source vertex out-of-range.");
+
+  // Transfer single source to the device for single source case
+  vertex_t* d_sources = sources;
+  rmm::device_uvector<vertex_t> d_sources_v(0, handle.get_stream());
+  if (aggregate_n_sources == 1 && n_sources) {
+    cudaPointerAttributes s_att;
+    CUDA_CHECK(cudaPointerGetAttributes(&s_att, sources));
+    if (s_att.devicePointer == nullptr) {
+      d_sources_v.resize(n_sources, handle.get_stream());
+      d_sources = d_sources_v.data();
+      raft::copy(d_sources, sources, n_sources, handle.get_stream());
+    }
+  }
 
   if (do_expensive_check) {
-    // nothing to do
+    vertex_partition_device_t<GraphViewType> vertex_partition(push_graph_view);
+    auto num_invalid_vertices =
+      count_if_v(handle,
+                 push_graph_view,
+                 d_sources,
+                 d_sources + n_sources,
+                 [vertex_partition] __device__(auto val) {
+                   return !(vertex_partition.is_valid_vertex(val) &&
+                            vertex_partition.is_local_vertex_nocheck(val));
+                 });
+    CUGRAPH_EXPECTS(num_invalid_vertices == 0,
+                    "Invalid input argument: sources have invalid vertex IDs.");
   }
 
   // 2. initialize distances and predecessors
 
   auto constexpr invalid_distance = std::numeric_limits<vertex_t>::max();
   auto constexpr invalid_vertex   = invalid_vertex_id<vertex_t>::value;
 
-  auto val_first = thrust::make_zip_iterator(thrust::make_tuple(distances, predecessor_first));
-  thrust::transform(handle.get_thrust_policy(),
-                    thrust::make_counting_iterator(push_graph_view.get_local_vertex_first()),
-                    thrust::make_counting_iterator(push_graph_view.get_local_vertex_last()),
-                    val_first,
-                    [source_vertex] __device__(auto val) {
-                      auto distance = invalid_distance;
-                      if (val == source_vertex) { distance = vertex_t{0}; }
-                      return thrust::make_tuple(distance, invalid_vertex);
-                    });
+  thrust::fill(rmm::exec_policy(handle.get_thrust_policy()),
+               distances,
+               distances + push_graph_view.get_number_of_local_vertices(),
+               invalid_distance);
+  thrust::fill(rmm::exec_policy(handle.get_thrust_policy()),
+               predecessor_first,
+               predecessor_first + push_graph_view.get_number_of_local_vertices(),
+               invalid_vertex);
+  auto vertex_partition = vertex_partition_device_view_t<vertex_t, GraphViewType::is_multi_gpu>(
+    push_graph_view.get_vertex_partition_view());
+  if (n_sources) {
+    thrust::for_each(
+      rmm::exec_policy(handle.get_thrust_policy()),
+      d_sources,
+      d_sources + n_sources,
+      [vertex_partition, distances, predecessor_first] __device__(auto v) {
+        *(distances + vertex_partition.get_local_vertex_offset_from_vertex_nocheck(v)) =
+          vertex_t{0};
+      });
+  }
 
   // 3. initialize BFS frontier
-
   enum class Bucket { cur, next, num_buckets };
   VertexFrontier<vertex_t,
                  void,
                  GraphViewType::is_multi_gpu,
                  static_cast<size_t>(Bucket::num_buckets)>
     vertex_frontier(handle);
 
-  if (push_graph_view.is_local_vertex_nocheck(source_vertex)) {
-    vertex_frontier.get_bucket(static_cast<size_t>(Bucket::cur)).insert(source_vertex);
+  // insert local source(s) in the bucket
+  if (aggregate_n_sources == 1) {
+    vertex_t src;
+    // FIXME: this (cheap) transfer could be skiped when is_local_vertex_nocheck accpets device mem
+    raft::copy(&src, sources, n_sources, handle.get_stream());
+    if (push_graph_view.is_local_vertex_nocheck(src)) {
+      vertex_frontier.get_bucket(static_cast<size_t>(Bucket::cur))
+        .insert(d_sources, d_sources + n_sources);
+    }
+  } else {
+    // pre-shuffled
+    vertex_frontier.get_bucket(static_cast<size_t>(Bucket::cur))
+      .insert(d_sources, d_sources + n_sources);
   }
-
   // 4. BFS iteration
-
   vertex_t depth{0};
   while (true) {
     if (direction_optimizing) {
@@ -167,7 +217,8 @@ void bfs(raft::handle_t const& handle,
          graph_view_t<vertex_t, edge_t, weight_t, false, multi_gpu> const& graph_view,
          vertex_t* distances,
          vertex_t* predecessors,
-         vertex_t source_vertex,
+         vertex_t* sources,
+         size_t n_sources,
          bool direction_optimizing,
          vertex_t depth_limit,
          bool do_expensive_check)
@@ -177,7 +228,8 @@ void bfs(raft::handle_t const& handle,
                 graph_view,
                 distances,
                 predecessors,
-                source_vertex,
+                sources,
+                n_sources,
                 direction_optimizing,
                 depth_limit,
                 do_expensive_check);
@@ -186,7 +238,8 @@ void bfs(raft::handle_t const& handle,
                 graph_view,
                 distances,
                 thrust::make_discard_iterator(),
-                source_vertex,
+                sources,
+                n_sources,
                 direction_optimizing,
                 depth_limit,
                 do_expensive_check);
@@ -199,7 +252,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int32_t, float, false, true> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -208,7 +262,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int32_t, double, false, true> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -217,7 +272,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int64_t, float, false, true> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -226,7 +282,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int64_t, double, false, true> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -235,7 +292,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int64_t, int64_t, float, false, true> const& graph_view,
                   int64_t* distances,
                   int64_t* predecessors,
-                  int64_t source_vertex,
+                  int64_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int64_t depth_limit,
                   bool do_expensive_check);
@@ -244,7 +302,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int64_t, int64_t, double, false, true> const& graph_view,
                   int64_t* distances,
                   int64_t* predecessors,
-                  int64_t source_vertex,
+                  int64_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int64_t depth_limit,
                   bool do_expensive_check);
@@ -253,7 +312,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int32_t, float, false, false> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -262,7 +322,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int32_t, double, false, false> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -271,7 +332,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int64_t, float, false, false> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -280,7 +342,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int32_t, int64_t, double, false, false> const& graph_view,
                   int32_t* distances,
                   int32_t* predecessors,
-                  int32_t source_vertex,
+                  int32_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int32_t depth_limit,
                   bool do_expensive_check);
@@ -289,7 +352,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int64_t, int64_t, float, false, false> const& graph_view,
                   int64_t* distances,
                   int64_t* predecessors,
-                  int64_t source_vertex,
+                  int64_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int64_t depth_limit,
                   bool do_expensive_check);
@@ -298,7 +362,8 @@ template void bfs(raft::handle_t const& handle,
                   graph_view_t<int64_t, int64_t, double, false, false> const& graph_view,
                   int64_t* distances,
                   int64_t* predecessors,
-                  int64_t source_vertex,
+                  int64_t* sources,
+                  size_t n_sources,
                   bool direction_optimizing,
                   int64_t depth_limit,
                   bool do_expensive_check);