[REVIEW] multi seeds BFS with one seed per component

cpp/include/algorithms.hpp

@@ -588,7 +588,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
@@ -1071,8 +1072,10 @@ namespace experimental {
  * @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
@@ -1086,7 +1089,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/utilities/cython.hpp

@@ -449,7 +449,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/egonet.cu

@@ -72,11 +72,8 @@ extract(
   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);
@@ -86,7 +83,6 @@ extract(
     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;
@@ -113,7 +109,8 @@ extract(
                                                                   csr_view,
                                                                   reached[i].data(),
                                                                   predecessors.data(),
-                                                                  h_source_vertex[i],
+                                                                  source_vertex + i,
+                                                                  1,
                                                                   direction_optimizing,
                                                                   radius);
 

cpp/src/experimental/bfs.cu

@@ -16,9 +16,12 @@
 
 #include <algorithms.hpp>
 #include <experimental/graph_view.hpp>
+#include <iostream>
+#include <patterns/count_if_v.cuh>
 #include <patterns/reduce_op.cuh>
 #include <patterns/update_frontier_v_push_if_out_nbr.cuh>
 #include <patterns/vertex_frontier.cuh>
+
 #include <utilities/error.hpp>
 #include <vertex_partition_device.cuh>
 
@@ -45,7 +48,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,47 +63,91 @@ 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 more than 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;
+  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());
+    }
+  } else {
+    d_sources = sources;
+  }
 
   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(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
-                    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_stream())->on(handle.get_stream()),
+               distances,
+               distances + push_graph_view.get_number_of_local_vertices(),
+               invalid_distance);
+  thrust::fill(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+               predecessor_first,
+               predecessor_first + push_graph_view.get_number_of_local_vertices(),
+               invalid_vertex);
+  vertex_partition_device_t<GraphViewType> vertex_partition(push_graph_view);
+  if (n_sources)
+    thrust::for_each(
+      rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+      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, 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, n_sources);
+  } else {
+    // pre-shuffled
+    vertex_frontier.get_bucket(static_cast<size_t>(Bucket::cur)).insert(d_sources, n_sources);
   }
-
   // 4. BFS iteration
-
   vertex_t depth{0};
   while (true) {
     if (direction_optimizing) {
@@ -160,7 +208,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)
@@ -170,7 +219,8 @@ void bfs(raft::handle_t const &handle,
                 graph_view,
                 distances,
                 predecessors,
-                source_vertex,
+                sources,
+                n_sources,
                 direction_optimizing,
                 depth_limit,
                 do_expensive_check);
@@ -179,7 +229,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);
@@ -192,7 +243,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);
@@ -201,7 +253,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);
@@ -210,7 +263,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);
@@ -219,7 +273,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);
@@ -228,7 +283,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);
@@ -237,7 +293,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);
@@ -246,7 +303,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);
@@ -255,7 +313,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);
@@ -264,7 +323,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);
@@ -273,7 +333,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);
@@ -282,7 +343,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);
@@ -291,7 +353,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);