Implement a vertex pair intersection primitive

cpp/include/cugraph/graph_functions.hpp

@@ -51,7 +51,8 @@ struct renumber_meta_t<vertex_t, edge_t, multi_gpu, std::enable_if_t<!multi_gpu>
  *
  * This function assumes that vertices are pre-shuffled to their target processes and edges are
  * pre-shuffled to their target processess and edge partitions using compute_gpu_id_from_vertex_t
- * and compute_gpu_id_from_edge_t & compute_partition_id_from_edge_t functors, respectively.
+ * and compute_gpu_id_from_ext_edge_endpoints_t & compute_partition_id_from_ext_edge_endpoints_t
+ * functors, respectively.
  *
  * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
  * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
@@ -65,12 +66,12 @@ struct renumber_meta_t<vertex_t, edge_t, multi_gpu, std::enable_if_t<!multi_gpu>
  * work (vertices should be pre-shuffled).
  * @param edgelist_srcs Pointers (one pointer per local edge partition assigned to this process) to
  * edge source vertex IDs. Source IDs are updated in-place ([INOUT] parameter). Applying the
- * compute_gpu_id_from_edge_t functor to every (destination ID, source ID) pair (if store_transposed
- * = true) or (source ID, destination ID) pair (if store_transposed = false) should return the local
- * GPU ID for this function to work (edges should be pre-shuffled). Applying the
- * compute_partition_id_from_edge_t to every (destination ID, source ID) pair (if store_transposed =
- * true) or (source ID, destination ID) pair (if store_transposed = false) should also return the
- * corresponding edge partition ID. The best way to enforce this is to use
+ * compute_gpu_id_from_ext_edge_endpoints_t functor to every (destination ID, source ID) pair (if
+ * store_transposed = true) or (source ID, destination ID) pair (if store_transposed = false) should
+ * return the local GPU ID for this function to work (edges should be pre-shuffled). Applying the
+ * compute_partition_id_from_ext_edge_endpoints_t to every (destination ID, source ID) pair (if
+ * store_transposed = true) or (source ID, destination ID) pair (if store_transposed = false) should
+ * also return the corresponding edge partition ID. The best way to enforce this is to use
  * shuffle_edgelist_by_gpu_id & groupby_and_count_edgelist_by_local_partition_id.
  * @param edgelist_dsts Pointers (one pointer per local edge partition assigned to this process) to
  * edge destination vertex IDs. Destination IDs are updated in-place ([INOUT] parameter).
@@ -347,8 +348,8 @@ void renumber_local_ext_vertices(raft::handle_t const& handle,
  * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
  * handles to various CUDA libraries) to run graph algorithms.
  * @param edgelist_srcs Vector of edge source vertex IDs. If multi-GPU, applying the
