Update API for node2vec and biased random walks

cpp/include/cugraph_c/sampling_algorithms.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -41,6 +41,7 @@ typedef struct {
  * @brief  Compute uniform random walks
  *
  * @param [in]  handle          Handle for accessing resources
+ * @param [in,out] rng_state    State of the random number generator, updated with each call
  * @param [in]  graph           Pointer to graph.  NOTE: Graph might be modified if the storage
  *                              needs to be transposed
  * @param [in]  start_vertices  Array of source vertices
@@ -52,6 +53,7 @@ typedef struct {
  */
 cugraph_error_code_t cugraph_uniform_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,
@@ -62,6 +64,7 @@ cugraph_error_code_t cugraph_uniform_random_walks(
  * @brief  Compute biased random walks
  *
  * @param [in]  handle          Handle for accessing resources
+ * @param [in,out] rng_state    State of the random number generator, updated with each call
  * @param [in]  graph           Pointer to graph.  NOTE: Graph might be modified if the storage
  *                              needs to be transposed
  * @param [in]  start_vertices  Array of source vertices
@@ -73,6 +76,7 @@ cugraph_error_code_t cugraph_uniform_random_walks(
  */
 cugraph_error_code_t cugraph_biased_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,
@@ -83,6 +87,7 @@ cugraph_error_code_t cugraph_biased_random_walks(
  * @brief  Compute random walks using the node2vec framework.
  *
  * @param [in]  handle          Handle for accessing resources
+ * @param [in,out] rng_state    State of the random number generator, updated with each call
  * @param [in]  graph           Pointer to graph.  NOTE: Graph might be modified if the storage
  *                              needs to be transposed
  * @param [in]  start_vertices  Array of source vertices
@@ -98,6 +103,7 @@ cugraph_error_code_t cugraph_biased_random_walks(
  */
 cugraph_error_code_t cugraph_node2vec_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,

cpp/src/c_api/random_walks.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -154,19 +154,20 @@ 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};
   cugraph::c_api::cugraph_random_walk_result_t* result_{nullptr};
 
   uniform_random_walks_functor(cugraph_resource_handle_t const* handle,
+                               cugraph_rng_state_t* rng_state,
                                cugraph_graph_t* graph,
                                cugraph_type_erased_device_array_view_t const* start_vertices,
                                size_t max_length)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
       start_vertices_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -224,10 +225,6 @@ 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_,
@@ -261,19 +258,20 @@ 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};
 
   biased_random_walks_functor(cugraph_resource_handle_t const* handle,
+                              cugraph_rng_state_t* rng_state,
                               cugraph_graph_t* graph,
                               cugraph_type_erased_device_array_view_t const* start_vertices,
                               size_t max_length)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
       start_vertices_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -293,8 +291,6 @@ struct biased_random_walks_functor : public cugraph::c_api::abstract_functor {
     // FIXME: Think about how to handle SG vice MG
     if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
       unsupported();
-    } else if constexpr (multi_gpu) {
-      unsupported();
     } else {
       // random walks expects store_transposed == false
       if constexpr (store_transposed) {
@@ -333,10 +329,6 @@ 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_,
@@ -348,8 +340,13 @@ struct biased_random_walks_functor : public cugraph::c_api::abstract_functor {
       //
       // Need to unrenumber the vertices in the resulting paths
       //
-      cugraph::unrenumber_local_int_vertices<vertex_t>(
-        handle_, paths.data(), paths.size(), number_map->data(), 0, paths.size() - 1, false);
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(
+        handle_,
+        paths.data(),
+        paths.size(),
+        number_map->data(),
+        graph_view.vertex_partition_range_lasts(),
+        false);
 
       result_ = new cugraph::c_api::cugraph_random_walk_result_t{
         false,
@@ -365,7 +362,6 @@ 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};
@@ -375,13 +371,15 @@ struct node2vec_random_walks_functor : public cugraph::c_api::abstract_functor {
   cugraph::c_api::cugraph_random_walk_result_t* result_{nullptr};
 
   node2vec_random_walks_functor(cugraph_resource_handle_t const* handle,
+                                cugraph_rng_state_t* rng_state,
                                 cugraph_graph_t* graph,
                                 cugraph_type_erased_device_array_view_t const* start_vertices,
                                 size_t max_length,
                                 double p,
                                 double q)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
       start_vertices_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -403,8 +401,6 @@ struct node2vec_random_walks_functor : public cugraph::c_api::abstract_functor {
     // FIXME: Think about how to handle SG vice MG
     if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
       unsupported();
-    } else if constexpr (multi_gpu) {
-      unsupported();
     } else {
       // random walks expects store_transposed == false
       if constexpr (store_transposed) {
@@ -443,10 +439,6 @@ 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_,
@@ -464,8 +456,13 @@ struct node2vec_random_walks_functor : public cugraph::c_api::abstract_functor {
       //
       // Need to unrenumber the vertices in the resulting paths
       //
-      cugraph::unrenumber_local_int_vertices<vertex_t>(
-        handle_, paths.data(), paths.size(), number_map->data(), 0, paths.size(), false);
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(
+        handle_,
+        paths.data(),
+        paths.size(),
+        number_map->data(),
+        graph_view.vertex_partition_range_lasts(),
+        false);
 
       result_ = new cugraph::c_api::cugraph_random_walk_result_t{
         false,
@@ -546,6 +543,7 @@ void cugraph_random_walk_result_free(cugraph_random_walk_result_t* result)
 
 cugraph_error_code_t cugraph_uniform_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,
@@ -560,13 +558,14 @@ cugraph_error_code_t cugraph_uniform_random_walks(
                "vertex type of graph and start_vertices must match",
                *error);
 
-  uniform_random_walks_functor functor(handle, graph, start_vertices, max_length);
+  uniform_random_walks_functor functor(handle, rng_state, graph, start_vertices, max_length);
 
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
 cugraph_error_code_t cugraph_biased_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,
@@ -581,13 +580,14 @@ cugraph_error_code_t cugraph_biased_random_walks(
                "vertex type of graph and start_vertices must match",
                *error);
 
-  biased_random_walks_functor functor(handle, graph, start_vertices, max_length);
+  biased_random_walks_functor functor(handle, rng_state, graph, start_vertices, max_length);
 
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
 cugraph_error_code_t cugraph_node2vec_random_walks(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   size_t max_length,
@@ -604,7 +604,7 @@ cugraph_error_code_t cugraph_node2vec_random_walks(
                "vertex type of graph and start_vertices must match",
                *error);
 
-  node2vec_random_walks_functor functor(handle, graph, start_vertices, max_length, p, q);
+  node2vec_random_walks_functor functor(handle, rng_state, graph, start_vertices, max_length, p, q);
 
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }