Replace raw pointers with device_span in induced subgraph (rapidsai#2348

)

Changed the interface of function extract_induced_subgraphs, and all corresponding test codes using this interface.

Authors:
  - Yang Hu (https://github.com/yang-hu-nv)

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

URL: rapidsai#2348

cpp/include/cugraph/graph_functions.hpp

@@ -19,6 +19,7 @@
 #include <cugraph/graph_view.hpp>
 
 #include <raft/handle.hpp>
+#include <raft/span.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <memory>
@@ -471,8 +472,8 @@ std::tuple<rmm::device_uvector<vertex_t>,
 extract_induced_subgraphs(
   raft::handle_t const& handle,
   graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
-  size_t const* subgraph_offsets /* size == num_subgraphs + 1 */,
-  vertex_t const* subgraph_vertices /* size == subgraph_offsets[num_subgraphs] */,
+  raft::device_span<size_t const> subgraph_offsets /* size == num_subgraphs + 1 */,
+  raft::device_span<vertex_t const> subgraph_vertices /* size == subgraph_offsets[num_subgraphs] */,
   size_t num_subgraphs,
   bool do_expensive_check = false);
 

cpp/src/community/legacy/egonet.cu

@@ -169,7 +169,12 @@ extract(raft::handle_t const& handle,
 
   // extract
   return cugraph::extract_induced_subgraphs(
-    handle, csr_view, neighbors_offsets.data().get(), neighbors.data().get(), n_subgraphs);
+    handle,
+    csr_view,
+    raft::device_span<size_t const>(neighbors_offsets.data().get(), neighbors_offsets.size()),
+    raft::device_span<vertex_t const>(neighbors.data().get(), neighbors.size()),
+    n_subgraphs,
+    false);
 }
 
 }  // namespace

cpp/src/structure/induced_subgraph_impl.cuh

