Update Uniform Neighborhood Sampling API

cpp/CMakeLists.txt

@@ -402,6 +402,7 @@ add_library(cugraph_c
         src/c_api/sssp.cpp
         src/c_api/extract_paths.cpp
         src/c_api/random_walks.cpp
+        src/c_api/random.cpp
         src/c_api/similarity.cpp
         src/c_api/louvain.cpp
         src/c_api/triangle_count.cpp

cpp/include/cugraph/algorithms.hpp

@@ -1780,6 +1780,9 @@ k_core(raft::handle_t const& handle,
 /**
  * @brief Uniform Neighborhood Sampling.
  *
+ * @deprecated This function should be replaced with uniform_neighbor_sample.  Input of the
+ * new function adds an optional parameter, output has a number of extra fields.
+ *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects from these outgoing neighbors to extract a subgraph.
  *
@@ -1822,6 +1825,68 @@ uniform_nbr_sample(raft::handle_t const& handle,
                    bool with_replacement = true,
                    uint64_t seed         = 0);
 
+/**
+ * @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 is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * label), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in, label (optional) identifies which vertex label this edge was derived from.
+ *
+ * @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.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_type_view Optional view object holding edge ids and types for @p graph_view.
+ * @param starting_vertices Device vector of starting vertex IDs for the sampling.
+ * @param starting_labels Optional device vector of starting vertex labels for the sampling.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level
+ * @param with_replacement boolean flag specifying if random sampling is done with replacement
+ * (true); or, without replacement (false); default = true;
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type, int32_t hop,
+ * optional int32_t label)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           rmm::device_uvector<int32_t>,
+           std::optional<rmm::device_uvector<int32_t>>>
+uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<
+    edge_property_view_t<edge_t,
+                         thrust::zip_iterator<thrust::tuple<edge_t const*, edge_type_t const*>>>>
+    edge_id_type_view,
+  rmm::device_uvector<vertex_t>&& starting_vertices,
+  std::optional<rmm::device_uvector<int32_t>>&& starting_labels,
+  raft::host_span<int32_t const> fan_out,
+  raft::random::RngState& rng_state,
+  bool with_replacement = true);
+
 /*
  * @brief Compute triangle counts.
  *

cpp/include/cugraph/detail/shuffle_wrappers.hpp

@@ -28,7 +28,9 @@ namespace detail {
  * their local GPUs based on edge partitioning.
  *
  * @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.
+ * @tparam edge_type_t Type of edge type identifiers. Needs to be an integral type.
  *
  * @param[in] handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
@@ -38,25 +40,32 @@ namespace detail {
  * partitioning to determine the local GPU.
  * @param[in] minors Vector of second elements in vertex pairs.
  * @param[in] weights Optional vector of vertex pair weight values.
+ * @param[in] edge_id_type_tuple Optional tuple of vectors of edge id and edge type values
  *
  * @return Tuple of vectors storing shuffled major vertices, minor vertices and optional weights.
  */
-template <typename vertex_t, typename weight_t>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>>
+template <typename vertex_t, typename edge_t, typename weight_t, typename edge_type_id_t>
+std::tuple<
+  rmm::device_uvector<vertex_t>,
+  rmm::device_uvector<vertex_t>,
+  std::optional<rmm::device_uvector<weight_t>>,
+  std::optional<std::tuple<rmm::device_uvector<edge_t>, rmm::device_uvector<edge_type_id_t>>>>
 shuffle_ext_vertex_pairs_to_local_gpu_by_edge_partitioning(
   raft::handle_t const& handle,
   rmm::device_uvector<vertex_t>&& majors,
   rmm::device_uvector<vertex_t>&& minors,
-  std::optional<rmm::device_uvector<weight_t>>&& weights);
+  std::optional<rmm::device_uvector<weight_t>>&& weights,
+  std::optional<std::tuple<rmm::device_uvector<edge_t>, rmm::device_uvector<edge_type_id_t>>>&&
+    edge_id_type_tuple);
 
 /**
  * @brief Shuffle internal (i.e. renumbered) vertex pairs (which can be edge end points) to their
  * local GPUs based on edge partitioning.
  *
  * @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.
+ * @tparam edge_type_t Type of edge type identifiers. Needs to be an integral type.
  *
  * @param[in] handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
@@ -66,20 +75,25 @@ shuffle_ext_vertex_pairs_to_local_gpu_by_edge_partitioning(
  * partitioning to determine the local GPU.
  * @param[in] minors Vector of second elements in vertex pairs.
  * @param[in] weights Optional vector of vertex pair weight values.
+ * @param[in] edge_id_type_tuple Optional tuple of vectors of edge id and edge type values
  * @param[in] vertex_partition_range_lasts Vector of each GPU's vertex partition range's last
  * (exclusive) vertex ID.
  *
  * @return Tuple of vectors storing shuffled major vertices, minor vertices and optional weights.
  */
-template <typename vertex_t, typename weight_t>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>>
+template <typename vertex_t, typename edge_t, typename weight_t, typename edge_type_id_t>
+std::tuple<
+  rmm::device_uvector<vertex_t>,
+  rmm::device_uvector<vertex_t>,
+  std::optional<rmm::device_uvector<weight_t>>,
+  std::optional<std::tuple<rmm::device_uvector<edge_t>, rmm::device_uvector<edge_type_id_t>>>>
 shuffle_int_vertex_pairs_to_local_gpu_by_edge_partitioning(
   raft::handle_t const& handle,
   rmm::device_uvector<vertex_t>&& majors,
   rmm::device_uvector<vertex_t>&& minors,
   std::optional<rmm::device_uvector<weight_t>>&& weights,
+  std::optional<std::tuple<rmm::device_uvector<edge_t>, rmm::device_uvector<edge_type_id_t>>>&&
+    edge_id_type_tuple,
   std::vector<vertex_t> const& vertex_partition_range_lasts);
 
 /**
@@ -139,6 +153,28 @@ rmm::device_uvector<vertex_t> shuffle_int_vertices_to_local_gpu_by_vertex_partit
   rmm::device_uvector<vertex_t>&& vertices,
   std::vector<vertex_t> const& vertex_partition_range_lasts);
 
+/**
+ * @brief Shuffle vertices using the internal vertex key function which returns the target GPU ID.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam value_t Type of vertex values. Needs to be an integral type.
+ *
+ * @param[in] handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
+ * @param[in] vertices Vertex IDs to shuffle
+ * @param[in] values Vertex Values to shuffle
+ * @param[in] vertex_partition_range_lasts From graph view, vector of last vertex id for each gpu
+ *
+ * @return tuple containing device vector of shuffled vertices and device vector of corresponding
+ *         values
+ */
+template <typename vertex_t, typename value_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value_t>>
+shuffle_int_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<vertex_t>&& vertices,
+  rmm::device_uvector<value_t>&& values,
+  std::vector<vertex_t> const& vertex_partition_range_lasts);
+
 /**
  * @brief Groupby and count edgelist using the key function which returns the target local partition
  * ID for an edge.

cpp/include/cugraph/edge_property.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cugraph/utilities/dataframe_buffer.hpp>
+#include <cugraph/utilities/thrust_tuple_utils.hpp>
 
 #include <raft/core/handle.hpp>
 
@@ -124,4 +125,22 @@ class edge_dummy_property_t {
   auto view() const { return edge_dummy_property_view_t{}; }
 };
 
+template <typename edge_t, typename... Ts>
+auto view_concat(edge_property_view_t<edge_t, Ts> const&... views)
+{
+  using concat_value_iterator = decltype(thrust::make_zip_iterator(
+    thrust_tuple_cat(detail::to_thrust_tuple(views.value_firsts()[0])...)));
+
+  std::vector<concat_value_iterator> edge_partition_concat_value_firsts{};
+  auto first_view = detail::get_first_of_pack(views...);
+  edge_partition_concat_value_firsts.resize(first_view.value_firsts().size());
+  for (size_t i = 0; i < edge_partition_concat_value_firsts.size(); ++i) {
+    edge_partition_concat_value_firsts[i] = thrust::make_zip_iterator(
+      thrust_tuple_cat(detail::to_thrust_tuple(views.value_firsts()[i])...));
+  }
+
+  return edge_property_view_t<edge_t, concat_value_iterator>(edge_partition_concat_value_firsts,
+                                                             first_view.edge_counts());
+}
+
 }  // namespace cugraph

cpp/include/cugraph/edge_src_dst_property.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -333,28 +333,6 @@ class edge_endpoint_dummy_property_view_t {
   using value_iterator = void*;
 };
 
-template <typename Iterator,
-          typename std::enable_if_t<std::is_arithmetic<
-            typename std::iterator_traits<Iterator>::value_type>::value>* = nullptr>
-auto to_thrust_tuple(Iterator iter)
-{
-  return thrust::make_tuple(iter);
-}
-
-template <typename Iterator,
-          typename std::enable_if_t<is_thrust_tuple_of_arithmetic<
-            typename std::iterator_traits<Iterator>::value_type>::value>* = nullptr>
-auto to_thrust_tuple(Iterator iter)
-{
-  return iter.get_iterator_tuple();
-}
-
-template <typename T, typename... Ts>
-decltype(auto) get_first_of_pack(T&& t, Ts&&...)
-{
-  return std::forward<T>(t);
-}
-
 }  // namespace detail
 
 template <typename GraphViewType, typename T>

cpp/include/cugraph/utilities/thrust_tuple_utils.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -239,4 +239,28 @@ struct thrust_tuple_get {
   }
 };
 
+namespace detail {
+template <typename Iterator,
+          typename std::enable_if_t<std::is_arithmetic<
+            typename std::iterator_traits<Iterator>::value_type>::value>* = nullptr>
+auto to_thrust_tuple(Iterator iter)
+{
+  return thrust::make_tuple(iter);
+}
+
+template <typename Iterator,
+          typename std::enable_if_t<is_thrust_tuple_of_arithmetic<
+            typename std::iterator_traits<Iterator>::value_type>::value>* = nullptr>
+auto to_thrust_tuple(Iterator iter)
+{
+  return iter.get_iterator_tuple();
+}
+
+template <typename T, typename... Ts>
+decltype(auto) get_first_of_pack(T&& t, Ts&&...)
+{
+  return std::forward<T>(t);
+}
+
+}  // namespace detail
 }  // namespace cugraph

cpp/include/cugraph_c/random.h

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cugraph_c/resource_handle.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct {
+  int32_t align_;
+} cugraph_rng_state_t;
+
+/**
+ * @brief     Create a Random Number Generator State
+ *
+ * @param [in]  seed        Initial value for seed.  In MG this should be different
+ *                          on each GPU
+ * @param [out] state       Pointer to the location to store the pointer to the RngState
+ * @param [out] error       Pointer to an error object storing details of any error.  Will
+ *                          be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_rng_state_create(const cugraph_resource_handle_t* handle,
+                                              uint64_t seed,
+                                              cugraph_rng_state_t** state,
+                                              cugraph_error_t** error);
+
+/**
+ * @brief    Destroy a Random Number Generator State
+ *
+ * @param [in]  p    Pointer to the Random Number Generator State
+ */
+void cugraph_rng_state_free(cugraph_rng_state_t* p);
+
+#ifdef __cplusplus
+}
+#endif