Fix bug found in how we handle empty frontiers in sampling

Author: ChuckHastings

cpp/include/cugraph/prims/detail/per_v_select_transform_e.cuh

@@ -585,6 +585,19 @@ per_v_select_transform_e(
   std::inclusive_scan(
     local_key_list_sizes.begin(), local_key_list_sizes.end(), local_key_list_offsets.begin() + 1);
 
+  // Empty frontier (after minor-comm size exchange): return immediately.
+  if (local_key_list_offsets.back() == 0) {
+    auto empty_results = allocate_dataframe_buffer<T>(0, handle.get_stream());
+    if (invalid_value) {
+      return std::make_tuple(std::nullopt, std::move(empty_results));
+    } else {
+      rmm::device_uvector<size_t> sample_offsets(size_t{1}, handle.get_stream());
+      sample_offsets.set_element_to_zero_async(size_t{0}, handle.get_stream());
+      return std::make_tuple(std::make_optional(std::move(sample_offsets)),
+                             std::move(empty_results));
+    }
+  }
+
   // 1. aggregate key_list
 
   std::optional<key_buffer_t> aggregate_local_key_list{std::nullopt};

cpp/include/cugraph/prims/detail/sample_and_compute_local_nbr_indices.cuh

@@ -5067,6 +5067,12 @@ homogeneous_uniform_sample_and_compute_local_nbr_indices(
   }
   assert(minor_comm_size == graph_view.number_of_local_edge_partitions());
 
+  if (local_frontier_offsets.back() == 0) {
+    return std::make_tuple(rmm::device_uvector<edge_t>(0, handle.get_stream()),
+                           std::optional<rmm::device_uvector<size_t>>{std::nullopt},
+                           std::vector<size_t>(local_frontier_offsets.size(), size_t{0}));
+  }
+
   auto aggregate_local_frontier_major_first =
     thrust_tuple_get_or_identity<KeyIterator, 0>(aggregate_local_frontier_key_first);
 
@@ -5191,6 +5197,12 @@ heterogeneous_uniform_sample_and_compute_local_nbr_indices(
   }
   assert(minor_comm_size == graph_view.number_of_local_edge_partitions());
 
+  if (local_frontier_offsets.back() == 0) {
+    return std::make_tuple(rmm::device_uvector<edge_t>(0, handle.get_stream()),
+                           std::optional<rmm::device_uvector<size_t>>{std::nullopt},
+                           std::vector<size_t>(local_frontier_offsets.size(), size_t{0}));
+  }
+
   auto num_edge_types = static_cast<edge_type_t>(Ks.size());
 
   auto edge_mask_view = graph_view.edge_mask_view();
@@ -5528,6 +5540,12 @@ homogeneous_biased_sample_and_compute_local_nbr_indices(
   }
   assert(minor_comm_size == graph_view.number_of_local_edge_partitions());
 
+  if (local_frontier_offsets.back() == 0) {
+    return std::make_tuple(rmm::device_uvector<edge_t>(0, handle.get_stream()),
+                           std::optional<rmm::device_uvector<size_t>>{std::nullopt},
+                           std::vector<size_t>(local_frontier_offsets.size(), size_t{0}));
+  }
+
   auto edge_mask_view = graph_view.edge_mask_view();
 
   // 1. compute biases for unique keys (to reduce memory footprint)
@@ -5684,6 +5702,12 @@ heterogeneous_biased_sample_and_compute_local_nbr_indices(
   }
   assert(minor_comm_size == graph_view.number_of_local_edge_partitions());
 
+  if (local_frontier_offsets.back() == 0) {
+    return std::make_tuple(rmm::device_uvector<edge_t>(0, handle.get_stream()),
+                           std::optional<rmm::device_uvector<size_t>>{std::nullopt},
+                           std::vector<size_t>(local_frontier_offsets.size(), size_t{0}));
+  }
+
   auto num_edge_types = static_cast<edge_type_t>(Ks.size());
 
   auto edge_mask_view = graph_view.edge_mask_view();

cpp/include/cugraph/prims/detail/transform_v_frontier_e.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -614,24 +614,26 @@ auto transform_v_frontier_e(raft::handle_t const& handle,
             get_dataframe_buffer_begin(aggregate_value_buffer));
       }
     } else {
-      raft::grid_1d_thread_t update_grid(
-        (local_frontier_offsets[i + 1] - local_frontier_offsets[i]),
-        detail::transform_v_frontier_e_kernel_block_size,
-        handle.get_device_properties().maxGridSize[0]);
+      auto frontier_size = local_frontier_offsets[i + 1] - local_frontier_offsets[i];
+      if (frontier_size > 0) {
+        raft::grid_1d_thread_t update_grid(frontier_size,
+                                           detail::transform_v_frontier_e_kernel_block_size,
+                                           handle.get_device_properties().maxGridSize[0]);
 
-      detail::transform_v_frontier_e_hypersparse_or_low_degree<false, GraphViewType>
-        <<<update_grid.num_blocks, update_grid.block_size, 0, handle.get_stream()>>>(
-          edge_partition,
-          edge_partition_frontier_key_first,
-          thrust::make_counting_iterator(size_t{0}),
-          thrust::make_counting_iterator(local_frontier_offsets[i + 1] - local_frontier_offsets[i]),
-          edge_partition_src_value_input,
-          edge_partition_dst_value_input,
-          edge_partition_e_value_input,
-          edge_partition_e_mask,
-          edge_partition_frontier_local_degree_offsets,
-          e_op,
-          get_dataframe_buffer_begin(aggregate_value_buffer));
+        detail::transform_v_frontier_e_hypersparse_or_low_degree<false, GraphViewType>
+          <<<update_grid.num_blocks, update_grid.block_size, 0, handle.get_stream()>>>(
+            edge_partition,
+            edge_partition_frontier_key_first,
+            thrust::make_counting_iterator(size_t{0}),
+            thrust::make_counting_iterator(frontier_size),
+            edge_partition_src_value_input,
+            edge_partition_dst_value_input,
+            edge_partition_e_value_input,
+            edge_partition_e_mask,
+            edge_partition_frontier_local_degree_offsets,
+            e_op,
+            get_dataframe_buffer_begin(aggregate_value_buffer));
+      }
     }
   }
 

