Biased Random Walks and Node2Vec implementation

cpp/include/cugraph/algorithms.hpp

@@ -1579,11 +1579,11 @@ std::
 template <typename vertex_t, typename edge_t, typename weight_t, bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>, std::optional<rmm::device_uvector<weight_t>>>
 uniform_random_walks(raft::handle_t const& handle,
+                     raft::random::RngState& rng_state,
                      graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view,
                      std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
                      raft::device_span<vertex_t const> start_vertices,
-                     size_t max_length,
-                     uint64_t seed = std::numeric_limits<uint64_t>::max());
+                     size_t max_length);
 
 /**
  * @brief returns biased random walks from starting sources, where each path is of given
@@ -1623,11 +1623,11 @@ uniform_random_walks(raft::handle_t const& handle,
 template <typename vertex_t, typename edge_t, typename weight_t, bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>, std::optional<rmm::device_uvector<weight_t>>>
 biased_random_walks(raft::handle_t const& handle,
+                    raft::random::RngState& rng_state,
                     graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view,
                     edge_property_view_t<edge_t, weight_t const*> edge_weight_view,
                     raft::device_span<vertex_t const> start_vertices,
-                    size_t max_length,
-                    uint64_t seed = std::numeric_limits<uint64_t>::max());
+                    size_t max_length);
 
 /**
  * @brief returns biased random walks with node2vec biases from starting sources,
@@ -1670,13 +1670,13 @@ biased_random_walks(raft::handle_t const& handle,
 template <typename vertex_t, typename edge_t, typename weight_t, bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>, std::optional<rmm::device_uvector<weight_t>>>
 node2vec_random_walks(raft::handle_t const& handle,
+                      raft::random::RngState& rng_state,
                       graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view,
                       std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
                       raft::device_span<vertex_t const> start_vertices,
                       size_t max_length,
                       weight_t p,
-                      weight_t q,
-                      uint64_t seed = std::numeric_limits<uint64_t>::max());
+                      weight_t q);
 
 #ifndef NO_CUGRAPH_OPS
 /**

cpp/src/c_api/random_walks.cpp

@@ -16,6 +16,7 @@
 
 #include "c_api/abstract_functor.hpp"
 #include "c_api/graph.hpp"
+#include "c_api/random.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
 
@@ -153,10 +154,11 @@ namespace {
 
 struct uniform_random_walks_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
+  //  FIXME: rng_state_ should be passed as a parameter
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   size_t max_length_{0};
-  size_t seed_{0};
   cugraph::c_api::cugraph_random_walk_result_t* result_{nullptr};
 
   uniform_random_walks_functor(cugraph_resource_handle_t const* handle,
@@ -222,13 +224,17 @@ struct uniform_random_walks_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_last(),
         false);
 
+      //  FIXME: remove once rng_state passed as parameter
+      rng_state_ = reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(
+        new cugraph::c_api::cugraph_rng_state_t{raft::random::RngState{0}});
+
       auto [paths, weights] = cugraph::uniform_random_walks(
         handle_,
+        rng_state_->rng_state_,
         graph_view,
         (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
         raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-        max_length_,
-        seed_);
+        max_length_);
 
       //
       // Need to unrenumber the vertices in the resulting paths
@@ -255,11 +261,12 @@ struct uniform_random_walks_functor : public cugraph::c_api::abstract_functor {
 
 struct biased_random_walks_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
+  //  FIXME: rng_state_ should be passed as a parameter
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   size_t max_length_{0};
   cugraph::c_api::cugraph_random_walk_result_t* result_{nullptr};
-  uint64_t seed_{0};
 
   biased_random_walks_functor(cugraph_resource_handle_t const* handle,
                               cugraph_graph_t* graph,
@@ -326,13 +333,17 @@ struct biased_random_walks_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_last(),
         false);
 
+      //  FIXME: remove once rng_state passed as parameter
+      rng_state_ = reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(
+        new cugraph::c_api::cugraph_rng_state_t{raft::random::RngState{0}});
+
       auto [paths, weights] = cugraph::biased_random_walks(
         handle_,
+        rng_state_->rng_state_,
         graph_view,
         edge_weights->view(),
         raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-        max_length_,
-        seed_);
+        max_length_);
 
       //
       // Need to unrenumber the vertices in the resulting paths
@@ -354,12 +365,13 @@ struct biased_random_walks_functor : public cugraph::c_api::abstract_functor {
 
 struct node2vec_random_walks_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
+  //  FIXME: rng_state_ should be passed as a parameter
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   size_t max_length_{0};
   double p_{0};
   double q_{0};
-  uint64_t seed_{0};
   cugraph::c_api::cugraph_random_walk_result_t* result_{nullptr};
 
   node2vec_random_walks_functor(cugraph_resource_handle_t const* handle,
@@ -431,15 +443,19 @@ struct node2vec_random_walks_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_last(),
         false);
 
+      //  FIXME: remove once rng_state passed as parameter
+      rng_state_ = reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(
+        new cugraph::c_api::cugraph_rng_state_t{raft::random::RngState{0}});
+
       auto [paths, weights] = cugraph::node2vec_random_walks(
         handle_,
+        rng_state_->rng_state_,
         graph_view,
         (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
         raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
         max_length_,
         static_cast<weight_t>(p_),
-        static_cast<weight_t>(q_),
-        seed_);
+        static_cast<weight_t>(q_));
 
       // FIXME:  Need to fix invalid_vtx issue here.  We can't unrenumber max_vertex_id+1
       // properly...

cpp/src/prims/detail/sample_and_compute_local_nbr_indices.cuh

@@ -392,11 +392,11 @@ compute_unique_keys(raft::handle_t const& handle,
       cuda::proclaim_return_type<size_t>(
         [unique_key_first = get_dataframe_buffer_begin(aggregate_local_frontier_unique_keys) +
                             local_frontier_unique_key_displacements[i],
-         num_unique_keys = local_frontier_unique_key_sizes[i]] __device__(key_t key) {
+         unique_key_last = get_dataframe_buffer_begin(aggregate_local_frontier_unique_keys) +
+                           local_frontier_unique_key_displacements[i] +
+                           local_frontier_unique_key_sizes[i]] __device__(key_t key) {
           return static_cast<size_t>(thrust::distance(
-            unique_key_first,
-            thrust::lower_bound(
-              thrust::seq, unique_key_first, unique_key_first + num_unique_keys, key)));
+            unique_key_first, thrust::find(thrust::seq, unique_key_first, unique_key_last, key)));
         }));
   }
 
@@ -1759,8 +1759,7 @@ biased_sample_and_compute_local_nbr_indices(
   std::optional<rmm::device_uvector<size_t>> key_indices{std::nullopt};
   std::vector<size_t> local_frontier_sample_offsets{};
   if (with_replacement) {
-    // computet segmented inclusive sums (one segment per seed)
-
+    // compute segmented inclusive sums (one segment per seed)
     auto unique_key_first = thrust::make_transform_iterator(
       thrust::make_counting_iterator(size_t{0}),
       cuda::proclaim_return_type<size_t>(

cpp/src/prims/per_v_random_select_transform_outgoing_e.cuh

@@ -351,7 +351,7 @@ per_v_random_select_transform_e(raft::handle_t const& handle,
       uniform_sample_and_compute_local_nbr_indices(
         handle,
         graph_view,
-        (minor_comm_size > 1) ? get_dataframe_buffer_begin(*aggregate_local_frontier)
+        (minor_comm_size > 1) ? get_dataframe_buffer_cbegin(*aggregate_local_frontier)
                               : frontier.begin(),
         local_frontier_displacements,
         local_frontier_sizes,
@@ -363,7 +363,7 @@ per_v_random_select_transform_e(raft::handle_t const& handle,
       biased_sample_and_compute_local_nbr_indices(
         handle,
         graph_view,
-        (minor_comm_size > 1) ? get_dataframe_buffer_begin(*aggregate_local_frontier)
+        (minor_comm_size > 1) ? get_dataframe_buffer_cbegin(*aggregate_local_frontier)
                               : frontier.begin(),
         edge_bias_src_value_input,
         edge_bias_dst_value_input,
@@ -392,7 +392,7 @@ per_v_random_select_transform_e(raft::handle_t const& handle,
         graph_view.local_edge_partition_view(i));
 
     auto edge_partition_frontier_key_first =
-      ((minor_comm_size > 1) ? get_dataframe_buffer_begin(*aggregate_local_frontier)
+      ((minor_comm_size > 1) ? get_dataframe_buffer_cbegin(*aggregate_local_frontier)
                              : frontier.begin()) +
       local_frontier_displacements[i];
     auto edge_partition_sample_local_nbr_index_first =