Move edge triangle count to the stable API

cpp/include/cugraph/algorithms.hpp

@@ -2007,6 +2007,10 @@ void triangle_count(raft::handle_t const& handle,
                     raft::device_span<edge_t> counts,
                     bool do_expensive_check = false);
 
+template <typename vertex_t, typename edge_t, bool multi_gpu>
+rmm::device_uvector<edge_t> edge_triangle_count(
+  raft::handle_t const& handle, graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view);
+
 /*
  * @brief Compute K-Truss.
  *

cpp/src/community/edge_triangle_count_impl.cuh

@@ -34,8 +34,9 @@ namespace detail {
 
 template <typename vertex_t, typename edge_t, typename EdgeIterator>
 struct update_edges_p_r_q_r_num_triangles {
-  size_t num_edges{};  // rename to num_edges
+  size_t num_edges{};
   const edge_t edge_first_or_second{};
+  size_t chunk_start{};
   raft::device_span<size_t const> intersection_offsets{};
   raft::device_span<vertex_t const> intersection_indices{};
   raft::device_span<edge_t> num_triangles{};
@@ -48,28 +49,22 @@ struct update_edges_p_r_q_r_num_triangles {
       thrust::seq, intersection_offsets.begin() + 1, intersection_offsets.end(), i);
     auto idx = thrust::distance(intersection_offsets.begin() + 1, itr);
     if (edge_first_or_second == 0) {
-      auto p_r_pair =
-        thrust::make_tuple(thrust::get<0>(*(edge_first + idx)), intersection_indices[i]);
+      auto p_r_pair = thrust::make_tuple(thrust::get<0>(*(edge_first + chunk_start + idx)),
+                                         intersection_indices[i]);
 
       // Find its position in 'edges'
       auto itr_p_r_p_q =
-        thrust::lower_bound(thrust::seq,
-                            edge_first,
-                            edge_first + num_edges,  // pass the number of vertex pairs
-                            p_r_pair);
+        thrust::lower_bound(thrust::seq, edge_first, edge_first + num_edges, p_r_pair);
 
       assert(*itr_p_r_p_q == p_r_pair);
       idx = thrust::distance(edge_first, itr_p_r_p_q);
     } else {
-      auto p_r_pair =
-        thrust::make_tuple(thrust::get<1>(*(edge_first + idx)), intersection_indices[i]);
+      auto p_r_pair = thrust::make_tuple(thrust::get<1>(*(edge_first + chunk_start + idx)),
+                                         intersection_indices[i]);
 
       // Find its position in 'edges'
       auto itr_p_r_p_q =
-        thrust::lower_bound(thrust::seq,
-                            edge_first,
-                            edge_first + num_edges,  // pass the number of vertex pairs
-                            p_r_pair);
+        thrust::lower_bound(thrust::seq, edge_first, edge_first + num_edges, p_r_pair);
       assert(*itr_p_r_p_q == p_r_pair);
       idx = thrust::distance(edge_first, itr_p_r_p_q);
     }
@@ -81,74 +76,103 @@ struct update_edges_p_r_q_r_num_triangles {
 template <typename vertex_t, typename edge_t, bool store_transposed, bool multi_gpu>
 std::enable_if_t<!multi_gpu, rmm::device_uvector<edge_t>> edge_triangle_count_impl(
   raft::handle_t const& handle,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  raft::device_span<vertex_t> edgelist_srcs,
-  raft::device_span<vertex_t> edgelist_dsts)
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view)
 {
+  using weight_t = float;
+  rmm::device_uvector<vertex_t> edgelist_srcs(0, handle.get_stream());
+  rmm::device_uvector<vertex_t> edgelist_dsts(0, handle.get_stream());
+  std::tie(edgelist_srcs, edgelist_dsts, std::ignore, std::ignore) = decompress_to_edgelist(
+    handle,
+    graph_view,
+    std::optional<edge_property_view_t<edge_t, weight_t const*>>{std::nullopt},
+    std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+    std::optional<raft::device_span<vertex_t const>>(std::nullopt));
+
   auto edge_first = thrust::make_zip_iterator(edgelist_srcs.begin(), edgelist_dsts.begin());
 
   thrust::sort(handle.get_thrust_policy(), edge_first, edge_first + edgelist_srcs.size());
 
-  // FIXME: Perform 'nbr_intersection' in chunks to reduce peak memory.
-  auto [intersection_offsets, intersection_indices] =
-    detail::nbr_intersection(handle,
-                             graph_view,
-                             cugraph::edge_dummy_property_t{}.view(),
-                             edge_first,
-                             edge_first + edgelist_srcs.size(),
-                             std::array<bool, 2>{true, true},
-                             false /*FIXME: pass 'do_expensive_check' as argument*/);
+  size_t approx_edges_to_intersect_per_iteration =
+    static_cast<size_t>(handle.get_device_properties().multiProcessorCount) * (1 << 17);
 
