Enable edge masking in additional primitives

cpp/include/cugraph/edge_partition_device_view.cuh

@@ -17,6 +17,7 @@
 
 #include <cugraph/edge_partition_view.hpp>
 #include <cugraph/utilities/error.hpp>
+#include <cugraph/utilities/mask_utils.cuh>
 #include <cugraph/utilities/misc_utils.cuh>
 
 #include <raft/core/device_span.hpp>
@@ -92,6 +93,54 @@ struct local_degree_op_t {
   }
 };
 
+template <typename vertex_t,
+          typename edge_t,
+          typename return_type_t,
+          bool multi_gpu,
+          bool use_dcs,
+          typename MaskIterator>
+struct local_degree_with_mask_op_t {
+  raft::device_span<edge_t const> offsets{};
+  std::conditional_t<multi_gpu, vertex_t, std::byte /* dummy */> major_range_first{};
+
+  std::conditional_t<use_dcs, raft::device_span<vertex_t const>, std::byte /* dummy */>
+    dcs_nzd_vertices{};
+  std::conditional_t<use_dcs, vertex_t, std::byte /* dummy */> major_hypersparse_first{};
+
+  MaskIterator mask_first{};
+
+  __device__ return_type_t operator()(vertex_t major) const
+  {
+    if constexpr (multi_gpu) {
+      vertex_t idx{};
+      if constexpr (use_dcs) {
+        if (major < major_hypersparse_first) {
+          idx = major - major_range_first;
+          return static_cast<return_type_t>(
+            count_set_bits(mask_first, offsets[idx], offsets[idx + 1] - offsets[idx]));
+        } else {
+          auto major_hypersparse_idx =
+            major_hypersparse_idx_from_major_nocheck_impl(dcs_nzd_vertices, major);
+          if (major_hypersparse_idx) {
+            idx = (major_hypersparse_first - major_range_first) + *major_hypersparse_idx;
+            return static_cast<return_type_t>(
+              count_set_bits(mask_first, offsets[idx], offsets[idx + 1] - offsets[idx]));
+          } else {
+            return return_type_t{0};
+          }
+        }
+      } else {
+        idx = major - major_range_first;
+        return static_cast<return_type_t>(
+          count_set_bits(mask_first, offsets[idx], offsets[idx + 1] - offsets[idx]));
+      }
+    } else {
+      return static_cast<return_type_t>(
+        count_set_bits(mask_first, offsets[major], offsets[major + 1] - offsets[major]));
+    }
+  }
+};
+
 template <typename vertex_t, typename edge_t>
 class edge_partition_device_view_base_t {
  public:
@@ -255,6 +304,122 @@ class edge_partition_device_view_t<vertex_t, edge_t, multi_gpu, std::enable_if_t
     return local_degrees;
   }
 
+  template <typename MaskIterator, typename MajorIterator>
+  size_t compute_number_of_edges_with_mask(MaskIterator mask_first,
+                                           MajorIterator major_first,
+                                           MajorIterator major_last,
+                                           rmm::cuda_stream_view stream) const
+  {
+    return dcs_nzd_vertices_ ? thrust::transform_reduce(
+                                 rmm::exec_policy(stream),
+                                 major_first,
+                                 major_last,
+                                 detail::local_degree_with_mask_op_t<
+                                   vertex_t,
+                                   edge_t,
+                                   size_t /* no limit on majors.size(), so edge_t can overflow */,
+                                   multi_gpu,
+                                   true,
+                                   MaskIterator>{this->offsets_,
+                                                 major_range_first_,
+                                                 *dcs_nzd_vertices_,
+                                                 *major_hypersparse_first_,
+                                                 mask_first},
+                                 size_t{0},
+                                 thrust::plus<size_t>())
+                             : thrust::transform_reduce(
+                                 rmm::exec_policy(stream),
+                                 major_first,
+                                 major_last,
+                                 detail::local_degree_with_mask_op_t<
+                                   vertex_t,
+                                   edge_t,
+                                   size_t /* no limit on majors.size(), so edge_t can overflow */,
+                                   multi_gpu,
+                                   false,
+                                   MaskIterator>{this->offsets_,
+                                                 major_range_first_,
+                                                 std::byte{0} /* dummy */,
+                                                 std::byte{0} /* dummy */,
+                                                 mask_first},
+                                 size_t{0},
+                                 thrust::plus<size_t>());
+  }
+
+  template <typename MaskIterator>
+  rmm::device_uvector<edge_t> compute_local_degrees_with_mask(MaskIterator mask_first,
+                                                              rmm::cuda_stream_view stream) const
+  {
+    rmm::device_uvector<edge_t> local_degrees(this->major_range_size(), stream);
+    if (dcs_nzd_vertices_) {
+      assert(major_hypersparse_first_);
+      thrust::transform(
+        rmm::exec_policy(stream),
+        thrust::make_counting_iterator(this->major_range_first()),
+        thrust::make_counting_iterator(this->major_range_last()),
+        local_degrees.begin(),
+        detail::
+          local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, true, MaskIterator>{
+            this->offsets_,
+            major_range_first_,
+            *dcs_nzd_vertices_,
+            major_hypersparse_first_.value_or(vertex_t{0}),
+            mask_first});
+    } else {
+      thrust::transform(
+        rmm::exec_policy(stream),
+        thrust::make_counting_iterator(this->major_range_first()),
+        thrust::make_counting_iterator(this->major_range_last()),
+        local_degrees.begin(),
+        detail::
+          local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, false, MaskIterator>{
+            this->offsets_,
+            major_range_first_,
+            std::byte{0} /* dummy */,
+            std::byte{0} /* dummy */,
+            mask_first});
+    }
+    return local_degrees;
+  }
+
+  template <typename MaskIterator, typename MajorIterator>
+  rmm::device_uvector<edge_t> compute_local_degrees_with_mask(MaskIterator mask_first,
+                                                              MajorIterator major_first,
+                                                              MajorIterator major_last,
+                                                              rmm::cuda_stream_view stream) const
+  {
+    rmm::device_uvector<edge_t> local_degrees(thrust::distance(major_first, major_last), stream);
+    if (dcs_nzd_vertices_) {
+      assert(major_hypersparse_first_);
+      thrust::transform(
+        rmm::exec_policy(stream),
+        major_first,
+        major_last,
+        local_degrees.begin(),
+        detail::
+          local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, true, MaskIterator>{
+            this->offsets_,
+            major_range_first_,
+            dcs_nzd_vertices_.value(),
+            major_hypersparse_first_.value_or(vertex_t{0}),
+            mask_first});
+    } else {
+      thrust::transform(
+        rmm::exec_policy(stream),
+        major_first,
+        major_last,
+        local_degrees.begin(),
+        detail::
+          local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, false, MaskIterator>{
+            this->offsets_,
+            major_range_first_,
+            std::byte{0} /* dummy */,
+            std::byte{0} /* dummy */,
+            mask_first});
+    }
+    return local_degrees;
+  }
+
   __host__ __device__ vertex_t major_value_start_offset() const
   {
     return major_value_start_offset_;
@@ -440,6 +605,71 @@ class edge_partition_device_view_t<vertex_t, edge_t, multi_gpu, std::enable_if_t
     return local_degrees;
   }
 
