Update primitives to better follow STL/thrust conventions in supporting predicate operations

cpp/src/centrality/betweenness_centrality_impl.cuh

@@ -17,13 +17,12 @@
 
 #include "prims/count_if_v.cuh"
 #include "prims/edge_bucket.cuh"
-#include "prims/extract_transform_e.cuh"
-#include "prims/extract_transform_v_frontier_outgoing_e.cuh"
+#include "prims/extract_transform_if_e.cuh"
 #include "prims/fill_edge_property.cuh"
 #include "prims/per_v_transform_reduce_incoming_outgoing_e.cuh"
 #include "prims/transform_e.cuh"
+#include "prims/transform_reduce_if_v_frontier_outgoing_e_by_dst.cuh"
 #include "prims/transform_reduce_v.cuh"
-#include "prims/transform_reduce_v_frontier_outgoing_e_by_dst.cuh"
 #include "prims/update_edge_src_dst_property.cuh"
 #include "prims/update_v_frontier.cuh"
 #include "prims/vertex_frontier.cuh"
@@ -49,32 +48,51 @@ namespace {
 
 template <typename vertex_t>
 struct brandes_e_op_t {
+  template <typename value_t, typename ignore_t>
+  __device__ value_t operator()(vertex_t, vertex_t, value_t src_sigma, vertex_t, ignore_t) const
+  {
+    return src_sigma;
+  }
+};
+
+template <typename vertex_t>
+struct brandes_pred_op_t {
   const vertex_t invalid_distance_{std::numeric_limits<vertex_t>::max()};
 
   template <typename value_t, typename ignore_t>
-  __device__ cuda::std::optional<value_t> operator()(
+  __device__ bool operator()(
     vertex_t, vertex_t, value_t src_sigma, vertex_t dst_distance, ignore_t) const
   {
-    return (dst_distance == invalid_distance_) ? cuda::std::make_optional(src_sigma)
-                                               : cuda::std::nullopt;
+    return (dst_distance == invalid_distance_);
   }
 };
 
 template <typename vertex_t>
 struct extract_edge_e_op_t {
-  vertex_t d{};
-
   template <typename edge_t, typename weight_t>
-  __device__ cuda::std::optional<thrust::tuple<vertex_t, vertex_t>> operator()(
+  __device__ thrust::tuple<vertex_t, vertex_t> operator()(
     vertex_t src,
     vertex_t dst,
     thrust::tuple<vertex_t, edge_t, weight_t> src_props,
     thrust::tuple<vertex_t, edge_t, weight_t> dst_props,
-    weight_t edge_centrality) const
+    weight_t) const
+  {
+    return thrust::make_tuple(src, dst);
+  }
+};
+
+template <typename vertex_t>
+struct extract_edge_pred_op_t {
+  vertex_t d{};
+
+  template <typename edge_t, typename weight_t>
+  __device__ bool operator()(vertex_t src,
+                             vertex_t dst,
+                             thrust::tuple<vertex_t, edge_t, weight_t> src_props,
+                             thrust::tuple<vertex_t, edge_t, weight_t> dst_props,
+                             weight_t) const
   {
-    return ((thrust::get<0>(dst_props) == d) && (thrust::get<0>(src_props) == (d - 1)))
-             ? cuda::std::optional<thrust::tuple<vertex_t, vertex_t>>{thrust::make_tuple(src, dst)}
-             : cuda::std::nullopt;
+    return ((thrust::get<0>(dst_props) == d) && (thrust::get<0>(src_props) == (d - 1)));
   }
 };
 
@@ -134,15 +152,16 @@ std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<edge_t>> brandes_b
     update_edge_src_property(handle, graph_view, sigmas.begin(), src_sigmas.mutable_view());
     update_edge_dst_property(handle, graph_view, distances.begin(), dst_distances.mutable_view());
 
-    auto [new_frontier, new_sigma] = cugraph::transform_reduce_v_frontier_outgoing_e_by_dst(
+    auto [new_frontier, new_sigma] = cugraph::transform_reduce_if_v_frontier_outgoing_e_by_dst(
       handle,
       graph_view,
       vertex_frontier.bucket(bucket_idx_cur),
       src_sigmas.view(),
       dst_distances.view(),
       cugraph::edge_dummy_property_t{}.view(),
       brandes_e_op_t<vertex_t>{},
-      reduce_op::plus<vertex_t>());
+      reduce_op::plus<vertex_t>(),
+      brandes_pred_op_t<vertex_t>{});
 
     auto next_frontier_bucket_indices = std::vector<size_t>{bucket_idx_next};
     update_v_frontier(handle,
@@ -348,13 +367,14 @@ void accumulate_edge_results(
     cugraph::edge_bucket_t<vertex_t, void, true, multi_gpu, true> edge_list(handle);
 
     {
-      auto [src, dst] = extract_transform_e(handle,
-                                            graph_view,
-                                            src_properties.view(),
-                                            dst_properties.view(),
-                                            centralities_view,
-                                            extract_edge_e_op_t<vertex_t>{d},
-                                            do_expensive_check);
+      auto [src, dst] = extract_transform_if_e(handle,
+                                               graph_view,
+                                               src_properties.view(),
+                                               dst_properties.view(),
+                                               centralities_view,
+                                               extract_edge_e_op_t<vertex_t>{},
+                                               extract_edge_pred_op_t<vertex_t>{d},
+                                               do_expensive_check);
 
       thrust::sort(handle.get_thrust_policy(),
                    thrust::make_zip_iterator(src.begin(), dst.begin()),

cpp/src/community/k_truss_impl.cuh

@@ -16,8 +16,7 @@
 #pragma once
 
 #include "prims/edge_bucket.cuh"
-#include "prims/extract_transform_e.cuh"
-#include "prims/extract_transform_v_frontier_outgoing_e.cuh"
+#include "prims/extract_transform_if_e.cuh"
 #include "prims/fill_edge_property.cuh"
 #include "prims/per_v_pair_dst_nbr_intersection.cuh"
 #include "prims/transform_e.cuh"
@@ -45,19 +44,6 @@
 
 namespace cugraph {
 
-template <typename vertex_t, typename edge_t>
-struct extract_weak_edges {
-  edge_t k{};
-  __device__ cuda::std::optional<thrust::tuple<vertex_t, vertex_t>> operator()(
-    vertex_t src, vertex_t dst, cuda::std::nullopt_t, cuda::std::nullopt_t, edge_t count) const
-  {
-    // No need to process edges with count == 0
-    return ((count < k - 2) && (count != 0))
-             ? cuda::std::optional<thrust::tuple<vertex_t, vertex_t>>{thrust::make_tuple(src, dst)}
-             : cuda::std::nullopt;
-  }
-};
-
 template <typename edge_t>
 struct is_k_or_greater_t {
   edge_t k{};
@@ -116,49 +102,50 @@ struct exclude_self_loop_t {
 };
 
 template <typename vertex_t, typename edge_t>
-struct extract_low_to_high_degree_edges_from_endpoints_t {
+struct extract_low_to_high_degree_edges_from_endpoints_e_op_t {
   raft::device_span<vertex_t const> srcs{};
   raft::device_span<vertex_t const> dsts{};
   raft::device_span<edge_t const> count{};
-  __device__ cuda::std::optional<thrust::tuple<vertex_t, vertex_t, edge_t>> operator()(
-    vertex_t src,
-    vertex_t dst,
-    edge_t src_out_degree,
-    edge_t dst_out_degree,
-    cuda::std::nullopt_t) const
+  __device__ thrust::tuple<vertex_t, vertex_t, edge_t> operator()(vertex_t src,
+                                                                  vertex_t dst,
+                                                                  edge_t src_out_degree,
+                                                                  edge_t dst_out_degree,
+                                                                  cuda::std::nullopt_t) const
   {
-    // FIXME: Not the most efficient way because the entire edgelist is scan just to find
-    // the direction of the edges
     auto itr = thrust::lower_bound(thrust::seq,
                                    thrust::make_zip_iterator(srcs.begin(), dsts.begin()),
                                    thrust::make_zip_iterator(srcs.end(), dsts.end()),
                                    thrust::make_tuple(src, dst));
 
-    if ((itr != thrust::make_zip_iterator(srcs.end(), dsts.end())) &&
-        (*itr == thrust::make_tuple(src, dst))) {
-      auto idx = thrust::distance(thrust::make_zip_iterator(srcs.begin(), dsts.begin()), itr);
-
-      if (src_out_degree < dst_out_degree) {
-        return cuda::std::optional<thrust::tuple<vertex_t, vertex_t, edge_t>>{
-          thrust::make_tuple(src, dst, count[idx])};
-      } else if (dst_out_degree < src_out_degree) {
-        return cuda::std::optional<thrust::tuple<vertex_t, vertex_t, edge_t>>{
-          thrust::make_tuple(dst, src, count[idx])};
-      } else {  // src_out_degree == dst_out_degree
-        if (src < dst /* tie-breaking using vertex ID */) {
-          return cuda::std::optional<thrust::tuple<vertex_t, vertex_t, edge_t>>{
-            thrust::make_tuple(src, dst, count[idx])};
-        } else {
-          return cuda::std::optional<thrust::tuple<vertex_t, vertex_t, edge_t>>{
-            thrust::make_tuple(dst, src, count[idx])};
-        }
+    auto idx = thrust::distance(thrust::make_zip_iterator(srcs.begin(), dsts.begin()), itr);
+
+    if (src_out_degree < dst_out_degree) {
+      return thrust::make_tuple(src, dst, count[idx]);
+    } else if (dst_out_degree < src_out_degree) {
+      return thrust::make_tuple(dst, src, count[idx]);
+    } else {  // src_out_degree == dst_out_degree
+      if (src < dst /* tie-breaking using vertex ID */) {
+        return thrust::make_tuple(src, dst, count[idx]);
+      } else {
+        return thrust::make_tuple(dst, src, count[idx]);
       }
-    } else {
-      return cuda::std::nullopt;
     }
   }
 };
 
+template <typename vertex_t, typename edge_t>
+struct extract_low_to_high_degree_edges_from_endpoints_pred_op_t {
+  raft::device_span<vertex_t const> srcs{};
+  raft::device_span<vertex_t const> dsts{};
+  __device__ bool operator()(vertex_t src, vertex_t dst, edge_t, edge_t, cuda::std::nullopt_t) const
+  {
+    return thrust::binary_search(thrust::seq,
+                                 thrust::make_zip_iterator(srcs.begin(), dsts.begin()),
+                                 thrust::make_zip_iterator(srcs.end(), dsts.end()),
+                                 thrust::make_tuple(src, dst));
+  }
+};
+
 }  // namespace
 
 template <typename vertex_t, typename edge_t, typename weight_t, bool multi_gpu>
@@ -322,13 +309,19 @@ k_truss(raft::handle_t const& handle,
 
     while (true) {
       // Extract weak edges
-      auto [weak_edgelist_srcs, weak_edgelist_dsts] =
-        extract_transform_e(handle,
-                            cur_graph_view,
-                            edge_src_dummy_property_t{}.view(),
-                            edge_dst_dummy_property_t{}.view(),
-                            edge_triangle_counts.view(),
-                            extract_weak_edges<vertex_t, edge_t>{k});
+      auto [weak_edgelist_srcs, weak_edgelist_dsts] = extract_transform_if_e(
+        handle,
+        cur_graph_view,
+        edge_src_dummy_property_t{}.view(),
+        edge_dst_dummy_property_t{}.view(),
+        edge_triangle_counts.view(),
+        cuda::proclaim_return_type<thrust::tuple<vertex_t, vertex_t>>(
+          [] __device__(vertex_t src, vertex_t dst, auto, auto, auto) {
+            return thrust::make_tuple(src, dst);
+          }),
+        cuda::proclaim_return_type<bool>([k] __device__(auto, auto, auto, auto, edge_t count) {
+          return ((count < k - 2) && (count != 0));
+        }));
 
       auto weak_edgelist_first =
         thrust::make_zip_iterator(weak_edgelist_srcs.begin(), weak_edgelist_dsts.begin());
@@ -518,18 +511,23 @@ k_truss(raft::handle_t const& handle,
       std::tie(std::get<0>(vertex_pair_buffer_unique),
                std::get<1>(vertex_pair_buffer_unique),
                decrease_count) =
-        extract_transform_e(
+        extract_transform_if_e(
           handle,
           cur_graph_view,
           edge_src_out_degrees.view(),
           edge_dst_out_degrees.view(),
           edge_dummy_property_t{}.view(),
-          extract_low_to_high_degree_edges_from_endpoints_t<vertex_t, edge_t>{
+          extract_low_to_high_degree_edges_from_endpoints_e_op_t<vertex_t, edge_t>{
             raft::device_span<vertex_t const>(std::get<0>(vertex_pair_buffer_unique).data(),
                                               std::get<0>(vertex_pair_buffer_unique).size()),
             raft::device_span<vertex_t const>(std::get<1>(vertex_pair_buffer_unique).data(),
                                               std::get<1>(vertex_pair_buffer_unique).size()),
-            raft::device_span<edge_t const>(decrease_count.data(), decrease_count.size())});
+            raft::device_span<edge_t const>(decrease_count.data(), decrease_count.size())},
+          extract_low_to_high_degree_edges_from_endpoints_pred_op_t<vertex_t, edge_t>{
+            raft::device_span<vertex_t const>(std::get<0>(vertex_pair_buffer_unique).data(),
+                                              std::get<0>(vertex_pair_buffer_unique).size()),
+            raft::device_span<vertex_t const>(std::get<1>(vertex_pair_buffer_unique).data(),
+                                              std::get<1>(vertex_pair_buffer_unique).size())});
 
       if constexpr (multi_gpu) {
         auto& comm           = handle.get_comms();

cpp/src/community/triangle_count_impl.cuh

@@ -15,7 +15,7 @@
  */
 #pragma once
 
-#include "prims/extract_transform_e.cuh"
+#include "prims/extract_transform_if_e.cuh"
 #include "prims/fill_edge_property.cuh"
 #include "prims/transform_e.cuh"
 #include "prims/transform_reduce_dst_nbr_intersection_of_e_endpoints_by_v.cuh"
@@ -64,21 +64,30 @@ struct is_two_or_greater_t {
 };
 
 template <typename vertex_t, typename edge_t>
-struct extract_low_to_high_degree_edges_t {
-  __device__ cuda::std::optional<thrust::tuple<vertex_t, vertex_t>> operator()(
-    vertex_t src,
-    vertex_t dst,
-    edge_t src_out_degree,
-    edge_t dst_out_degree,
-    cuda::std::nullopt_t) const
+struct extract_low_to_high_degree_edges_e_op_t {
+  __device__ thrust::tuple<vertex_t, vertex_t> operator()(vertex_t src,
+                                                          vertex_t dst,
+                                                          edge_t src_out_degree,
+                                                          edge_t dst_out_degree,
+                                                          cuda::std::nullopt_t) const
   {
-    return (src_out_degree < dst_out_degree)
-             ? cuda::std::optional<thrust::tuple<vertex_t, vertex_t>>{thrust::make_tuple(src, dst)}
-             : (((src_out_degree == dst_out_degree) &&
-                 (src < dst) /* tie-breaking using vertex ID */)
-                  ? cuda::std::optional<thrust::tuple<vertex_t, vertex_t>>{thrust::make_tuple(src,
-                                                                                              dst)}
-                  : cuda::std::nullopt);
+    return thrust::make_tuple(src, dst);
+  }
+};
+
+template <typename vertex_t, typename edge_t>
+struct extract_low_to_high_degree_edges_pred_op_t {
+  __device__ bool operator()(vertex_t src,
+                             vertex_t dst,
+                             edge_t src_out_degree,
+                             edge_t dst_out_degree,
+                             cuda::std::nullopt_t) const
+  {
+    return (src_out_degree < dst_out_degree) ? true
+                                             : (((src_out_degree == dst_out_degree) &&
+                                                 (src < dst) /* tie-breaking using vertex ID */)
+                                                  ? true
+                                                  : false);
   }
 };
 
@@ -433,12 +442,14 @@ void triangle_count(raft::handle_t const& handle,
       handle, cur_graph_view, out_degrees.begin(), edge_src_out_degrees.mutable_view());
     update_edge_dst_property(
       handle, cur_graph_view, out_degrees.begin(), edge_dst_out_degrees.mutable_view());
-    auto [srcs, dsts] = extract_transform_e(handle,
-                                            cur_graph_view,
-                                            edge_src_out_degrees.view(),
-                                            edge_dst_out_degrees.view(),
-                                            edge_dummy_property_t{}.view(),
-                                            extract_low_to_high_degree_edges_t<vertex_t, edge_t>{});
+    auto [srcs, dsts] =
+      extract_transform_if_e(handle,
+                             cur_graph_view,
+                             edge_src_out_degrees.view(),
+                             edge_dst_out_degrees.view(),
+                             edge_dummy_property_t{}.view(),
+                             extract_low_to_high_degree_edges_e_op_t<vertex_t, edge_t>{},
+                             extract_low_to_high_degree_edges_pred_op_t<vertex_t, edge_t>{});
 
     if constexpr (multi_gpu) {
       std::tie(