Merge branch 'branch-24.06' into basic_tutorial

benchmarks/cugraph/standalone/cugraph_dask_funcs.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -44,7 +44,7 @@ def read_csv(input_csv_file, scale):
     chunksize = cugraph.dask.get_chunksize(input_csv_file)
     return dask_cudf.read_csv(
         input_csv_file,
-        chunksize=chunksize,
+        blocksize=chunksize,
         delimiter=" ",
         # names=names,
         dtype=dtypes,

benchmarks/cugraph/standalone/cugraph_funcs.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -75,10 +75,8 @@ def read_csv(input_csv_file, scale):
     dtypes = [vertex_t, vertex_t, "float32"]
     names = (["src", "dst", "weight"],)
 
-    chunksize = cugraph.dask.get_chunksize(input_csv_file)
     return cudf.read_csv(
         input_csv_file,
-        chunksize=chunksize,
         delimiter=" ",
         # names=names,
         dtype=dtypes,

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/include/cugraph_c/sampling_algorithms.h

@@ -236,6 +236,15 @@ typedef enum cugraph_compression_type_t {
 cugraph_error_code_t cugraph_sampling_options_create(cugraph_sampling_options_t** options,
                                                      cugraph_error_t** error);
 
+/**
+ * @ingroup samplingC
+ * @brief   Set flag to retain seeds (original sources)
+ *
+ * @param options - opaque pointer to the sampling options
+ * @param value - Boolean value to assign to the option
+ */
+void cugraph_sampling_set_retain_seeds(cugraph_sampling_options_t* options, bool_t value);
+
 /**
  * @ingroup samplingC
  * @brief   Set flag to renumber results
@@ -335,6 +344,8 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
  * output.  If specified then the all data from @p label_list[i] will be shuffled to rank @p.  This
  * cannot be specified unless @p start_vertex_labels is also specified
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
+ * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
  * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
  *                           We only support fanout values of type INT32
  * @param [in/out] rng_state State of the random number generator, updated with each call
@@ -354,6 +365,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,

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/c_api/uniform_neighbor_sampling.cpp

@@ -40,6 +40,7 @@ struct cugraph_sampling_options_t {
   bool_t renumber_results_{FALSE};
   cugraph_compression_type_t compression_type_{cugraph_compression_type_t::COO};
   bool_t compress_per_hop_{FALSE};
+  bool_t retain_seeds_{FALSE};
 };
 
 struct cugraph_sample_result_t {
@@ -68,6 +69,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
@@ -81,6 +83,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
     cugraph_type_erased_device_array_view_t const* start_vertex_labels,
     cugraph_type_erased_device_array_view_t const* label_list,
     cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
+    cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
     cugraph_rng_state_t* rng_state,
     cugraph::c_api::cugraph_sampling_options_t options,
@@ -99,6 +102,9 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
       label_to_comm_rank_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           label_to_comm_rank)),
+      label_offsets_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          label_offsets)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
@@ -267,8 +273,13 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
               std::move(edge_id),
               std::move(edge_type),
               std::move(hop),
-              std::nullopt,
-              std::nullopt,
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
               offsets ? std::make_optional(
                           raft::device_span<size_t const>{offsets->data(), offsets->size()})
                       : std::nullopt,
@@ -304,8 +315,13 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
               std::move(edge_id),
               std::move(edge_type),
               std::move(hop),
-              std::nullopt,
-              std::nullopt,
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
               offsets ? std::make_optional(
                           raft::device_span<size_t const>{offsets->data(), offsets->size()})
                       : std::nullopt,
@@ -402,6 +418,12 @@ extern "C" cugraph_error_code_t cugraph_sampling_options_create(
   return CUGRAPH_SUCCESS;
 }
 
+extern "C" void cugraph_sampling_set_retain_seeds(cugraph_sampling_options_t* options, bool_t value)
+{
+  auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t*>(options);
+  internal_pointer->retain_seeds_ = value;
+}
+
 extern "C" void cugraph_sampling_set_renumber_results(cugraph_sampling_options_t* options,
                                                       bool_t value)
 {
@@ -871,13 +893,21 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
+  auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
+
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "must specify label_offsets if retain_seeds is true",
+               *error);
+
   CAPI_EXPECTS((start_vertex_labels == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                     start_vertex_labels)
@@ -911,16 +941,16 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
-  uniform_neighbor_sampling_functor functor{
-    handle,
-    graph,
-    start_vertices,
-    start_vertex_labels,
-    label_list,
-    label_to_comm_rank,
-    fan_out,
-    rng_state,
-    *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options),
-    do_expensive_check};
+  uniform_neighbor_sampling_functor functor{handle,
+                                            graph,
+                                            start_vertices,
+                                            start_vertex_labels,
+                                            label_list,
+                                            label_to_comm_rank,
+                                            label_offsets,
+                                            fan_out,
+                                            rng_state,
+                                            std::move(options_cpp),
+                                            do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }

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 =