From 79b06fa8228feeff371f4e3b53e0e92cbb90f527 Mon Sep 17 00:00:00 2001
From: Naim <naim@uib.no>
Date: Tue, 9 Jan 2024 23:27:21 +0100
Subject: [PATCH] Change parameter order, add do_expensive_check falg parameter

---
 .../cugraph/detail/shuffle_wrappers.hpp       |  5 ++--
 cpp/src/community/louvain_impl.cuh            |  2 +-
 cpp/src/detail/permute_range.cu               | 30 ++++++++++++-------
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/cpp/include/cugraph/detail/shuffle_wrappers.hpp b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
index 4e08542911f..c77ecb7aa01 100644
--- a/cpp/include/cugraph/detail/shuffle_wrappers.hpp
+++ b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
@@ -156,9 +156,10 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
 template <typename vertex_t>
 rmm::device_uvector<vertex_t> permute_range(raft::handle_t const& handle,
                                             raft::random::RngState& rng_state,
-                                            vertex_t local_range_size,
                                             vertex_t local_start,
-                                            bool multi_gpu = false);
+                                            vertex_t local_range_size,
+                                            bool multi_gpu          = false,
+                                            bool do_expensive_check = false);
 
 /**
  * @brief Shuffle internal (i.e. renumbered) vertices to their local GPUs based on vertex
diff --git a/cpp/src/community/louvain_impl.cuh b/cpp/src/community/louvain_impl.cuh
index 99c5f0617d0..4919dda5a75 100644
--- a/cpp/src/community/louvain_impl.cuh
+++ b/cpp/src/community/louvain_impl.cuh
@@ -90,8 +90,8 @@ std::pair<std::unique_ptr<Dendrogram<vertex_t>>, weight_t> louvain(
       auto random_cluster_assignments = cugraph::detail::permute_range<vertex_t>(
         handle,
         *rng_state,
-        current_graph_view.local_vertex_partition_range_size(),
         current_graph_view.local_vertex_partition_range_first(),
+        current_graph_view.local_vertex_partition_range_size(),
         multi_gpu);
 
       raft::copy(dendrogram->current_level_begin(),
diff --git a/cpp/src/detail/permute_range.cu b/cpp/src/detail/permute_range.cu
index 009f3578189..069272ae7a1 100644
--- a/cpp/src/detail/permute_range.cu
+++ b/cpp/src/detail/permute_range.cu
@@ -38,11 +38,12 @@ namespace detail {
 template <typename vertex_t>
 rmm::device_uvector<vertex_t> permute_range(raft::handle_t const& handle,
                                             raft::random::RngState& rng_state,
-                                            vertex_t local_range_size,
                                             vertex_t local_range_start,
-                                            bool multi_gpu)
+                                            vertex_t local_range_size,
+                                            bool multi_gpu,
+                                            bool do_expensive_check)
 {
-  if (multi_gpu) {
+  if (do_expensive_check && multi_gpu) {
     auto& comm           = handle.get_comms();
     auto const comm_size = comm.get_size();
     auto const comm_rank = comm.get_rank();
@@ -69,8 +70,7 @@ rmm::device_uvector<vertex_t> permute_range(raft::handle_t const& handle,
     auto const comm_size = comm.get_size();
     auto const comm_rank = comm.get_rank();
 
-    std::vector<size_t> tx_value_counts(comm_size);
-    std::fill(tx_value_counts.begin(), tx_value_counts.end(), 0);
+    std::vector<size_t> tx_value_counts(comm_size, 0)
 
     {
       rmm::device_uvector<vertex_t> d_target_ranks(permuted_integers.size(), handle.get_stream());
@@ -138,8 +138,14 @@ rmm::device_uvector<vertex_t> permute_range(raft::handle_t const& handle,
     // take care of deficits and extras numbers
     auto& comm           = handle.get_comms();
     auto const comm_rank = comm.get_rank();
-    int nr_extras = static_cast<int>(permuted_integers.size()) - static_cast<int>(local_range_size);
-    int nr_deficits = nr_extras >= 0 ? 0 : -nr_extras;
+
+    size_t nr_extras{0};
+    size_t nr_deficits{0};
+    if (permuted_integers.size() > static_cast<size_t>(local_range_size)) {
+      nr_extras = permuted_integers.size() - static_cast<size_t>(local_range_size);
+    } else {
+      nr_deficits = static_cast<size_t>(local_range_size) - permuted_integers.size();
+    }
 
     auto extra_cluster_ids = cugraph::detail::device_allgatherv(
       handle,
@@ -165,15 +171,17 @@ rmm::device_uvector<vertex_t> permute_range(raft::handle_t const& handle,
 
 template rmm::device_uvector<int32_t> permute_range(raft::handle_t const& handle,
                                                     raft::random::RngState& rng_state,
-                                                    int32_t local_range_size,
                                                     int32_t local_range_start,
-                                                    bool multi_gpu);
+                                                    int32_t local_range_size,
+                                                    bool multi_gpu,
+                                                    bool do_expensive_check);
 
 template rmm::device_uvector<int64_t> permute_range(raft::handle_t const& handle,
                                                     raft::random::RngState& rng_state,
-                                                    int64_t local_range_size,
                                                     int64_t local_range_start,
-                                                    bool multi_gpu);
+                                                    int64_t local_range_size,
+                                                    bool multi_gpu,
+                                                    bool do_expensive_check);
 
 }  // namespace detail
 }  // namespace cugraph