Compile Optimizations - batch 3

cpp/include/cugraph/mtmg/detail/per_device_edgelist.hpp

@@ -16,8 +16,8 @@
 
 #pragma once
 
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/mtmg/handle.hpp>
+#include <cugraph/shuffle_functions.hpp>
 
 // FIXME: Could use std::span once compiler supports C++20
 #include <raft/core/host_span.hpp>
@@ -301,38 +301,35 @@ class per_device_edgelist_t {
       if (edge_end_time_) resize_and_copy_buffers(*edge_end_time_, total_size, stream);
     }
 
-    auto tmp_wgt     = wgt_ ? std::make_optional(std::move((*wgt_)[0])) : std::nullopt;
-    auto tmp_edge_id = edge_id_ ? std::make_optional(std::move((*edge_id_)[0])) : std::nullopt;
-    auto tmp_edge_type =
-      edge_type_ ? std::make_optional(std::move((*edge_type_)[0])) : std::nullopt;
-    auto tmp_edge_start_time =
-      edge_start_time_ ? std::make_optional(std::move((*edge_start_time_)[0])) : std::nullopt;
-    auto tmp_edge_end_time =
-      edge_end_time_ ? std::make_optional(std::move((*edge_end_time_)[0])) : std::nullopt;
-
-    std::tie(store_transposed ? dst_[0] : src_[0],
-             store_transposed ? src_[0] : dst_[0],
-             tmp_wgt,
-             tmp_edge_id,
-             tmp_edge_type,
-             tmp_edge_start_time,
-             tmp_edge_end_time,
-             std::ignore) =
-      cugraph::detail::shuffle_ext_vertex_pairs_with_values_to_local_gpu_by_edge_partitioning(
-        handle.raft_handle(),
-        store_transposed ? std::move(dst_[0]) : std::move(src_[0]),
-        store_transposed ? std::move(src_[0]) : std::move(dst_[0]),
-        std::move(tmp_wgt),
-        std::move(tmp_edge_id),
-        std::move(tmp_edge_type),
-        std::move(tmp_edge_start_time),
-        std::move(tmp_edge_end_time));
-
-    if (tmp_wgt) ((*wgt_)[0]) = std::move(*tmp_wgt);
-    if (tmp_edge_id) ((*edge_id_)[0]) = std::move(*tmp_edge_id);
-    if (tmp_edge_type) ((*edge_type_)[0]) = std::move(*tmp_edge_type);
-    if (tmp_edge_start_time) ((*edge_start_time_)[0]) = std::move(*tmp_edge_start_time);
-    if (tmp_edge_end_time) ((*edge_end_time_)[0]) = std::move(*tmp_edge_end_time);
+    std::vector<cugraph::arithmetic_device_uvector_t> tmp_edge_properties{};
+    if (wgt_) tmp_edge_properties.push_back(std::move((*wgt_)[0]));
+    if (edge_id_) tmp_edge_properties.push_back(std::move((*edge_id_)[0]));
+    if (edge_type_) tmp_edge_properties.push_back(std::move((*edge_type_)[0]));
+    if (edge_start_time_) tmp_edge_properties.push_back(std::move((*edge_start_time_)[0]));
+    if (edge_start_time_) tmp_edge_properties.push_back(std::move((*edge_start_time_)[0]));
+
+    std::tie(src_[0], dst_[0], tmp_edge_properties, std::ignore) =
+      cugraph::shuffle_ext_edges(handle.raft_handle(),
+                                 std::move(src_[0]),
+                                 std::move(dst_[0]),
+                                 std::move(tmp_edge_properties),
+                                 store_transposed);
+
+    size_t pos{0};
+    if (wgt_)
+      ((*wgt_)[0]) = std::move(std::get<rmm::device_uvector<weight_t>>(tmp_edge_properties[pos++]));
+    if (edge_id_)
+      ((*edge_id_)[0]) =
+        std::move(std::get<rmm::device_uvector<edge_t>>(tmp_edge_properties[pos++]));
+    if (edge_type_)
+      ((*edge_type_)[0]) =
+        std::move(std::get<rmm::device_uvector<edge_type_t>>(tmp_edge_properties[pos++]));
+    if (edge_start_time_)
+      ((*edge_start_time_)[0]) =
+        std::move(std::get<rmm::device_uvector<edge_time_t>>(tmp_edge_properties[pos++]));
+    if (edge_end_time_)
+      ((*edge_end_time_)[0]) =
+        std::move(std::get<rmm::device_uvector<edge_time_t>>(tmp_edge_properties[pos++]));
   }
 
  private:

