Change parameter order, add do_expensive_check falg parameter

rapidsai · Jan 9, 2024 · 79b06fa · 79b06fa
1 parent 50c852a
commit 79b06fa
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git 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
@@ -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
@@ -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