+  auto num_chunks        = ((edgelist_srcs.size() % approx_edges_to_intersect_per_iteration) == 0)
+                             ? (edgelist_srcs.size() / approx_edges_to_intersect_per_iteration)
+                             : (edgelist_srcs.size() / approx_edges_to_intersect_per_iteration) + 1;
+  size_t prev_chunk_size = 0;
+  auto num_edges         = edgelist_srcs.size();
   rmm::device_uvector<edge_t> num_triangles(edgelist_srcs.size(), handle.get_stream());
 
-  // Update the number of triangles of each (p, q) edges by looking at their intersection
-  // size
-  thrust::adjacent_difference(handle.get_thrust_policy(),
-                              intersection_offsets.begin() + 1,
-                              intersection_offsets.end(),
-                              num_triangles.begin());
-
-  // Given intersection offsets and indices that are used to update the number of
-  // triangles of (p, q) edges where `r`s are the intersection indices, update
-  // the number of triangles of the pairs (p, r) and (q, r).
-
-  thrust::for_each(
-    handle.get_thrust_policy(),
-    thrust::make_counting_iterator<edge_t>(0),
-    thrust::make_counting_iterator<edge_t>(intersection_indices.size()),
-    update_edges_p_r_q_r_num_triangles<vertex_t, edge_t, decltype(edge_first)>{
-      edgelist_srcs.size(),
-      0,
-      raft::device_span<size_t const>(intersection_offsets.data(), intersection_offsets.size()),
-      raft::device_span<vertex_t const>(intersection_indices.data(), intersection_indices.size()),
-      raft::device_span<edge_t>(num_triangles.data(), num_triangles.size()),
-      edge_first});
-
-  thrust::for_each(
-    handle.get_thrust_policy(),
-    thrust::make_counting_iterator<edge_t>(0),
-    thrust::make_counting_iterator<edge_t>(intersection_indices.size()),
-    update_edges_p_r_q_r_num_triangles<vertex_t, edge_t, decltype(edge_first)>{
-      edgelist_srcs.size(),
-      1,
-      raft::device_span<size_t const>(intersection_offsets.data(), intersection_offsets.size()),
-      raft::device_span<vertex_t const>(intersection_indices.data(), intersection_indices.size()),
-      raft::device_span<edge_t>(num_triangles.data(), num_triangles.size()),
-      edge_first});
+  // Need to ensure that the vector has its values initialized to 0 before incrementing
+  thrust::fill(handle.get_thrust_policy(), num_triangles.begin(), num_triangles.end(), 0);
+
+  for (size_t i = 0; i < num_chunks; ++i) {
+    auto chunk_size = std::min(approx_edges_to_intersect_per_iteration, num_edges);
+    num_edges -= chunk_size;
+    // Perform 'nbr_intersection' in chunks to reduce peak memory.
+    auto [intersection_offsets, intersection_indices] =
+      detail::nbr_intersection(handle,
+                               graph_view,
+                               cugraph::edge_dummy_property_t{}.view(),
+                               edge_first + prev_chunk_size,
+                               edge_first + prev_chunk_size + chunk_size,
+                               std::array<bool, 2>{true, true},
+                               false /*FIXME: pass 'do_expensive_check' as argument*/);
+
+    // Update the number of triangles of each (p, q) edges by looking at their intersection
+    // size
+    thrust::for_each(
+      handle.get_thrust_policy(),
+      thrust::make_counting_iterator<edge_t>(0),
+      thrust::make_counting_iterator<edge_t>(chunk_size),
+      [chunk_start          = prev_chunk_size,
+       num_triangles        = raft::device_span<edge_t>(num_triangles.data(), num_triangles.size()),
+       intersection_offsets = raft::device_span<size_t const>(
+         intersection_offsets.data(), intersection_offsets.size())] __device__(auto i) {
+        num_triangles[chunk_start + i] += (intersection_offsets[i + 1] - intersection_offsets[i]);
+      });
+
+    // Given intersection offsets and indices that are used to update the number of
+    // triangles of (p, q) edges where `r`s are the intersection indices, update
+    // the number of triangles of the pairs (p, r) and (q, r).
+    thrust::for_each(
+      handle.get_thrust_policy(),
+      thrust::make_counting_iterator<edge_t>(0),
+      thrust::make_counting_iterator<edge_t>(intersection_indices.size()),
+      update_edges_p_r_q_r_num_triangles<vertex_t, edge_t, decltype(edge_first)>{
+        edgelist_srcs.size(),
+        0,
+        prev_chunk_size,
+        raft::device_span<size_t const>(intersection_offsets.data(), intersection_offsets.size()),
+        raft::device_span<vertex_t const>(intersection_indices.data(), intersection_indices.size()),
+        raft::device_span<edge_t>(num_triangles.data(), num_triangles.size()),
+        edge_first});
+
+    thrust::for_each(
+      handle.get_thrust_policy(),
+      thrust::make_counting_iterator<edge_t>(0),
+      thrust::make_counting_iterator<edge_t>(intersection_indices.size()),
+      update_edges_p_r_q_r_num_triangles<vertex_t, edge_t, decltype(edge_first)>{
+        edgelist_srcs.size(),
+        1,
+        prev_chunk_size,
+        raft::device_span<size_t const>(intersection_offsets.data(), intersection_offsets.size()),
+        raft::device_span<vertex_t const>(intersection_indices.data(), intersection_indices.size()),
+        raft::device_span<edge_t>(num_triangles.data(), num_triangles.size()),
+        edge_first});
+
+    prev_chunk_size += chunk_size;
+  }
 
   return num_triangles;
 }
 
 }  // namespace detail
 
