From c568fd0b440f8afaf805ceaa492a59d5a2917be7 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 01:00:26 +0100
Subject: [PATCH 01/13] Update select_random_vertices to select as many random
 vetices local vertex partition range size on each GPU

---
 cpp/include/cugraph/graph_functions.hpp       |  6 ++-
 .../structure/select_random_vertices_impl.hpp | 45 ++++++++++++++++++-
 .../structure/select_random_vertices_mg.cu    |  6 +++
 .../structure/select_random_vertices_sg.cu    |  6 +++
 4 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/cpp/include/cugraph/graph_functions.hpp b/cpp/include/cugraph/graph_functions.hpp
index 6a75a420bf8..453f8b0bd33 100644
--- a/cpp/include/cugraph/graph_functions.hpp
+++ b/cpp/include/cugraph/graph_functions.hpp
@@ -903,6 +903,9 @@ weight_t compute_total_edge_weight(
  * @param  select_count The number of vertices to select from the graph
  * @param  with_replacement If true, select with replacement, if false select without replacement
  * @param  sort_vertices If true, return the sorted vertices (in the ascending order).
+ * @param  shuffle_int_to_local If true and If @p given_set is not specified
+ * then shuffle internal (i.e. renumbered) vertices to their local GPUs based on vertex
+ * partitioning, otherwise shuffle as many vertices as local vertex partition size to each GPU.
  * @return Device vector of selected vertices.
  */
 template <typename vertex_t, typename edge_t, bool store_transposed, bool multi_gpu>
@@ -914,7 +917,8 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool do_expensive_check = false);
+  bool shuffle_int_to_local = true,
+  bool do_expensive_check   = false);
 
 /**
  * @brief renumber sampling output
diff --git a/cpp/src/structure/select_random_vertices_impl.hpp b/cpp/src/structure/select_random_vertices_impl.hpp
index b6a0c364848..c142bfab101 100644
--- a/cpp/src/structure/select_random_vertices_impl.hpp
+++ b/cpp/src/structure/select_random_vertices_impl.hpp
@@ -52,6 +52,7 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check)
 {
   size_t num_of_elements_in_given_set{0};
@@ -232,8 +233,48 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   }
 
   if constexpr (multi_gpu) {
-    mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
-      handle, std::move(mg_sample_buffer), partition_range_lasts);
+    if (given_set) {
+      mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        handle, std::move(mg_sample_buffer), partition_range_lasts);
+    } else {
+      if (shuffle_int_to_local) {
+        mg_sample_buffer =
+          cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+            handle, std::move(mg_sample_buffer), partition_range_lasts);
+
+      } else {
+        // shuffle as many vertices as local vertex partition size to each GPU.
+
+        auto& comm           = handle.get_comms();
+        auto const comm_size = comm.get_size();
+        auto const comm_rank = comm.get_rank();
+        std::vector<size_t> tx_value_counts(comm_size, 0);
+        auto sample_buffer_sizes = cugraph::host_scalar_allgather(
+          handle.get_comms(), mg_sample_buffer.size(), handle.get_stream());
+
+        auto expected_sample_buffer_sizes = cugraph::host_scalar_allgather(
+          handle.get_comms(), graph_view.local_vertex_partition_range_size(), handle.get_stream());
+
+        std::vector<size_t> nr_smaples(comm_size, 0);
+
+        // find out how many elements current GPU needs to send to other GPUs
+        for (int i = 0; i < comm_size; i++) {
+          size_t nr_samples_ith_gpu = sample_buffer_sizes[i];
+          for (int j = 0; nr_samples_ith_gpu > 0 && j < comm_size; j++) {
+            if (expected_sample_buffer_sizes[j] > static_cast<vertex_t>(nr_smaples[j])) {
+              size_t delta =
+                std::min(nr_samples_ith_gpu, expected_sample_buffer_sizes[j] - nr_smaples[j]);
+              if (comm_rank == i) { tx_value_counts[j] = delta; }
+              nr_smaples[j] += delta;
+              nr_samples_ith_gpu -= delta;
+            }
+          }
+        }
+
+        std::tie(mg_sample_buffer, std::ignore) = cugraph::shuffle_values(
+          handle.get_comms(), mg_sample_buffer.begin(), tx_value_counts, handle.get_stream());
+      }
+    }
   }
 
   if (given_set) {
diff --git a/cpp/src/structure/select_random_vertices_mg.cu b/cpp/src/structure/select_random_vertices_mg.cu
index 595da12f678..1c80f441d76 100644
--- a/cpp/src/structure/select_random_vertices_mg.cu
+++ b/cpp/src/structure/select_random_vertices_mg.cu
@@ -26,6 +26,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -36,6 +37,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -46,6 +48,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -56,6 +59,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -66,6 +70,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -76,6 +81,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/structure/select_random_vertices_sg.cu b/cpp/src/structure/select_random_vertices_sg.cu
index 1ca1878c9db..1cf6002d729 100644
--- a/cpp/src/structure/select_random_vertices_sg.cu
+++ b/cpp/src/structure/select_random_vertices_sg.cu
@@ -26,6 +26,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -36,6 +37,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -46,6 +48,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -56,6 +59,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -66,6 +70,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -76,6 +81,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
+  bool shuffle_int_to_local,
   bool do_expensive_check);
 
 }  // namespace cugraph

From b2dd3072faaf7408c2027304b260bdf4b5bb08c6 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 01:47:15 +0100
Subject: [PATCH 02/13] Update select_random_vertices to select as many random
 vetices local vertex partition range size on each GPU

---
 cpp/src/structure/select_random_vertices_impl.hpp  | 14 +++++++-------
 .../structure/mg_select_random_vertices_test.cpp   |  2 +-
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/cpp/src/structure/select_random_vertices_impl.hpp b/cpp/src/structure/select_random_vertices_impl.hpp
index c142bfab101..5220b890c98 100644
--- a/cpp/src/structure/select_random_vertices_impl.hpp
+++ b/cpp/src/structure/select_random_vertices_impl.hpp
@@ -236,15 +236,11 @@ rmm::device_uvector<vertex_t> select_random_vertices(
     if (given_set) {
       mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
         handle, std::move(mg_sample_buffer), partition_range_lasts);
-    } else {
-      if (shuffle_int_to_local) {
-        mg_sample_buffer =
-          cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
-            handle, std::move(mg_sample_buffer), partition_range_lasts);
 
-      } else {
+    } else {
+      if (!shuffle_int_to_local &&
+          select_count == static_cast<size_t>(graph_view.number_of_vertices())) {
         // shuffle as many vertices as local vertex partition size to each GPU.
-
         auto& comm           = handle.get_comms();
         auto const comm_size = comm.get_size();
         auto const comm_rank = comm.get_rank();
@@ -273,6 +269,10 @@ rmm::device_uvector<vertex_t> select_random_vertices(
 
         std::tie(mg_sample_buffer, std::ignore) = cugraph::shuffle_values(
           handle.get_comms(), mg_sample_buffer.begin(), tx_value_counts, handle.get_stream());
+      } else {
+        mg_sample_buffer =
+          cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+            handle, std::move(mg_sample_buffer), partition_range_lasts);
       }
     }
   }
diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index e49e1ebcb99..d000a6107fb 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -148,7 +148,7 @@ class Tests_MGSelectRandomVertices
     //
 
     for (int idx = 0; idx < with_replacement_flags.size(); idx++) {
-      bool with_replacement   = false;
+      bool with_replacement   = with_replacement_flags[idx];
       auto d_sampled_vertices = cugraph::select_random_vertices(
         *handle_,
         mg_graph_view,

From 77b06663de0f1351fb28ccdcea2f0af03cdfc592 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 03:00:49 +0100
Subject: [PATCH 03/13] Update select_random_vertices to select as many random
 vetices local vertex partition range size on each GPU

---
 .../structure/select_random_vertices_impl.hpp | 12 ++--
 .../mg_select_random_vertices_test.cpp        | 68 ++++++++++++++++++-
 2 files changed, 71 insertions(+), 9 deletions(-)

diff --git a/cpp/src/structure/select_random_vertices_impl.hpp b/cpp/src/structure/select_random_vertices_impl.hpp
index 5220b890c98..989ce74585d 100644
--- a/cpp/src/structure/select_random_vertices_impl.hpp
+++ b/cpp/src/structure/select_random_vertices_impl.hpp
@@ -236,9 +236,8 @@ rmm::device_uvector<vertex_t> select_random_vertices(
     if (given_set) {
       mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
         handle, std::move(mg_sample_buffer), partition_range_lasts);
-
-    } else {
-      if (!shuffle_int_to_local &&
+    } else if (!shuffle_int_to_local) {
+      if (!with_replacement &&
           select_count == static_cast<size_t>(graph_view.number_of_vertices())) {
         // shuffle as many vertices as local vertex partition size to each GPU.
         auto& comm           = handle.get_comms();
@@ -269,11 +268,10 @@ rmm::device_uvector<vertex_t> select_random_vertices(
 
         std::tie(mg_sample_buffer, std::ignore) = cugraph::shuffle_values(
           handle.get_comms(), mg_sample_buffer.begin(), tx_value_counts, handle.get_stream());
-      } else {
-        mg_sample_buffer =
-          cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
-            handle, std::move(mg_sample_buffer), partition_range_lasts);
       }
+    } else {
+      mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        handle, std::move(mg_sample_buffer), partition_range_lasts);
     }
   }
 
diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index d000a6107fb..20a5643a825 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -183,6 +183,70 @@ class Tests_MGSelectRandomVertices
                       });
       }
     }
+
+    std::vector<bool> sort_vertices_flags        = {true, false};
+    std::vector<bool> shuffle_int_to_local_flags = {true, false};
+    std::vector<vertex_t> select_counts          = {mg_graph_view.number_of_vertices(),
+                                                    mg_graph_view.number_of_vertices() / 4};
+
+    for (int i = 0; i < with_replacement_flags.size(); i++) {
+      for (int j = 0; j < sort_vertices_flags.size(); j++) {
+        for (int k = 0; k < shuffle_int_to_local_flags.size(); k++) {
+          for (int l = 0; l < select_counts.size(); l++) {
+            bool with_replacement     = with_replacement_flags[i];
+            bool sort_vertices        = sort_vertices_flags[j];
+            bool shuffle_int_to_local = shuffle_int_to_local_flags[k];
+            auto select_count         = static_cast<size_t>(select_counts[l]);
+
+            auto d_sampled_vertices = cugraph::select_random_vertices(
+              *handle_,
+              mg_graph_view,
+              std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+              rng_state,
+              select_count,
+              with_replacement,
+              sort_vertices,
+              shuffle_int_to_local);
+
+            RAFT_CUDA_TRY(cudaDeviceSynchronize());
+
+            auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+
+            if (select_random_vertices_usecase.check_correctness) {
+              if (!with_replacement) {
+                std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+
+                auto nr_duplicates =
+                  std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                                h_sampled_vertices.end());
+
+                ASSERT_EQ(nr_duplicates, 0);
+              }
+
+              if (shuffle_int_to_local) {
+                auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
+                auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
+
+                std::for_each(h_sampled_vertices.begin(),
+                              h_sampled_vertices.end(),
+                              [vertex_first, vertex_last](vertex_t v) {
+                                ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
+                              });
+              } else {
+                if (!with_replacement &&
+                    select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
+                  ASSERT_EQ(h_sampled_vertices.size(),
+                            mg_graph_view.local_vertex_partition_range_size());
+                }
+
+                std::cout << "silv: " << shuffle_int_to_local << " sc: " << select_count
+                          << "  got: " << h_sampled_vertices.size() << std::endl;
+              }
+            }
+          }
+        }
+      }
+    }
   }
 
  private:
@@ -242,8 +306,8 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_MGSelectRandomVertices_Rmat,
   ::testing::Combine(
-    ::testing::Values(SelectRandomVertices_Usecase{500, false},
-                      SelectRandomVertices_Usecase{500, false}),
+    ::testing::Values(SelectRandomVertices_Usecase{500, true},
+                      SelectRandomVertices_Usecase{500, true}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()

From aad93722a63227673e291bfafe2c2264a845983c Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 03:30:34 +0100
Subject: [PATCH 04/13] Update select_random_vertices to select as many random
 vetices local vertex partition range size on each GPU

---
 cpp/include/cugraph/graph_functions.hpp       |  9 +++---
 .../structure/select_random_vertices_impl.hpp | 10 ++-----
 .../structure/select_random_vertices_mg.cu    | 12 ++++----
 .../structure/select_random_vertices_sg.cu    | 12 ++++----
 .../mg_select_random_vertices_test.cpp        | 28 ++++++++++---------
 5 files changed, 34 insertions(+), 37 deletions(-)

diff --git a/cpp/include/cugraph/graph_functions.hpp b/cpp/include/cugraph/graph_functions.hpp
index 453f8b0bd33..f64b6e19f59 100644
--- a/cpp/include/cugraph/graph_functions.hpp
+++ b/cpp/include/cugraph/graph_functions.hpp
@@ -903,9 +903,8 @@ weight_t compute_total_edge_weight(
  * @param  select_count The number of vertices to select from the graph
  * @param  with_replacement If true, select with replacement, if false select without replacement
  * @param  sort_vertices If true, return the sorted vertices (in the ascending order).
- * @param  shuffle_int_to_local If true and If @p given_set is not specified
- * then shuffle internal (i.e. renumbered) vertices to their local GPUs based on vertex
- * partitioning, otherwise shuffle as many vertices as local vertex partition size to each GPU.
+ * @param  shuffle_random_vertices_using_vertex_partition If true, shuffle random vertices to GPUs
+ * based on vertex partitioning.
  * @return Device vector of selected vertices.
  */
 template <typename vertex_t, typename edge_t, bool store_transposed, bool multi_gpu>
