From e4b58ec04e4add07578ce779977835cfe28f516a Mon Sep 17 00:00:00 2001
From: Naim <110031745+naimnv@users.noreply.github.com>
Date: Tue, 12 Mar 2024 19:00:31 +0100
Subject: [PATCH] Make external vertex and edge shuffling function public
 (#4227)

Make vertex and edge shuffling function public

Authors:
  - Naim (https://github.com/naimnv)
  - Ralph Liu (https://github.com/nv-rliu)

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

URL: https://github.com/rapidsai/cugraph/pull/4227
---
 cpp/CMakeLists.txt                            |   4 +-
 cpp/include/cugraph/graph_functions.hpp       |  63 +++++++++++
 .../shuffle_vertex_pairs.cu                   | 102 ++++++++++++++++++
 .../{detail => utilities}/shuffle_vertices.cu |  70 ++++++++++++
 4 files changed, 237 insertions(+), 2 deletions(-)
 rename cpp/src/{detail => utilities}/shuffle_vertex_pairs.cu (80%)
 rename cpp/src/{detail => utilities}/shuffle_vertices.cu (73%)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index a3392627fb8..3131404712f 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -186,9 +186,9 @@ endif()
 #       which should give us a better parallel schedule.
 
 set(CUGRAPH_SOURCES
-    src/detail/shuffle_vertices.cu
+    src/utilities/shuffle_vertices.cu
     src/detail/permute_range.cu
-    src/detail/shuffle_vertex_pairs.cu
+    src/utilities/shuffle_vertex_pairs.cu
     src/detail/collect_local_vertex_values.cu
     src/detail/groupby_and_count.cu
     src/detail/collect_comm_wrapper.cu
diff --git a/cpp/include/cugraph/graph_functions.hpp b/cpp/include/cugraph/graph_functions.hpp
index 90425f86bef..6d4470e8251 100644
--- a/cpp/include/cugraph/graph_functions.hpp
+++ b/cpp/include/cugraph/graph_functions.hpp
@@ -1052,4 +1052,67 @@ remove_multi_edges(raft::handle_t const& handle,
                    std::optional<rmm::device_uvector<edge_type_t>>&& edgelist_edge_types,
                    bool keep_min_value_edge = false);
 
+/**
+ * @brief Shuffle external vertex ids to the proper GPU.
+ *
+ * @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 vertices  List of vertex ids
+ * @return Vector of vertex ids mapped to this GPU.
+ */
+template <typename vertex_t>
+rmm::device_uvector<vertex_t> shuffle_external_vertices(raft::handle_t const& handle,
+                                                        rmm::device_uvector<vertex_t>&& vertices);
+
+/**
+ * @brief Shuffle external vertex ids and values to the proper GPU.
+ *
+ * @tparam vertex_t   Type of vertex identifiers. Needs to be an integral type.
+ * @tparam value_t    Type of values. currently supported types are int32_t,
+ * int64_t, size_t, float and double.
+ *
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param vertices  List of vertex ids
+ * @param values List of values
+ * @return Tuple of vectors storing vertex ids and values mapped to this GPU.
+ */
+template <typename vertex_t, typename value_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<vertex_t>&& vertices,
+                                    rmm::device_uvector<value_t>&& values);
+
+/**
+ * @brief Shuffle external edges to the proper GPU.
+ *
+ * @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.
+ *
+ * @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  List of source vertex ids
+ * @param edge_dsts  List of destination vertex ids
+ * @param edge_weights  Optional list of edge weights
+ * @param edge_ids  Optional list of edge ids
+ * @param edge_types Optional list of edge types
+ * @return Tuple of vectors storing edge sources, destinations, optional weights,
+ *          optional edge ids, optional edge types mapped to this GPU.
+ */
+template <typename vertex_t, typename edge_t, typename weight_t, typename edge_type_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>>>
+shuffle_external_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);
+
 }  // namespace cugraph
diff --git a/cpp/src/detail/shuffle_vertex_pairs.cu b/cpp/src/utilities/shuffle_vertex_pairs.cu
similarity index 80%
rename from cpp/src/detail/shuffle_vertex_pairs.cu
rename to cpp/src/utilities/shuffle_vertex_pairs.cu
index 33a7834f5ff..b473796aa9d 100644
--- a/cpp/src/detail/shuffle_vertex_pairs.cu
+++ b/cpp/src/utilities/shuffle_vertex_pairs.cu
@@ -519,4 +519,106 @@ shuffle_int_vertex_pairs_with_values_to_local_gpu_by_edge_partitioning(
   std::vector<int64_t> const& vertex_partition_range_lasts);
 
 }  // namespace detail
