Implement two_hop_neighbors C API

cpp/include/cugraph_c/graph_functions.h

@@ -87,15 +87,18 @@ void cugraph_vertex_pairs_free(cugraph_vertex_pairs_t* vertex_pairs);
  * @param [in]  start_vertices Optional type erased array of starting vertices
  *                             If NULL use all, if specified compute two-hop
  *                             neighbors for these starting vertices
+ * @param [in]  do_expensive_check
+ *                             A flag to run expensive checks for input arguments (if set to true)
  * @param [out] result         Opaque pointer to resulting vertex pairs
  * @param [out] error          Pointer to an error object storing details of any error.  Will
  *                             be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
 cugraph_error_code_t cugraph_two_hop_neighbors(
   const cugraph_resource_handle_t* handle,
-  const cugraph_graph_t* graph,
+  cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
+  bool_t do_expensive_check,
   cugraph_vertex_pairs_t** result,
   cugraph_error_t** error);
 

cpp/src/c_api/graph_functions.cpp

@@ -19,9 +19,11 @@
 #include <c_api/abstract_functor.hpp>
 #include <c_api/graph.hpp>
 #include <c_api/graph_functions.hpp>
+#include <c_api/graph_helper.hpp>
 #include <c_api/resource_handle.hpp>
 #include <c_api/utils.hpp>
 
+#include <cugraph/algorithms.hpp>
 #include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/graph_functions.hpp>
@@ -80,15 +82,23 @@ struct create_vertex_pairs_functor : public cugraph::c_api::abstract_functor {
 };
 
 struct two_hop_neighbors_functor : public cugraph::c_api::abstract_functor {
-  raft::handle_t const& handle_;
-  cugraph::c_api::cugraph_graph_t const* graph_;
+  raft::handle_t const& handle_{};
+  cugraph::c_api::cugraph_graph_t* graph_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_vertex_pairs_t* result_{};
+  bool do_expensive_check_{false};
 
   two_hop_neighbors_functor(::cugraph_resource_handle_t const* handle,
-                            ::cugraph_graph_t const* graph)
+                            ::cugraph_graph_t* graph,
+                            ::cugraph_type_erased_device_array_view_t const* start_vertices,
+                            bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
-      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t const*>(graph))
+      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      start_vertices_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertices)),
+      do_expensive_check_(do_expensive_check)
   {
   }
 
@@ -103,12 +113,91 @@ struct two_hop_neighbors_functor : public cugraph::c_api::abstract_functor {
     if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
       unsupported();
     } else {
-      auto graph = reinterpret_cast<
-        cugraph::graph_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu>*>(graph_->graph_);
+      // k_hop_nbrs expects store_transposed == false
+      if constexpr (store_transposed) {
+        error_code_ = cugraph::c_api::
+          transpose_storage<vertex_t, edge_t, weight_t, store_transposed, multi_gpu>(
+            handle_, graph_, error_.get());
+        if (error_code_ != CUGRAPH_SUCCESS) return;
+      }
+
+      auto graph =
+        reinterpret_cast<cugraph::graph_t<vertex_t, edge_t, weight_t, false, multi_gpu>*>(
+          graph_->graph_);
 
       auto graph_view = graph->view();
+      auto number_map = reinterpret_cast<rmm::device_uvector<vertex_t>*>(graph_->number_map_);
+
+      rmm::device_uvector<vertex_t> start_vertices(0, handle_.get_stream());
+
+      if (start_vertices_ != nullptr) {
+        start_vertices.resize(start_vertices_->size_, handle_.get_stream());
+        raft::copy(start_vertices.data(),
+                   start_vertices_->as_type<vertex_t const>(),
+                   start_vertices_->size_,
+                   handle_.get_stream());
+
+        cugraph::renumber_ext_vertices<vertex_t, multi_gpu>(
+          handle_,
+          start_vertices.data(),
+          start_vertices.size(),
+          number_map->data(),
+          graph_view.local_vertex_partition_range_first(),
+          graph_view.local_vertex_partition_range_last(),
+          do_expensive_check_);
+
+        if constexpr (multi_gpu) {
+          start_vertices =
+            cugraph::detail::shuffle_ext_vertices_by_gpu_id(handle_, std::move(start_vertices));
+        }
+      } else {
+        start_vertices.resize(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
+        cugraph::detail::sequence_fill(handle_.get_stream(),
+                                       start_vertices.data(),
+                                       start_vertices.size(),
+                                       graph_view.local_vertex_partition_range_first());
+      }
+
+      auto [offsets, dst] = cugraph::k_hop_nbrs(
+        handle_,
+        graph_view,
+        raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+        size_t{2},
+        do_expensive_check_);
+
+      auto src = cugraph::c_api::expand_sparse_offsets(
+        raft::device_span<size_t const>{offsets.data(), offsets.size()},
+        vertex_t{0},
+        handle_.get_stream());
+
+      // convert ids back to srcs:  src[i] = start_vertices[src[i]]
+      cugraph::unrenumber_local_int_vertices(handle_,
+                                             src.data(),
+                                             src.size(),
+                                             start_vertices.data(),
+                                             vertex_t{0},
+                                             graph_view.local_vertex_partition_range_size(),
+                                             do_expensive_check_);
 
-      CUGRAPH_FAIL("Not implemented");
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(
+        handle_,
+        src.data(),
+        src.size(),
+        number_map->data(),
+        graph_view.vertex_partition_range_lasts(),
+        do_expensive_check_);
+
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(
+        handle_,
+        dst.data(),
+        dst.size(),
+        number_map->data(),
+        graph_view.vertex_partition_range_lasts(),
+        do_expensive_check_);
+
+      result_ = new cugraph::c_api::cugraph_vertex_pairs_t{
+        new cugraph::c_api::cugraph_type_erased_device_array_t(src, graph_->vertex_type_),
+        new cugraph::c_api::cugraph_type_erased_device_array_t(dst, graph_->vertex_type_)};
     }
   }
 };