cpp/src/sampling/neighbor_sampling_impl.cuh

@@ -15,6 +15,7 @@
 #include <cugraph/graph.hpp>
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/host_scalar_comm.hpp>
 #include <cugraph/utilities/thrust_wrappers/fill.hpp>
 #include <cugraph/vertex_partition_view.hpp>
 
@@ -189,6 +190,16 @@ neighbor_sample_impl(raft::handle_t const& handle,
   }
 
   for (size_t hop = 0; hop < num_hops; ++hop) {
+    {
+      size_t local_frontier_size = (hop == 0) ? starting_vertices.size() : frontier_vertices.size();
+      size_t frontier_size       = local_frontier_size;
+      if constexpr (multi_gpu) {
+        frontier_size = host_scalar_allreduce(
+          handle.get_comms(), local_frontier_size, raft::comms::op_t::SUM, handle.get_stream());
+      }
+      if (frontier_size == 0) { break; }
+    }
+
     std::optional<std::vector<size_t>> level_Ks{std::nullopt};
     std::unique_ptr<bool[]> gather_flags{};
     std::vector<raft::device_span<vertex_t const>> next_frontier_vertex_spans{};

cpp/src/sampling/temporal_sampling_impl.cuh

@@ -19,6 +19,7 @@
 #include <cugraph/sampling_functions.hpp>
 #include <cugraph/utilities/device_functors.cuh>
 #include <cugraph/utilities/error.hpp>
+#include <cugraph/utilities/host_scalar_comm.hpp>
 #include <cugraph/utilities/thrust_wrappers/fill.hpp>
 #include <cugraph/utilities/thrust_wrappers/sort.hpp>
 #include <cugraph/vertex_partition_view.hpp>
@@ -615,6 +616,15 @@ temporal_neighbor_sample_impl(
   produced_edge_lists.reserve(num_hops * 2);  // at most a biased/uniform + a gather list per hop
 
   for (size_t hop = 0; hop < num_hops; ++hop) {
+    {
+      size_t frontier_size = frontier_vertices.size();
+      if constexpr (multi_gpu) {
+        frontier_size = host_scalar_allreduce(
+          handle.get_comms(), frontier_size, raft::comms::op_t::SUM, handle.get_stream());
+      }
+      if (frontier_size == 0) { break; }
+    }
+
     std::optional<std::vector<size_t>> level_Ks{std::nullopt};
     std::unique_ptr<bool[]> gather_flags{};
     std::vector<raft::device_span<vertex_t const>> next_frontier_vertex_spans{};