Bipartite R-mat graph generation.

cpp/include/cugraph/graph_generators.hpp

@@ -127,6 +127,45 @@ std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>> generat
   double c           = 0.19,
   bool clip_and_flip = false);
 
+/**
+ * @brief generate an edge list for a bi-partite R-mat graph.
+ *
+ * The source vertex IDs will be in the range of [0, 2^src_scale) and the destination vertex IDs
+ * will be in the range of [0, 2^dst_scale). This function allows multi-edges.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param rng_state RAFT RNG state, updated with each call
+ * @param src_scale Scale factor to set the range of source vertex IDs (or the first vertex set) in
+ * the bi-partite graph. Vertex IDs have values in [0, V_src), where V_src = 1 << @p src_scale.
+ * @param dst_scale Scale factor to set the range of destination vertex IDs (or the second vertex
+ * set) in the bi-partite graph. Vertex IDs have values in [0, V_dst), where V_dst = 1 << @p
+ * dst_scale.
+ * @param num_edges Number of edges to generate.
+ * @param a a, b, c, d (= 1.0 - (a + b + c)) in the R-mat graph generator (vist https://graph500.org
+ * for additional details). a, b, c, d should be non-negative and a + b + c should be no larger
+ * than 1.0.
+ * @param b a, b, c, d (= 1.0 - (a + b + c)) in the R-mat graph generator (vist https://graph500.org
+ * for additional details). a, b, c, d should be non-negative and a + b + c should be no larger
+ * than 1.0.
+ * @param c a, b, c, d (= 1.0 - (a + b + c)) in the R-mat graph generator (vist https://graph500.org
+ * for additional details). a, b, c, d should be non-negative and a + b + c should be no larger
+ * than 1.0.
+ * @return std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>> A tuple of
+ * rmm::device_uvector objects for edge source vertex IDs and edge destination vertex IDs.
+ */
+template <typename vertex_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>>
+generate_bipartite_rmat_edgelist(raft::handle_t const& handle,
+                                 raft::random::RngState& rng_state,
+                                 size_t src_scale,
+                                 size_t dst_scale,
+                                 size_t num_edges,
+                                 double a = 0.57,
+                                 double b = 0.19,
+                                 double c = 0.19);
+
 enum class generator_distribution_t { POWER_LAW = 0, UNIFORM };
 
 /**

cpp/src/generators/generate_rmat_edgelist.cu

@@ -44,6 +44,7 @@ std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>> generat
   double c,
   bool clip_and_flip)
 {
+  CUGRAPH_EXPECTS(scale >= 1, "Invalid input argument; scale should be at least 1.");
   CUGRAPH_EXPECTS((size_t{1} << scale) <= static_cast<size_t>(std::numeric_limits<vertex_t>::max()),
                   "Invalid input argument: scale too large for vertex_t.");
   CUGRAPH_EXPECTS((a >= 0.0) && (b >= 0.0) && (c >= 0.0) && (a + b + c <= 1.0),
@@ -212,6 +213,83 @@ generate_rmat_edgelists(raft::handle_t const& handle,
                                            clip_and_flip);
 }
 
+template <typename vertex_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<vertex_t>> generate_bipartite_rmat_edgelist(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  size_t src_scale,
+  size_t dst_scale,
+  size_t num_edges,
+  double a,
+  double b,
+  double c)
+{
+  CUGRAPH_EXPECTS(src_scale >= 1 && dst_scale >= 1, "Invalid input argument; src_scale & dst_scale should be at least 1.");
+  CUGRAPH_EXPECTS((size_t{1} << src_scale) <= static_cast<size_t>(std::numeric_limits<vertex_t>::max()),
+                  "Invalid input argument: src_scale too large for vertex_t.");
+  CUGRAPH_EXPECTS((size_t{1} << dst_scale) <= static_cast<size_t>(std::numeric_limits<vertex_t>::max()),
+                  "Invalid input argument: dst_scale too large for vertex_t.");
+  CUGRAPH_EXPECTS((a >= 0.0) && (b >= 0.0) && (c >= 0.0) && (a + b + c <= 1.0),
+                  "Invalid input argument: a, b, c should be non-negative and a + b + c should not "
+                  "be larger than 1.0.");
+
+  // to limit memory footprint (1024 is a tuning parameter)
+  auto max_edges_to_generate_per_iteration =
+    static_cast<size_t>(handle.get_device_properties().multiProcessorCount) * 1024;
+  rmm::device_uvector<float> rands(
+    std::min(num_edges, max_edges_to_generate_per_iteration) * (src_scale + dst_scale), handle.get_stream());
+
+  rmm::device_uvector<vertex_t> srcs(num_edges, handle.get_stream());
+  rmm::device_uvector<vertex_t> dsts(num_edges, handle.get_stream());
+
+  size_t num_edges_generated{0};
+  while (num_edges_generated < num_edges) {
+    auto num_edges_to_generate =
+      std::min(num_edges - num_edges_generated, max_edges_to_generate_per_iteration);
+    auto pair_first = thrust::make_zip_iterator(thrust::make_tuple(srcs.begin(), dsts.begin())) +
+                      num_edges_generated;
+
+    detail::uniform_random_fill(
+      handle.get_stream(), rands.data(), num_edges_to_generate * (src_scale + dst_scale), 0.0f, 1.0f, rng_state);
+
+    thrust::transform(
+      handle.get_thrust_policy(),
+      thrust::make_counting_iterator(size_t{0}),
+      thrust::make_counting_iterator(num_edges_to_generate),
+      pair_first,
+      // if a + b == 0.0, a_norm is irrelevant, if (1.0 - (a+b)) == 0.0, c_norm is irrelevant
+      [src_scale,
+       dst_scale,
+       rands    = rands.data(),
+       a_plus_b = a + b,
+       a_plus_c = a + c,
+       a_norm   = (a + b) > 0.0 ? a / (a + b) : 0.0,
+       c_norm   = (1.0 - (a + b)) > 0.0 ? c / (1.0 - (a + b)) : 0.0] __device__(auto i) {
+        vertex_t src{0};
+        vertex_t dst{0};
+        size_t rand_offset = i * (src_scale + dst_scale);
+        for (int level = 0; level < static_cast<int>(std::max(src_scale, dst_scale)); ++level) {
+          auto dst_threshold = a_plus_c;
+          if (level < src_scale) {
+            auto r = rands[rand_offset++];
+            auto src_bit_set = r > a_plus_b;
+            src += src_bit_set ? static_cast<vertex_t>(vertex_t{1} << (src_scale - (level + 1))) : 0;
+            dst_threshold = src_bit_set ? c_nrom : a_norm;
+          }
+          if (level < dst_scale) {
+            auto r = rands[rand_offset++];
+            auto dst_bit_set = r > dst_threshold;
+            dst += dst_bit_set ? static_cast<vertex_t>(vertex_t{1} << (dst_scale - (level + 1))) : 0;
+          }
+        }
+        return thrust::make_tuple(src, dst);
+      });
+    num_edges_generated += num_edges_to_generate;
+  }
+
+  return std::make_tuple(std::move(srcs), std::move(dsts));
+}
+
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
 generate_rmat_edgelist<int32_t>(raft::handle_t const& handle,
                                 raft::random::RngState& rng_state,
@@ -254,6 +332,28 @@ generate_rmat_edgelists<int64_t>(raft::handle_t const& handle,
                                  generator_distribution_t edge_distribution,
                                  bool clip_and_flip);
 
+template <typename int32_t>
+std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>> generate_bipartite_rmat_edgelist(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  size_t src_scale,
+  size_t dst_scale,
+  size_t num_edges,
+  double a,
+  double b,
+  double c);
+
+template <typename int64_t>
+std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int64_t>> generate_bipartite_rmat_edgelist(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  size_t src_scale,
+  size_t dst_scale,
+  size_t num_edges,
+  double a,
+  double b,
+  double c);
+
 /* DEPRECATED */
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
 generate_rmat_edgelist<int32_t>(raft::handle_t const& handle,