@@ -170,12 +259,13 @@ extern "C" void cugraph_vertex_pairs_free(cugraph_vertex_pairs_t* vertex_pairs)
 
 extern "C" cugraph_error_code_t cugraph_two_hop_neighbors(
   const cugraph_resource_handle_t* handle,
-  const cugraph_graph_t* graph,
+  cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
+  bool_t do_expensive_check,
   cugraph_vertex_pairs_t** result,
   cugraph_error_t** error)
 {
-  two_hop_neighbors_functor functor(handle, graph);
+  two_hop_neighbors_functor functor(handle, graph, start_vertices, do_expensive_check);
 
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }

cpp/src/c_api/graph_functions.hpp

@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <c_api/array.hpp>
+
 namespace cugraph {
 namespace c_api {
 

cpp/src/c_api/graph_helper.cu

@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#pragma once
 
 #include <structure/detail/structure_utils.cuh>
 
@@ -43,5 +42,13 @@ template rmm::device_uvector<int64_t> expand_sparse_offsets(
   int64_t base_vertex_id,
   rmm::cuda_stream_view const& stream);
 
+template rmm::device_uvector<int32_t> expand_sparse_offsets(raft::device_span<size_t const> offsets,
+                                                            int32_t base_vertex_id,
+                                                            rmm::cuda_stream_view const& stream);
+
+template rmm::device_uvector<int64_t> expand_sparse_offsets(raft::device_span<size_t const> offsets,
+                                                            int64_t base_vertex_id,
+                                                            rmm::cuda_stream_view const& stream);
+
 }  // namespace c_api
 }  // namespace cugraph

cpp/tests/CMakeLists.txt

@@ -591,6 +591,7 @@ if(BUILD_CUGRAPH_MG_TESTS)
     ConfigureCTestMG(MG_CAPI_SIMILARITY_TEST c_api/mg_similarity_test.c c_api/mg_test_utils.cpp)
     ConfigureCTestMG(MG_CAPI_K_CORE_TEST c_api/mg_k_core_test.c c_api/mg_test_utils.cpp)
     ConfigureCTestMG(MG_CAPI_INDUCED_SUBGRAPH_TEST c_api/mg_induced_subgraph_test.c c_api/mg_test_utils.cpp)
+    ConfigureCTestMG(MG_CAPI_TWO_HOP_NEIGHBORS_TEST c_api/mg_two_hop_neighbors_test.c c_api/mg_test_utils.cpp)
 endif()
 
 ###################################################################################################
@@ -645,6 +646,7 @@ ConfigureCTest(CAPI_CORE_NUMBER_TEST c_api/core_number_test.c)
 ConfigureCTest(CAPI_SIMILARITY_TEST c_api/similarity_test.c)
 ConfigureCTest(CAPI_K_CORE_TEST c_api/k_core_test.c)
 ConfigureCTest(CAPI_INDUCED_SUBGRAPH_TEST c_api/induced_subgraph_test.c)
+ConfigureCTest(CAPI_TWO_HOP_NEIGHBORS_TEST c_api/two_hop_neighbors_test.c)
 
 ###################################################################################################
 ### enable testing ################################################################################