cpp/include/cugraph/shuffle_functions.hpp

@@ -15,7 +15,9 @@
  */
 #pragma once
 
+#include <cugraph/arithmetic_variant_types.hpp>
 #include <cugraph/large_buffer_manager.hpp>
+#include <cugraph/utilities/dataframe_buffer.hpp>
 
 #include <raft/core/handle.hpp>
 #include <raft/core/host_span.hpp>
@@ -75,51 +77,56 @@ shuffle_ext_vertex_value_pairs(raft::handle_t const& handle,
  * @brief Shuffle external edges to the owning GPUs (by edge partitioning)
  *
  * @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 weight. Currently float and double are supported.
- * @tparam edge_type_t Type of edge type. Needs to be an integral type, currently only int32_t is
- * supported.
- * @tparam edge_time_t Type of edge time.  Needs to be an integral type.
  *
  * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
  * handles to various CUDA libraries) to run graph algorithms.
  * @param edge_srcs  Vector of source vertex ids
  * @param edge_dsts  Vector of destination vertex ids
- * @param edge_weights  Optional vector of edge weights
- * @param edge_ids  Optional vector of edge ids
- * @param edge_types Optional vector of edge types
- * @param edge_start_times Optional vector of edge start times
- * @param edge_end_times Optional vector of edge end times
+ * @param edge_properties  Vector of edge properties, each element is an arithmetic device vector
  * @param store_transposed Should be true if shuffled edges will be used with a cugraph::graph_t
  * object with store_tranposed = true. Should be false otherwise.
- * @return Tuple of vectors storing edge sources, destinations, optional weights,
- *          optional edge ids, optional edge types, optional edge start times, optional edge end
- * times mapped to this GPU and a vector storing the number of edges received from each GPU.
+ * @return Tuple of vectors storing edge sources, destinations, and edge properties
  */
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename edge_time_t>
+template <typename vertex_t>
 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_type_t>>,
-           std::optional<rmm::device_uvector<edge_time_t>>,
-           std::optional<rmm::device_uvector<edge_time_t>>,
+           std::vector<cugraph::arithmetic_device_uvector_t>,
            std::vector<size_t>>
 shuffle_ext_edges(raft::handle_t const& handle,
                   rmm::device_uvector<vertex_t>&& edge_srcs,
                   rmm::device_uvector<vertex_t>&& edge_dsts,
-                  std::optional<rmm::device_uvector<weight_t>>&& edge_weights,
-                  std::optional<rmm::device_uvector<edge_t>>&& edge_ids,
-                  std::optional<rmm::device_uvector<edge_type_t>>&& edge_types,
-                  std::optional<rmm::device_uvector<edge_time_t>>&& edge_start_times,
-                  std::optional<rmm::device_uvector<edge_time_t>>&& edge_end_times,
+                  std::vector<cugraph::arithmetic_device_uvector_t>&& edge_properties,
                   bool store_transposed,
                   std::optional<large_buffer_type_t> large_buffer_type = std::nullopt);
 