+
+template <typename vertex_t, typename edge_t, typename weight_t, typename edge_type_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>>>
+shuffle_external_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)
+{
+  auto& comm                 = handle.get_comms();
+  auto const comm_size       = comm.get_size();
+  auto& major_comm           = handle.get_subcomm(cugraph::partition_manager::major_comm_name());
+  auto const major_comm_size = major_comm.get_size();
+  auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
+  auto const minor_comm_size = minor_comm.get_size();
+
+  return detail::shuffle_ext_vertex_pairs_with_values_to_local_gpu_by_edge_partitioning(
+    handle,
+    std::move(edge_srcs),
+    std::move(edge_dsts),
+    std::move(edge_weights),
+    std::move(edge_ids),
+    std::move(edge_types));
+}
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int32_t>&& majors,
+                       rmm::device_uvector<int32_t>&& minors,
+                       std::optional<rmm::device_uvector<float>>&& weights,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int32_t>&& majors,
+                       rmm::device_uvector<int32_t>&& minors,
+                       std::optional<rmm::device_uvector<double>>&& weights,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int32_t>&& majors,
+                       rmm::device_uvector<int32_t>&& minors,
+                       std::optional<rmm::device_uvector<float>>&& weights,
+                       std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int32_t>&& majors,
+                       rmm::device_uvector<int32_t>&& minors,
+                       std::optional<rmm::device_uvector<double>>&& weights,
+                       std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int64_t>&& majors,
+                       rmm::device_uvector<int64_t>&& minors,
+                       std::optional<rmm::device_uvector<float>>&& weights,
+                       std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>>
+shuffle_external_edges(raft::handle_t const& handle,
+                       rmm::device_uvector<int64_t>&& majors,
+                       rmm::device_uvector<int64_t>&& minors,
+                       std::optional<rmm::device_uvector<double>>&& weights,
+                       std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+                       std::optional<rmm::device_uvector<int32_t>>&& edge_types);
+
 }  // namespace cugraph
diff --git a/cpp/src/detail/shuffle_vertices.cu b/cpp/src/utilities/shuffle_vertices.cu
similarity index 73%
rename from cpp/src/detail/shuffle_vertices.cu
rename to cpp/src/utilities/shuffle_vertices.cu
index be6875f1073..b396201f509 100644
--- a/cpp/src/detail/shuffle_vertices.cu
+++ b/cpp/src/utilities/shuffle_vertices.cu
@@ -249,4 +249,74 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<double>&& values);
 
 }  // namespace detail
+
+template <typename vertex_t, typename value_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<vertex_t>&& vertices,
+                                    rmm::device_uvector<value_t>&& values)
+{
+  return detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
+    handle, std::move(vertices), std::move(values));
+}
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int32_t>&& vertices,
+                                    rmm::device_uvector<int32_t>&& values);
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int32_t>&& vertices,
+                                    rmm::device_uvector<size_t>&& values);
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int32_t>&& vertices,
+                                    rmm::device_uvector<float>&& values);
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<double>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int32_t>&& vertices,
+                                    rmm::device_uvector<double>&& values);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int32_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int64_t>&& vertices,
+                                    rmm::device_uvector<int32_t>&& values);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int64_t>&& vertices,
+                                    rmm::device_uvector<int64_t>&& values);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int64_t>&& vertices,
+                                    rmm::device_uvector<size_t>&& values);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int64_t>&& vertices,
+                                    rmm::device_uvector<float>&& values);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<double>>
+shuffle_external_vertex_value_pairs(raft::handle_t const& handle,
+                                    rmm::device_uvector<int64_t>&& vertices,
+                                    rmm::device_uvector<double>&& values);
+
+template <typename vertex_t>
+rmm::device_uvector<vertex_t> shuffle_external_vertices(raft::handle_t const& handle,
+                                                        rmm::device_uvector<vertex_t>&& vertices)
+{
+  return detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(handle,
+                                                                          std::move(vertices));
+}
+
+template rmm::device_uvector<int32_t> shuffle_external_vertices(
+  raft::handle_t const& handle, rmm::device_uvector<int32_t>&& d_vertices);
+
+template rmm::device_uvector<int64_t> shuffle_external_vertices(
+  raft::handle_t const& handle, rmm::device_uvector<int64_t>&& d_vertices);
+
 }  // namespace cugraph