- * compute_gpu_id_from_edge_t to every edge should return the local GPU ID for this function to work
- * (edges should be pre-shuffled).
+ * compute_gpu_id_from_ext_edge_endpoints_t to every edge should return the local GPU ID for this
+ * function to work (edges should be pre-shuffled).
  * @param edgelist_dsts Vector of edge destination vertex IDs.
  * @param edgelist_weights Vector of edge weights.
  * @param reciprocal Flag indicating whether to keep (if set to `false`) or discard (if set to
@@ -500,8 +501,8 @@ extract_induced_subgraphs(
  * compute_gpu_id_from_vertex_t to every vertex should return the local GPU ID for this function to
  * work (vertices should be pre-shuffled).
  * @param edgelist_srcs Vector of edge source vertex IDs. If multi-GPU, applying the
- * compute_gpu_id_from_edge_t to every edge should return the local GPU ID for this function to work
- * (edges should be pre-shuffled).
+ * compute_gpu_id_from_ext_edge_endpoints_t to every edge should return the local GPU ID for this
+ * function to work (edges should be pre-shuffled).
  * @param edgelist_dsts Vector of edge destination vertex IDs.
  * @param edgelist_weights Vector of edge weights.
  * @param edgelist_id_type_pairs Vector of edge ID and type pairs.

cpp/include/cugraph/utilities/device_functors.cuh

@@ -37,6 +37,16 @@ struct typecast_t {
   __device__ output_t operator()(input_t val) const { return static_cast<output_t>(val); }
 };
 
+template <typename Iterator>
+struct indirection_t {
+  Iterator first{};
+
+  __device__ typename thrust::iterator_traits<Iterator>::value_type operator()(size_t i) const
+  {
+    return *(first + i);
+  }
+};
+
 template <typename T>
 struct not_equal_t {
   T compare{};

cpp/src/components/weakly_connected_components_impl.cuh

@@ -711,7 +711,7 @@ void weakly_connected_components_impl(raft::handle_t const& handle,
           get_dataframe_buffer_begin(edge_buffer),
           get_dataframe_buffer_end(edge_buffer),
           [key_func =
-             cugraph::detail::compute_gpu_id_from_edge_t<vertex_t>{
+             cugraph::detail::compute_gpu_id_from_ext_edge_endpoints_t<vertex_t>{
                comm_size, row_comm_size, col_comm_size}] __device__(auto val) {
             return key_func(thrust::get<0>(val), thrust::get<1>(val));
           },

cpp/src/detail/graph_utils.cuh

@@ -54,21 +54,19 @@ struct compute_gpu_id_from_ext_vertex_t {
 
 template <typename vertex_t>
 struct compute_gpu_id_from_int_vertex_t {
-  raft::device_span<vertex_t> vertex_partition_range_lasts_span;
+  raft::device_span<vertex_t const> vertex_partition_range_lasts{};
 
   __device__ int operator()(vertex_t v) const
   {
-    return static_cast<int>(
-      thrust::distance(vertex_partition_range_lasts_span.begin(),
-                       thrust::upper_bound(thrust::seq,
-                                           vertex_partition_range_lasts_span.begin(),
-                                           vertex_partition_range_lasts_span.end(),
-                                           v)));
+    return static_cast<int>(thrust::distance(
+      vertex_partition_range_lasts.begin(),
+      thrust::upper_bound(
+        thrust::seq, vertex_partition_range_lasts.begin(), vertex_partition_range_lasts.end(), v)));
   }
 };
 
 template <typename vertex_t>
-struct compute_gpu_id_from_edge_t {
+struct compute_gpu_id_from_ext_edge_endpoints_t {
   int comm_size{0};
   int row_comm_size{0};
   int col_comm_size{0};
@@ -83,7 +81,49 @@ struct compute_gpu_id_from_edge_t {
 };
 
 template <typename vertex_t>
-struct compute_partition_id_from_edge_t {
+struct compute_gpu_id_from_int_edge_endpoints_t {
+  raft::device_span<vertex_t const> vertex_partition_range_lasts{};
+  int comm_size{0};
+  int row_comm_size{0};
+  int col_comm_size{0};
+
+  __device__ int operator()(vertex_t major, vertex_t minor) const
+  {
+    auto major_comm_rank =
+      static_cast<int>(thrust::distance(vertex_partition_range_lasts.begin(),
+                                        thrust::upper_bound(thrust::seq,
+                                                            vertex_partition_range_lasts.begin(),
+                                                            vertex_partition_range_lasts.end(),
+                                                            major)));
+    auto minor_comm_rank =
+      static_cast<int>(thrust::distance(vertex_partition_range_lasts.begin(),
+                                        thrust::upper_bound(thrust::seq,
+                                                            vertex_partition_range_lasts.begin(),
+                                                            vertex_partition_range_lasts.end(),
+                                                            minor)));
+    return (minor_comm_rank / row_comm_size) * row_comm_size + (major_comm_rank % row_comm_size);
+  }
+
+  __device__ int operator()(thrust::tuple<vertex_t, vertex_t> pair /* major, minor */) const
+  {
+    auto major_comm_rank =
+      static_cast<int>(thrust::distance(vertex_partition_range_lasts.begin(),
+                                        thrust::upper_bound(thrust::seq,
+                                                            vertex_partition_range_lasts.begin(),
+                                                            vertex_partition_range_lasts.end(),
+                                                            thrust::get<0>(pair))));
+    auto minor_comm_rank =
+      static_cast<int>(thrust::distance(vertex_partition_range_lasts.begin(),
+                                        thrust::upper_bound(thrust::seq,
+                                                            vertex_partition_range_lasts.begin(),
+                                                            vertex_partition_range_lasts.end(),
+                                                            thrust::get<1>(pair))));
+    return (minor_comm_rank / row_comm_size) * row_comm_size + (major_comm_rank % row_comm_size);
+  }
+};
+
+template <typename vertex_t>
+struct compute_partition_id_from_ext_edge_endpoints_t {
   int comm_size{0};
   int row_comm_size{0};
   int col_comm_size{0};

cpp/src/detail/shuffle_wrappers.cu

@@ -84,7 +84,7 @@ shuffle_edgelist_by_gpu_id(raft::handle_t const& handle,
       edge_first,
       edge_first + d_edgelist_majors.size(),
       [key_func =
-         cugraph::detail::compute_gpu_id_from_edge_t<vertex_t>{
+         cugraph::detail::compute_gpu_id_from_ext_edge_endpoints_t<vertex_t>{
            comm_size, row_comm_size, col_comm_size}] __device__(auto val) {
         return key_func(thrust::get<0>(val), thrust::get<1>(val));
       },
@@ -133,7 +133,7 @@ shuffle_edgelist_by_gpu_id(raft::handle_t const& handle,
       edge_first,
       edge_first + d_edgelist_majors.size(),
       [key_func =
-         cugraph::detail::compute_gpu_id_from_edge_t<vertex_t>{
+         cugraph::detail::compute_gpu_id_from_ext_edge_endpoints_t<vertex_t>{
            comm_size, row_comm_size, col_comm_size}] __device__(auto val) {
         return key_func(thrust::get<0>(val), thrust::get<1>(val));
       },
@@ -364,7 +364,7 @@ rmm::device_uvector<size_t> groupby_and_count_edgelist_by_local_partition_id(
       [comm_size,
        row_comm_size,
        partition_id_key_func =
-         cugraph::detail::compute_partition_id_from_edge_t<vertex_t>{
+         cugraph::detail::compute_partition_id_from_ext_edge_endpoints_t<vertex_t>{
            comm_size, row_comm_size, col_comm_size},
        gpu_id_key_func = cugraph::detail::compute_gpu_id_from_ext_vertex_t<vertex_t>{
          comm_size}] __device__(auto pair) {
@@ -412,7 +412,7 @@ rmm::device_uvector<size_t> groupby_and_count_edgelist_by_local_partition_id(
   } else {
     auto local_partition_id_op =
       [comm_size,
-       key_func = cugraph::detail::compute_partition_id_from_edge_t<vertex_t>{
+       key_func = cugraph::detail::compute_partition_id_from_ext_edge_endpoints_t<vertex_t>{
          comm_size, row_comm_size, col_comm_size}] __device__(auto pair) {
         return key_func(thrust::get<0>(pair), thrust::get<1>(pair)) /
                comm_size;  // global partition id to local partition id