perf experiment

rapidsai · Aug 23, 2024 · 6d8c7ef · 6d8c7ef
1 parent ac33784
commit 6d8c7ef
Showing 1 changed file with 23 additions and 19 deletions.
diff --git a/cpp/src/prims/detail/extract_transform_v_frontier_e.cuh b/cpp/src/prims/detail/extract_transform_v_frontier_e.cuh
@@ -613,24 +613,11 @@ extract_transform_v_frontier_e(raft::handle_t const& handle,
     }
   }
 
-  std::vector<size_t> local_frontier_sizes{};
-  if constexpr (GraphViewType::is_multi_gpu) {
-    auto& minor_comm     = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
-    local_frontier_sizes = host_scalar_allgather(
-      minor_comm,
-      static_cast<size_t>(thrust::distance(frontier_key_first, frontier_key_last)),
-      handle.get_stream());
-  } else {
-    local_frontier_sizes = std::vector<size_t>{static_cast<size_t>(
-      static_cast<vertex_t>(thrust::distance(frontier_key_first, frontier_key_last)))};
-  }
-
   // update frontier bitmap (used to reduce broadcast bandwidth size)
 
   std::
     conditional_t<try_bitmap, std::optional<rmm::device_uvector<uint32_t>>, std::byte /* dummy */>
       frontier_bitmap{};
-  std::conditional_t<try_bitmap, std::vector<bool>, std::byte /* dummy */> use_bitmap_flags{};
   if constexpr (try_bitmap) {
     auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
     auto const minor_comm_size = minor_comm.get_size();
@@ -648,12 +635,6 @@ extract_transform_v_frontier_e(raft::handle_t const& handle,
                                         graph_view.local_vertex_partition_range_first() + bool_size,
                                         handle.get_stream());
     }
-    auto tmp_flags = host_scalar_allgather(
-      minor_comm, frontier_bitmap ? uint8_t{1} : uint8_t{0}, handle.get_stream());
-    use_bitmap_flags.resize(tmp_flags.size());
-    std::transform(tmp_flags.begin(), tmp_flags.end(), use_bitmap_flags.begin(), [](auto flag) {
-      return flag == uint8_t{1};
-    });
   }
 
   // compute max_pushes
@@ -679,6 +660,29 @@ extract_transform_v_frontier_e(raft::handle_t const& handle,
       frontier_major_first, frontier_major_last, handle.get_stream());
   }
 
+  // communication over minor_comm
+
+  std::vector<size_t> local_frontier_sizes{};
+  std::conditional_t<try_bitmap, std::vector<bool>, std::byte /* dummy */> use_bitmap_flags{};
+  if constexpr (GraphViewType::is_multi_gpu) {
+    auto& minor_comm     = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
+    local_frontier_sizes = host_scalar_allgather(
+      minor_comm,
+      static_cast<size_t>(thrust::distance(frontier_key_first, frontier_key_last)),
+      handle.get_stream());
+    if constexpr (try_bitmap) {
+      auto tmp_flags = host_scalar_allgather(
+        minor_comm, frontier_bitmap ? uint8_t{1} : uint8_t{0}, handle.get_stream());
+      use_bitmap_flags.resize(tmp_flags.size());
+      std::transform(tmp_flags.begin(), tmp_flags.end(), use_bitmap_flags.begin(), [](auto flag) {
+        return flag == uint8_t{1};
+      });
+    }
+  } else {
+    local_frontier_sizes = std::vector<size_t>{static_cast<size_t>(
+      static_cast<vertex_t>(thrust::distance(frontier_key_first, frontier_key_last)))};
+  }
+
   // set-up stream ppol
 
   std::optional<std::vector<size_t>> stream_pool_indices{std::nullopt};