Lookup edge src dst using edge id and type

cpp/CMakeLists.txt

@@ -286,6 +286,12 @@ set(CUGRAPH_SOURCES
     src/community/k_truss_sg_v64_e64.cu
     src/community/k_truss_sg_v32_e32.cu
     src/community/k_truss_sg_v32_e64.cu
+    src/lookup/lookup_src_dst_mg_v32_e32.cu
+    src/lookup/lookup_src_dst_mg_v32_e64.cu
+    src/lookup/lookup_src_dst_mg_v64_e64.cu
+    src/lookup/lookup_src_dst_sg_v32_e32.cu
+    src/lookup/lookup_src_dst_sg_v32_e64.cu
+    src/lookup/lookup_src_dst_sg_v64_e64.cu
     src/sampling/random_walks_old_sg_v32_e32.cu
     src/sampling/random_walks_old_sg_v32_e64.cu
     src/sampling/random_walks_old_sg_v64_e64.cu

cpp/include/cugraph/sampling_functions.hpp

@@ -15,6 +15,8 @@
  */
 #pragma once
 
+#include <cugraph/src_dst_lookup_container.hpp>
+
 #include <raft/core/device_span.hpp>
 #include <raft/core/handle.hpp>
 
@@ -449,5 +451,83 @@ sort_sampled_edgelist(raft::handle_t const& handle,
                       size_t num_hops,
                       bool src_is_major       = true,
                       bool do_expensive_check = false);
+/*
+ * @brief Build map to lookup source and destination using edge id and type
+ *
+ * @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 edge_type_t Type of edge types. Needs to be an integral type.
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @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.
+ * @param edge_id_view View object holding edge ids of the edges of the graph pointed @p graph_view
+ * @param edge_type_view View object holding edge types of the edges of the graph pointed @p
+ * graph_view
+ * @return An object of type cugraph::lookup_container_t that encapsulates edge id and type to
+ * source and destination lookup map.
+ */
+template <typename vertex_t, typename edge_t, typename edge_type_t, bool multi_gpu>
+lookup_container_t<edge_t, edge_type_t, vertex_t> build_edge_id_and_type_to_src_dst_lookup_map(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view,
+  edge_property_view_t<edge_t, edge_t const*> edge_id_view,
+  edge_property_view_t<edge_t, edge_type_t const*> edge_type_view);
+
+/*
+ * @brief Lookup edge sources and destinations using edge ids and a single edge type.
+ * Use this function to lookup endpoints of edges belonging to the same edge type.
+ *
+ * @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 edge_type_t Type of edge types. Needs to be an integral type.
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param lookup_container Object of type cugraph::lookup_container_t that encapsulates edge id and
+ * type to source and destination lookup map.
+ * @param edge_ids_to_lookup Device span of edge ids to lookup
+ * @param edge_type_to_lookup Type of the edges corresponding to edge ids in @p edge_ids_to_lookup
+ * @return A tuple of device vector containing edge sources and destinations for edge ids
+ * in @p edge_ids_to_lookup and edge type @. If an edge id in @p edge_ids_to_lookup or
+ * edge type @edge_type_to_lookup is not found, the corresponding entry in the device vectors of
+ * the returned tuple will contain cugraph::invalid_vertex_id<vertex_t>.
+ */
+template <typename vertex_t, typename edge_t, typename edge_type_t, bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>>
+lookup_endpoints_from_edge_ids_and_single_type(
+  raft::handle_t const& handle,
+  lookup_container_t<edge_t, edge_type_t, vertex_t> const& lookup_container,
+  raft::device_span<edge_t const> edge_ids_to_lookup,
+  edge_type_t edge_type_to_lookup);
+
+/*
+ * @brief Lookup edge sources and destinations using edge ids and edge types.
+ * Use this function to lookup endpoints of edges belonging to different edge types.
+ *
+ * @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 edge_type_t Type of edge types. Needs to be an integral type.
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param lookup_container Object of type cugraph::lookup_container_t that encapsulates edge id and
+ * type to source and destination lookup map.
+ * @param edge_ids_to_lookup Device span of edge ids to lookup
+ * @param edge_types_to_lookup Device span of edge types corresponding to the edge ids
+ * in @p edge_ids_to_lookup
+ * @return A tuple of device vector containing edge sources and destinations for the edge ids
+ * in @p edge_ids_to_lookup and the edge types in @p edge_types_to_lookup. If an edge id in
+ * @p edge_ids_to_lookup or edge type in @p edge_types_to_lookup is not found, the
+ * corresponding entry in the device vectors of the returned tuple will contain
+ * cugraph::invalid_vertex_id<vertex_t>.
+ */
+template <typename vertex_t, typename edge_t, typename edge_type_t, bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>>
+lookup_endpoints_from_edge_ids_and_types(
+  raft::handle_t const& handle,
+  lookup_container_t<edge_t, edge_type_t, vertex_t> const& lookup_container,
+  raft::device_span<edge_t const> edge_ids_to_lookup,
+  raft::device_span<edge_type_t const> edge_types_to_lookup);
 
 }  // namespace cugraph