@@ -917,8 +916,8 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local = true,
-  bool do_expensive_check   = false);
+  bool shuffle_random_vertices_using_vertex_partition = true,
+  bool do_expensive_check                             = false);
 
 /**
  * @brief renumber sampling output
diff --git a/cpp/src/structure/select_random_vertices_impl.hpp b/cpp/src/structure/select_random_vertices_impl.hpp
index 989ce74585d..ba9bc640646 100644
--- a/cpp/src/structure/select_random_vertices_impl.hpp
+++ b/cpp/src/structure/select_random_vertices_impl.hpp
@@ -52,7 +52,7 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check)
 {
   size_t num_of_elements_in_given_set{0};
@@ -233,12 +233,8 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   }
 
   if constexpr (multi_gpu) {
-    if (given_set) {
-      mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
-        handle, std::move(mg_sample_buffer), partition_range_lasts);
-    } else if (!shuffle_int_to_local) {
-      if (!with_replacement &&
-          select_count == static_cast<size_t>(graph_view.number_of_vertices())) {
+    if (!shuffle_random_vertices_using_vertex_partition) {
+      if (select_count == static_cast<size_t>(graph_view.number_of_vertices())) {
         // shuffle as many vertices as local vertex partition size to each GPU.
         auto& comm           = handle.get_comms();
         auto const comm_size = comm.get_size();
diff --git a/cpp/src/structure/select_random_vertices_mg.cu b/cpp/src/structure/select_random_vertices_mg.cu
index 1c80f441d76..4dc364c86b8 100644
--- a/cpp/src/structure/select_random_vertices_mg.cu
+++ b/cpp/src/structure/select_random_vertices_mg.cu
@@ -26,7 +26,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -37,7 +37,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -48,7 +48,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -59,7 +59,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -70,7 +70,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -81,7 +81,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/structure/select_random_vertices_sg.cu b/cpp/src/structure/select_random_vertices_sg.cu
index 1cf6002d729..e628d421a90 100644
--- a/cpp/src/structure/select_random_vertices_sg.cu
+++ b/cpp/src/structure/select_random_vertices_sg.cu
@@ -26,7 +26,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -37,7 +37,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -48,7 +48,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -59,7 +59,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -70,7 +70,7 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -81,7 +81,7 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_int_to_local,
+  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 20a5643a825..41773763948 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -184,19 +184,20 @@ class Tests_MGSelectRandomVertices
       }
     }
 
-    std::vector<bool> sort_vertices_flags        = {true, false};
-    std::vector<bool> shuffle_int_to_local_flags = {true, false};
-    std::vector<vertex_t> select_counts          = {mg_graph_view.number_of_vertices(),
-                                                    mg_graph_view.number_of_vertices() / 4};
+    std::vector<bool> sort_vertices_flags                                  = {true, false};
+    std::vector<bool> shuffle_random_vertices_using_vertex_partition_flags = {true, false};
+    std::vector<vertex_t> select_counts = {mg_graph_view.number_of_vertices(),
+                                           mg_graph_view.number_of_vertices() / 4};
 
     for (int i = 0; i < with_replacement_flags.size(); i++) {
       for (int j = 0; j < sort_vertices_flags.size(); j++) {
-        for (int k = 0; k < shuffle_int_to_local_flags.size(); k++) {
+        for (int k = 0; k < shuffle_random_vertices_using_vertex_partition_flags.size(); k++) {
           for (int l = 0; l < select_counts.size(); l++) {
-            bool with_replacement     = with_replacement_flags[i];
-            bool sort_vertices        = sort_vertices_flags[j];
-            bool shuffle_int_to_local = shuffle_int_to_local_flags[k];
-            auto select_count         = static_cast<size_t>(select_counts[l]);
+            bool with_replacement = with_replacement_flags[i];
+            bool sort_vertices    = sort_vertices_flags[j];
+            bool shuffle_random_vertices_using_vertex_partition =
+              shuffle_random_vertices_using_vertex_partition_flags[k];
+            auto select_count = static_cast<size_t>(select_counts[l]);
 
             auto d_sampled_vertices = cugraph::select_random_vertices(
               *handle_,
@@ -206,7 +207,7 @@ class Tests_MGSelectRandomVertices
               select_count,
               with_replacement,
               sort_vertices,
-              shuffle_int_to_local);
+              shuffle_random_vertices_using_vertex_partition);
 
             RAFT_CUDA_TRY(cudaDeviceSynchronize());
 
@@ -223,7 +224,7 @@ class Tests_MGSelectRandomVertices
                 ASSERT_EQ(nr_duplicates, 0);
               }
 
-              if (shuffle_int_to_local) {
+              if (shuffle_random_vertices_using_vertex_partition) {
                 auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
                 auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
 
@@ -239,8 +240,9 @@ class Tests_MGSelectRandomVertices
                             mg_graph_view.local_vertex_partition_range_size());
                 }
 
-                std::cout << "silv: " << shuffle_int_to_local << " sc: " << select_count
-                          << "  got: " << h_sampled_vertices.size() << std::endl;
+                std::cout << "silv: " << shuffle_random_vertices_using_vertex_partition
+                          << " sc: " << select_count << "  got: " << h_sampled_vertices.size()
+                          << std::endl;
               }
             }
           }

From 028610d103997a16ef4137d868579184b075d422 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 03:31:57 +0100
Subject: [PATCH 05/13] Update comments and fix logic

---
 cpp/tests/structure/mg_select_random_vertices_test.cpp | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 41773763948..db28c5b8e4d 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -239,10 +239,6 @@ class Tests_MGSelectRandomVertices
                   ASSERT_EQ(h_sampled_vertices.size(),
                             mg_graph_view.local_vertex_partition_range_size());
                 }
-
-                std::cout << "silv: " << shuffle_random_vertices_using_vertex_partition
-                          << " sc: " << select_count << "  got: " << h_sampled_vertices.size()
-                          << std::endl;
               }
             }
           }

From 10d859993ee68736c73bd4f40ff10fb42ea2301e Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 03:44:10 +0100
Subject: [PATCH 06/13] Update test

---
 cpp/tests/structure/mg_select_random_vertices_test.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index db28c5b8e4d..00c37e4ead1 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -234,8 +234,7 @@ class Tests_MGSelectRandomVertices
                                 ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
                               });
               } else {
-                if (!with_replacement &&
-                    select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
+                if (select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
                   ASSERT_EQ(h_sampled_vertices.size(),
                             mg_graph_view.local_vertex_partition_range_size());
                 }

From 70da14498db04ba65dec5b998c2ac744549d7e38 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 04:28:31 +0100
Subject: [PATCH 07/13] Update test

---
 .../mg_select_random_vertices_test.cpp        | 98 +++++++++++--------
 1 file changed, 57 insertions(+), 41 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 00c37e4ead1..66ba5ab1e58 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -79,6 +79,8 @@ class Tests_MGSelectRandomVertices
     //
 
     std::vector<bool> with_replacement_flags = {true, false};
+    std::vector<bool> shuffle_flags          = {true, false};
+
     {
       // Generate distributed vertex set to sample from
       std::srand((unsigned)std::chrono::duration_cast<std::chrono::milliseconds>(
@@ -146,58 +148,72 @@ class Tests_MGSelectRandomVertices
     //
     // Test sampling from [0, V)
     //
+    std::vector<size_t> select_counts = {select_random_vertices_usecase.select_count,
+                                         static_cast<size_t>(mg_graph_view.number_of_vertices())};
 
     for (int idx = 0; idx < with_replacement_flags.size(); idx++) {
-      bool with_replacement   = with_replacement_flags[idx];
-      auto d_sampled_vertices = cugraph::select_random_vertices(
-        *handle_,
-        mg_graph_view,
-        std::optional<raft::device_span<vertex_t const>>{std::nullopt},
-        rng_state,
-        select_random_vertices_usecase.select_count,
-        with_replacement,
-        true);
-
-      RAFT_CUDA_TRY(cudaDeviceSynchronize());
-
-      auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+      for (int k = 0; k < shuffle_flags.size(); k++) {
+        for (int l = 0; l < select_counts.size(); l++) {
+          bool with_replacement               = with_replacement_flags[idx];
+          bool shuffle_using_vertex_partition = shuffle_flags[k];
+          auto select_count                   = select_counts[l];
+
+          auto d_sampled_vertices = cugraph::select_random_vertices(
+            *handle_,
+            mg_graph_view,
+            std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+            rng_state,
+            select_count,
+            with_replacement,
+            true,
+            shuffle_using_vertex_partition);
+
+          RAFT_CUDA_TRY(cudaDeviceSynchronize());
+
+          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+
+          if (select_random_vertices_usecase.check_correctness) {
+            if (!with_replacement) {
+              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+
+              auto nr_duplicates =
+                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                              h_sampled_vertices.end());
+
+              ASSERT_EQ(nr_duplicates, 0);
+            }
 
-      if (select_random_vertices_usecase.check_correctness) {
-        if (!with_replacement) {
-          std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+            if (shuffle_using_vertex_partition) {
+              auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
+              auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
 
-          auto nr_duplicates =
-            std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                          h_sampled_vertices.end());
+              std::for_each(h_sampled_vertices.begin(),
+                            h_sampled_vertices.end(),
+                            [vertex_first, vertex_last](vertex_t v) {
+                              ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
+                            });
 
-          ASSERT_EQ(nr_duplicates, 0);
+            } else {
+              if (select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
+                ASSERT_EQ(h_sampled_vertices.size(),
+                          mg_graph_view.local_vertex_partition_range_size());
+              }
+            }
+          }
         }
-
-        auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
-        auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
-
-        std::for_each(h_sampled_vertices.begin(),
-                      h_sampled_vertices.end(),
-                      [vertex_first, vertex_last](vertex_t v) {
-                        ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
-                      });
       }
     }
 
-    std::vector<bool> sort_vertices_flags                                  = {true, false};
-    std::vector<bool> shuffle_random_vertices_using_vertex_partition_flags = {true, false};
-    std::vector<vertex_t> select_counts = {mg_graph_view.number_of_vertices(),
-                                           mg_graph_view.number_of_vertices() / 4};
+    std::vector<bool> sort_vertices_flags = {true, false};
 
     for (int i = 0; i < with_replacement_flags.size(); i++) {
       for (int j = 0; j < sort_vertices_flags.size(); j++) {
-        for (int k = 0; k < shuffle_random_vertices_using_vertex_partition_flags.size(); k++) {
+        for (int k = 0; k < shuffle_flags.size(); k++) {
           for (int l = 0; l < select_counts.size(); l++) {
-            bool with_replacement = with_replacement_flags[i];
-            bool sort_vertices    = sort_vertices_flags[j];
-            bool shuffle_random_vertices_using_vertex_partition =
-              shuffle_random_vertices_using_vertex_partition_flags[k];
-            auto select_count = static_cast<size_t>(select_counts[l]);
+            bool with_replacement               = with_replacement_flags[i];
+            bool sort_vertices                  = sort_vertices_flags[j];
+            bool shuffle_using_vertex_partition = shuffle_flags[k];
+            auto select_count                   = static_cast<size_t>(select_counts[l]);
 
             auto d_sampled_vertices = cugraph::select_random_vertices(
               *handle_,
@@ -207,7 +223,7 @@ class Tests_MGSelectRandomVertices
               select_count,
               with_replacement,
               sort_vertices,
-              shuffle_random_vertices_using_vertex_partition);
+              shuffle_using_vertex_partition);
 
             RAFT_CUDA_TRY(cudaDeviceSynchronize());
 
@@ -224,7 +240,7 @@ class Tests_MGSelectRandomVertices
                 ASSERT_EQ(nr_duplicates, 0);
               }
 
-              if (shuffle_random_vertices_using_vertex_partition) {
+              if (shuffle_using_vertex_partition) {
                 auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
                 auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
 

From 47998030cfc7cf9f5dbd00113b8668d2a6e75e26 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 5 Dec 2023 04:30:07 +0100
Subject: [PATCH 08/13] Update test

---
 .../mg_select_random_vertices_test.cpp        | 53 -------------------
 1 file changed, 53 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 66ba5ab1e58..2689b7f3cd6 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -151,59 +151,6 @@ class Tests_MGSelectRandomVertices
     std::vector<size_t> select_counts = {select_random_vertices_usecase.select_count,
                                          static_cast<size_t>(mg_graph_view.number_of_vertices())};
 
-    for (int idx = 0; idx < with_replacement_flags.size(); idx++) {
-      for (int k = 0; k < shuffle_flags.size(); k++) {
-        for (int l = 0; l < select_counts.size(); l++) {
-          bool with_replacement               = with_replacement_flags[idx];
-          bool shuffle_using_vertex_partition = shuffle_flags[k];
-          auto select_count                   = select_counts[l];
-
-          auto d_sampled_vertices = cugraph::select_random_vertices(
-            *handle_,
-            mg_graph_view,
-            std::optional<raft::device_span<vertex_t const>>{std::nullopt},
-            rng_state,
-            select_count,
-            with_replacement,
-            true,
-            shuffle_using_vertex_partition);
-
-          RAFT_CUDA_TRY(cudaDeviceSynchronize());
-
-          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
-
-          if (select_random_vertices_usecase.check_correctness) {
-            if (!with_replacement) {
-              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
-
-              auto nr_duplicates =
-                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                              h_sampled_vertices.end());
-
-              ASSERT_EQ(nr_duplicates, 0);
-            }
-
-            if (shuffle_using_vertex_partition) {
-              auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
-              auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
-
-              std::for_each(h_sampled_vertices.begin(),
-                            h_sampled_vertices.end(),
-                            [vertex_first, vertex_last](vertex_t v) {
-                              ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
-                            });
-
-            } else {
-              if (select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
-                ASSERT_EQ(h_sampled_vertices.size(),
-                          mg_graph_view.local_vertex_partition_range_size());
-              }
-            }
-          }
-        }
-      }
-    }
-
     std::vector<bool> sort_vertices_flags = {true, false};
 
     for (int i = 0; i < with_replacement_flags.size(); i++) {

From 3fb397359242b889564412c2a8844c3aa612b302 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 12 Dec 2023 02:00:17 +0100
Subject: [PATCH 09/13] Undo changes

---
 cpp/include/cugraph/graph_functions.hpp       |  5 +-
 .../structure/select_random_vertices_impl.hpp | 39 +---------
 .../structure/select_random_vertices_mg.cu    |  6 --
 .../structure/select_random_vertices_sg.cu    |  6 --
 .../mg_select_random_vertices_test.cpp        | 76 +++++++------------
 5 files changed, 31 insertions(+), 101 deletions(-)

diff --git a/cpp/include/cugraph/graph_functions.hpp b/cpp/include/cugraph/graph_functions.hpp
index f64b6e19f59..6a75a420bf8 100644
--- a/cpp/include/cugraph/graph_functions.hpp
+++ b/cpp/include/cugraph/graph_functions.hpp
@@ -903,8 +903,6 @@ weight_t compute_total_edge_weight(
  * @param  select_count The number of vertices to select from the graph
  * @param  with_replacement If true, select with replacement, if false select without replacement
  * @param  sort_vertices If true, return the sorted vertices (in the ascending order).
- * @param  shuffle_random_vertices_using_vertex_partition If true, shuffle random vertices to GPUs
- * based on vertex partitioning.
  * @return Device vector of selected vertices.
  */
 template <typename vertex_t, typename edge_t, bool store_transposed, bool multi_gpu>
