Skip to content

Commit

Permalink
mdspan-ify rmat_rectangular_gen
Browse files Browse the repository at this point in the history 
I added two mdspan overloads of rmat_rectangular_gen,
each corresponding to one of the existing overloads.
I also added tests for the two overloads.
Both tests build and pass.
  • Loading branch information
@mhoemmen
mhoemmen committed Sep 20, 2022
1 parent 0741542 commit 039a3b1
Show file tree
Hide file tree
Showing 2 changed files with 209 additions and 11 deletions.
118 changes: 107 additions & 11 deletions cpp/include/raft/random/rmat_rectangular_generator.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,51 +18,117 @@

#include "detail/rmat_rectangular_generator.cuh"

#include <raft/mdarray.hpp>

namespace raft::random {

/**
* @brief Generate RMAT for a rectangular shaped adjacency matrices (useful when
* @brief Generate RMAT for a rectangular adjacency matrix (useful when
* graphs to be generated are bipartite)
*
* @tparam IdxT node indices type
* @tparam IdxT type of each node index
* @tparam ProbT data type used for probability distributions (either fp32 or fp64)
*
* @param[out] out generated edgelist [on device] [dim = n_edges x 2]. On each row
* the first element corresponds to the source node id while the
* second, the destination node id. If you don't need this output
* @param[in] handle RAFT handle, containing the CUDA stream on which to schedule work
* @param[in] r underlying state of the random generator. Especially useful when
* one wants to call this API for multiple times in order to generate
* a larger graph. For that case, just create this object with the
* initial seed once and after every call continue to pass the same
* object for the successive calls.
* @param[out] out generated edgelist [on device] [dim = n_edges x 2]. In each row
* the first element is the source node id, and the second element
* is the destination node id. If you don't need this output
* then pass a `nullptr` in its place.
* @param[out] out_src list of source node id's [on device] [len = n_edges]. If you
* don't need this output then pass a `nullptr` in its place.
* @param[out] out_dst list of destination node id's [on device] [len = n_edges]. If
* you don't need this output then pass a `nullptr` in its place.
* @param[in] theta distribution of each quadrant at each level of resolution.
* Since these are probabilities, each of the 2x2 matrix for
* Since these are probabilities, each of the 2x2 matrices for
* each level of the RMAT must sum to one. [on device]
* [dim = max(r_scale, c_scale) x 2 x 2]. Of course, it is assumed
* that each of the group of 2 x 2 numbers all sum up to 1.
* @param[in] r_scale 2^r_scale represents the number of source nodes
* @param[in] c_scale 2^c_scale represents the number of destination nodes
* @param[in] n_edges number of edges to generate
* @param[in] stream cuda stream to schedule the work on
*
* We call the `r_scale != c_scale` case the "rectangular adjacency matrix" case (IOW generating
* bipartite graphs). In this case, at `depth >= r_scale`, the distribution is assumed to be:
* `[theta[4 * depth] + theta[4 * depth + 2], theta[4 * depth + 1] + theta[4 * depth + 3]; 0, 0]`.
* Then for `depth >= c_scale`, the distribution is assumed to be:
* `[theta[4 * depth] + theta[4 * depth + 1], 0; theta[4 * depth + 2] + theta[4 * depth + 3], 0]`.
*
* @note This can generate duplicate edges and self-loops. It is the responsibility of the
* caller to clean them up accordingly.
* @note This also only generates directed graphs. If undirected graphs are needed, then a
* separate post-processing step is expected to be done by the caller.
*
* @{
*/
template <typename IdxT, typename ProbT>
void rmat_rectangular_gen(const raft::handle_t& handle,
raft::random::RngState& r,
raft::device_vector_view<IdxT, IdxT> out,
raft::device_vector_view<IdxT, IdxT> out_src,
raft::device_vector_view<IdxT, IdxT> out_dst,
raft::device_vector_view<const ProbT, IdxT> theta,
IdxT r_scale,
IdxT c_scale,
IdxT n_edges)
{
detail::rmat_rectangular_gen_caller(out.data_handle(),
out_src.data_handle(),
out_dst.data_handle(),
theta.data_handle(),
r_scale,
c_scale,
n_edges,
handle.get_stream(),
r);
}

/**
* @brief Generate RMAT for a rectangular adjacency matrix (useful when
* graphs to be generated are bipartite)
*
* @tparam IdxT type of each node index
* @tparam ProbT data type used for probability distributions (either fp32 or fp64)
*
* @param[out] out generated edgelist [on device] [dim = n_edges x 2]. In each row
* the first element is the source node id, and the second element
* is the destination node id. If you don't need this output
* then pass a `nullptr` in its place.
* @param[out] out_src list of source node id's [on device] [len = n_edges]. If you
* don't need this output then pass a `nullptr` in its place.
* @param[out] out_dst list of destination node id's [on device] [len = n_edges]. If
* you don't need this output then pass a `nullptr` in its place.
* @param[in] theta distribution of each quadrant at each level of resolution.
* Since these are probabilities, each of the 2x2 matrices for
* each level of the RMAT must sum to one. [on device]
* [dim = max(r_scale, c_scale) x 2 x 2]. Of course, it is assumed
* that each of the group of 2 x 2 numbers all sum up to 1.
* @param[in] r_scale 2^r_scale represents the number of source nodes
* @param[in] c_scale 2^c_scale represents the number of destination nodes
* @param[in] n_edges number of edges to generate
* @param[in] stream cuda stream on which to schedule the work
* @param[in] r underlying state of the random generator. Especially useful when
* one wants to call this API for multiple times in order to generate
* a larger graph. For that case, just create this object with the
* initial seed once and after every call continue to pass the same
* object for the successive calls.
*
* When `r_scale != c_scale` it is referred to as rectangular adjacency matrix case (IOW generating
* We call the `r_scale != c_scale` case the "rectangular adjacency matrix" case (IOW generating
* bipartite graphs). In this case, at `depth >= r_scale`, the distribution is assumed to be:
* `[theta[4 * depth] + theta[4 * depth + 2], theta[4 * depth + 1] + theta[4 * depth + 3]; 0, 0]`.
* Then for the `depth >= c_scale`, the distribution is assumed to be:
* Then for `depth >= c_scale`, the distribution is assumed to be:
* `[theta[4 * depth] + theta[4 * depth + 1], 0; theta[4 * depth + 2] + theta[4 * depth + 3], 0]`.
*
* @note This can generate duplicate edges and self-loops. It is the responsibility of the
* caller to clean them up accordingly.
* @note This also only generates directed graphs. If undirected graphs are needed, then a
* separate post-processing step is expected to be done by the caller.
*
* @{
*/
template <typename IdxT, typename ProbT>
void rmat_rectangular_gen(IdxT* out,
Expand Down Expand Up @@ -99,6 +165,36 @@ void rmat_rectangular_gen(IdxT* out,
detail::rmat_rectangular_gen_caller(
out, out_src, out_dst, a, b, c, r_scale, c_scale, n_edges, stream, r);
}

/**
* This is the same as the previous method but assumes the same a, b, c, d probability
* distributions across all the scales
*/
template <typename IdxT, typename ProbT>
void rmat_rectangular_gen(const raft::handle_t& handle,
raft::random::RngState& r,
raft::device_vector_view<IdxT, IdxT> out,
raft::device_vector_view<IdxT, IdxT> out_src,
raft::device_vector_view<IdxT, IdxT> out_dst,
ProbT a,
ProbT b,
ProbT c,
IdxT r_scale,
IdxT c_scale,
IdxT n_edges)
{
detail::rmat_rectangular_gen_caller(out.data_handle(),
out_src.data_handle(),
out_dst.data_handle(),
a,
b,
c,
r_scale,
c_scale,
n_edges,
handle.get_stream(),
r);
}
/** @} */

} // end namespace raft::random
102 changes: 102 additions & 0 deletions cpp/test/random/rmat_rectangular_generator.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -253,6 +253,105 @@ class RmatGenTest : public ::testing::TestWithParam<RmatInputs> {
size_t max_scale;
};

class RmatGenMdspanTest : public ::testing::TestWithParam<RmatInputs> {
public:
RmatGenMdspanTest()
: handle{},
stream{handle.get_stream()},
params{::testing::TestWithParam<RmatInputs>::GetParam()},
out{params.n_edges * 2, stream},
out_src{params.n_edges, stream},
out_dst{params.n_edges, stream},
theta{0, stream},
h_theta{},
state{params.seed, GeneratorType::GenPC},
max_scale{std::max(params.r_scale, params.c_scale)}
{
theta.resize(4 * max_scale, stream);
uniform<float>(state, theta.data(), theta.size(), 0.0f, 1.0f, stream);
normalize<float, float>(theta.data(),
theta.data(),
max_scale,
params.r_scale,
params.c_scale,
params.r_scale != params.c_scale,
params.theta_array,
stream);
h_theta.resize(theta.size());
RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
raft::update_host(h_theta.data(), theta.data(), theta.size(), stream);
RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
}

protected:
void SetUp() override
{
using index_type = size_t;
raft::device_vector_view<size_t, index_type> out_view(out.data(), out.size());
raft::device_vector_view<size_t, index_type> out_src_view(out_src.data(), out_src.size());
raft::device_vector_view<size_t, index_type> out_dst_view(out_dst.data(), out_dst.size());

if (params.theta_array) {
raft::device_vector_view<const float, index_type> theta_view(theta.data(), theta.size());
rmat_rectangular_gen(handle,
state,
out_view,
out_src_view,
out_dst_view,
theta_view,
params.r_scale,
params.c_scale,
params.n_edges);
} else {
rmat_rectangular_gen(handle,
state,
out_view,
out_src_view,
out_dst_view,
h_theta[0],
h_theta[1],
h_theta[2],
params.r_scale,
params.c_scale,
params.n_edges);
}
RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
}

void validate()
{
rmm::device_uvector<int> hist{theta.size(), stream};
RAFT_CUDA_TRY(cudaMemsetAsync(hist.data(), 0, hist.size() * sizeof(int), stream));
compute_hist<<<raft::ceildiv<size_t>(out.size() / 2, 256), 256, 0, stream>>>(
hist.data(), out.data(), out.size(), max_scale, params.r_scale, params.c_scale);
RAFT_CUDA_TRY(cudaGetLastError());
rmm::device_uvector<float> computed_theta{theta.size(), stream};
normalize<float, int>(computed_theta.data(),
hist.data(),
max_scale,
params.r_scale,
params.c_scale,
false,
true,
stream);
RAFT_CUDA_TRY(cudaGetLastError());
RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
ASSERT_TRUE(devArrMatchHost(
h_theta.data(), computed_theta.data(), theta.size(), CompareApprox<float>(params.eps)));
}

protected:
raft::handle_t handle;
cudaStream_t stream;

RmatInputs params;
rmm::device_uvector<size_t> out, out_src, out_dst;
rmm::device_uvector<float> theta;
std::vector<float> h_theta;
RngState state;
size_t max_scale;
};

static const float TOLERANCE = 0.01f;

const std::vector<RmatInputs> inputs = {
Expand Down Expand Up @@ -287,5 +386,8 @@ const std::vector<RmatInputs> inputs = {
TEST_P(RmatGenTest, Result) { validate(); }
INSTANTIATE_TEST_SUITE_P(RmatGenTests, RmatGenTest, ::testing::ValuesIn(inputs));

TEST_P(RmatGenMdspanTest, Result) { validate(); }
INSTANTIATE_TEST_SUITE_P(RmatGenMdspanTests, RmatGenMdspanTest, ::testing::ValuesIn(inputs));

} // namespace random
} // namespace raft

0 comments on commit 039a3b1

Please sign in to comment.