cpp/include/cugraph/src_dst_lookup_container.hpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cugraph/graph.hpp>
+#include <cugraph/utilities/dataframe_buffer.hpp>
+#include <cugraph/utilities/thrust_tuple_utils.hpp>
+
+#include <raft/core/device_span.hpp>
+#include <raft/core/handle.hpp>
+
+#include <rmm/exec_policy.hpp>
+
+#include <thrust/tuple.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace cugraph {
+
+template <typename edge_id_t,
+          typename edge_type_t,
+          typename vertex_t,
+          typename value_t = thrust::tuple<vertex_t, vertex_t>>
+class lookup_container_t {
+  template <typename _edge_id_t, typename _edge_type_t, typename _vertex_t, typename _value_t>
+  struct lookup_container_impl;
+  std::unique_ptr<lookup_container_impl<edge_id_t, edge_type_t, vertex_t, value_t>> pimpl;
+
+ public:
+  using edge_id_type   = edge_id_t;
+  using edge_type_type = edge_type_t;
+  using value_type     = value_t;
+
+  static_assert(std::is_integral_v<edge_id_t>);
+  static_assert(std::is_integral_v<edge_type_t>);
+  static_assert(is_thrust_tuple_of_integral<value_t>::value);
+
+  ~lookup_container_t();
+  lookup_container_t();
+  lookup_container_t(raft::handle_t const& handle,
+                     std::vector<edge_type_t> types,
+                     std::vector<edge_id_t> type_counts);
+  lookup_container_t(const lookup_container_t&);
+
+  void insert(raft::handle_t const& handle,
+              edge_type_t typ,
+              raft::device_span<edge_id_t const> edge_ids_to_insert,
+              dataframe_buffer_type_t<value_t>&& values_to_insert);
+
+  dataframe_buffer_type_t<value_t> lookup_from_edge_ids_and_single_type(
+    raft::handle_t const& handle,
+    raft::device_span<edge_id_t const> edge_ids_to_lookup,
+    edge_type_t edge_type_to_lookup,
+    bool multi_gpu) const;
+
+  dataframe_buffer_type_t<value_t> lookup_from_edge_ids_and_types(
+    raft::handle_t const& handle,
+    raft::device_span<edge_id_t const> edge_ids_to_lookup,
+    raft::device_span<edge_type_t const> edge_types_to_lookup,
+    bool multi_gpu) const;
+};
+
+}  // namespace cugraph

cpp/include/cugraph/utilities/dataframe_buffer.hpp

@@ -207,6 +207,14 @@ auto get_dataframe_buffer_end(BufferType& buffer)
     std::make_index_sequence<std::tuple_size<BufferType>::value>(), buffer);
 }
 
+template <typename T>
+struct dataframe_buffer_type {
+  using type = decltype(allocate_dataframe_buffer<T>(size_t{0}, rmm::cuda_stream_view{}));
+};
+
+template <typename T>
+using dataframe_buffer_type_t = typename dataframe_buffer_type<T>::type;
+
 template <typename BufferType,
           typename std::enable_if_t<is_arithmetic_vector<std::remove_cv_t<BufferType>,
                                                          rmm::device_uvector>::value>* = nullptr>

cpp/include/cugraph/utilities/thrust_tuple_utils.hpp

@@ -113,6 +113,19 @@ struct is_thrust_tuple_of_arithmetic<thrust::tuple<Ts...>> {
   static constexpr bool value = (... && is_valid<Ts>);
 };
 
+template <typename TupleType>
+struct is_thrust_tuple_of_integral : std::false_type {};
+
+template <typename... Ts>
+struct is_thrust_tuple_of_integral<thrust::tuple<Ts...>> {
+ private:
+  template <typename T>
+  static constexpr bool is_valid = std::is_integral_v<T> || std::is_same_v<T, thrust::null_type>;
+
+ public:
+  static constexpr bool value = (... && is_valid<Ts>);
+};
+
 template <typename T>
 struct is_std_tuple : std::false_type {};
 

cpp/src/detail/graph_partition_utils.cuh

@@ -50,6 +50,21 @@ struct compute_gpu_id_from_ext_vertex_t {
   }
 };
 
+template <typename edge_t>
+struct compute_gpu_id_from_ext_edge_id_t {
+  int comm_size{0};
+  int major_comm_size{0};
+  int minor_comm_size{0};
+
+  __host__ __device__ int operator()(edge_t e) const
+  {
+    cuco::detail::MurmurHash3_32<edge_t> hash_func{};
+    auto vertex_partition_id = static_cast<int>(hash_func(e) % comm_size);
+    return partition_manager::compute_global_comm_rank_from_vertex_partition_id(
+      major_comm_size, minor_comm_size, vertex_partition_id);
+  }
+};
+
 template <typename vertex_t>
 struct compute_gpu_id_from_int_vertex_t {
   raft::device_span<vertex_t const> vertex_partition_range_lasts{};