@@ -916,8 +914,7 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition = true,
-  bool do_expensive_check                             = false);
+  bool do_expensive_check = false);
 
 /**
  * @brief renumber sampling output
diff --git a/cpp/src/structure/select_random_vertices_impl.hpp b/cpp/src/structure/select_random_vertices_impl.hpp
index ba9bc640646..b6a0c364848 100644
--- a/cpp/src/structure/select_random_vertices_impl.hpp
+++ b/cpp/src/structure/select_random_vertices_impl.hpp
@@ -52,7 +52,6 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check)
 {
   size_t num_of_elements_in_given_set{0};
@@ -233,42 +232,8 @@ rmm::device_uvector<vertex_t> select_random_vertices(
   }
 
   if constexpr (multi_gpu) {
-    if (!shuffle_random_vertices_using_vertex_partition) {
-      if (select_count == static_cast<size_t>(graph_view.number_of_vertices())) {
-        // shuffle as many vertices as local vertex partition size to each GPU.
-        auto& comm           = handle.get_comms();
-        auto const comm_size = comm.get_size();
-        auto const comm_rank = comm.get_rank();
-        std::vector<size_t> tx_value_counts(comm_size, 0);
-        auto sample_buffer_sizes = cugraph::host_scalar_allgather(
-          handle.get_comms(), mg_sample_buffer.size(), handle.get_stream());
-
-        auto expected_sample_buffer_sizes = cugraph::host_scalar_allgather(
-          handle.get_comms(), graph_view.local_vertex_partition_range_size(), handle.get_stream());
-
-        std::vector<size_t> nr_smaples(comm_size, 0);
-
-        // find out how many elements current GPU needs to send to other GPUs
-        for (int i = 0; i < comm_size; i++) {
-          size_t nr_samples_ith_gpu = sample_buffer_sizes[i];
-          for (int j = 0; nr_samples_ith_gpu > 0 && j < comm_size; j++) {
-            if (expected_sample_buffer_sizes[j] > static_cast<vertex_t>(nr_smaples[j])) {
-              size_t delta =
-                std::min(nr_samples_ith_gpu, expected_sample_buffer_sizes[j] - nr_smaples[j]);
-              if (comm_rank == i) { tx_value_counts[j] = delta; }
-              nr_smaples[j] += delta;
-              nr_samples_ith_gpu -= delta;
-            }
-          }
-        }
-
-        std::tie(mg_sample_buffer, std::ignore) = cugraph::shuffle_values(
-          handle.get_comms(), mg_sample_buffer.begin(), tx_value_counts, handle.get_stream());
-      }
-    } else {
-      mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
-        handle, std::move(mg_sample_buffer), partition_range_lasts);
-    }
+    mg_sample_buffer = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+      handle, std::move(mg_sample_buffer), partition_range_lasts);
   }
 
   if (given_set) {
diff --git a/cpp/src/structure/select_random_vertices_mg.cu b/cpp/src/structure/select_random_vertices_mg.cu
index 4dc364c86b8..595da12f678 100644
--- a/cpp/src/structure/select_random_vertices_mg.cu
+++ b/cpp/src/structure/select_random_vertices_mg.cu
@@ -26,7 +26,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -37,7 +36,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -48,7 +46,6 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -59,7 +56,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -70,7 +66,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -81,7 +76,6 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/structure/select_random_vertices_sg.cu b/cpp/src/structure/select_random_vertices_sg.cu
index e628d421a90..1ca1878c9db 100644
--- a/cpp/src/structure/select_random_vertices_sg.cu
+++ b/cpp/src/structure/select_random_vertices_sg.cu
@@ -26,7 +26,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -37,7 +36,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -48,7 +46,6 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -59,7 +56,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int32_t> select_random_vertices(
@@ -70,7 +66,6 @@ template rmm::device_uvector<int32_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> select_random_vertices(
@@ -81,7 +76,6 @@ template rmm::device_uvector<int64_t> select_random_vertices(
   size_t select_count,
   bool with_replacement,
   bool sort_vertices,
-  bool shuffle_random_vertices_using_vertex_partition,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 2689b7f3cd6..ff7fa067bc4 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -79,7 +79,6 @@ class Tests_MGSelectRandomVertices
     //
 
     std::vector<bool> with_replacement_flags = {true, false};
-    std::vector<bool> shuffle_flags          = {true, false};
 
     {
       // Generate distributed vertex set to sample from
@@ -155,53 +154,34 @@ class Tests_MGSelectRandomVertices
 
     for (int i = 0; i < with_replacement_flags.size(); i++) {
       for (int j = 0; j < sort_vertices_flags.size(); j++) {
-        for (int k = 0; k < shuffle_flags.size(); k++) {
-          for (int l = 0; l < select_counts.size(); l++) {
-            bool with_replacement               = with_replacement_flags[i];
-            bool sort_vertices                  = sort_vertices_flags[j];
-            bool shuffle_using_vertex_partition = shuffle_flags[k];
-            auto select_count                   = static_cast<size_t>(select_counts[l]);
-
-            auto d_sampled_vertices = cugraph::select_random_vertices(
-              *handle_,
-              mg_graph_view,
-              std::optional<raft::device_span<vertex_t const>>{std::nullopt},
-              rng_state,
-              select_count,
-              with_replacement,
-              sort_vertices,
-              shuffle_using_vertex_partition);
-
-            RAFT_CUDA_TRY(cudaDeviceSynchronize());
-
-            auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
-
-            if (select_random_vertices_usecase.check_correctness) {
-              if (!with_replacement) {
-                std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
-
-                auto nr_duplicates =
-                  std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                                h_sampled_vertices.end());
-
-                ASSERT_EQ(nr_duplicates, 0);
-              }
-
-              if (shuffle_using_vertex_partition) {
-                auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
-                auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
-
-                std::for_each(h_sampled_vertices.begin(),
-                              h_sampled_vertices.end(),
-                              [vertex_first, vertex_last](vertex_t v) {
-                                ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
-                              });
-              } else {
-                if (select_count == static_cast<size_t>(mg_graph_view.number_of_vertices())) {
-                  ASSERT_EQ(h_sampled_vertices.size(),
-                            mg_graph_view.local_vertex_partition_range_size());
-                }
-              }
+        for (int l = 0; l < select_counts.size(); l++) {
+          bool with_replacement = with_replacement_flags[i];
+          bool sort_vertices    = sort_vertices_flags[j];
+
+          auto select_count = static_cast<size_t>(select_counts[l]);
+
+          auto d_sampled_vertices = cugraph::select_random_vertices(
+            *handle_,
+            mg_graph_view,
+            std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+            rng_state,
+            select_count,
+            with_replacement,
+            sort_vertices);
+
+          RAFT_CUDA_TRY(cudaDeviceSynchronize());
+
+          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+
+          if (select_random_vertices_usecase.check_correctness) {
+            if (!with_replacement) {
+              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+
+              auto nr_duplicates =
+                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                              h_sampled_vertices.end());
+
+              ASSERT_EQ(nr_duplicates, 0);
             }
           }
         }

From 7692477580863b4d3492dabfd68b75797534a1d0 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 12 Dec 2023 02:03:56 +0100
Subject: [PATCH 10/13] Undo changes

---
 .../mg_select_random_vertices_test.cpp        | 46 ++++++++-----------
 1 file changed, 20 insertions(+), 26 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index ff7fa067bc4..b15b120bd37 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -147,42 +147,36 @@ class Tests_MGSelectRandomVertices
     //
     // Test sampling from [0, V)
     //
-    std::vector<size_t> select_counts = {select_random_vertices_usecase.select_count,
-                                         static_cast<size_t>(mg_graph_view.number_of_vertices())};
 
     std::vector<bool> sort_vertices_flags = {true, false};
 
     for (int i = 0; i < with_replacement_flags.size(); i++) {
       for (int j = 0; j < sort_vertices_flags.size(); j++) {
-        for (int l = 0; l < select_counts.size(); l++) {
-          bool with_replacement = with_replacement_flags[i];
-          bool sort_vertices    = sort_vertices_flags[j];
+        bool with_replacement = with_replacement_flags[i];
+        bool sort_vertices    = sort_vertices_flags[j];
+
+        auto d_sampled_vertices = cugraph::select_random_vertices(
+          *handle_,
+          mg_graph_view,
+          std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+          rng_state,
+          select_random_vertices_usecase.select_count,
+          with_replacement,
+          sort_vertices);
 
-          auto select_count = static_cast<size_t>(select_counts[l]);
-
-          auto d_sampled_vertices = cugraph::select_random_vertices(
-            *handle_,
-            mg_graph_view,
-            std::optional<raft::device_span<vertex_t const>>{std::nullopt},
-            rng_state,
-            select_count,
-            with_replacement,
-            sort_vertices);
-
-          RAFT_CUDA_TRY(cudaDeviceSynchronize());
+        RAFT_CUDA_TRY(cudaDeviceSynchronize());
 
-          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+        auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
 
-          if (select_random_vertices_usecase.check_correctness) {
-            if (!with_replacement) {
-              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+        if (select_random_vertices_usecase.check_correctness) {
+          if (!with_replacement) {
+            std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
 
-              auto nr_duplicates =
-                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                              h_sampled_vertices.end());
+            auto nr_duplicates =
+              std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                            h_sampled_vertices.end());
 
-              ASSERT_EQ(nr_duplicates, 0);
-            }
+            ASSERT_EQ(nr_duplicates, 0);
           }
         }
       }

From 0573351353c03a0614f5a5be85622bff75400ace Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 12 Dec 2023 02:40:41 +0100
Subject: [PATCH 11/13] Undo changes

---
 .../mg_select_random_vertices_test.cpp        | 149 ++++++++++--------
 1 file changed, 84 insertions(+), 65 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index b15b120bd37..74e95435356 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -79,6 +79,7 @@ class Tests_MGSelectRandomVertices
     //
 
     std::vector<bool> with_replacement_flags = {true, false};
+    std::vector<bool> sort_vertices_flags    = {true, false};
 
     {
       // Generate distributed vertex set to sample from
@@ -108,75 +109,93 @@ class Tests_MGSelectRandomVertices
           ? select_random_vertices_usecase.select_count
           : std::rand() % (num_of_elements_in_given_set + 1);
 
-      for (int idx = 0; idx < with_replacement_flags.size(); idx++) {
-        bool with_replacement = with_replacement_flags[idx];
-        auto d_sampled_vertices =
-          cugraph::select_random_vertices(*handle_,
-                                          mg_graph_view,
-                                          std::make_optional(raft::device_span<vertex_t const>{
-                                            d_given_set.data(), d_given_set.size()}),
-                                          rng_state,
-                                          select_count,
-                                          with_replacement,
-                                          true);
-
-        RAFT_CUDA_TRY(cudaDeviceSynchronize());
-
-        auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
-
-        if (select_random_vertices_usecase.check_correctness) {
-          if (!with_replacement) {
-            std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
-
-            auto nr_duplicates =
-              std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                            h_sampled_vertices.end());
-
-            ASSERT_EQ(nr_duplicates, 0);
+      for (int i = 0; i < with_replacement_flags.size(); i++) {
+        for (int j = 0; j < sort_vertices_flags.size(); j++) {
+          bool with_replacement = with_replacement_flags[i];
+          bool sort_vertices    = sort_vertices_flags[j];
+
+          auto d_sampled_vertices =
+            cugraph::select_random_vertices(*handle_,
+                                            mg_graph_view,
+                                            std::make_optional(raft::device_span<vertex_t const>{
+                                              d_given_set.data(), d_given_set.size()}),
+                                            rng_state,
+                                            select_count,
+                                            with_replacement,
+                                            sort_vertices);
+
+          RAFT_CUDA_TRY(cudaDeviceSynchronize());
+
+          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+
+          if (select_random_vertices_usecase.check_correctness) {
+            if (!with_replacement) {
+              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+
+              auto nr_duplicates =
+                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                              h_sampled_vertices.end());
+
+              ASSERT_EQ(nr_duplicates, 0);
+            }
+
+            std::sort(h_given_set.begin(), h_given_set.end());
+            if (sort_vertices) {
+              std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end());
+            } else {
+              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+            }
+            std::for_each(
+              h_sampled_vertices.begin(), h_sampled_vertices.end(), [&h_given_set](vertex_t v) {
+                ASSERT_TRUE(std::binary_search(h_given_set.begin(), h_given_set.end(), v));
+              });
           }
-
-          std::sort(h_given_set.begin(), h_given_set.end());
-          std::for_each(
-            h_sampled_vertices.begin(), h_sampled_vertices.end(), [&h_given_set](vertex_t v) {
-              ASSERT_TRUE(std::binary_search(h_given_set.begin(), h_given_set.end(), v));
-            });
         }
       }
-    }
-
-    //
-    // Test sampling from [0, V)
-    //
-
-    std::vector<bool> sort_vertices_flags = {true, false};
-
-    for (int i = 0; i < with_replacement_flags.size(); i++) {
-      for (int j = 0; j < sort_vertices_flags.size(); j++) {
-        bool with_replacement = with_replacement_flags[i];
-        bool sort_vertices    = sort_vertices_flags[j];
-
-        auto d_sampled_vertices = cugraph::select_random_vertices(
-          *handle_,
-          mg_graph_view,
-          std::optional<raft::device_span<vertex_t const>>{std::nullopt},
-          rng_state,
-          select_random_vertices_usecase.select_count,
-          with_replacement,
-          sort_vertices);
-
-        RAFT_CUDA_TRY(cudaDeviceSynchronize());
-
-        auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
-
-        if (select_random_vertices_usecase.check_correctness) {
-          if (!with_replacement) {
-            std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
-
-            auto nr_duplicates =
-              std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
-                            h_sampled_vertices.end());
 
-            ASSERT_EQ(nr_duplicates, 0);
+      //
+      // Test sampling from [0, V)
+      //
+
+      for (int i = 0; i < with_replacement_flags.size(); i++) {
+        for (int j = 0; j < sort_vertices_flags.size(); j++) {
+          bool with_replacement = with_replacement_flags[i];
+          bool sort_vertices    = sort_vertices_flags[j];
+
+          auto d_sampled_vertices = cugraph::select_random_vertices(
+            *handle_,
+            mg_graph_view,
+            std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+            rng_state,
+            select_random_vertices_usecase.select_count,
+            with_replacement,
+            sort_vertices);
+
+          RAFT_CUDA_TRY(cudaDeviceSynchronize());
+
+          auto h_sampled_vertices = cugraph::test::to_host(*handle_, d_sampled_vertices);
+
+          if (select_random_vertices_usecase.check_correctness) {
+            if (!with_replacement) {
+              std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
+
+              auto nr_duplicates =
+                std::distance(std::unique(h_sampled_vertices.begin(), h_sampled_vertices.end()),
+                              h_sampled_vertices.end());
+
+              ASSERT_EQ(nr_duplicates, 0);
+            }
+            if (sort_vertices) {
+              std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end());
+            }
+
+            auto vertex_first = mg_graph_view.local_vertex_partition_range_first();
+            auto vertex_last  = mg_graph_view.local_vertex_partition_range_last();
+            std::for_each(h_sampled_vertices.begin(),
+                          h_sampled_vertices.end(),
+                          [vertex_first, vertex_last](vertex_t v) {
+                            ASSERT_TRUE((v >= vertex_first) && (v < vertex_last));
+                          });
           }
         }
       }

From e8665609c0471543d926e3200c8e8050c0bb5ee0 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 12 Dec 2023 02:42:53 +0100
Subject: [PATCH 12/13] Check all possible flag values

---
 cpp/tests/structure/mg_select_random_vertices_test.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index 74e95435356..ba4c322103b 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -141,7 +141,7 @@ class Tests_MGSelectRandomVertices
 
             std::sort(h_given_set.begin(), h_given_set.end());
             if (sort_vertices) {
-              std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end());
+              assert(std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end()));
             } else {
               std::sort(h_sampled_vertices.begin(), h_sampled_vertices.end());
             }
@@ -186,7 +186,7 @@ class Tests_MGSelectRandomVertices
               ASSERT_EQ(nr_duplicates, 0);
             }
             if (sort_vertices) {
-              std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end());
+              assert(std::is_sorted(h_sampled_vertices.begin(), h_sampled_vertices.end()));
             }
 
             auto vertex_first = mg_graph_view.local_vertex_partition_range_first();

From e827bd7e7102c5a3780a91d02b0a884735c031a8 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 12 Dec 2023 02:49:03 +0100
Subject: [PATCH 13/13] Set check_correctness to false for benchmark test

---
 cpp/tests/structure/mg_select_random_vertices_test.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/tests/structure/mg_select_random_vertices_test.cpp b/cpp/tests/structure/mg_select_random_vertices_test.cpp
index ba4c322103b..8392a6831ca 100644
--- a/cpp/tests/structure/mg_select_random_vertices_test.cpp
+++ b/cpp/tests/structure/mg_select_random_vertices_test.cpp
@@ -259,8 +259,8 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_MGSelectRandomVertices_Rmat,
   ::testing::Combine(
-    ::testing::Values(SelectRandomVertices_Usecase{500, true},
-                      SelectRandomVertices_Usecase{500, true}),
+    ::testing::Values(SelectRandomVertices_Usecase{500, false},
+                      SelectRandomVertices_Usecase{500, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()