+/**
+ * @ingroup graph_functions_cpp
+ * @brief Shuffle internal edges to the owning GPUs (by edge partitioning)
+ *
+ * @tparam vertex_t    Type of vertex identifiers. Needs to be an integral type.
+ *
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param edge_srcs  Vector of source vertex ids
+ * @param edge_dsts  Vector of destination vertex ids
+ * @param edge_properties  Vector of edge properties, each element is an arithmetic device vector
+ * @param store_transposed Should be true if shuffled edges will be used with a cugraph::graph_t
+ * object with store_tranposed = true. Should be false otherwise.
+ * @return Tuple of vectors storing edge sources, destinations, and edge properties
+ */
+template <typename vertex_t>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::vector<cugraph::arithmetic_device_uvector_t>,
+           std::vector<size_t>>
+shuffle_int_edges(raft::handle_t const& handle,
+                  rmm::device_uvector<vertex_t>&& majors,
+                  rmm::device_uvector<vertex_t>&& minors,
+                  std::vector<cugraph::arithmetic_device_uvector_t>&& edge_properties,
+                  bool store_transposed,
+                  raft::host_span<vertex_t const> vertex_partition_range_lasts,
+                  std::optional<large_buffer_type_t> large_buffer_type = std::nullopt);
+
 /**
  * @brief Shuffle local edge sources (already placed by edge partitioning) to the owning GPUs (by
  * vertex partitioning).
@@ -129,7 +136,8 @@ shuffle_ext_edges(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 edge_srcs  Vector of local edge source IDs
- * @param vertex_partition_range_lasts  Span of vertex partition range lasts (size = number of GPUs)
+ * @param vertex_partition_range_lasts  Span of vertex partition range lasts (size = number of
+ * GPUs)
  * @param store_transposed Should be true if shuffled edges will be used with a cugraph::graph_t
  * object with store_tranposed = true. Should be false otherwise.
  * @return Vector of shuffled edge source vertex IDs (shuffled by vertex partitioning).

cpp/src/c_api/allgather.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,7 +25,6 @@
 
 #include <cugraph/algorithms.hpp>
 #include <cugraph/detail/collect_comm_wrapper.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/decompress_to_edgelist.cpp

@@ -24,7 +24,6 @@
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/ecg.cpp

@@ -17,7 +17,6 @@
 #include <cugraph_c/community_algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/extract_vertex_list.cpp

@@ -24,7 +24,6 @@
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/graph_functions.cpp

@@ -25,7 +25,6 @@
 #include <cugraph_c/graph_functions.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/shuffle_functions.hpp>
@@ -73,27 +72,14 @@ struct create_vertex_pairs_functor : public cugraph::c_api::abstract_functor {
         second_copy.data(), second_->as_type<vertex_t>(), second_->size_, handle_.get_stream());
 
       if constexpr (multi_gpu) {
-        std::tie(first_copy,
-                 second_copy,
-                 std::ignore,
-                 std::ignore,
-                 std::ignore,
-                 std::ignore,
-                 std::ignore,
-                 std::ignore) =
-          cugraph::detail::shuffle_ext_vertex_pairs_with_values_to_local_gpu_by_edge_partitioning<
-            vertex_t,
-            edge_t,
-            weight_t,
-            edge_type_type_t,
-            int32_t>(handle_,
-                     std::move(first_copy),
-                     std::move(second_copy),
-                     std::nullopt,
-                     std::nullopt,
-                     std::nullopt,
-                     std::nullopt,
-                     std::nullopt);
+        std::vector<cugraph::arithmetic_device_uvector_t> edge_properties{};
+
+        std::tie(first_copy, second_copy, std::ignore, std::ignore) =
+          cugraph::shuffle_ext_edges(handle_,
+                                     std::move(first_copy),
+                                     std::move(second_copy),
+                                     std::move(edge_properties),
+                                     false);
       }
       // FIXME: use std::tuple (template) instead.
       result_ = new cugraph::c_api::cugraph_vertex_pairs_t{

cpp/src/c_api/graph_mg.cpp

@@ -146,39 +146,40 @@ struct create_graph_functor : public cugraph::c_api::abstract_functor {
       rmm::device_uvector<vertex_t> edgelist_dsts =
         concatenate<vertex_t>(handle_, dst_, num_arrays_);
 
-      std::optional<rmm::device_uvector<weight_t>> edgelist_weights =
-        weights_ ? std::make_optional(concatenate<weight_t>(handle_, weights_, num_arrays_))
+      std::vector<cugraph::arithmetic_device_uvector_t> edgelist_edge_properties{};
+
+      if (weights_)
+        edgelist_edge_properties.push_back(concatenate<weight_t>(handle_, weights_, num_arrays_));
+      if (edge_ids_)
+        edgelist_edge_properties.push_back(concatenate<edge_t>(handle_, edge_ids_, num_arrays_));
+      if (edge_type_ids_)
+        edgelist_edge_properties.push_back(
+          concatenate<edge_type_t>(handle_, edge_type_ids_, num_arrays_));
+
+      std::tie(edgelist_srcs, edgelist_dsts, edgelist_edge_properties, std::ignore) =
+        cugraph::shuffle_ext_edges(handle_,
+                                   std::move(edgelist_srcs),
+                                   std::move(edgelist_dsts),
+                                   std::move(edgelist_edge_properties),
+                                   store_transposed);
+
+      size_t pos{0};
+      auto edgelist_weights =
+        weights_ ? std::make_optional(std::move(
+                     std::get<rmm::device_uvector<weight_t>>(edgelist_edge_properties[pos++])))
                  : std::nullopt;
-
-      std::optional<rmm::device_uvector<edge_t>> edgelist_edge_ids =
-        edge_ids_ ? std::make_optional(concatenate<edge_t>(handle_, edge_ids_, num_arrays_))
+      auto edgelist_edge_ids =
+        edge_ids_ ? std::make_optional(std::move(
+                      std::get<rmm::device_uvector<edge_t>>(edgelist_edge_properties[pos++])))
                   : std::nullopt;
-
-      std::optional<rmm::device_uvector<edge_type_t>> edgelist_edge_types =
-        edge_type_ids_
-          ? std::make_optional(concatenate<edge_type_t>(handle_, edge_type_ids_, num_arrays_))
-          : std::nullopt;
+      auto edgelist_edge_types =
+        edge_type_ids_ ? std::make_optional(std::move(std::get<rmm::device_uvector<edge_type_t>>(
+                           edgelist_edge_properties[pos++])))
+                       : std::nullopt;
 
       std::optional<rmm::device_uvector<edge_time_t>> edgelist_edge_start_times{std::nullopt};
       std::optional<rmm::device_uvector<edge_time_t>> edgelist_edge_end_times{std::nullopt};
 
-      std::tie(edgelist_srcs,
-               edgelist_dsts,
-               edgelist_weights,
-               edgelist_edge_ids,
-               edgelist_edge_types,
-               edgelist_edge_start_times,
-               edgelist_edge_end_times,
-               std::ignore) = cugraph::shuffle_ext_edges(handle_,
-                                                         std::move(edgelist_srcs),
-                                                         std::move(edgelist_dsts),
-                                                         std::move(edgelist_weights),
-                                                         std::move(edgelist_edge_ids),
-                                                         std::move(edgelist_edge_types),
-                                                         std::move(edgelist_edge_start_times),
-                                                         std::move(edgelist_edge_end_times),
-                                                         store_transposed);
-
       if (vertex_list) {
         vertex_list = cugraph::shuffle_ext_vertices(handle_, std::move(*vertex_list));
       }

cpp/src/c_api/hits.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,11 +18,11 @@
 #include "c_api/graph.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
+#include "detail/shuffle_wrappers.hpp"
 
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/induced_subgraph.cpp

@@ -24,7 +24,6 @@
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/k_core.cpp

@@ -19,11 +19,11 @@
 #include "c_api/graph.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
+#include "detail/shuffle_wrappers.hpp"
 
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/k_truss.cpp

@@ -23,7 +23,6 @@
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
 

cpp/src/c_api/katz.cpp

@@ -19,11 +19,11 @@
 #include "c_api/graph.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
+#include "detail/shuffle_wrappers.hpp"
 
 #include <cugraph_c/algorithms.h>
 
 #include <cugraph/algorithms.hpp>
-#include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>