@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 #pragma once
-
 #include <cugraph/edge_partition_device_view.cuh>
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/graph_view.hpp>
@@ -48,8 +47,8 @@ std::tuple<rmm::device_uvector<vertex_t>,
 extract_induced_subgraphs(
   raft::handle_t const& handle,
   graph_view_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu> const& graph_view,
-  size_t const* subgraph_offsets /* size == num_subgraphs + 1 */,
-  vertex_t const* subgraph_vertices /* size == subgraph_offsets[num_subgraphs] */,
+  raft::device_span<size_t const> subgraph_offsets /* size == num_subgraphs + 1 */,
+  raft::device_span<vertex_t const> subgraph_vertices /* size == subgraph_offsets[num_subgraphs] */,
   size_t num_subgraphs,
   bool do_expensive_check)
 {
@@ -69,24 +68,25 @@ extract_induced_subgraphs(
   if (do_expensive_check) {
     size_t should_be_zero{std::numeric_limits<size_t>::max()};
     size_t num_aggregate_subgraph_vertices{};
-    raft::update_host(&should_be_zero, subgraph_offsets, 1, handle.get_stream());
-    raft::update_host(
-      &num_aggregate_subgraph_vertices, subgraph_offsets + num_subgraphs, 1, handle.get_stream());
+    raft::update_host(&should_be_zero, subgraph_offsets.data(), 1, handle.get_stream());
+    raft::update_host(&num_aggregate_subgraph_vertices,
+                      subgraph_offsets.data() + num_subgraphs,
+                      1,
+                      handle.get_stream());
     handle.sync_stream();
     CUGRAPH_EXPECTS(should_be_zero == 0,
                     "Invalid input argument: subgraph_offsets[0] should be 0.");
 
-    CUGRAPH_EXPECTS(
-      thrust::is_sorted(
-        handle.get_thrust_policy(), subgraph_offsets, subgraph_offsets + (num_subgraphs + 1)),
-      "Invalid input argument: subgraph_offsets is not sorted.");
+    CUGRAPH_EXPECTS(thrust::is_sorted(
+                      handle.get_thrust_policy(), subgraph_offsets.begin(), subgraph_offsets.end()),
+                    "Invalid input argument: subgraph_offsets is not sorted.");
     auto vertex_partition =
       vertex_partition_device_view_t<vertex_t, multi_gpu>(graph_view.local_vertex_partition_view());
 
     CUGRAPH_EXPECTS(
       thrust::count_if(handle.get_thrust_policy(),
-                       subgraph_vertices,
-                       subgraph_vertices + num_aggregate_subgraph_vertices,
+                       subgraph_vertices.begin(),
+                       subgraph_vertices.end(),
                        [vertex_partition] __device__(auto v) {
                          return !vertex_partition.is_valid_vertex(v) ||
                                 !vertex_partition.in_local_vertex_partition_range_nocheck(v);
@@ -101,8 +101,8 @@ extract_induced_subgraphs(
         [subgraph_offsets, subgraph_vertices] __device__(auto i) {
           // vertices are sorted and unique
           return !thrust::is_sorted(thrust::seq,
-                                    subgraph_vertices + subgraph_offsets[i],
-                                    subgraph_vertices + subgraph_offsets[i + 1]) ||
+                                    subgraph_vertices.begin() + subgraph_offsets[i],
+                                    subgraph_vertices.begin() + subgraph_offsets[i + 1]) ||
                  (thrust::count_if(
                     thrust::seq,
                     thrust::make_counting_iterator(subgraph_offsets[i]),
@@ -127,8 +127,10 @@ extract_induced_subgraphs(
     // 2-1. Phase 1: calculate memory requirements
 
     size_t num_aggregate_subgraph_vertices{};
-    raft::update_host(
-      &num_aggregate_subgraph_vertices, subgraph_offsets + num_subgraphs, 1, handle.get_stream());
+    raft::update_host(&num_aggregate_subgraph_vertices,
+                      subgraph_offsets.data() + num_subgraphs,
+                      1,
+                      handle.get_stream());
     handle.sync_stream();
 
     rmm::device_uvector<size_t> subgraph_vertex_output_offsets(
@@ -145,9 +147,10 @@ extract_induced_subgraphs(
       thrust::make_counting_iterator(num_aggregate_subgraph_vertices),
       subgraph_vertex_output_offsets.begin(),
       [subgraph_offsets, subgraph_vertices, num_subgraphs, edge_partition] __device__(auto i) {
-        auto subgraph_idx = thrust::distance(
-          subgraph_offsets + 1,
-          thrust::upper_bound(thrust::seq, subgraph_offsets, subgraph_offsets + num_subgraphs, i));
+        auto subgraph_idx =
+          thrust::distance(subgraph_offsets.begin() + 1,
+                           thrust::upper_bound(
+                             thrust::seq, subgraph_offsets.begin(), subgraph_offsets.end() - 1, i));
         vertex_t const* indices{nullptr};
         thrust::optional<weight_t const*> weights{thrust::nullopt};
         edge_t local_degree{};
@@ -158,9 +161,9 @@ extract_induced_subgraphs(
           thrust::seq,
           indices,
           indices + local_degree,
-          [vertex_first = subgraph_vertices + subgraph_offsets[subgraph_idx],
+          [vertex_first = subgraph_vertices.begin() + subgraph_offsets[subgraph_idx],
            vertex_last =
-             subgraph_vertices + subgraph_offsets[subgraph_idx + 1]] __device__(auto nbr) {
+             subgraph_vertices.begin() + subgraph_offsets[subgraph_idx + 1]] __device__(auto nbr) {
             return thrust::binary_search(thrust::seq, vertex_first, vertex_last, nbr);
           });
       });
@@ -201,9 +204,9 @@ extract_induced_subgraphs(
        edge_weights = edge_weights ? thrust::optional<weight_t*>{(*edge_weights).data()}
                                    : thrust::nullopt] __device__(auto i) {
         auto subgraph_idx = thrust::distance(
-          subgraph_offsets + 1,
+          subgraph_offsets.begin() + 1,
           thrust::upper_bound(
-            thrust::seq, subgraph_offsets, subgraph_offsets + num_subgraphs, size_t{i}));
+            thrust::seq, subgraph_offsets.begin(), subgraph_offsets.end() - 1, size_t{i}));
         vertex_t const* indices{nullptr};
         thrust::optional<weight_t const*> weights{thrust::nullopt};
         edge_t local_degree{};
@@ -219,34 +222,35 @@ extract_induced_subgraphs(
             triplet_first + local_degree,
             thrust::make_zip_iterator(thrust::make_tuple(edge_majors, edge_minors, *edge_weights)) +
               subgraph_vertex_output_offsets[i],
-            [vertex_first = subgraph_vertices + subgraph_offsets[subgraph_idx],
+            [vertex_first = subgraph_vertices.begin() + subgraph_offsets[subgraph_idx],
              vertex_last =
-               subgraph_vertices + subgraph_offsets[subgraph_idx + 1]] __device__(auto t) {
+               subgraph_vertices.begin() + subgraph_offsets[subgraph_idx + 1]] __device__(auto t) {
               return thrust::binary_search(
                 thrust::seq, vertex_first, vertex_last, thrust::get<1>(t));
             });
         } else {
           auto pair_first = thrust::make_zip_iterator(
             thrust::make_tuple(thrust::make_constant_iterator(subgraph_vertices[i]), indices));
           // FIXME: this is inefficient for high local degree vertices
-          thrust::copy_if(thrust::seq,
-                          pair_first,
-                          pair_first + local_degree,
-                          thrust::make_zip_iterator(thrust::make_tuple(edge_majors, edge_minors)) +
-                            subgraph_vertex_output_offsets[i],
-                          [vertex_first = subgraph_vertices + subgraph_offsets[subgraph_idx],
-                           vertex_last  = subgraph_vertices +
-                                         subgraph_offsets[subgraph_idx + 1]] __device__(auto t) {
-                            return thrust::binary_search(
-                              thrust::seq, vertex_first, vertex_last, thrust::get<1>(t));
-                          });
+          thrust::copy_if(
+            thrust::seq,
+            pair_first,
+            pair_first + local_degree,
+            thrust::make_zip_iterator(thrust::make_tuple(edge_majors, edge_minors)) +
+              subgraph_vertex_output_offsets[i],
+            [vertex_first = subgraph_vertices.begin() + subgraph_offsets[subgraph_idx],
+             vertex_last =
+               subgraph_vertices.begin() + subgraph_offsets[subgraph_idx + 1]] __device__(auto t) {
+              return thrust::binary_search(
+                thrust::seq, vertex_first, vertex_last, thrust::get<1>(t));
+            });
         }
       });
 
     rmm::device_uvector<size_t> subgraph_edge_offsets(num_subgraphs + 1, handle.get_stream());
     thrust::gather(handle.get_thrust_policy(),
-                   subgraph_offsets,
-                   subgraph_offsets + (num_subgraphs + 1),
+                   subgraph_offsets.begin(),
+                   subgraph_offsets.end(),
                    subgraph_vertex_output_offsets.begin(),
                    subgraph_edge_offsets.begin());
 #ifdef TIMING

cpp/src/structure/induced_subgraph_mg.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, 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,8 +25,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int32_t, float, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -36,8 +36,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int32_t, float, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -47,8 +47,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int32_t, double, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -58,8 +58,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int32_t, double, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -69,8 +69,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int64_t, float, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -80,8 +80,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int64_t, float, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -91,8 +91,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int64_t, double, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -102,8 +102,8 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int32_t, int64_t, double, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int32_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int32_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -113,8 +113,8 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int64_t, int64_t, float, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int64_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int64_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -124,8 +124,8 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int64_t, int64_t, float, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int64_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int64_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -135,8 +135,8 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int64_t, int64_t, double, true, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int64_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int64_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);
 
@@ -146,8 +146,8 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<size_t>>
 extract_induced_subgraphs(raft::handle_t const& handle,
                           graph_view_t<int64_t, int64_t, double, false, true> const& graph_view,
-                          size_t const* subgraph_offsets,
-                          int64_t const* subgraph_vertices,
+                          raft::device_span<size_t const> subgraph_offsets,
+                          raft::device_span<int64_t const> subgraph_vertices,
                           size_t num_subgraphs,
                           bool do_expensive_check);