Merge branch 'branch-22.06' of github.com:rapidsai/cugraph into enh_u…

…pdate_frontier_split

cpp/include/cugraph/algorithms.hpp

@@ -1574,6 +1574,14 @@ uniform_nbr_sample(raft::handle_t const& handle,
 /**
  * @brief Uniform Neighborhood Sampling.
  *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function a set of tuples (src, dst, weight, count), identifying the randomly
+ * selected edges.  src is the source vertex, dst is the destination vertex, weight is the weight
+ * of the edge and count identifies the number of times this edge was encountered during the
+ * sampling of this graph (so it is >= 1).
+ *
  * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
  * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
  * @tparam weight_t Type of edge weights. Needs to be a floating point type.
@@ -1587,22 +1595,25 @@ uniform_nbr_sample(raft::handle_t const& handle,
  * @param with_replacement boolean flag specifying if random sampling is done with replacement
  * (true); or, without replacement (false); default = true;
  * @param seed A seed to initialize the random number generator
- * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex, weight_t wgt)
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex, weight_t
+ * weight, edge_t count)
  */
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
           bool store_transposed,
           bool multi_gpu>
-std::
-  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>, rmm::device_uvector<weight_t>>
-  uniform_nbr_sample(
-    raft::handle_t const& handle,
-    graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
-    raft::device_span<vertex_t> starting_vertices,
-    raft::host_span<const int> fan_out,
-    bool with_replacement = true,
-    uint64_t seed         = 0);
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<weight_t>,
+           rmm::device_uvector<edge_t>>
+uniform_nbr_sample(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
+  raft::device_span<vertex_t> starting_vertices,
+  raft::host_span<const int> fan_out,
+  bool with_replacement = true,
+  uint64_t seed         = 0);
 
 /*
  * @brief Compute triangle counts.

cpp/src/c_api/uniform_neighbor_sampling.cpp

@@ -35,11 +35,13 @@ struct cugraph_sample_result_t {
   bool experimental_{true};
   cugraph_type_erased_device_array_t* src_{nullptr};
   cugraph_type_erased_device_array_t* dst_{nullptr};
-  // FIXME: Will be deleted once experimental replaces curren
+  // FIXME: Will be deleted once experimental replaces current
   cugraph_type_erased_device_array_t* label_{nullptr};
   cugraph_type_erased_device_array_t* index_{nullptr};
-  // FIXME: Will be deleted once experimental replaces curren
+  // FIXME: Will be deleted once experimental replaces current
   cugraph_type_erased_host_array_t* count_{nullptr};
+  // FIXME: Rename to count_ once experimental replaces current
+  cugraph_type_erased_device_array_t* experimental_count_{nullptr};
 };
 
 }  // namespace c_api
@@ -233,7 +235,7 @@ struct experimental_uniform_neighbor_sampling_functor : public cugraph::c_api::a
         graph_view.local_vertex_partition_range_last(),
         false);
 
-      auto&& [srcs, dsts, weights] = cugraph::uniform_nbr_sample(
+      auto&& [srcs, dsts, weights, counts] = cugraph::uniform_nbr_sample(
         handle_,
         graph_view,
         raft::device_span<vertex_t>(start.data(), start.size()),
@@ -262,7 +264,8 @@ struct experimental_uniform_neighbor_sampling_functor : public cugraph::c_api::a
         new cugraph::c_api::cugraph_type_erased_device_array_t(dsts, graph_->vertex_type_),
         nullptr,
         new cugraph::c_api::cugraph_type_erased_device_array_t(weights, graph_->weight_type_),
-        nullptr};
+        nullptr,
+        new cugraph::c_api::cugraph_type_erased_device_array_t(counts, graph_->edge_type_)};
     }
   }
 };

cpp/src/sampling/detail/graph_functions.hpp

@@ -175,6 +175,16 @@ gather_one_hop_edgelist(
   GraphViewType const& graph_view,
   const rmm::device_uvector<typename GraphViewType::vertex_type>& active_majors);
 
+template <typename vertex_t, typename edge_t, typename weight_t>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<weight_t>,
+           rmm::device_uvector<edge_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<vertex_t>&& src,
+                            rmm::device_uvector<vertex_t>&& dst,
+                            rmm::device_uvector<weight_t>&& wgt);
+
 }  // namespace detail
 
 }  // namespace cugraph

cpp/src/sampling/detail/sampling_utils_impl.cuh

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cugraph/detail/decompress_edge_partition.cuh>
+#include <cugraph/detail/graph_utils.cuh>
 #include <cugraph/edge_partition_device_view.cuh>
 #include <cugraph/partition_manager.hpp>
 #include <cugraph/utilities/device_comm.cuh>
@@ -907,5 +908,46 @@ gather_one_hop_edgelist(
   return std::make_tuple(std::move(majors), std::move(minors), std::move(weights));
 }
 
+template <typename vertex_t, typename edge_t, typename weight_t>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<weight_t>,
+           rmm::device_uvector<edge_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<vertex_t>&& src,
+                            rmm::device_uvector<vertex_t>&& dst,
+                            rmm::device_uvector<weight_t>&& wgt)
+{
+  auto tuple_iter_begin =
+    thrust::make_zip_iterator(thrust::make_tuple(src.begin(), dst.begin(), wgt.begin()));
+
+  thrust::sort(handle.get_thrust_policy(), tuple_iter_begin, tuple_iter_begin + src.size());
+
+  auto num_uniques =
+    thrust::count_if(handle.get_thrust_policy(),
+                     thrust::make_counting_iterator(size_t{0}),
+                     thrust::make_counting_iterator(src.size()),
+                     detail::is_first_in_run_pair_t<vertex_t>{src.data(), dst.data()});
+
+  rmm::device_uvector<vertex_t> result_src(num_uniques, handle.get_stream());
+  rmm::device_uvector<vertex_t> result_dst(num_uniques, handle.get_stream());
+  rmm::device_uvector<weight_t> result_wgt(num_uniques, handle.get_stream());
+  rmm::device_uvector<edge_t> result_count(num_uniques, handle.get_stream());
+
+  rmm::device_uvector<edge_t> count(src.size(), handle.get_stream());
+  thrust::fill(handle.get_thrust_policy(), count.begin(), count.end(), edge_t{1});
+
+  thrust::reduce_by_key(handle.get_thrust_policy(),
+                        tuple_iter_begin,
+                        tuple_iter_begin + src.size(),
+                        count.begin(),
+                        thrust::make_zip_iterator(thrust::make_tuple(
+                          result_src.begin(), result_dst.begin(), result_wgt.begin())),
+                        result_count.begin());
+
+  return std::make_tuple(
+    std::move(result_src), std::move(result_dst), std::move(result_wgt), std::move(result_count));
+}
+
 }  // namespace detail
 }  // namespace cugraph

cpp/src/sampling/detail/sampling_utils_sg.cu

@@ -229,5 +229,60 @@ gather_one_hop_edgelist(raft::handle_t const& handle,
                         graph_view_t<int64_t, int64_t, double, false, false> const& graph_view,
                         rmm::device_uvector<int64_t> const& active_majors);
 
+//  Only need to build once, not separately for SG/MG
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int32_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int32_t>&& src,
+                            rmm::device_uvector<int32_t>&& dst,
+                            rmm::device_uvector<float>&& wgt);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int64_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int32_t>&& src,
+                            rmm::device_uvector<int32_t>&& dst,
+                            rmm::device_uvector<float>&& wgt);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int64_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int64_t>&& src,
+                            rmm::device_uvector<int64_t>&& dst,
+                            rmm::device_uvector<float>&& wgt);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int32_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int32_t>&& src,
+                            rmm::device_uvector<int32_t>&& dst,
+                            rmm::device_uvector<double>&& wgt);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int64_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int32_t>&& src,
+                            rmm::device_uvector<int32_t>&& dst,
+                            rmm::device_uvector<double>&& wgt);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int64_t>>
+count_and_remove_duplicates(raft::handle_t const& handle,
+                            rmm::device_uvector<int64_t>&& src,
+                            rmm::device_uvector<int64_t>&& dst,
+                            rmm::device_uvector<double>&& wgt);
+
 }  // namespace detail
 }  // namespace cugraph

cpp/src/sampling/uniform_neighbor_sampling_impl.hpp

@@ -41,7 +41,8 @@ namespace detail {
 template <typename graph_view_t>
 std::tuple<rmm::device_uvector<typename graph_view_t::vertex_type>,
            rmm::device_uvector<typename graph_view_t::vertex_type>,
-           rmm::device_uvector<typename graph_view_t::weight_type>>
+           rmm::device_uvector<typename graph_view_t::weight_type>,
+           rmm::device_uvector<typename graph_view_t::edge_type>>
 uniform_nbr_sample_impl(
   raft::handle_t const& handle,
   graph_view_t const& graph_view,
@@ -150,8 +151,8 @@ uniform_nbr_sample_impl(
     ++level;
   }
 
-  return std::make_tuple(
-    std::move(d_result_src), std::move(d_result_dst), std::move(*d_result_indices));
+  return count_and_remove_duplicates<vertex_t, edge_t, weight_t>(
+    handle, std::move(d_result_src), std::move(d_result_dst), std::move(*d_result_indices));
 }
 }  // namespace detail
 
@@ -160,15 +161,17 @@ template <typename vertex_t,
           typename weight_t,
           bool store_transposed,
           bool multi_gpu>
-std::
-  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>, rmm::device_uvector<weight_t>>
-  uniform_nbr_sample(
-    raft::handle_t const& handle,
-    graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
-    raft::device_span<vertex_t> starting_vertices,
-    raft::host_span<const int> fan_out,
-    bool with_replacement,
-    uint64_t seed)
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<weight_t>,
+           rmm::device_uvector<edge_t>>
+uniform_nbr_sample(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
+  raft::device_span<vertex_t> starting_vertices,
+  raft::host_span<const int> fan_out,
+  bool with_replacement,
+  uint64_t seed)
 {
   rmm::device_uvector<vertex_t> d_start_vs(starting_vertices.size(), handle.get_stream());
   raft::copy(

cpp/src/sampling/uniform_neighbor_sampling_mg.cpp

@@ -20,58 +20,70 @@
 
 namespace cugraph {
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int32_t, int32_t, float, false, true> const& graph_view,
-                     raft::device_span<int32_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int32_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int32_t, int32_t, float, false, true> const& graph_view,
+                   raft::device_span<int32_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int32_t, int64_t, float, false, true> const& graph_view,
-                     raft::device_span<int32_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int64_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int32_t, int64_t, float, false, true> const& graph_view,
+                   raft::device_span<int32_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>, rmm::device_uvector<float>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int64_t, int64_t, float, false, true> const& graph_view,
-                     raft::device_span<int64_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<float>,
+                    rmm::device_uvector<int64_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int64_t, int64_t, float, false, true> const& graph_view,
+                   raft::device_span<int64_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<double>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int32_t, int32_t, double, false, true> const& graph_view,
-                     raft::device_span<int32_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int32_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int32_t, int32_t, double, false, true> const& graph_view,
+                   raft::device_span<int32_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<double>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int32_t, int64_t, double, false, true> const& graph_view,
-                     raft::device_span<int32_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int64_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int32_t, int64_t, double, false, true> const& graph_view,
+                   raft::device_span<int32_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>, rmm::device_uvector<double>>
-  uniform_nbr_sample(raft::handle_t const& handle,
-                     graph_view_t<int64_t, int64_t, double, false, true> const& graph_view,
-                     raft::device_span<int64_t> starting_vertices,
-                     raft::host_span<const int> fan_out,
-                     bool with_replacement,
-                     uint64_t seed);
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<double>,
+                    rmm::device_uvector<int64_t>>
+uniform_nbr_sample(raft::handle_t const& handle,
+                   graph_view_t<int64_t, int64_t, double, false, true> const& graph_view,
+                   raft::device_span<int64_t> starting_vertices,
+                   raft::host_span<const int> fan_out,
+                   bool with_replacement,
+                   uint64_t seed);
 
 }  // namespace cugraph