+  template <typename MaskIterator, typename MajorIterator>
+  size_t compute_number_of_edges_with_mask(MaskIterator mask_first,
+                                           MajorIterator major_first,
+                                           MajorIterator major_last,
+                                           rmm::cuda_stream_view stream) const
+  {
+    return thrust::transform_reduce(
+      rmm::exec_policy(stream),
+      major_first,
+      major_last,
+      detail::local_degree_with_mask_op_t<
+        vertex_t,
+        edge_t,
+        size_t /* no limit on majors.size(), so edge_t can overflow */,
+        multi_gpu,
+        false,
+        MaskIterator>{this->offsets_,
+                      std::byte{0} /* dummy */,
+                      std::byte{0} /* dummy */,
+                      std::byte{0} /* dummy */,
+                      mask_first},
+      size_t{0},
+      thrust::plus<size_t>());
+  }
+
+  template <typename MaskIterator>
+  rmm::device_uvector<edge_t> compute_local_degrees_with_mask(MaskIterator mask_first,
+                                                              rmm::cuda_stream_view stream) const
+  {
+    rmm::device_uvector<edge_t> local_degrees(this->major_range_size(), stream);
+    thrust::transform(
+      rmm::exec_policy(stream),
+      thrust::make_counting_iterator(this->major_range_first()),
+      thrust::make_counting_iterator(this->major_range_last()),
+      local_degrees.begin(),
+      detail::local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, false, MaskIterator>{
+        this->offsets_,
+        std::byte{0} /* dummy */,
+        std::byte{0} /* dummy */,
+        std::byte{0} /* dummy */,
+        mask_first});
+    return local_degrees;
+  }
+
+  template <typename MaskIterator, typename MajorIterator>
+  rmm::device_uvector<edge_t> compute_local_degrees_with_mask(MaskIterator mask_first,
+                                                              MajorIterator major_first,
+                                                              MajorIterator major_last,
+                                                              rmm::cuda_stream_view stream) const
+  {
+    rmm::device_uvector<edge_t> local_degrees(thrust::distance(major_first, major_last), stream);
+    thrust::transform(
+      rmm::exec_policy(stream),
+      major_first,
+      major_last,
+      local_degrees.begin(),
+      detail::local_degree_with_mask_op_t<vertex_t, edge_t, edge_t, multi_gpu, false, MaskIterator>{
+        this->offsets_,
+        std::byte{0} /* dummy */,
+        std::byte{0} /* dummy */,
+        std::byte{0} /* dummy */,
+        mask_first});
+    return local_degrees;
+  }
+
   __host__ __device__ vertex_t major_value_start_offset() const { return vertex_t{0}; }
 
   __host__ __device__ thrust::optional<vertex_t> major_hypersparse_first() const noexcept

cpp/include/cugraph/utilities/mask_utils.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-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.
@@ -25,6 +25,7 @@
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/transform.h>
+#include <thrust/transform_reduce.h>
 
 namespace cugraph {
 

cpp/src/centrality/betweenness_centrality_impl.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-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.
@@ -69,7 +69,7 @@ struct extract_edge_e_op_t {
     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)
+    weight_t edge_centrality) const
   {
     return ((thrust::get<0>(dst_props) == d) && (thrust::get<0>(src_props) == (d - 1)))
              ? thrust::optional<thrust::tuple<vertex_t, vertex_t>>{thrust::make_tuple(src, dst)}