Update decompress_to_edgelist to handle edge types (#4397)

Update decompress_to_edgelist to handle edge types

Authors:
  - Naim (https://github.com/naimnv)

Approvers:
  - Seunghwa Kang (https://github.com/seunghwak)
  - Joseph Nke (https://github.com/jnke2016)
  - Chuck Hastings (https://github.com/ChuckHastings)

URL: #4397

Author: naimnv

cpp/include/cugraph/detail/decompress_edge_partition.cuh

@@ -209,20 +209,27 @@ void decompress_edge_partition_to_fill_edgelist_majors(
   }
 }
 
-template <typename vertex_t, typename edge_t, typename weight_t, bool multi_gpu>
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool multi_gpu>
 void decompress_edge_partition_to_edgelist(
   raft::handle_t const& handle,
   edge_partition_device_view_t<vertex_t, edge_t, multi_gpu> edge_partition,
   std::optional<edge_partition_edge_property_device_view_t<edge_t, weight_t const*>>
     edge_partition_weight_view,
   std::optional<edge_partition_edge_property_device_view_t<edge_t, edge_t const*>>
     edge_partition_id_view,
+  std::optional<edge_partition_edge_property_device_view_t<edge_t, edge_type_t const*>>
+    edge_partition_type_view,
   std::optional<edge_partition_edge_property_device_view_t<edge_t, uint32_t const*, bool>>
     edge_partition_mask_view,
   raft::device_span<vertex_t> edgelist_majors /* [OUT] */,
   raft::device_span<vertex_t> edgelist_minors /* [OUT] */,
   std::optional<raft::device_span<weight_t>> edgelist_weights /* [OUT] */,
   std::optional<raft::device_span<edge_t>> edgelist_ids /* [OUT] */,
+  std::optional<raft::device_span<edge_type_t>> edgelist_types /* [OUT] */,
   std::optional<std::vector<vertex_t>> const& segment_offsets)
 {
   auto number_of_edges = edge_partition.number_of_edges();
@@ -271,6 +278,22 @@ void decompress_edge_partition_to_edgelist(
                    (*edgelist_ids).begin());
     }
   }
+
+  if (edge_partition_type_view) {
+    assert(edgelist_types.has_value());
+    if (edge_partition_mask_view) {
+      copy_if_mask_set(handle,
+                       (*edge_partition_type_view).value_first(),
+                       (*edge_partition_type_view).value_first() + number_of_edges,
+                       (*edge_partition_mask_view).value_first(),
+                       (*edgelist_types).begin());
+    } else {
+      thrust::copy(handle.get_thrust_policy(),
+                   (*edge_partition_type_view).value_first(),
+                   (*edge_partition_type_view).value_first() + number_of_edges,
+                   (*edgelist_types).begin());
+    }
+  }
 }
 
 }  // namespace detail

cpp/include/cugraph/graph_functions.hpp

@@ -344,15 +344,17 @@ void renumber_local_ext_vertices(raft::handle_t const& handle,
  * @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.
  * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge types. Needs to be an integral type.
  * @tparam store_transposed Flag indicating whether to use sources (if false) or destinations (if
  * true) as major indices in storing edges using a 2D sparse matrix. transposed.
  * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
  * or multi-GPU (true).
  * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
  * handles to various CUDA libraries) to run graph algorithms.
  * @param graph_view Graph view object of the graph to be decompressed.
- * @param edge_id_view Optional view object holding edge ids for @p graph_view.
  * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param renumber_map If valid, return the renumbered edge list based on the provided @p
  * renumber_map
  * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
@@ -363,17 +365,20 @@ void renumber_local_ext_vertices(raft::handle_t const& handle,
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
+          typename edge_type_t,
           bool store_transposed,
           bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>,
            rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>>
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>>
 decompress_to_edgelist(
   raft::handle_t const& handle,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   std::optional<raft::device_span<vertex_t const>> renumber_map,
   bool do_expensive_check = false);
 

cpp/src/c_api/betweenness_centrality.cpp

@@ -234,12 +234,13 @@ struct edge_betweenness_centrality_functor : public cugraph::c_api::abstract_fun
           normalized_,
           do_expensive_check_);
 
-      auto [src_ids, dst_ids, output_centralities, output_edge_ids] =
-        cugraph::decompress_to_edgelist(
+      auto [src_ids, dst_ids, output_centralities, output_edge_ids, output_edge_types] =
+        cugraph::decompress_to_edgelist<vertex_t, edge_t, weight_t, int32_t, false, multi_gpu>(
           handle_,
           graph_view,
           std::make_optional(centralities.view()),
           (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+          std::nullopt,
           (number_map != nullptr) ? std::make_optional(raft::device_span<vertex_t const>{
                                       number_map->data(), number_map->size()})
                                   : std::nullopt);

cpp/src/community/k_truss_impl.cuh

@@ -671,12 +671,14 @@ k_truss(raft::handle_t const& handle,
     edge_weight_view =
       edge_weight ? std::make_optional((*edge_weight).view())
                   : std::optional<edge_property_view_t<edge_t, weight_t const*>>{std::nullopt};
-    std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts, std::ignore) = decompress_to_edgelist(
-      handle,
-      cur_graph_view,
-      edge_weight_view,
-      std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-      std::optional<raft::device_span<vertex_t const>>(std::nullopt));
+    std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts, std::ignore, std::ignore) =
+      decompress_to_edgelist(
+        handle,
+        cur_graph_view,
+        edge_weight_view,
+        std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        std::optional<raft::device_span<vertex_t const>>(std::nullopt));
 
     auto num_triangles = edge_triangle_count<vertex_t, edge_t, false, false>(
       handle,
@@ -894,12 +896,14 @@ k_truss(raft::handle_t const& handle,
       num_triangles.resize(num_edges_with_triangles, handle.get_stream());
     }
 
-    std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts, std::ignore) = decompress_to_edgelist(
-      handle,
-      cur_graph_view,
-      edge_weight_view ? std::make_optional(*edge_weight_view) : std::nullopt,
-      std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-      std::optional<raft::device_span<vertex_t const>>(std::nullopt));
+    std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts, std::ignore, std::ignore) =
+      decompress_to_edgelist(
+        handle,
+        cur_graph_view,
+        edge_weight_view ? std::make_optional(*edge_weight_view) : std::nullopt,
+        std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        std::optional<raft::device_span<vertex_t const>>(std::nullopt));
 
     std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts) =
       symmetrize_edgelist<vertex_t, weight_t, false, multi_gpu>(handle,

cpp/src/prims/transform_reduce_dst_nbr_intersection_of_e_endpoints_by_v.cuh

@@ -309,16 +309,19 @@ void transform_reduce_dst_nbr_intersection_of_e_endpoints_by_v(
     detail::decompress_edge_partition_to_edgelist<vertex_t,
                                                   edge_t,
                                                   weight_t,
+                                                  int32_t,
                                                   GraphViewType::is_multi_gpu>(
       handle,
       edge_partition,
       std::nullopt,
       std::nullopt,
+      std::nullopt,
       edge_partition_e_mask,
       raft::device_span<vertex_t>(majors.data(), majors.size()),
       raft::device_span<vertex_t>(minors.data(), minors.size()),
       std::nullopt,
       std::nullopt,
+      std::nullopt,
       segment_offsets);
 
     auto vertex_pair_first =

cpp/src/structure/coarsen_graph_impl.cuh

@@ -171,18 +171,20 @@ decompress_edge_partition_to_relabeled_and_grouped_and_coarsened_edgelist(
                             ? std::make_optional<rmm::device_uvector<weight_t>>(
                                 edgelist_majors.size(), handle.get_stream())
                             : std::nullopt;
-  detail::decompress_edge_partition_to_edgelist<vertex_t, edge_t, weight_t, multi_gpu>(
+  detail::decompress_edge_partition_to_edgelist<vertex_t, edge_t, weight_t, int32_t, multi_gpu>(
     handle,
     edge_partition,
     edge_partition_weight_view,
     std::nullopt,
+    std::nullopt,
     edge_partition_e_mask,
     raft::device_span<vertex_t>(edgelist_majors.data(), edgelist_majors.size()),
     raft::device_span<vertex_t>(edgelist_minors.data(), edgelist_minors.size()),
     edgelist_weights ? std::make_optional<raft::device_span<weight_t>>((*edgelist_weights).data(),
                                                                        (*edgelist_weights).size())
                      : std::nullopt,
     std::nullopt,
+    std::nullopt,
     segment_offsets);
 
   auto pair_first =