-template <typename vertex_t, typename edge_t, bool store_transposed, bool multi_gpu>
+template <typename vertex_t, typename edge_t, bool multi_gpu>
 rmm::device_uvector<edge_t> edge_triangle_count(
-  raft::handle_t const& handle,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  raft::device_span<vertex_t> edgelist_srcs,
-  raft::device_span<vertex_t> edgelist_dsts)
+  raft::handle_t const& handle, graph_view_t<vertex_t, edge_t, false, multi_gpu> const& graph_view)
 {
-  return detail::edge_triangle_count_impl(handle, graph_view, edgelist_srcs, edgelist_dsts);
+  return detail::edge_triangle_count_impl(handle, graph_view);
 }
 
 }  // namespace cugraph

cpp/src/community/edge_triangle_count_sg.cu

@@ -20,20 +20,14 @@ namespace cugraph {
 // SG instantiation
 template rmm::device_uvector<int32_t> edge_triangle_count(
   raft::handle_t const& handle,
-  cugraph::graph_view_t<int32_t, int32_t, false, false> const& graph_view,
-  raft::device_span<int32_t> edgelist_srcs,
-  raft::device_span<int32_t> edgelist_dsts);
+  cugraph::graph_view_t<int32_t, int32_t, false, false> const& graph_view);
 
 template rmm::device_uvector<int64_t> edge_triangle_count(
   raft::handle_t const& handle,
-  cugraph::graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  raft::device_span<int32_t> edgelist_srcs,
-  raft::device_span<int32_t> edgelist_dsts);
+  cugraph::graph_view_t<int32_t, int64_t, false, false> const& graph_view);
 
 template rmm::device_uvector<int64_t> edge_triangle_count(
   raft::handle_t const& handle,
-  cugraph::graph_view_t<int64_t, int64_t, false, false> const& graph_view,
-  raft::device_span<int64_t> edgelist_srcs,
-  raft::device_span<int64_t> edgelist_dsts);
+  cugraph::graph_view_t<int64_t, int64_t, false, false> const& graph_view);
 
 }  // namespace cugraph