Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "77cfea689d41c6fce0e7714e46758e2866507fc2"
Unverified Commit 81831111 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] clang-format auto fix. (#4805) 



* [Misc] clang-format auto fix.

* fix
Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent 16e771c0
Show whitespace changes
Inline Side-by-side
Showing with 3263 additions and 3237 deletions
src/geometry/cpu/geometry_op_impl.cc
View file @ 81831111
......@@ -4,9 +4,11 @@
* \brief Geometry operator CPU implementation
*/
#include <dgl/random.h>
#include <numeric>
#include <vector>
#include <utility>
#include <vector>
#include "../geometry_op.h"
namespace dgl {
......@@ -25,8 +27,9 @@ void IndexShuffle(IdType *idxs, int64_t num_elems) {
template void IndexShuffle<int32_t>(int32_t *idxs, int64_t num_elems);
template void IndexShuffle<int64_t>(int64_t *idxs, int64_t num_elems);
/*! \brief Groupwise index shuffle algorithm. This function will perform shuffle in subarrays
* indicated by group index. The group index is similar to indptr in CSRMatrix.
/*! \brief Groupwise index shuffle algorithm. This function will perform shuffle
* in subarrays indicated by group index. The group index is similar to indptr
* in CSRMatrix.
*
* \param group_idxs group index array.
* \param idxs index array for shuffle.
......@@ -34,64 +37,73 @@ template void IndexShuffle<int64_t>(int64_t *idxs, int64_t num_elems);
* \param num_elems length of idxs
*/
template <typename IdType>
void GroupIndexShuffle(const IdType *group_idxs, IdType *idxs,
int64_t num_groups_idxs, int64_t num_elems) {
void GroupIndexShuffle(
const IdType *group_idxs, IdType *idxs, int64_t num_groups_idxs,
int64_t num_elems) {
if (num_groups_idxs < 2) return; // empty idxs array
CHECK_LE(group_idxs[num_groups_idxs - 1], num_elems) << "group_idxs out of range";
CHECK_LE(group_idxs[num_groups_idxs - 1], num_elems)
<< "group_idxs out of range";
for (int64_t i = 0; i < num_groups_idxs - 1; ++i) {
auto subarray_len = group_idxs[i + 1] - group_idxs[i];
IndexShuffle(idxs + group_idxs[i], subarray_len);
}
}
template void GroupIndexShuffle<int32_t>(
const int32_t *group_idxs, int32_t *idxs, int64_t num_groups_idxs, int64_t num_elems);
const int32_t *group_idxs, int32_t *idxs, int64_t num_groups_idxs,
int64_t num_elems);
template void GroupIndexShuffle<int64_t>(
const int64_t *group_idxs, int64_t *idxs, int64_t num_groups_idxs, int64_t num_elems);
const int64_t *group_idxs, int64_t *idxs, int64_t num_groups_idxs,
int64_t num_elems);
template <typename IdType>
IdArray RandomPerm(int64_t num_nodes) {
IdArray perm = aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
IdType* perm_data = static_cast<IdType*>(perm->data);
IdArray perm =
aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
IdType *perm_data = static_cast<IdType *>(perm->data);
std::iota(perm_data, perm_data + num_nodes, 0);
IndexShuffle(perm_data, num_nodes);
return perm;
}
template <typename IdType>
IdArray GroupRandomPerm(const IdType *group_idxs, int64_t num_group_idxs, int64_t num_nodes) {
IdArray perm = aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
IdType* perm_data = static_cast<IdType*>(perm->data);
IdArray GroupRandomPerm(
const IdType *group_idxs, int64_t num_group_idxs, int64_t num_nodes) {
IdArray perm =
aten::NewIdArray(num_nodes, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
IdType *perm_data = static_cast<IdType *>(perm->data);
std::iota(perm_data, perm_data + num_nodes, 0);
GroupIndexShuffle(group_idxs, perm_data, num_group_idxs, num_nodes);
return perm;
}
/*!
* \brief Farthest Point Sampler without the need to compute all pairs of distance.
* \brief Farthest Point Sampler without the need to compute all pairs of
* distance.
*
* The input array has shape (N, d), where N is the number of points, and d is the dimension.
* It consists of a (flatten) batch of point clouds.
* The input array has shape (N, d), where N is the number of points, and d is
* the dimension. It consists of a (flatten) batch of point clouds.
*
* In each batch, the algorithm starts with the sample index specified by ``start_idx``.
* Then for each point, we maintain the minimum to-sample distance.
* Finally, we pick the point with the maximum such distance.
* This process will be repeated for ``sample_points`` - 1 times.
* In each batch, the algorithm starts with the sample index specified by
* ``start_idx``. Then for each point, we maintain the minimum to-sample
* distance. Finally, we pick the point with the maximum such distance. This
* process will be repeated for ``sample_points`` - 1 times.
*/
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result) {
const FloatType* array_data = static_cast<FloatType*>(array->data);
void FarthestPointSampler(
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result) {
const FloatType *array_data = static_cast<FloatType *>(array->data);
const int64_t point_in_batch = array->shape[0] / batch_size;
const int64_t dim = array->shape[1];
// distance
FloatType* dist_data = static_cast<FloatType*>(dist->data);
FloatType *dist_data = static_cast<FloatType *>(dist->data);
// sample for each cloud in the batch
IdType* start_idx_data = static_cast<IdType*>(start_idx->data);
IdType *start_idx_data = static_cast<IdType *>(start_idx->data);
// return value
IdType* ret_data = static_cast<IdType*>(result->data);
IdType *ret_data = static_cast<IdType *>(result->data);
int64_t array_start = 0, ret_start = 0;
// loop for each point cloud sample in this batch
......@@ -136,29 +148,30 @@ void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_poin
}
}
template void FarthestPointSampler<kDGLCPU, float, int32_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCPU, float, int64_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCPU, double, int32_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCPU, double, int64_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void WeightedNeighborMatching(const aten::CSRMatrix &csr, const NDArray weight, IdArray result) {
void WeightedNeighborMatching(
const aten::CSRMatrix &csr, const NDArray weight, IdArray result) {
const int64_t num_nodes = result->shape[0];
const IdType *indptr_data = static_cast<IdType*>(csr.indptr->data);
const IdType *indices_data = static_cast<IdType*>(csr.indices->data);
IdType *result_data = static_cast<IdType*>(result->data);
FloatType *weight_data = static_cast<FloatType*>(weight->data);
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
IdType *result_data = static_cast<IdType *>(result->data);
FloatType *weight_data = static_cast<FloatType *>(weight->data);
// build node visiting order
IdArray vis_order = RandomPerm<IdType>(num_nodes);
IdType *vis_order_data = static_cast<IdType*>(vis_order->data);
IdType *vis_order_data = static_cast<IdType *>(vis_order->data);
for (int64_t n = 0; n < num_nodes; ++n) {
auto u = vis_order_data[n];
......@@ -193,16 +206,16 @@ template void WeightedNeighborMatching<kDGLCPU, double, int64_t>(
template <DGLDeviceType XPU, typename IdType>
void NeighborMatching(const aten::CSRMatrix &csr, IdArray result) {
const int64_t num_nodes = result->shape[0];
const IdType *indptr_data = static_cast<IdType*>(csr.indptr->data);
const IdType *indices_data = static_cast<IdType*>(csr.indices->data);
IdType *result_data = static_cast<IdType*>(result->data);
const IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
const IdType *indices_data = static_cast<IdType *>(csr.indices->data);
IdType *result_data = static_cast<IdType *>(result->data);
// build vis order
IdArray u_vis_order = RandomPerm<IdType>(num_nodes);
IdType *u_vis_order_data = static_cast<IdType*>(u_vis_order->data);
IdType *u_vis_order_data = static_cast<IdType *>(u_vis_order->data);
IdArray v_vis_order = GroupRandomPerm<IdType>(
indptr_data, csr.indptr->shape[0], csr.indices->shape[0]);
IdType *v_vis_order_data = static_cast<IdType*>(v_vis_order->data);
IdType *v_vis_order_data = static_cast<IdType *>(v_vis_order->data);
for (int64_t n = 0; n < num_nodes; ++n) {
auto u = u_vis_order_data[n];
......@@ -221,8 +234,10 @@ void NeighborMatching(const aten::CSRMatrix &csr, IdArray result) {
}
}
}
template void NeighborMatching<kDGLCPU, int32_t>(const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCPU, int64_t>(const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCPU, int32_t>(
const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCPU, int64_t>(
const aten::CSRMatrix &csr, IdArray result);
} // namespace impl
} // namespace geometry
......
src/geometry/cuda/edge_coarsening_impl.cu
View file @ 81831111
......@@ -3,14 +3,16 @@
* \file geometry/cuda/edge_coarsening_impl.cu
* \brief Edge coarsening CUDA implementation
*/
#include <curand_kernel.h>
#include <dgl/array.h>
#include <dgl/random.h>
#include <dmlc/thread_local.h>
#include <curand_kernel.h>
#include <cstdint>
#include "../geometry_op.h"
#include "../../runtime/cuda/cuda_common.h"
#include "../../array/cuda/utils.h"
#include "../../runtime/cuda/cuda_common.h"
#include "../geometry_op.h"
#define BLOCKS(N, T) (N + T - 1) / T
......@@ -26,7 +28,8 @@ constexpr int EMPTY_IDX = -1;
__device__ bool done_d;
__global__ void init_done_kernel() { done_d = true; }
__global__ void generate_uniform_kernel(float *ret_values, size_t num, uint64_t seed) {
__global__ void generate_uniform_kernel(
float *ret_values, size_t num, uint64_t seed) {
size_t id = blockIdx.x * blockDim.x + threadIdx.x;
if (id < num) {
curandState state;
......@@ -36,7 +39,8 @@ __global__ void generate_uniform_kernel(float *ret_values, size_t num, uint64_t
}
template <typename IdType>
__global__ void colorize_kernel(const float *prop, int64_t num_elem, IdType *result) {
__global__ void colorize_kernel(
const float *prop, int64_t num_elem, IdType *result) {
const IdType idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < num_elem) {
if (result[idx] < 0) { // if unmatched
......@@ -47,9 +51,9 @@ __global__ void colorize_kernel(const float *prop, int64_t num_elem, IdType *res
}
template <typename FloatType, typename IdType>
__global__ void weighted_propose_kernel(const IdType *indptr, const IdType *indices,
const FloatType *weights, int64_t num_elem,
IdType *proposal, IdType *result) {
__global__ void weighted_propose_kernel(
const IdType *indptr, const IdType *indices, const FloatType *weights,
int64_t num_elem, IdType *proposal, IdType *result) {
const IdType idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < num_elem) {
if (result[idx] != BLUE) return;
......@@ -61,8 +65,7 @@ __global__ void weighted_propose_kernel(const IdType *indptr, const IdType *indi
for (IdType i = indptr[idx]; i < indptr[idx + 1]; ++i) {
auto v = indices[i];
if (result[v] < 0)
has_unmatched_neighbor = true;
if (result[v] < 0) has_unmatched_neighbor = true;
if (result[v] == RED && weights[i] >= weight_max) {
v_max = v;
weight_max = weights[i];
......@@ -70,15 +73,14 @@ __global__ void weighted_propose_kernel(const IdType *indptr, const IdType *indi
}
proposal[idx] = v_max;
if (!has_unmatched_neighbor)
result[idx] = idx;
if (!has_unmatched_neighbor) result[idx] = idx;
}
}
template <typename FloatType, typename IdType>
__global__ void weighted_respond_kernel(const IdType *indptr, const IdType *indices,
const FloatType *weights, int64_t num_elem,
IdType *proposal, IdType *result) {
__global__ void weighted_respond_kernel(
const IdType *indptr, const IdType *indices, const FloatType *weights,
int64_t num_elem, IdType *proposal, IdType *result) {
const IdType idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < num_elem) {
if (result[idx] != RED) return;
......@@ -93,9 +95,7 @@ __global__ void weighted_respond_kernel(const IdType *indptr, const IdType *indi
if (result[v] < 0) {
has_unmatched_neighbors = true;
}
if (result[v] == BLUE
&& proposal[v] == idx
&& weights[i] >= weight_max) {
if (result[v] == BLUE && proposal[v] == idx && weights[i] >= weight_max) {
v_max = v;
weight_max = weights[i];
}
......@@ -105,8 +105,7 @@ __global__ void weighted_respond_kernel(const IdType *indptr, const IdType *indi
result[idx] = min(idx, v_max);
}
if (!has_unmatched_neighbors)
result[idx] = idx;
if (!has_unmatched_neighbors) result[idx] = idx;
}
}
......@@ -114,8 +113,8 @@ __global__ void weighted_respond_kernel(const IdType *indptr, const IdType *indi
* nodes with BLUE(-1) and RED(-2) and checks whether the node matching
* process has finished.
*/
template<typename IdType>
bool Colorize(IdType * result_data, int64_t num_nodes, float * const prop) {
template <typename IdType>
bool Colorize(IdType *result_data, int64_t num_nodes, float *const prop) {
// initial done signal
cudaStream_t stream = runtime::getCurrentCUDAStream();
CUDA_KERNEL_CALL(init_done_kernel, 1, 1, 0, stream);
......@@ -124,14 +123,17 @@ bool Colorize(IdType * result_data, int64_t num_nodes, float * const prop) {
uint64_t seed = dgl::RandomEngine::ThreadLocal()->RandInt(UINT64_MAX);
auto num_threads = cuda::FindNumThreads(num_nodes);
auto num_blocks = cuda::FindNumBlocks<'x'>(BLOCKS(num_nodes, num_threads));
CUDA_KERNEL_CALL(generate_uniform_kernel, num_blocks, num_threads, 0, stream,
prop, num_nodes, seed);
CUDA_KERNEL_CALL(
generate_uniform_kernel, num_blocks, num_threads, 0, stream, prop,
num_nodes, seed);
// call kernel
CUDA_KERNEL_CALL(colorize_kernel, num_blocks, num_threads, 0, stream,
prop, num_nodes, result_data);
CUDA_KERNEL_CALL(
colorize_kernel, num_blocks, num_threads, 0, stream, prop, num_nodes,
result_data);
bool done_h = false;
CUDA_CALL(cudaMemcpyFromSymbol(&done_h, done_d, sizeof(done_h), 0, cudaMemcpyDeviceToHost));
CUDA_CALL(cudaMemcpyFromSymbol(
&done_h, done_d, sizeof(done_h), 0, cudaMemcpyDeviceToHost));
return done_h;
}
......@@ -151,9 +153,10 @@ bool Colorize(IdType * result_data, int64_t num_nodes, float * const prop) {
* pair and mark them with the smaller id between them.
*/
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void WeightedNeighborMatching(const aten::CSRMatrix &csr, const NDArray weight, IdArray result) {
void WeightedNeighborMatching(
const aten::CSRMatrix &csr, const NDArray weight, IdArray result) {
cudaStream_t stream = runtime::getCurrentCUDAStream();
const auto& ctx = result->ctx;
const auto &ctx = result->ctx;
auto device = runtime::DeviceAPI::Get(ctx);
device->SetDevice(ctx);
......@@ -162,23 +165,27 @@ void WeightedNeighborMatching(const aten::CSRMatrix &csr, const NDArray weight,
IdArray proposal = aten::Full(-1, num_nodes, sizeof(IdType) * 8, ctx);
// get data ptrs
IdType *indptr_data = static_cast<IdType*>(csr.indptr->data);
IdType *indices_data = static_cast<IdType*>(csr.indices->data);
IdType *result_data = static_cast<IdType*>(result->data);
IdType *proposal_data = static_cast<IdType*>(proposal->data);
FloatType *weight_data = static_cast<FloatType*>(weight->data);
IdType *indptr_data = static_cast<IdType *>(csr.indptr->data);
IdType *indices_data = static_cast<IdType *>(csr.indices->data);
IdType *result_data = static_cast<IdType *>(result->data);
IdType *proposal_data = static_cast<IdType *>(proposal->data);
FloatType *weight_data = static_cast<FloatType *>(weight->data);
// allocate workspace for prop used in Colorize()
float *prop = static_cast<float*>(
float *prop = static_cast<float *>(
device->AllocWorkspace(ctx, num_nodes * sizeof(float)));
auto num_threads = cuda::FindNumThreads(num_nodes);
auto num_blocks = cuda::FindNumBlocks<'x'>(BLOCKS(num_nodes, num_threads));
while (!Colorize<IdType>(result_data, num_nodes, prop)) {
CUDA_KERNEL_CALL(weighted_propose_kernel, num_blocks, num_threads, 0, stream,
indptr_data, indices_data, weight_data, num_nodes, proposal_data, result_data);
CUDA_KERNEL_CALL(weighted_respond_kernel, num_blocks, num_threads, 0, stream,
indptr_data, indices_data, weight_data, num_nodes, proposal_data, result_data);
CUDA_KERNEL_CALL(
weighted_propose_kernel, num_blocks, num_threads, 0, stream,
indptr_data, indices_data, weight_data, num_nodes, proposal_data,
result_data);
CUDA_KERNEL_CALL(
weighted_respond_kernel, num_blocks, num_threads, 0, stream,
indptr_data, indices_data, weight_data, num_nodes, proposal_data,
result_data);
}
device->FreeWorkspace(ctx, prop);
}
......@@ -204,7 +211,7 @@ template void WeightedNeighborMatching<kDGLCUDA, double, int64_t>(
template <DGLDeviceType XPU, typename IdType>
void NeighborMatching(const aten::CSRMatrix &csr, IdArray result) {
const int64_t num_edges = csr.indices->shape[0];
const auto& ctx = result->ctx;
const auto &ctx = result->ctx;
auto device = runtime::DeviceAPI::Get(ctx);
device->SetDevice(ctx);
......@@ -212,17 +219,20 @@ void NeighborMatching(const aten::CSRMatrix &csr, IdArray result) {
cudaStream_t stream = runtime::getCurrentCUDAStream();
NDArray weight = NDArray::Empty(
{num_edges}, DGLDataType{kDGLFloat, sizeof(float) * 8, 1}, ctx);
float *weight_data = static_cast<float*>(weight->data);
float *weight_data = static_cast<float *>(weight->data);
uint64_t seed = dgl::RandomEngine::ThreadLocal()->RandInt(UINT64_MAX);
auto num_threads = cuda::FindNumThreads(num_edges);
auto num_blocks = cuda::FindNumBlocks<'x'>(BLOCKS(num_edges, num_threads));
CUDA_KERNEL_CALL(generate_uniform_kernel, num_blocks, num_threads, 0, stream,
weight_data, num_edges, seed);
CUDA_KERNEL_CALL(
generate_uniform_kernel, num_blocks, num_threads, 0, stream, weight_data,
num_edges, seed);
WeightedNeighborMatching<XPU, float, IdType>(csr, weight, result);
}
template void NeighborMatching<kDGLCUDA, int32_t>(const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCUDA, int64_t>(const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCUDA, int32_t>(
const aten::CSRMatrix &csr, IdArray result);
template void NeighborMatching<kDGLCUDA, int64_t>(
const aten::CSRMatrix &csr, IdArray result);
} // namespace impl
} // namespace geometry
......
src/geometry/cuda/geometry_op_impl.cu
View file @ 81831111
......@@ -5,8 +5,8 @@
*/
#include <dgl/array.h>
#include "../../runtime/cuda/cuda_common.h"
#include "../../c_api_common.h"
#include "../../runtime/cuda/cuda_common.h"
#include "../geometry_op.h"
#define THREADS 1024
......@@ -16,21 +16,23 @@ namespace geometry {
namespace impl {
/*!
* \brief Farthest Point Sampler without the need to compute all pairs of distance.
* \brief Farthest Point Sampler without the need to compute all pairs of
* distance.
*
* The input array has shape (N, d), where N is the number of points, and d is the dimension.
* It consists of a (flatten) batch of point clouds.
* The input array has shape (N, d), where N is the number of points, and d is
* the dimension. It consists of a (flatten) batch of point clouds.
*
* In each batch, the algorithm starts with the sample index specified by ``start_idx``.
* Then for each point, we maintain the minimum to-sample distance.
* Finally, we pick the point with the maximum such distance.
* This process will be repeated for ``sample_points`` - 1 times.
* In each batch, the algorithm starts with the sample index specified by
* ``start_idx``. Then for each point, we maintain the minimum to-sample
* distance. Finally, we pick the point with the maximum such distance. This
* process will be repeated for ``sample_points`` - 1 times.
*/
template <typename FloatType, typename IdType>
__global__ void fps_kernel(const FloatType *array_data, const int64_t batch_size,
__global__ void fps_kernel(
const FloatType* array_data, const int64_t batch_size,
const int64_t sample_points, const int64_t point_in_batch,
const int64_t dim, const IdType *start_idx,
FloatType *dist_data, IdType *ret_data) {
const int64_t dim, const IdType* start_idx, FloatType* dist_data,
IdType* ret_data) {
const int64_t thread_idx = threadIdx.x;
const int64_t batch_idx = blockIdx.x;
......@@ -90,8 +92,9 @@ __global__ void fps_kernel(const FloatType *array_data, const int64_t batch_size
}
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result) {
void FarthestPointSampler(
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result) {
cudaStream_t stream = runtime::getCurrentCUDAStream();
const FloatType* array_data = static_cast<FloatType*>(array->data);
......@@ -109,24 +112,23 @@ void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_poin
IdType* start_idx_data = static_cast<IdType*>(start_idx->data);
CUDA_CALL(cudaSetDevice(array->ctx.device_id));
CUDA_KERNEL_CALL(fps_kernel,
batch_size, THREADS, 0, stream,
array_data, batch_size, sample_points,
point_in_batch, dim, start_idx_data, dist_data, ret_data);
CUDA_KERNEL_CALL(
fps_kernel, batch_size, THREADS, 0, stream, array_data, batch_size,
sample_points, point_in_batch, dim, start_idx_data, dist_data, ret_data);
}
template void FarthestPointSampler<kDGLCUDA, float, int32_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCUDA, float, int64_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCUDA, double, int32_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
template void FarthestPointSampler<kDGLCUDA, double, int64_t>(
NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
} // namespace impl
} // namespace geometry
......
src/geometry/geometry.cc
View file @ 81831111
......@@ -4,28 +4,34 @@
* \brief DGL geometry utilities implementation
*/
#include <dgl/array.h>
#include <dgl/runtime/ndarray.h>
#include <dgl/base_heterograph.h>
#include <dgl/runtime/ndarray.h>
#include "../array/check.h"
#include "../c_api_common.h"
#include "./geometry_op.h"
#include "../array/check.h"
using namespace dgl::runtime;
namespace dgl {
namespace geometry {
void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result) {
CHECK_EQ(array->ctx, result->ctx) << "Array and the result should be on the same device.";
CHECK_EQ(array->shape[0], dist->shape[0]) << "Shape of array and dist mismatch";
CHECK_EQ(start_idx->shape[0], batch_size) << "Shape of start_idx and batch_size mismatch";
CHECK_EQ(result->shape[0], batch_size * sample_points) << "Invalid shape of result";
void FarthestPointSampler(
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result) {
CHECK_EQ(array->ctx, result->ctx)
<< "Array and the result should be on the same device.";
CHECK_EQ(array->shape[0], dist->shape[0])
<< "Shape of array and dist mismatch";
CHECK_EQ(start_idx->shape[0], batch_size)
<< "Shape of start_idx and batch_size mismatch";
CHECK_EQ(result->shape[0], batch_size * sample_points)
<< "Invalid shape of result";
ATEN_FLOAT_TYPE_SWITCH(array->dtype, FloatType, "values", {
ATEN_ID_TYPE_SWITCH(result->dtype, IdType, {
ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "FarthestPointSampler", {
ATEN_XPU_SWITCH_CUDA(
array->ctx.device_type, XPU, "FarthestPointSampler", {
impl::FarthestPointSampler<XPU, FloatType, IdType>(
array, batch_size, sample_points, dist, start_idx, result);
});
......@@ -33,9 +39,11 @@ void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_poin
});
}
void NeighborMatching(HeteroGraphPtr graph, const NDArray weight, IdArray result) {
void NeighborMatching(
HeteroGraphPtr graph, const NDArray weight, IdArray result) {
if (!aten::IsNullArray(weight)) {
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "NeighborMatching", {
ATEN_XPU_SWITCH_CUDA(
graph->Context().device_type, XPU, "NeighborMatching", {
ATEN_FLOAT_TYPE_SWITCH(weight->dtype, FloatType, "weight", {
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
impl::WeightedNeighborMatching<XPU, FloatType, IdType>(
......@@ -44,10 +52,10 @@ void NeighborMatching(HeteroGraphPtr graph, const NDArray weight, IdArray result
});
});
} else {
ATEN_XPU_SWITCH_CUDA(graph->Context().device_type, XPU, "NeighborMatching", {
ATEN_XPU_SWITCH_CUDA(
graph->Context().device_type, XPU, "NeighborMatching", {
ATEN_ID_TYPE_SWITCH(graph->DataType(), IdType, {
impl::NeighborMatching<XPU, IdType>(
graph->GetCSRMatrix(0), result);
impl::NeighborMatching<XPU, IdType>(graph->GetCSRMatrix(0), result);
});
});
}
......@@ -56,7 +64,7 @@ void NeighborMatching(HeteroGraphPtr graph, const NDArray weight, IdArray result
///////////////////////// C APIs /////////////////////////
DGL_REGISTER_GLOBAL("geometry._CAPI_FarthestPointSampler")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const NDArray data = args[0];
const int64_t batch_size = args[1];
const int64_t sample_points = args[2];
......@@ -64,24 +72,28 @@ DGL_REGISTER_GLOBAL("geometry._CAPI_FarthestPointSampler")
IdArray start_idx = args[4];
IdArray result = args[5];
FarthestPointSampler(data, batch_size, sample_points, dist, start_idx, result);
FarthestPointSampler(
data, batch_size, sample_points, dist, start_idx, result);
});
DGL_REGISTER_GLOBAL("geometry._CAPI_NeighborMatching")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef graph = args[0];
const NDArray weight = args[1];
IdArray result = args[2];
// sanity check
aten::CheckCtx(graph->Context(), {weight, result}, {"edge_weight, result"});
aten::CheckCtx(
graph->Context(), {weight, result}, {"edge_weight, result"});
aten::CheckContiguous({weight, result}, {"edge_weight", "result"});
CHECK_EQ(graph->NumEdgeTypes(), 1) << "homogeneous graph has only one edge type";
CHECK_EQ(graph->NumEdgeTypes(), 1)
<< "homogeneous graph has only one edge type";
CHECK_EQ(result->ndim, 1) << "result should be an 1D tensor.";
auto pair = graph->meta_graph()->FindEdge(0);
const dgl_type_t node_type = pair.first;
CHECK_EQ(graph->NumVertices(node_type), result->shape[0])
<< "The number of nodes should be the same as the length of result tensor.";
<< "The number of nodes should be the same as the length of result "
"tensor.";
if (!aten::IsNullArray(weight)) {
CHECK_EQ(weight->ndim, 1) << "weight should be an 1D tensor.";
CHECK_EQ(graph->NumEdges(0), weight->shape[0])
......
src/geometry/geometry_op.h
View file @ 81831111
......@@ -13,16 +13,19 @@ namespace geometry {
namespace impl {
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void FarthestPointSampler(NDArray array, int64_t batch_size, int64_t sample_points,
NDArray dist, IdArray start_idx, IdArray result);
void FarthestPointSampler(
NDArray array, int64_t batch_size, int64_t sample_points, NDArray dist,
IdArray start_idx, IdArray result);
/*! \brief Implementation of weighted neighbor matching process of edge coarsening used
* in Metis and Graclus for homogeneous graph coarsening. This procedure keeps
* picking an unmarked vertex and matching it with one its unmarked neighbors
* (that maximizes its edge weight) until no match can be done.
/*! \brief Implementation of weighted neighbor matching process of edge
* coarsening used in Metis and Graclus for homogeneous graph coarsening. This
* procedure keeps picking an unmarked vertex and matching it with one its
* unmarked neighbors (that maximizes its edge weight) until no match can be
* done.
*/
template <DGLDeviceType XPU, typename FloatType, typename IdType>
void WeightedNeighborMatching(const aten::CSRMatrix &csr, const NDArray weight, IdArray result);
void WeightedNeighborMatching(
const aten::CSRMatrix &csr, const NDArray weight, IdArray result);
/*! \brief Implementation of neighbor matching process of edge coarsening used
* in Metis and Graclus for homogeneous graph coarsening. This procedure keeps
......
src/graph/creators.cc
View file @ 81831111
......@@ -10,60 +10,51 @@ namespace dgl {
// creator implementation
HeteroGraphPtr CreateHeteroGraph(
GraphPtr meta_graph,
const std::vector<HeteroGraphPtr>& rel_graphs,
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs,
const std::vector<int64_t>& num_nodes_per_type) {
return HeteroGraphPtr(new HeteroGraph(meta_graph, rel_graphs, num_nodes_per_type));
return HeteroGraphPtr(
new HeteroGraph(meta_graph, rel_graphs, num_nodes_per_type));
}
HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray row, IdArray col,
bool row_sorted, bool col_sorted, dgl_format_code_t formats) {
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray row,
IdArray col, bool row_sorted, bool col_sorted, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCOO(
num_vtypes, num_src, num_dst, row, col, row_sorted, col_sorted, formats);
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, const aten::COOMatrix& mat,
dgl_format_code_t formats) {
int64_t num_vtypes, const aten::COOMatrix& mat, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCOO(num_vtypes, mat, formats);
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids,
dgl_format_code_t formats) {
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray indptr,
IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCSR(
num_vtypes, num_src, num_dst, indptr, indices, edge_ids, formats);
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, const aten::CSRMatrix& mat,
dgl_format_code_t formats) {
int64_t num_vtypes, const aten::CSRMatrix& mat, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCSR(num_vtypes, mat, formats);
auto ret = HeteroGraphPtr(new HeteroGraph(
unit_g->meta_graph(),
{unit_g}));
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(),
{unit_g}));
auto ret = HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, int64_t num_src, int64_t num_dst,
IdArray indptr, IdArray indices, IdArray edge_ids,
dgl_format_code_t formats) {
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray indptr,
IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCSC(
num_vtypes, num_src, num_dst, indptr, indices, edge_ids, formats);
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, const aten::CSRMatrix& mat,
dgl_format_code_t formats) {
int64_t num_vtypes, const aten::CSRMatrix& mat, dgl_format_code_t formats) {
auto unit_g = UnitGraph::CreateFromCSC(num_vtypes, mat, formats);
return HeteroGraphPtr(new HeteroGraph(unit_g->meta_graph(), {unit_g}));
}
......
src/graph/gk_ops.cc
View file @ 81831111
......@@ -37,16 +37,18 @@ gk_csr_t *Convert2GKCsr(const aten::CSRMatrix mat, bool is_row) {
size_t num_ptrs;
if (is_row) {
num_ptrs = gk_csr->nrows + 1;
gk_indptr = gk_csr->rowptr = gk_zmalloc(gk_csr->nrows+1,
const_cast<char*>("gk_csr_ExtractPartition: rowptr"));
gk_indices = gk_csr->rowind = gk_imalloc(nnz,
const_cast<char*>("gk_csr_ExtractPartition: rowind"));
gk_indptr = gk_csr->rowptr = gk_zmalloc(
gk_csr->nrows + 1,
const_cast<char *>("gk_csr_ExtractPartition: rowptr"));
gk_indices = gk_csr->rowind =
gk_imalloc(nnz, const_cast<char *>("gk_csr_ExtractPartition: rowind"));
} else {
num_ptrs = gk_csr->ncols + 1;
gk_indptr = gk_csr->colptr = gk_zmalloc(gk_csr->ncols+1,
const_cast<char*>("gk_csr_ExtractPartition: colptr"));
gk_indices = gk_csr->colind = gk_imalloc(nnz,
const_cast<char*>("gk_csr_ExtractPartition: colind"));
gk_indptr = gk_csr->colptr = gk_zmalloc(
gk_csr->ncols + 1,
const_cast<char *>("gk_csr_ExtractPartition: colptr"));
gk_indices = gk_csr->colind =
gk_imalloc(nnz, const_cast<char *>("gk_csr_ExtractPartition: colind"));
}
for (size_t i = 0; i < num_ptrs; i++) {
......@@ -98,7 +100,8 @@ aten::CSRMatrix Convert2DGLCsr(gk_csr_t *gk_csr, bool is_row) {
eids[i] = i;
}
return aten::CSRMatrix(gk_csr->nrows, gk_csr->ncols, indptr_arr, indices_arr, eids_arr);
return aten::CSRMatrix(
gk_csr->nrows, gk_csr->ncols, indptr_arr, indices_arr, eids_arr);
}
#endif // !defined(_WIN32)
......
src/graph/graph.cc
View file @ 81831111
......@@ -5,11 +5,13 @@
*/
#include <dgl/graph.h>
#include <dgl/sampler.h>
#include <algorithm>
#include <unordered_map>
#include <set>
#include <functional>
#include <set>
#include <tuple>
#include <unordered_map>
#include "../c_api_common.h"
namespace dgl {
......@@ -21,8 +23,8 @@ Graph::Graph(IdArray src_ids, IdArray dst_ids, size_t num_nodes) {
num_edges_ = src_ids->shape[0];
CHECK(static_cast<int64_t>(num_edges_) == dst_ids->shape[0])
<< "vectors in COO must have the same length";
const dgl_id_t *src_data = static_cast<dgl_id_t*>(src_ids->data);
const dgl_id_t *dst_data = static_cast<dgl_id_t*>(dst_ids->data);
const dgl_id_t* src_data = static_cast<dgl_id_t*>(src_ids->data);
const dgl_id_t* dst_data = static_cast<dgl_id_t*>(dst_ids->data);
all_edges_src_.reserve(num_edges_);
all_edges_dst_.reserve(num_edges_);
for (uint64_t i = 0; i < num_edges_; i++) {
......@@ -53,15 +55,15 @@ bool Graph::IsMultigraph() const {
pairs.emplace_back(all_edges_src_[eid], all_edges_dst_[eid]);
}
// sort according to src and dst ids
std::sort(pairs.begin(), pairs.end(),
[] (const Pair& t1, const Pair& t2) {
return std::get<0>(t1) < std::get<0>(t2)
|| (std::get<0>(t1) == std::get<0>(t2) && std::get<1>(t1) < std::get<1>(t2));
std::sort(pairs.begin(), pairs.end(), [](const Pair& t1, const Pair& t2) {
return std::get<0>(t1) < std::get<0>(t2) ||
(std::get<0>(t1) == std::get<0>(t2) &&
std::get<1>(t1) < std::get<1>(t2));
});
for (uint64_t eid = 0; eid < num_edges_-1; ++eid) {
for (uint64_t eid = 0; eid < num_edges_ - 1; ++eid) {
// As src and dst are all sorted, we only need to compare i and i+1
if (std::get<0>(pairs[eid]) == std::get<0>(pairs[eid+1]) &&
std::get<1>(pairs[eid]) == std::get<1>(pairs[eid+1]))
if (std::get<0>(pairs[eid]) == std::get<0>(pairs[eid + 1]) &&
std::get<1>(pairs[eid]) == std::get<1>(pairs[eid + 1]))
return true;
}
......@@ -125,7 +127,7 @@ BoolArray Graph::HasVertices(IdArray vids) const {
int64_t* rst_data = static_cast<int64_t*>(rst->data);
const int64_t nverts = NumVertices();
for (int64_t i = 0; i < len; ++i) {
rst_data[i] = (vid_data[i] < nverts && vid_data[i] >= 0)? 1 : 0;
rst_data[i] = (vid_data[i] < nverts && vid_data[i] >= 0) ? 1 : 0;
}
return rst;
}
......@@ -151,18 +153,18 @@ BoolArray Graph::HasEdgesBetween(IdArray src_ids, IdArray dst_ids) const {
if (srclen == 1) {
// one-many
for (int64_t i = 0; i < dstlen; ++i) {
rst_data[i] = HasEdgeBetween(src_data[0], dst_data[i])? 1 : 0;
rst_data[i] = HasEdgeBetween(src_data[0], dst_data[i]) ? 1 : 0;
}
} else if (dstlen == 1) {
// many-one
for (int64_t i = 0; i < srclen; ++i) {
rst_data[i] = HasEdgeBetween(src_data[i], dst_data[0])? 1 : 0;
rst_data[i] = HasEdgeBetween(src_data[i], dst_data[0]) ? 1 : 0;
}
} else {
// many-many
CHECK(srclen == dstlen) << "Invalid src and dst id array.";
for (int64_t i = 0; i < srclen; ++i) {
rst_data[i] = HasEdgeBetween(src_data[i], dst_data[i])? 1 : 0;
rst_data[i] = HasEdgeBetween(src_data[i], dst_data[i]) ? 1 : 0;
}
}
return rst;
......@@ -174,11 +176,11 @@ IdArray Graph::Predecessors(dgl_id_t vid, uint64_t radius) const {
CHECK(radius >= 1) << "invalid radius: " << radius;
std::set<dgl_id_t> vset;
for (auto& it : reverse_adjlist_[vid].succ)
vset.insert(it);
for (auto& it : reverse_adjlist_[vid].succ) vset.insert(it);
const int64_t len = vset.size();
IdArray rst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray rst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t* rst_data = static_cast<int64_t*>(rst->data);
std::copy(vset.begin(), vset.end(), rst_data);
......@@ -191,11 +193,11 @@ IdArray Graph::Successors(dgl_id_t vid, uint64_t radius) const {
CHECK(radius >= 1) << "invalid radius: " << radius;
std::set<dgl_id_t> vset;
for (auto& it : adjlist_[vid].succ)
vset.insert(it);
for (auto& it : adjlist_[vid].succ) vset.insert(it);
const int64_t len = vset.size();
IdArray rst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray rst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t* rst_data = static_cast<int64_t*>(rst->data);
std::copy(vset.begin(), vset.end(), rst_data);
......@@ -204,19 +206,20 @@ IdArray Graph::Successors(dgl_id_t vid, uint64_t radius) const {
// O(E)
IdArray Graph::EdgeId(dgl_id_t src, dgl_id_t dst) const {
CHECK(HasVertex(src) && HasVertex(dst)) << "invalid edge: " << src << " -> " << dst;
CHECK(HasVertex(src) && HasVertex(dst))
<< "invalid edge: " << src << " -> " << dst;
const auto& succ = adjlist_[src].succ;
std::vector<dgl_id_t> edgelist;
for (size_t i = 0; i < succ.size(); ++i) {
if (succ[i] == dst)
edgelist.push_back(adjlist_[src].edge_id[i]);
if (succ[i] == dst) edgelist.push_back(adjlist_[src].edge_id[i]);
}
// FIXME: signed? Also it seems that we are using int64_t everywhere...
const int64_t len = edgelist.size();
IdArray rst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray rst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
// FIXME: signed?
int64_t* rst_data = static_cast<int64_t*>(rst->data);
......@@ -243,10 +246,11 @@ EdgeArray Graph::EdgeIds(IdArray src_ids, IdArray dst_ids) const {
std::vector<dgl_id_t> src, dst, eid;
for (i = 0, j = 0; i < srclen && j < dstlen; i += src_stride, j += dst_stride) {
for (i = 0, j = 0; i < srclen && j < dstlen;
i += src_stride, j += dst_stride) {
const dgl_id_t src_id = src_data[i], dst_id = dst_data[j];
CHECK(HasVertex(src_id) && HasVertex(dst_id)) <<
"invalid edge: " << src_id << " -> " << dst_id;
CHECK(HasVertex(src_id) && HasVertex(dst_id))
<< "invalid edge: " << src_id << " -> " << dst_id;
const auto& succ = adjlist_[src_id].succ;
for (size_t k = 0; k < succ.size(); ++k) {
if (succ[k] == dst_id) {
......@@ -286,8 +290,7 @@ EdgeArray Graph::FindEdges(IdArray eids) const {
for (uint64_t i = 0; i < (uint64_t)len; ++i) {
dgl_id_t eid = eid_data[i];
if (eid >= num_edges_)
LOG(FATAL) << "invalid edge id:" << eid;
if (eid >= num_edges_) LOG(FATAL) << "invalid edge id:" << eid;
rst_src_data[i] = all_edges_src_[eid];
rst_dst_data[i] = all_edges_dst_[eid];
......@@ -301,9 +304,12 @@ EdgeArray Graph::FindEdges(IdArray eids) const {
EdgeArray Graph::InEdges(dgl_id_t vid) const {
CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
const int64_t len = reverse_adjlist_[vid].succ.size();
IdArray src = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray src = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t* src_data = static_cast<int64_t*>(src->data);
int64_t* dst_data = static_cast<int64_t*>(dst->data);
int64_t* eid_data = static_cast<int64_t*>(eid->data);
......@@ -347,9 +353,12 @@ EdgeArray Graph::InEdges(IdArray vids) const {
EdgeArray Graph::OutEdges(dgl_id_t vid) const {
CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
const int64_t len = adjlist_[vid].succ.size();
IdArray src = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray src = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t* src_data = static_cast<int64_t*>(src->data);
int64_t* dst_data = static_cast<int64_t*>(dst->data);
int64_t* eid_data = static_cast<int64_t*>(eid->data);
......@@ -390,11 +399,14 @@ EdgeArray Graph::OutEdges(IdArray vids) const {
}
// O(E*log(E)) if sort is required; otherwise, O(E)
EdgeArray Graph::Edges(const std::string &order) const {
EdgeArray Graph::Edges(const std::string& order) const {
const int64_t len = num_edges_;
IdArray src = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray src = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray dst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
if (order == "srcdst") {
typedef std::tuple<int64_t, int64_t, int64_t> Tuple;
......@@ -404,10 +416,11 @@ EdgeArray Graph::Edges(const std::string &order) const {
tuples.emplace_back(all_edges_src_[eid], all_edges_dst_[eid], eid);
}
// sort according to src and dst ids
std::sort(tuples.begin(), tuples.end(),
[] (const Tuple& t1, const Tuple& t2) {
return std::get<0>(t1) < std::get<0>(t2)
|| (std::get<0>(t1) == std::get<0>(t2) && std::get<1>(t1) < std::get<1>(t2));
std::sort(
tuples.begin(), tuples.end(), [](const Tuple& t1, const Tuple& t2) {
return std::get<0>(t1) < std::get<0>(t2) ||
(std::get<0>(t1) == std::get<0>(t2) &&
std::get<1>(t1) < std::get<1>(t2));
});
// make return arrays
......@@ -488,8 +501,11 @@ Subgraph Graph::VertexSubgraph(IdArray vids) const {
}
}
}
rst.induced_edges = IdArray::Empty({static_cast<int64_t>(edges.size())}, vids->dtype, vids->ctx);
std::copy(edges.begin(), edges.end(), static_cast<int64_t*>(rst.induced_edges->data));
rst.induced_edges = IdArray::Empty(
{static_cast<int64_t>(edges.size())}, vids->dtype, vids->ctx);
std::copy(
edges.begin(), edges.end(),
static_cast<int64_t*>(rst.induced_edges->data));
return rst;
}
......@@ -524,7 +540,9 @@ Subgraph Graph::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
rst.induced_vertices = IdArray::Empty(
{static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
std::copy(nodes.begin(), nodes.end(), static_cast<int64_t*>(rst.induced_vertices->data));
std::copy(
nodes.begin(), nodes.end(),
static_cast<int64_t*>(rst.induced_vertices->data));
} else {
rst.graph = std::make_shared<Graph>();
rst.induced_edges = eids;
......@@ -536,59 +554,58 @@ Subgraph Graph::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
rst.graph->AddEdge(src_id, dst_id);
}
for (uint64_t i = 0; i < NumVertices(); ++i)
nodes.push_back(i);
for (uint64_t i = 0; i < NumVertices(); ++i) nodes.push_back(i);
rst.induced_vertices = IdArray::Empty(
{static_cast<int64_t>(nodes.size())}, eids->dtype, eids->ctx);
std::copy(nodes.begin(), nodes.end(), static_cast<int64_t*>(rst.induced_vertices->data));
std::copy(
nodes.begin(), nodes.end(),
static_cast<int64_t*>(rst.induced_vertices->data));
}
return rst;
}
std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const {
std::vector<IdArray> Graph::GetAdj(
bool transpose, const std::string& fmt) const {
uint64_t num_edges = NumEdges();
uint64_t num_nodes = NumVertices();
if (fmt == "coo") {
IdArray idx = IdArray::Empty(
{2 * static_cast<int64_t>(num_edges)},
DGLDataType{kDGLInt, 64, 1},
{2 * static_cast<int64_t>(num_edges)}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
int64_t *idx_data = static_cast<int64_t*>(idx->data);
int64_t* idx_data = static_cast<int64_t*>(idx->data);
if (transpose) {
std::copy(all_edges_src_.begin(), all_edges_src_.end(), idx_data);
std::copy(all_edges_dst_.begin(), all_edges_dst_.end(), idx_data + num_edges);
std::copy(
all_edges_dst_.begin(), all_edges_dst_.end(), idx_data + num_edges);
} else {
std::copy(all_edges_dst_.begin(), all_edges_dst_.end(), idx_data);
std::copy(all_edges_src_.begin(), all_edges_src_.end(), idx_data + num_edges);
std::copy(
all_edges_src_.begin(), all_edges_src_.end(), idx_data + num_edges);
}
IdArray eid = IdArray::Empty(
{static_cast<int64_t>(num_edges)},
DGLDataType{kDGLInt, 64, 1},
{static_cast<int64_t>(num_edges)}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
int64_t *eid_data = static_cast<int64_t*>(eid->data);
int64_t* eid_data = static_cast<int64_t*>(eid->data);
for (uint64_t eid = 0; eid < num_edges; ++eid) {
eid_data[eid] = eid;
}
return std::vector<IdArray>{idx, eid};
} else if (fmt == "csr") {
IdArray indptr = IdArray::Empty(
{static_cast<int64_t>(num_nodes) + 1},
DGLDataType{kDGLInt, 64, 1},
{static_cast<int64_t>(num_nodes) + 1}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
IdArray indices = IdArray::Empty(
{static_cast<int64_t>(num_edges)},
DGLDataType{kDGLInt, 64, 1},
{static_cast<int64_t>(num_edges)}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
IdArray eid = IdArray::Empty(
{static_cast<int64_t>(num_edges)},
DGLDataType{kDGLInt, 64, 1},
{static_cast<int64_t>(num_edges)}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
int64_t *indptr_data = static_cast<int64_t*>(indptr->data);
int64_t *indices_data = static_cast<int64_t*>(indices->data);
int64_t *eid_data = static_cast<int64_t*>(eid->data);
const AdjacencyList *adjlist;
int64_t* indptr_data = static_cast<int64_t*>(indptr->data);
int64_t* indices_data = static_cast<int64_t*>(indices->data);
int64_t* eid_data = static_cast<int64_t*>(eid->data);
const AdjacencyList* adjlist;
if (transpose) {
// Out-edges.
adjlist = &adjlist_;
......@@ -599,9 +616,11 @@ std::vector<IdArray> Graph::GetAdj(bool transpose, const std::string &fmt) const
indptr_data[0] = 0;
for (size_t i = 0; i < adjlist->size(); i++) {
indptr_data[i + 1] = indptr_data[i] + adjlist->at(i).succ.size();
std::copy(adjlist->at(i).succ.begin(), adjlist->at(i).succ.end(),
std::copy(
adjlist->at(i).succ.begin(), adjlist->at(i).succ.end(),
indices_data + indptr_data[i]);
std::copy(adjlist->at(i).edge_id.begin(), adjlist->at(i).edge_id.end(),
std::copy(
adjlist->at(i).edge_id.begin(), adjlist->at(i).edge_id.end(),
eid_data + indptr_data[i]);
}
return std::vector<IdArray>{indptr, indices, eid};
......
src/graph/graph_apis.cc
View file @ 81831111
......@@ -3,72 +3,74 @@
* \file graph/graph.cc
* \brief DGL graph index APIs
*/
#include <dgl/packed_func_ext.h>
#include <dgl/graph.h>
#include <dgl/immutable_graph.h>
#include <dgl/graph_op.h>
#include <dgl/sampler.h>
#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>
#include <dgl/packed_func_ext.h>
#include <dgl/sampler.h>
#include "../c_api_common.h"
using dgl::runtime::DGLArgs;
using dgl::runtime::DGLArgValue;
using dgl::runtime::DGLRetValue;
using dgl::runtime::PackedFunc;
using dgl::runtime::NDArray;
using dgl::runtime::PackedFunc;
namespace dgl {
///////////////////////////// Graph API ///////////////////////////////////
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreateMutable")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
*rv = GraphRef(Graph::Create());
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const IdArray src_ids = args[0];
const IdArray dst_ids = args[1];
const int64_t num_nodes = args[2];
const bool readonly = args[3];
if (readonly) {
*rv = GraphRef(ImmutableGraph::CreateFromCOO(num_nodes, src_ids, dst_ids));
*rv = GraphRef(
ImmutableGraph::CreateFromCOO(num_nodes, src_ids, dst_ids));
} else {
*rv = GraphRef(Graph::CreateFromCOO(num_nodes, src_ids, dst_ids));
}
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const IdArray indptr = args[0];
const IdArray indices = args[1];
const std::string edge_dir = args[2];
IdArray edge_ids = IdArray::Empty({indices->shape[0]},
DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t *edge_data = static_cast<int64_t *>(edge_ids->data);
for (int64_t i = 0; i < edge_ids->shape[0]; i++)
edge_data[i] = i;
*rv = GraphRef(ImmutableGraph::CreateFromCSR(indptr, indices, edge_ids, edge_dir));
IdArray edge_ids = IdArray::Empty(
{indices->shape[0]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
int64_t* edge_data = static_cast<int64_t*>(edge_ids->data);
for (int64_t i = 0; i < edge_ids->shape[0]; i++) edge_data[i] = i;
*rv = GraphRef(
ImmutableGraph::CreateFromCSR(indptr, indices, edge_ids, edge_dir));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreateMMap")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const std::string shared_mem_name = args[0];
*rv = GraphRef(ImmutableGraph::CreateFromCSR(shared_mem_name));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
uint64_t num_vertices = args[1];
g->AddVertices(num_vertices);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t src = args[1];
const dgl_id_t dst = args[2];
......@@ -76,7 +78,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray src = args[1];
const IdArray dst = args[2];
......@@ -84,51 +86,51 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
g->Clear();
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = g->IsMultigraph();
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsReadonly")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = g->IsReadonly();
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = static_cast<int64_t>(g->NumVertices());
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = static_cast<int64_t>(g->NumEdges());
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
*rv = g->HasVertex(vid);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
*rv = g->HasVertices(vids);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t src = args[1];
const dgl_id_t dst = args[2];
......@@ -136,7 +138,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray src = args[1];
const IdArray dst = args[2];
......@@ -144,7 +146,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
const uint64_t radius = args[2];
......@@ -152,7 +154,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
const uint64_t radius = args[2];
......@@ -160,7 +162,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t src = args[1];
const dgl_id_t dst = args[2];
......@@ -168,7 +170,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray src = args[1];
const IdArray dst = args[2];
......@@ -176,89 +178,89 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t eid = args[1];
const auto& pair = g->FindEdge(eid);
*rv = PackedFunc([pair] (DGLArgs args, DGLRetValue* rv) {
*rv = PackedFunc([pair](DGLArgs args, DGLRetValue* rv) {
const int choice = args[0];
const int64_t ret = (choice == 0? pair.first : pair.second);
const int64_t ret = (choice == 0 ? pair.first : pair.second);
*rv = ret;
});
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray eids = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->FindEdges(eids));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->InEdges(vid));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->InEdges(vids));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->OutEdges(vid));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->OutEdges(vids));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
std::string order = args[1];
*rv = ConvertEdgeArrayToPackedFunc(g->Edges(order));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
*rv = static_cast<int64_t>(g->InDegree(vid));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
*rv = g->InDegrees(vids);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const dgl_id_t vid = args[1];
*rv = static_cast<int64_t>(g->OutDegree(vid));
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
*rv = g->OutDegrees(vids);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray vids = args[1];
std::shared_ptr<Subgraph> subg(new Subgraph(g->VertexSubgraph(vids)));
......@@ -266,7 +268,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray eids = args[1];
bool preserve_nodes = args[2];
......@@ -276,7 +278,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
bool transpose = args[1];
std::string format = args[2];
......@@ -285,13 +287,13 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphContext")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = g->Context();
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = g->NumBits();
});
......@@ -299,25 +301,25 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumBits")
// Subgraph C APIs
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
SubgraphRef subg = args[0];
*rv = GraphRef(subg->graph);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetInducedVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
SubgraphRef subg = args[0];
*rv = subg->induced_vertices;
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSubgraphGetInducedEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
SubgraphRef subg = args[0];
*rv = subg->induced_edges;
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLSortAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
g->SortCSR();
});
......
src/graph/graph_op.cc
View file @ 81831111
......@@ -3,13 +3,15 @@
* \file graph/graph.cc
* \brief Graph operation implementation
*/
#include <dgl/graph_op.h>
#include <dgl/array.h>
#include <dgl/graph_op.h>
#include <dgl/immutable_graph.h>
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
#include <dgl/runtime/parallel_for.h>
#include <algorithm>
#include "../c_api_common.h"
using namespace dgl::runtime;
......@@ -19,7 +21,7 @@ namespace {
// generate consecutive dgl ids
class RangeIter : public std::iterator<std::input_iterator_tag, dgl_id_t> {
public:
explicit RangeIter(dgl_id_t from): cur_(from) {}
explicit RangeIter(dgl_id_t from) : cur_(from) {}
RangeIter& operator++() {
++cur_;
......@@ -31,15 +33,9 @@ class RangeIter : public std::iterator<std::input_iterator_tag, dgl_id_t> {
++cur_;
return retval;
}
bool operator==(RangeIter other) const {
return cur_ == other.cur_;
}
bool operator!=(RangeIter other) const {
return cur_ != other.cur_;
}
dgl_id_t operator*() const {
return cur_;
}
bool operator==(RangeIter other) const { return cur_ == other.cur_; }
bool operator!=(RangeIter other) const { return cur_ != other.cur_; }
dgl_id_t operator*() const { return cur_; }
private:
dgl_id_t cur_;
......@@ -88,7 +84,8 @@ GraphPtr GraphOp::DisjointUnion(std::vector<GraphPtr> graphs) {
rst->AddVertices(gr->NumVertices());
for (uint64_t i = 0; i < gr->NumEdges(); ++i) {
// TODO(minjie): quite ugly to expose internal members
rst->AddEdge(mg->all_edges_src_[i] + cumsum, mg->all_edges_dst_[i] + cumsum);
rst->AddEdge(
mg->all_edges_src_[i] + cumsum, mg->all_edges_dst_[i] + cumsum);
}
cumsum += gr->NumVertices();
}
......@@ -136,14 +133,17 @@ GraphPtr GraphOp::DisjointUnion(std::vector<GraphPtr> graphs) {
cum_num_edges += g_num_edges;
}
return ImmutableGraph::CreateFromCSR(indptr_arr, indices_arr, edge_ids_arr, "in");
return ImmutableGraph::CreateFromCSR(
indptr_arr, indices_arr, edge_ids_arr, "in");
}
}
std::vector<GraphPtr> GraphOp::DisjointPartitionByNum(GraphPtr graph, int64_t num) {
std::vector<GraphPtr> GraphOp::DisjointPartitionByNum(
GraphPtr graph, int64_t num) {
CHECK(num != 0 && graph->NumVertices() % num == 0)
<< "Number of partitions must evenly divide the number of nodes.";
IdArray sizes = IdArray::Empty({num}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray sizes = IdArray::Empty(
{num}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
int64_t* sizes_data = static_cast<int64_t*>(sizes->data);
std::fill(sizes_data, sizes_data + num, graph->NumVertices() / num);
return DisjointPartitionBySizes(graph, sizes);
......@@ -170,10 +170,11 @@ std::vector<GraphPtr> GraphOp::DisjointPartitionBySizes(
MutableGraphPtr mg = Graph::Create();
// TODO(minjie): quite ugly to expose internal members
// copy adj
mg->adjlist_.insert(mg->adjlist_.end(),
graph->adjlist_.begin() + node_offset,
mg->adjlist_.insert(
mg->adjlist_.end(), graph->adjlist_.begin() + node_offset,
graph->adjlist_.begin() + node_offset + sizes_data[i]);
mg->reverse_adjlist_.insert(mg->reverse_adjlist_.end(),
mg->reverse_adjlist_.insert(
mg->reverse_adjlist_.end(),
graph->reverse_adjlist_.begin() + node_offset,
graph->reverse_adjlist_.begin() + node_offset + sizes_data[i]);
// relabel adjs
......@@ -209,12 +210,15 @@ std::vector<GraphPtr> GraphOp::DisjointPartitionBySizes(
}
} else {
// Input is an immutable graph. Partition it into several multiple graphs.
ImmutableGraphPtr graph = std::dynamic_pointer_cast<ImmutableGraph>(batched_graph);
ImmutableGraphPtr graph =
std::dynamic_pointer_cast<ImmutableGraph>(batched_graph);
// TODO(minjie): why in csr?
CSRPtr in_csr_ptr = graph->GetInCSR();
const dgl_id_t* indptr = static_cast<dgl_id_t*>(in_csr_ptr->indptr()->data);
const dgl_id_t* indices = static_cast<dgl_id_t*>(in_csr_ptr->indices()->data);
const dgl_id_t* edge_ids = static_cast<dgl_id_t*>(in_csr_ptr->edge_ids()->data);
const dgl_id_t* indices =
static_cast<dgl_id_t*>(in_csr_ptr->indices()->data);
const dgl_id_t* edge_ids =
static_cast<dgl_id_t*>(in_csr_ptr->edge_ids()->data);
dgl_id_t cum_sum_edges = 0;
for (int64_t i = 0; i < len; ++i) {
const int64_t start_pos = cumsum[i];
......@@ -257,7 +261,8 @@ IdArray GraphOp::MapParentIdToSubgraphId(IdArray parent_vids, IdArray query) {
const auto query_len = query->shape[0];
const dgl_id_t* parent_data = static_cast<dgl_id_t*>(parent_vids->data);
const dgl_id_t* query_data = static_cast<dgl_id_t*>(query->data);
IdArray rst = IdArray::Empty({query_len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray rst = IdArray::Empty(
{query_len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
dgl_id_t* rst_data = static_cast<dgl_id_t*>(rst->data);
const bool is_sorted = std::is_sorted(parent_data, parent_data + parent_len);
......@@ -300,11 +305,12 @@ IdArray GraphOp::ExpandIds(IdArray ids, IdArray offset) {
const auto id_len = ids->shape[0];
const auto off_len = offset->shape[0];
CHECK_EQ(id_len + 1, off_len);
const dgl_id_t *id_data = static_cast<dgl_id_t*>(ids->data);
const dgl_id_t *off_data = static_cast<dgl_id_t*>(offset->data);
const dgl_id_t* id_data = static_cast<dgl_id_t*>(ids->data);
const dgl_id_t* off_data = static_cast<dgl_id_t*>(offset->data);
const int64_t len = off_data[off_len - 1];
IdArray rst = IdArray::Empty({len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
dgl_id_t *rst_data = static_cast<dgl_id_t*>(rst->data);
IdArray rst = IdArray::Empty(
{len}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
dgl_id_t* rst_data = static_cast<dgl_id_t*>(rst->data);
for (int64_t i = 0; i < id_len; i++) {
const int64_t local_len = off_data[i + 1] - off_data[i];
for (int64_t j = 0; j < local_len; j++) {
......@@ -325,10 +331,11 @@ GraphPtr GraphOp::ToSimpleGraph(GraphPtr graph) {
hashmap.insert(dst);
}
}
indptr[src+1] = indices.size();
indptr[src + 1] = indices.size();
}
CSRPtr csr(new CSR(graph->NumVertices(), indices.size(),
indptr.begin(), indices.begin(), RangeIter(0)));
CSRPtr csr(new CSR(
graph->NumVertices(), indices.size(), indptr.begin(), indices.begin(),
RangeIter(0)));
return std::make_shared<ImmutableGraph>(csr);
}
......@@ -403,8 +410,9 @@ GraphPtr GraphOp::ToBidirectedImmutableGraph(GraphPtr g) {
g->NumVertices(), srcs_array, dsts_array);
}
HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hops) {
const dgl_id_t *nid = static_cast<dgl_id_t *>(nodes->data);
HaloSubgraph GraphOp::GetSubgraphWithHalo(
GraphPtr g, IdArray nodes, int num_hops) {
const dgl_id_t* nid = static_cast<dgl_id_t*>(nodes->data);
const auto id_len = nodes->shape[0];
// A map contains all nodes in the subgraph.
// The key is the old node Ids, the value indicates whether a node is a inner
......@@ -414,8 +422,7 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
// vector. The first few nodes are the inner nodes in the subgraph.
std::vector<dgl_id_t> old_node_ids(nid, nid + id_len);
std::vector<std::vector<dgl_id_t>> outer_nodes(num_hops);
for (int64_t i = 0; i < id_len; i++)
all_nodes[nid[i]] = true;
for (int64_t i = 0; i < id_len; i++) all_nodes[nid[i]] = true;
auto orig_nodes = all_nodes;
std::vector<dgl_id_t> edge_src, edge_dst, edge_eid;
......@@ -428,9 +435,9 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
auto dst = in_edges.dst;
auto eid = in_edges.id;
auto num_edges = eid->shape[0];
const dgl_id_t *src_data = static_cast<dgl_id_t *>(src->data);
const dgl_id_t *dst_data = static_cast<dgl_id_t *>(dst->data);
const dgl_id_t *eid_data = static_cast<dgl_id_t *>(eid->data);
const dgl_id_t* src_data = static_cast<dgl_id_t*>(src->data);
const dgl_id_t* dst_data = static_cast<dgl_id_t*>(dst->data);
const dgl_id_t* eid_data = static_cast<dgl_id_t*>(eid->data);
for (int64_t i = 0; i < num_edges; i++) {
// We check if the source node is in the original node.
auto it1 = orig_nodes.find(src_data[i]);
......@@ -451,15 +458,15 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
// Now we need to traverse the graph with the in-edges to access nodes
// and edges more hops away.
for (int k = 1; k < num_hops; k++) {
const std::vector<dgl_id_t> &nodes = outer_nodes[k-1];
const std::vector<dgl_id_t>& nodes = outer_nodes[k - 1];
EdgeArray in_edges = g->InEdges(aten::VecToIdArray(nodes));
auto src = in_edges.src;
auto dst = in_edges.dst;
auto eid = in_edges.id;
auto num_edges = eid->shape[0];
const dgl_id_t *src_data = static_cast<dgl_id_t *>(src->data);
const dgl_id_t *dst_data = static_cast<dgl_id_t *>(dst->data);
const dgl_id_t *eid_data = static_cast<dgl_id_t *>(eid->data);
const dgl_id_t* src_data = static_cast<dgl_id_t*>(src->data);
const dgl_id_t* dst_data = static_cast<dgl_id_t*>(dst->data);
const dgl_id_t* eid_data = static_cast<dgl_id_t*>(eid->data);
for (int64_t i = 0; i < num_edges; i++) {
edge_src.push_back(src_data[i]);
edge_dst.push_back(dst_data[i]);
......@@ -482,12 +489,12 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
}
num_edges = edge_src.size();
IdArray new_src = IdArray::Empty({num_edges}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
IdArray new_dst = IdArray::Empty({num_edges}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0});
dgl_id_t *new_src_data = static_cast<dgl_id_t *>(new_src->data);
dgl_id_t *new_dst_data = static_cast<dgl_id_t *>(new_dst->data);
IdArray new_src = IdArray::Empty(
{num_edges}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
IdArray new_dst = IdArray::Empty(
{num_edges}, DGLDataType{kDGLInt, 64, 1}, DGLContext{kDGLCPU, 0});
dgl_id_t* new_src_data = static_cast<dgl_id_t*>(new_src->data);
dgl_id_t* new_dst_data = static_cast<dgl_id_t*>(new_dst->data);
for (size_t i = 0; i < edge_src.size(); i++) {
new_src_data[i] = old2new[edge_src[i]];
new_dst_data[i] = old2new[edge_dst[i]];
......@@ -499,7 +506,8 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
inner_nodes[i] = all_nodes[old_nid];
}
GraphPtr subg = ImmutableGraph::CreateFromCOO(old_node_ids.size(), new_src, new_dst);
GraphPtr subg =
ImmutableGraph::CreateFromCOO(old_node_ids.size(), new_src, new_dst);
HaloSubgraph halo_subg;
halo_subg.graph = subg;
halo_subg.induced_vertices = aten::VecToIdArray(old_node_ids);
......@@ -509,7 +517,8 @@ HaloSubgraph GraphOp::GetSubgraphWithHalo(GraphPtr g, IdArray nodes, int num_hop
return halo_subg;
}
GraphPtr GraphOp::ReorderImmutableGraph(ImmutableGraphPtr ig, IdArray new_order) {
GraphPtr GraphOp::ReorderImmutableGraph(
ImmutableGraphPtr ig, IdArray new_order) {
CSRPtr in_csr, out_csr;
COOPtr coo;
// We only need to reorder one of the graph structure.
......@@ -517,12 +526,14 @@ GraphPtr GraphOp::ReorderImmutableGraph(ImmutableGraphPtr ig, IdArray new_order)
in_csr = ig->GetInCSR();
auto csrmat = in_csr->ToCSRMatrix();
auto new_csrmat = aten::CSRReorder(csrmat, new_order, new_order);
in_csr = CSRPtr(new CSR(new_csrmat.indptr, new_csrmat.indices, new_csrmat.data));
in_csr =
CSRPtr(new CSR(new_csrmat.indptr, new_csrmat.indices, new_csrmat.data));
} else if (ig->HasOutCSR()) {
out_csr = ig->GetOutCSR();
auto csrmat = out_csr->ToCSRMatrix();
auto new_csrmat = aten::CSRReorder(csrmat, new_order, new_order);
out_csr = CSRPtr(new CSR(new_csrmat.indptr, new_csrmat.indices, new_csrmat.data));
out_csr =
CSRPtr(new CSR(new_csrmat.indptr, new_csrmat.indices, new_csrmat.data));
} else {
coo = ig->GetCOO();
auto coomat = coo->ToCOOMatrix();
......@@ -536,14 +547,14 @@ GraphPtr GraphOp::ReorderImmutableGraph(ImmutableGraphPtr ig, IdArray new_order)
}
DGL_REGISTER_GLOBAL("transform._CAPI_DGLPartitionWithHalo")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef graph = args[0];
IdArray node_parts = args[1];
int num_hops = args[2];
const dgl_id_t *part_data = static_cast<dgl_id_t *>(node_parts->data);
const dgl_id_t* part_data = static_cast<dgl_id_t*>(node_parts->data);
int64_t num_nodes = node_parts->shape[0];
std::unordered_map<int, std::vector<dgl_id_t> > part_map;
std::unordered_map<int, std::vector<dgl_id_t>> part_map;
for (int64_t i = 0; i < num_nodes; i++) {
dgl_id_t part_id = part_data[i];
auto it = part_map.find(part_id);
......@@ -556,7 +567,7 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLPartitionWithHalo")
}
}
std::vector<int> part_ids;
std::vector<std::vector<dgl_id_t> > part_nodes;
std::vector<std::vector<dgl_id_t>> part_nodes;
int max_part_id = 0;
for (auto it = part_map.begin(); it != part_map.end(); it++) {
max_part_id = std::max(it->first, max_part_id);
......@@ -570,12 +581,13 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLPartitionWithHalo")
// try to construct in-CSR in openmp for loop, which will lead
// to some unexpected results.
graph_ptr->GetInCSR();
std::vector<std::shared_ptr<HaloSubgraph> > subgs(max_part_id + 1);
std::vector<std::shared_ptr<HaloSubgraph>> subgs(max_part_id + 1);
int num_partitions = part_nodes.size();
runtime::parallel_for(0, num_partitions, [&](size_t b, size_t e) {
for (auto i = b; i < e; ++i) {
auto nodes = aten::VecToIdArray(part_nodes[i]);
HaloSubgraph subg = GraphOp::GetSubgraphWithHalo(graph_ptr, nodes, num_hops);
HaloSubgraph subg =
GraphOp::GetSubgraphWithHalo(graph_ptr, nodes, num_hops);
std::shared_ptr<HaloSubgraph> subg_ptr(new HaloSubgraph(subg));
int part_id = part_ids[i];
subgs[part_id] = subg_ptr;
......@@ -589,25 +601,26 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLPartitionWithHalo")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGetSubgraphWithHalo")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef graph = args[0];
IdArray nodes = args[1];
int num_hops = args[2];
HaloSubgraph subg = GraphOp::GetSubgraphWithHalo(graph.sptr(), nodes, num_hops);
HaloSubgraph subg =
GraphOp::GetSubgraphWithHalo(graph.sptr(), nodes, num_hops);
std::shared_ptr<HaloSubgraph> subg_ptr(new HaloSubgraph(subg));
*rv = SubgraphRef(subg_ptr);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_GetHaloSubgraphInnerNodes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
SubgraphRef g = args[0];
auto gptr = std::dynamic_pointer_cast<HaloSubgraph>(g.sptr());
CHECK(gptr) << "The input graph has to be immutable graph";
*rv = gptr->inner_nodes;
});
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
List<GraphRef> graphs = args[0];
std::vector<GraphPtr> ptrs(graphs.size());
for (size_t i = 0; i < graphs.size(); ++i) {
......@@ -617,7 +630,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
int64_t num = args[1];
const auto& ret = GraphOp::DisjointPartitionByNum(g.sptr(), num);
......@@ -629,7 +642,7 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray sizes = args[1];
const auto& ret = GraphOp::DisjointPartitionBySizes(g.sptr(), sizes);
......@@ -638,35 +651,35 @@ DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
ret_list.push_back(GraphRef(gp));
}
*rv = ret_list;
});
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
bool backtracking = args[1];
*rv = GraphOp::LineGraph(g.sptr(), backtracking);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLToImmutable")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = ImmutableGraph::ToImmutable(g.sptr());
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLToSimpleGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = GraphOp::ToSimpleGraph(g.sptr());
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLToBidirectedMutableGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
*rv = GraphOp::ToBidirectedMutableGraph(g.sptr());
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLReorderGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
const IdArray new_order = args[1];
auto gptr = std::dynamic_pointer_cast<ImmutableGraph>(g.sptr());
......@@ -675,7 +688,7 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLReorderGraph")
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLReassignEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef graph = args[0];
bool is_incsr = args[1];
auto gptr = std::dynamic_pointer_cast<ImmutableGraph>(graph.sptr());
......@@ -683,7 +696,8 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLReassignEdges")
CSRPtr csr = is_incsr ? gptr->GetInCSR() : gptr->GetOutCSR();
auto csrmat = csr->ToCSRMatrix();
int64_t num_edges = csrmat.data->shape[0];
IdArray new_data = IdArray::Empty({num_edges}, csrmat.data->dtype, csrmat.data->ctx);
IdArray new_data =
IdArray::Empty({num_edges}, csrmat.data->dtype, csrmat.data->ctx);
// Return the original edge Ids.
*rv = new_data;
// TODO(zhengda) I need to invalidate out-CSR and COO.
......@@ -692,8 +706,8 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLReassignEdges")
// TODO(zhengda) after assignment, we actually don't need to store them
// physically.
ATEN_ID_TYPE_SWITCH(new_data->dtype, IdType, {
IdType *typed_new_data = static_cast<IdType*>(new_data->data);
IdType *typed_data = static_cast<IdType*>(csrmat.data->data);
IdType* typed_new_data = static_cast<IdType*>(new_data->data);
IdType* typed_data = static_cast<IdType*>(csrmat.data->data);
for (int64_t i = 0; i < num_edges; i++) {
typed_new_data[i] = typed_data[i];
typed_data[i] = i;
......@@ -702,7 +716,7 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLReassignEdges")
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLToBidirectedImmutableGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef g = args[0];
auto gptr = g.sptr();
auto immutable_g = std::dynamic_pointer_cast<ImmutableGraph>(gptr);
......@@ -719,29 +733,31 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLToBidirectedImmutableGraph")
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const IdArray parent_vids = args[0];
const IdArray query = args[1];
*rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
});
template<class IdType>
IdArray MapIds(IdArray ids, IdArray range_starts, IdArray range_ends, IdArray typed_map,
template <class IdType>
IdArray MapIds(
IdArray ids, IdArray range_starts, IdArray range_ends, IdArray typed_map,
int num_parts, int num_types) {
int64_t num_ids = ids->shape[0];
int64_t num_ranges = range_starts->shape[0];
IdArray ret = IdArray::Empty({num_ids * 2}, ids->dtype, ids->ctx);
const IdType *range_start_data = static_cast<IdType *>(range_starts->data);
const IdType *range_end_data = static_cast<IdType *>(range_ends->data);
const IdType *ids_data = static_cast<IdType *>(ids->data);
const IdType *typed_map_data = static_cast<IdType *>(typed_map->data);
IdType *types_data = static_cast<IdType *>(ret->data);
IdType *per_type_ids_data = static_cast<IdType *>(ret->data) + num_ids;
const IdType* range_start_data = static_cast<IdType*>(range_starts->data);
const IdType* range_end_data = static_cast<IdType*>(range_ends->data);
const IdType* ids_data = static_cast<IdType*>(ids->data);
const IdType* typed_map_data = static_cast<IdType*>(typed_map->data);
IdType* types_data = static_cast<IdType*>(ret->data);
IdType* per_type_ids_data = static_cast<IdType*>(ret->data) + num_ids;
runtime::parallel_for(0, ids->shape[0], [&](size_t b, size_t e) {
for (auto i = b; i < e; ++i) {
IdType id = ids_data[i];
auto it = std::lower_bound(range_end_data, range_end_data + num_ranges, id);
auto it =
std::lower_bound(range_end_data, range_end_data + num_ranges, id);
// The range must exist.
BUG_IF_FAIL(it != range_end_data + num_ranges);
size_t range_id = it - range_end_data;
......@@ -752,8 +768,9 @@ IdArray MapIds(IdArray ids, IdArray range_starts, IdArray range_ends, IdArray ty
if (part_id == 0) {
per_type_ids_data[i] = id - range_start_data[range_id];
} else {
per_type_ids_data[i] = id - range_start_data[range_id]
+ typed_map_data[num_parts * type_id + part_id - 1];
per_type_ids_data[i] =
id - range_start_data[range_id] +
typed_map_data[num_parts * type_id + part_id - 1];
}
}
});
......@@ -761,7 +778,7 @@ IdArray MapIds(IdArray ids, IdArray range_starts, IdArray range_ends, IdArray ty
}
DGL_REGISTER_GLOBAL("distributed.id_map._CAPI_DGLHeteroMapIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
const IdArray ids = args[0];
const IdArray range_starts = args[1];
const IdArray range_ends = args[2];
......@@ -778,7 +795,8 @@ DGL_REGISTER_GLOBAL("distributed.id_map._CAPI_DGLHeteroMapIds")
IdArray ret;
ATEN_ID_TYPE_SWITCH(ids->dtype, IdType, {
ret = MapIds<IdType>(ids, range_starts, range_ends, typed_map, num_parts, num_types);
ret = MapIds<IdType>(
ids, range_starts, range_ends, typed_map, num_parts, num_types);
});
*rv = ret;
});
......
src/graph/graph_traversal.cc
View file @ 81831111
......@@ -5,6 +5,7 @@
*/
#include <dgl/graph_traversal.h>
#include <dgl/packed_func_ext.h>
#include "../c_api_common.h"
using namespace dgl::runtime;
......@@ -13,7 +14,7 @@ namespace dgl {
namespace traverse {
DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSNodes_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef g = args[0];
const IdArray src = args[1];
bool reversed = args[2];
......@@ -28,7 +29,7 @@ DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSNodes_v2")
});
DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSEdges_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef g = args[0];
const IdArray src = args[1];
bool reversed = args[2];
......@@ -44,7 +45,7 @@ DGL_REGISTER_GLOBAL("traversal._CAPI_DGLBFSEdges_v2")
});
DGL_REGISTER_GLOBAL("traversal._CAPI_DGLTopologicalNodes_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef g = args[0];
bool reversed = args[1];
aten::CSRMatrix csr;
......@@ -58,9 +59,8 @@ DGL_REGISTER_GLOBAL("traversal._CAPI_DGLTopologicalNodes_v2")
*rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
});
DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSEdges_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef g = args[0];
const IdArray source = args[1];
const bool reversed = args[2];
......@@ -76,7 +76,7 @@ DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSEdges_v2")
});
DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSLabeledEdges_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef g = args[0];
const IdArray source = args[1];
const bool reversed = args[2];
......@@ -90,14 +90,12 @@ DGL_REGISTER_GLOBAL("traversal._CAPI_DGLDFSLabeledEdges_v2")
csr = g.sptr()->GetCSRMatrix(0);
}
const auto& front = aten::DGLDFSLabeledEdges(csr,
source,
has_reverse_edge,
has_nontree_edge,
return_labels);
const auto& front = aten::DGLDFSLabeledEdges(
csr, source, has_reverse_edge, has_nontree_edge, return_labels);
if (return_labels) {
*rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.tags, front.sections});
*rv = ConvertNDArrayVectorToPackedFunc(
{front.ids, front.tags, front.sections});
} else {
*rv = ConvertNDArrayVectorToPackedFunc({front.ids, front.sections});
}
......
src/graph/heterograph.cc
View file @ 81831111
......@@ -4,14 +4,16 @@
* \brief Heterograph implementation
*/
#include "./heterograph.h"
#include <dgl/array.h>
#include <dgl/immutable_graph.h>
#include <dgl/graph_serializer.h>
#include <dgl/immutable_graph.h>
#include <dmlc/memory_io.h>
#include <memory>
#include <vector>
#include <tuple>
#include <utility>
#include <vector>
using namespace dgl::runtime;
......@@ -23,7 +25,8 @@ using dgl::ImmutableGraph;
HeteroSubgraph EdgeSubgraphPreserveNodes(
const HeteroGraph* hg, const std::vector<IdArray>& eids) {
CHECK_EQ(eids.size(), hg->NumEdgeTypes())
<< "Invalid input: the input list size must be the same as the number of edge type.";
<< "Invalid input: the input list size must be the same as the number of "
"edge type.";
HeteroSubgraph ret;
ret.induced_vertices.resize(hg->NumVertexTypes());
ret.induced_edges = eids;
......@@ -33,14 +36,14 @@ HeteroSubgraph EdgeSubgraphPreserveNodes(
auto pair = hg->meta_graph()->FindEdge(etype);
const dgl_type_t src_vtype = pair.first;
const dgl_type_t dst_vtype = pair.second;
const auto& rel_vsg = hg->GetRelationGraph(etype)->EdgeSubgraph(
{eids[etype]}, true);
const auto& rel_vsg =
hg->GetRelationGraph(etype)->EdgeSubgraph({eids[etype]}, true);
subrels[etype] = rel_vsg.graph;
ret.induced_vertices[src_vtype] = rel_vsg.induced_vertices[0];
ret.induced_vertices[dst_vtype] = rel_vsg.induced_vertices[1];
}
ret.graph = HeteroGraphPtr(new HeteroGraph(
hg->meta_graph(), subrels, hg->NumVerticesPerType()));
ret.graph = HeteroGraphPtr(
new HeteroGraph(hg->meta_graph(), subrels, hg->NumVerticesPerType()));
return ret;
}
......@@ -49,34 +52,36 @@ HeteroSubgraph EdgeSubgraphNoPreserveNodes(
// TODO(minjie): In general, all relabeling should be separated with subgraph
// operations.
CHECK_EQ(eids.size(), hg->NumEdgeTypes())
<< "Invalid input: the input list size must be the same as the number of edge type.";
<< "Invalid input: the input list size must be the same as the number of "
"edge type.";
HeteroSubgraph ret;
ret.induced_vertices.resize(hg->NumVertexTypes());
ret.induced_edges = eids;
// NOTE(minjie): EdgeSubgraph when preserve_nodes is false is quite complicated in
// heterograph. This is because we need to make sure bipartite graphs that incident
// on the same vertex type must have the same ID space. For example, suppose we have
// following heterograph:
// NOTE(minjie): EdgeSubgraph when preserve_nodes is false is quite
// complicated in heterograph. This is because we need to make sure bipartite
// graphs that incident on the same vertex type must have the same ID space.
// For example, suppose we have following heterograph:
//
// Meta graph: A -> B -> C
// UnitGraph graphs:
// * A -> B: (0, 0), (0, 1)
// * B -> C: (1, 0), (1, 1)
//
// Suppose for A->B, we only keep edge (0, 0), while for B->C we only keep (1, 0). We need
// to make sure that in the result subgraph, node type B still has two nodes. This means
// we cannot simply compute EdgeSubgraph for B->C which will relabel node#1 of type B to be
// node #0.
// Suppose for A->B, we only keep edge (0, 0), while for B->C we only keep (1,
// 0). We need to make sure that in the result subgraph, node type B still has
// two nodes. This means we cannot simply compute EdgeSubgraph for B->C which
// will relabel node#1 of type B to be node #0.
//
// One implementation is as follows:
// (1) For each bipartite graph, slice out the edges using the given eids.
// (2) Make a dictionary map<vtype, vector<IdArray>>, where the key is the vertex type
// and the value is the incident nodes from the bipartite graphs that has the vertex
// type as either srctype or dsttype.
// (2) Make a dictionary map<vtype, vector<IdArray>>, where the key is the
// vertex type
// and the value is the incident nodes from the bipartite graphs that has
// the vertex type as either srctype or dsttype.
// (3) Then for each vertex type, use aten::Relabel_ on its vector<IdArray>.
// aten::Relabel_ computes the union of the vertex sets and relabel
// the unique elements from zero. The returned mapping array is the final induced
// vertex set for that vertex type.
// the unique elements from zero. The returned mapping array is the final
// induced vertex set for that vertex type.
// (4) Use the relabeled edges to construct the bipartite graph.
// step (1) & (2)
std::vector<EdgeArray> subedges(hg->NumEdgeTypes());
......@@ -103,10 +108,9 @@ HeteroSubgraph EdgeSubgraphNoPreserveNodes(
const dgl_type_t src_vtype = pair.first;
const dgl_type_t dst_vtype = pair.second;
subrels[etype] = UnitGraph::CreateFromCOO(
(src_vtype == dst_vtype)? 1 : 2,
(src_vtype == dst_vtype) ? 1 : 2,
ret.induced_vertices[src_vtype]->shape[0],
ret.induced_vertices[dst_vtype]->shape[0],
subedges[etype].src,
ret.induced_vertices[dst_vtype]->shape[0], subedges[etype].src,
subedges[etype].dst);
}
ret.graph = HeteroGraphPtr(new HeteroGraph(
......@@ -114,36 +118,39 @@ HeteroSubgraph EdgeSubgraphNoPreserveNodes(
return ret;
}
void HeteroGraphSanityCheck(GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs) {
void HeteroGraphSanityCheck(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs) {
// Sanity check
CHECK_EQ(meta_graph->NumEdges(), rel_graphs.size());
CHECK(!rel_graphs.empty()) << "Empty heterograph is not allowed.";
// all relation graphs must have only one edge type
for (const auto &rg : rel_graphs) {
CHECK_EQ(rg->NumEdgeTypes(), 1) << "Each relation graph must have only one edge type.";
for (const auto& rg : rel_graphs) {
CHECK_EQ(rg->NumEdgeTypes(), 1)
<< "Each relation graph must have only one edge type.";
}
auto ctx = rel_graphs[0]->Context();
for (const auto &rg : rel_graphs) {
CHECK_EQ(rg->Context(), ctx) << "Each relation graph must have the same context.";
for (const auto& rg : rel_graphs) {
CHECK_EQ(rg->Context(), ctx)
<< "Each relation graph must have the same context.";
}
}
std::vector<int64_t>
InferNumVerticesPerType(GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs) {
std::vector<int64_t> InferNumVerticesPerType(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs) {
// create num verts per type
std::vector<int64_t> num_verts_per_type(meta_graph->NumVertices(), -1);
EdgeArray etype_array = meta_graph->Edges();
dgl_type_t *srctypes = static_cast<dgl_type_t *>(etype_array.src->data);
dgl_type_t *dsttypes = static_cast<dgl_type_t *>(etype_array.dst->data);
dgl_type_t *etypes = static_cast<dgl_type_t *>(etype_array.id->data);
dgl_type_t* srctypes = static_cast<dgl_type_t*>(etype_array.src->data);
dgl_type_t* dsttypes = static_cast<dgl_type_t*>(etype_array.dst->data);
dgl_type_t* etypes = static_cast<dgl_type_t*>(etype_array.id->data);
for (size_t i = 0; i < meta_graph->NumEdges(); ++i) {
dgl_type_t srctype = srctypes[i];
dgl_type_t dsttype = dsttypes[i];
dgl_type_t etype = etypes[i];
const auto& rg = rel_graphs[etype];
const auto sty = 0;
const auto dty = rg->NumVertexTypes() == 1? 0 : 1;
const auto dty = rg->NumVertexTypes() == 1 ? 0 : 1;
size_t nv;
// # nodes of source type
......@@ -164,7 +171,8 @@ InferNumVerticesPerType(GraphPtr meta_graph, const std::vector<HeteroGraphPtr>&
return num_verts_per_type;
}
std::vector<UnitGraphPtr> CastToUnitGraphs(const std::vector<HeteroGraphPtr>& rel_graphs) {
std::vector<UnitGraphPtr> CastToUnitGraphs(
const std::vector<HeteroGraphPtr>& rel_graphs) {
std::vector<UnitGraphPtr> relation_graphs(rel_graphs.size());
for (size_t i = 0; i < rel_graphs.size(); ++i) {
HeteroGraphPtr relg = rel_graphs[i];
......@@ -181,9 +189,9 @@ std::vector<UnitGraphPtr> CastToUnitGraphs(const std::vector<HeteroGraphPtr>& re
} // namespace
HeteroGraph::HeteroGraph(
GraphPtr meta_graph,
const std::vector<HeteroGraphPtr>& rel_graphs,
const std::vector<int64_t>& num_nodes_per_type) : BaseHeteroGraph(meta_graph) {
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs,
const std::vector<int64_t>& num_nodes_per_type)
: BaseHeteroGraph(meta_graph) {
if (num_nodes_per_type.size() == 0)
num_verts_per_type_ = InferNumVerticesPerType(meta_graph, rel_graphs);
else
......@@ -193,7 +201,7 @@ HeteroGraph::HeteroGraph(
}
bool HeteroGraph::IsMultigraph() const {
for (const auto &hg : relation_graphs_) {
for (const auto& hg : relation_graphs_) {
if (hg->IsMultigraph()) {
return true;
}
......@@ -206,9 +214,11 @@ BoolArray HeteroGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
return aten::LT(vids, NumVertices(vtype));
}
HeteroSubgraph HeteroGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
HeteroSubgraph HeteroGraph::VertexSubgraph(
const std::vector<IdArray>& vids) const {
CHECK_EQ(vids.size(), NumVertexTypes())
<< "Invalid input: the input list size must be the same as the number of vertex types.";
<< "Invalid input: the input list size must be the same as the number of "
"vertex types.";
HeteroSubgraph ret;
ret.induced_vertices = vids;
std::vector<int64_t> num_vertices_per_type(NumVertexTypes());
......@@ -220,15 +230,16 @@ HeteroSubgraph HeteroGraph::VertexSubgraph(const std::vector<IdArray>& vids) con
auto pair = meta_graph_->FindEdge(etype);
const dgl_type_t src_vtype = pair.first;
const dgl_type_t dst_vtype = pair.second;
const std::vector<IdArray> rel_vids = (src_vtype == dst_vtype) ?
std::vector<IdArray>({vids[src_vtype]}) :
std::vector<IdArray>({vids[src_vtype], vids[dst_vtype]});
const std::vector<IdArray> rel_vids =
(src_vtype == dst_vtype)
? std::vector<IdArray>({vids[src_vtype]})
: std::vector<IdArray>({vids[src_vtype], vids[dst_vtype]});
const auto& rel_vsg = GetRelationGraph(etype)->VertexSubgraph(rel_vids);
subrels[etype] = rel_vsg.graph;
ret.induced_edges[etype] = rel_vsg.induced_edges[0];
}
ret.graph = HeteroGraphPtr(new HeteroGraph(
meta_graph_, subrels, std::move(num_vertices_per_type)));
ret.graph = HeteroGraphPtr(
new HeteroGraph(meta_graph_, subrels, std::move(num_vertices_per_type)));
return ret;
}
......@@ -248,11 +259,11 @@ HeteroGraphPtr HeteroGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
for (auto g : hgindex->relation_graphs_) {
rel_graphs.push_back(UnitGraph::AsNumBits(g, bits));
}
return HeteroGraphPtr(new HeteroGraph(hgindex->meta_graph_, rel_graphs,
hgindex->num_verts_per_type_));
return HeteroGraphPtr(new HeteroGraph(
hgindex->meta_graph_, rel_graphs, hgindex->num_verts_per_type_));
}
HeteroGraphPtr HeteroGraph::CopyTo(HeteroGraphPtr g, const DGLContext &ctx) {
HeteroGraphPtr HeteroGraph::CopyTo(HeteroGraphPtr g, const DGLContext& ctx) {
if (ctx == g->Context()) {
return g;
}
......@@ -262,23 +273,20 @@ HeteroGraphPtr HeteroGraph::CopyTo(HeteroGraphPtr g, const DGLContext &ctx) {
for (auto g : hgindex->relation_graphs_) {
rel_graphs.push_back(UnitGraph::CopyTo(g, ctx));
}
return HeteroGraphPtr(new HeteroGraph(hgindex->meta_graph_, rel_graphs,
hgindex->num_verts_per_type_));
return HeteroGraphPtr(new HeteroGraph(
hgindex->meta_graph_, rel_graphs, hgindex->num_verts_per_type_));
}
void HeteroGraph::PinMemory_() {
for (auto g : relation_graphs_)
g->PinMemory_();
for (auto g : relation_graphs_) g->PinMemory_();
}
void HeteroGraph::UnpinMemory_() {
for (auto g : relation_graphs_)
g->UnpinMemory_();
for (auto g : relation_graphs_) g->UnpinMemory_();
}
void HeteroGraph::RecordStream(DGLStreamHandle stream) {
for (auto g : relation_graphs_)
g->RecordStream(stream);
for (auto g : relation_graphs_) g->RecordStream(stream);
}
std::string HeteroGraph::SharedMemName() const {
......@@ -286,13 +294,13 @@ std::string HeteroGraph::SharedMemName() const {
}
HeteroGraphPtr HeteroGraph::CopyToSharedMem(
HeteroGraphPtr g, const std::string& name, const std::vector<std::string>& ntypes,
HeteroGraphPtr g, const std::string& name,
const std::vector<std::string>& ntypes,
const std::vector<std::string>& etypes, const std::set<std::string>& fmts) {
// TODO(JJ): Raise error when calling shared_memory if graph index is on gpu
auto hg = std::dynamic_pointer_cast<HeteroGraph>(g);
CHECK_NOTNULL(hg);
if (hg->SharedMemName() == name)
return g;
if (hg->SharedMemName() == name) return g;
// Copy buffer to share memory
auto mem = std::make_shared<SharedMemory>(name);
......@@ -312,7 +320,7 @@ HeteroGraphPtr HeteroGraph::CopyToSharedMem(
std::vector<HeteroGraphPtr> relgraphs(g->NumEdgeTypes());
for (dgl_type_t etype = 0 ; etype < g->NumEdgeTypes() ; ++etype) {
for (dgl_type_t etype = 0; etype < g->NumEdgeTypes(); ++etype) {
auto src_dst_type = g->GetEndpointTypes(etype);
int num_vtypes = (src_dst_type.first == src_dst_type.second ? 1 : 2);
aten::COOMatrix coo;
......@@ -341,10 +349,11 @@ HeteroGraphPtr HeteroGraph::CopyToSharedMem(
}
std::tuple<HeteroGraphPtr, std::vector<std::string>, std::vector<std::string>>
HeteroGraph::CreateFromSharedMem(const std::string &name) {
HeteroGraph::CreateFromSharedMem(const std::string& name) {
bool exist = SharedMemory::Exist(name);
if (!exist) {
return std::make_tuple(nullptr, std::vector<std::string>(), std::vector<std::string>());
return std::make_tuple(
nullptr, std::vector<std::string>(), std::vector<std::string>());
}
auto mem = std::make_shared<SharedMemory>(name);
auto mem_buf = mem->Open(SHARED_MEM_METAINFO_SIZE_MAX);
......@@ -367,7 +376,7 @@ std::tuple<HeteroGraphPtr, std::vector<std::string>, std::vector<std::string>>
CHECK(shm.Read(&num_verts_per_type)) << "Invalid number of vertices per type";
std::vector<HeteroGraphPtr> relgraphs(metagraph->NumEdges());
for (dgl_type_t etype = 0 ; etype < metagraph->NumEdges() ; ++etype) {
for (dgl_type_t etype = 0; etype < metagraph->NumEdges(); ++etype) {
auto src_dst = metagraph->FindEdge(etype);
int num_vtypes = (src_dst.first == src_dst.second) ? 1 : 2;
aten::COOMatrix coo;
......@@ -387,7 +396,8 @@ std::tuple<HeteroGraphPtr, std::vector<std::string>, std::vector<std::string>>
num_vtypes, csc, csr, coo, has_csc, has_csr, has_coo);
}
auto ret = std::make_shared<HeteroGraph>(metagraph, relgraphs, num_verts_per_type);
auto ret =
std::make_shared<HeteroGraph>(metagraph, relgraphs, num_verts_per_type);
ret->shared_mem_ = mem;
std::vector<std::string> ntypes;
......@@ -400,11 +410,12 @@ std::tuple<HeteroGraphPtr, std::vector<std::string>, std::vector<std::string>>
HeteroGraphPtr HeteroGraph::GetGraphInFormat(dgl_format_code_t formats) const {
std::vector<HeteroGraphPtr> format_rels(NumEdgeTypes());
for (dgl_type_t etype = 0; etype < NumEdgeTypes(); ++etype) {
auto relgraph = std::dynamic_pointer_cast<UnitGraph>(GetRelationGraph(etype));
auto relgraph =
std::dynamic_pointer_cast<UnitGraph>(GetRelationGraph(etype));
format_rels[etype] = relgraph->GetGraphInFormat(formats);
}
return HeteroGraphPtr(new HeteroGraph(
meta_graph_, format_rels, NumVerticesPerType()));
return HeteroGraphPtr(
new HeteroGraph(meta_graph_, format_rels, NumVerticesPerType()));
}
FlattenedHeteroGraphPtr HeteroGraph::Flatten(
......@@ -418,15 +429,16 @@ FlattenedHeteroGraphPtr HeteroGraph::Flatten(
}
template <class IdType>
FlattenedHeteroGraphPtr HeteroGraph::FlattenImpl(const std::vector<dgl_type_t>& etypes) const {
FlattenedHeteroGraphPtr HeteroGraph::FlattenImpl(
const std::vector<dgl_type_t>& etypes) const {
std::unordered_map<dgl_type_t, size_t> srctype_offsets, dsttype_offsets;
size_t src_nodes = 0, dst_nodes = 0;
std::vector<dgl_type_t> induced_srctype, induced_dsttype;
std::vector<IdType> induced_srcid, induced_dstid;
std::vector<dgl_type_t> srctype_set, dsttype_set;
// XXXtype_offsets contain the mapping from node type and number of nodes after this
// loop.
// XXXtype_offsets contain the mapping from node type and number of nodes
// after this loop.
for (dgl_type_t etype : etypes) {
auto src_dsttype = meta_graph_->FindEdge(etype);
dgl_type_t srctype = src_dsttype.first;
......@@ -443,15 +455,16 @@ FlattenedHeteroGraphPtr HeteroGraph::FlattenImpl(const std::vector<dgl_type_t>&
dsttype_set.push_back(dsttype);
}
}
// Sort the node types so that we can compare the sets and decide whether a homogeneous graph
// should be returned.
// Sort the node types so that we can compare the sets and decide whether a
// homogeneous graph should be returned.
std::sort(srctype_set.begin(), srctype_set.end());
std::sort(dsttype_set.begin(), dsttype_set.end());
bool homograph = (srctype_set.size() == dsttype_set.size()) &&
bool homograph =
(srctype_set.size() == dsttype_set.size()) &&
std::equal(srctype_set.begin(), srctype_set.end(), dsttype_set.begin());
// XXXtype_offsets contain the mapping from node type to node ID offsets after these
// two loops.
// XXXtype_offsets contain the mapping from node type to node ID offsets after
// these two loops.
for (size_t i = 0; i < srctype_set.size(); ++i) {
dgl_type_t ntype = srctype_set[i];
size_t num_nodes = srctype_offsets[ntype];
......@@ -492,14 +505,12 @@ FlattenedHeteroGraphPtr HeteroGraph::FlattenImpl(const std::vector<dgl_type_t>&
src_arrs.push_back(edges.src + srctype_offset);
dst_arrs.push_back(edges.dst + dsttype_offset);
eid_arrs.push_back(edges.id);
induced_etypes.push_back(aten::Full(etype, num_edges, NumBits(), Context()));
induced_etypes.push_back(
aten::Full(etype, num_edges, NumBits(), Context()));
}
HeteroGraphPtr gptr = UnitGraph::CreateFromCOO(
homograph ? 1 : 2,
src_nodes,
dst_nodes,
aten::Concat(src_arrs),
homograph ? 1 : 2, src_nodes, dst_nodes, aten::Concat(src_arrs),
aten::Concat(dst_arrs));
// Sanity check
......@@ -507,15 +518,20 @@ FlattenedHeteroGraphPtr HeteroGraph::FlattenImpl(const std::vector<dgl_type_t>&
CHECK_EQ(gptr->NumBits(), NumBits());
FlattenedHeteroGraph* result = new FlattenedHeteroGraph;
result->graph = HeteroGraphRef(HeteroGraphPtr(new HeteroGraph(gptr->meta_graph(), {gptr})));
result->induced_srctype = aten::VecToIdArray(induced_srctype).CopyTo(Context());
result->induced_srctype_set = aten::VecToIdArray(srctype_set).CopyTo(Context());
result->graph = HeteroGraphRef(
HeteroGraphPtr(new HeteroGraph(gptr->meta_graph(), {gptr})));
result->induced_srctype =
aten::VecToIdArray(induced_srctype).CopyTo(Context());
result->induced_srctype_set =
aten::VecToIdArray(srctype_set).CopyTo(Context());
result->induced_srcid = aten::VecToIdArray(induced_srcid).CopyTo(Context());
result->induced_etype = aten::Concat(induced_etypes);
result->induced_etype_set = aten::VecToIdArray(etypes).CopyTo(Context());
result->induced_eid = aten::Concat(eid_arrs);
result->induced_dsttype = aten::VecToIdArray(induced_dsttype).CopyTo(Context());
result->induced_dsttype_set = aten::VecToIdArray(dsttype_set).CopyTo(Context());
result->induced_dsttype =
aten::VecToIdArray(induced_dsttype).CopyTo(Context());
result->induced_dsttype_set =
aten::VecToIdArray(dsttype_set).CopyTo(Context());
result->induced_dstid = aten::VecToIdArray(induced_dstid).CopyTo(Context());
return FlattenedHeteroGraphPtr(result);
}
......@@ -545,21 +561,25 @@ void HeteroGraph::Save(dmlc::Stream* fs) const {
GraphPtr HeteroGraph::AsImmutableGraph() const {
CHECK(NumVertexTypes() == 1) << "graph has more than one node types";
CHECK(NumEdgeTypes() == 1) << "graph has more than one edge types";
auto unit_graph = CHECK_NOTNULL(
std::dynamic_pointer_cast<UnitGraph>(GetRelationGraph(0)));
auto unit_graph =
CHECK_NOTNULL(std::dynamic_pointer_cast<UnitGraph>(GetRelationGraph(0)));
return unit_graph->AsImmutableGraph();
}
HeteroGraphPtr HeteroGraph::LineGraph(bool backtracking) const {
CHECK_EQ(1, meta_graph_->NumEdges()) << "Only support Homogeneous graph now (one edge type)";
CHECK_EQ(1, meta_graph_->NumVertices()) << "Only support Homogeneous graph now (one node type)";
CHECK_EQ(1, meta_graph_->NumEdges())
<< "Only support Homogeneous graph now (one edge type)";
CHECK_EQ(1, meta_graph_->NumVertices())
<< "Only support Homogeneous graph now (one node type)";
CHECK_EQ(1, relation_graphs_.size()) << "Only support Homogeneous graph now";
UnitGraphPtr ug = relation_graphs_[0];
const auto &ulg = ug->LineGraph(backtracking);
const auto& ulg = ug->LineGraph(backtracking);
std::vector<HeteroGraphPtr> rel_graph = {ulg};
std::vector<int64_t> num_nodes_per_type = {static_cast<int64_t>(ulg->NumVertices(0))};
return HeteroGraphPtr(new HeteroGraph(meta_graph_, rel_graph, std::move(num_nodes_per_type)));
std::vector<int64_t> num_nodes_per_type = {
static_cast<int64_t>(ulg->NumVertices(0))};
return HeteroGraphPtr(
new HeteroGraph(meta_graph_, rel_graph, std::move(num_nodes_per_type)));
}
} // namespace dgl
src/graph/heterograph_capi.cc
View file @ 81831111
......@@ -5,11 +5,12 @@
*/
#include <dgl/array.h>
#include <dgl/aten/coo.h>
#include <dgl/packed_func_ext.h>
#include <dgl/immutable_graph.h>
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/c_runtime_api.h>
#include <dgl/runtime/container.h>
#include <dgl/runtime/parallel_for.h>
#include <dgl/runtime/c_runtime_api.h>
#include <set>
#include "../c_api_common.h"
......@@ -25,7 +26,7 @@ namespace dgl {
// XXX(minjie): Ideally, Unitgraph should be invisible to python side
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCOO")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
int64_t nvtypes = args[0];
int64_t num_src = args[1];
int64_t num_dst = args[2];
......@@ -40,13 +41,13 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCOO")
formats_vec.push_back(ParseSparseFormat(fmt));
}
const auto code = SparseFormatsToCode(formats_vec);
auto hgptr = CreateFromCOO(nvtypes, num_src, num_dst, row, col,
row_sorted, col_sorted, code);
auto hgptr = CreateFromCOO(
nvtypes, num_src, num_dst, row, col, row_sorted, col_sorted, code);
*rv = HeteroGraphRef(hgptr);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCSR")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
int64_t nvtypes = args[0];
int64_t num_src = args[1];
int64_t num_dst = args[2];
......@@ -62,16 +63,18 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCSR")
}
const auto code = SparseFormatsToCode(formats_vec);
if (!transpose) {
auto hgptr = CreateFromCSR(nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
auto hgptr = CreateFromCSR(
nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
*rv = HeteroGraphRef(hgptr);
} else {
auto hgptr = CreateFromCSC(nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
auto hgptr = CreateFromCSC(
nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
*rv = HeteroGraphRef(hgptr);
}
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateHeteroGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef meta_graph = args[0];
List<HeteroGraphRef> rel_graphs = args[1];
std::vector<HeteroGraphPtr> rel_ptrs;
......@@ -83,8 +86,9 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateHeteroGraph")
*rv = HeteroGraphRef(hgptr);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateHeteroGraphWithNumNodes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
DGL_REGISTER_GLOBAL(
"heterograph_index._CAPI_DGLHeteroCreateHeteroGraphWithNumNodes")
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef meta_graph = args[0];
List<HeteroGraphRef> rel_graphs = args[1];
IdArray num_nodes_per_type = args[2];
......@@ -101,20 +105,20 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateHeteroGraphWithNumNo
///////////////////////// HeteroGraph member functions /////////////////////////
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetMetaGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->meta_graph();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsMetaGraphUniBipartite")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
GraphPtr mg = hg->meta_graph();
*rv = mg->IsUniBipartite();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetRelationGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
CHECK_LE(etype, hg->NumEdgeTypes()) << "invalid edge type " << etype;
......@@ -126,12 +130,13 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetRelationGraph")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFlattenedGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> etypes = args[1];
std::vector<dgl_id_t> etypes_vec;
for (Value val : etypes) {
// (gq) have to decompose it into two statements because of a weird MSVC internal error
// (gq) have to decompose it into two statements because of a weird MSVC
// internal error
dgl_id_t id = val->data;
etypes_vec.push_back(id);
}
......@@ -140,7 +145,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFlattenedGraph")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t vtype = args[1];
int64_t num = args[2];
......@@ -148,7 +153,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddVertices")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t src = args[2];
......@@ -157,7 +162,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdge")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray src = args[2];
......@@ -166,63 +171,63 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdges")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
hg->Clear();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDataType")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->DataType();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroContext")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->Context();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsPinned")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->IsPinned();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->NumBits();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->IsMultigraph();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsReadonly")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = hg->IsReadonly();
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t vtype = args[1];
*rv = static_cast<int64_t>(hg->NumVertices(vtype));
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
*rv = static_cast<int64_t>(hg->NumEdges(etype));
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t vtype = args[1];
dgl_id_t vid = args[2];
......@@ -230,7 +235,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertex")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t vtype = args[1];
IdArray vids = args[2];
......@@ -238,7 +243,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertices")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t src = args[2];
......@@ -247,7 +252,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgeBetween")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray src = args[2];
......@@ -256,7 +261,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgesBetween")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t dst = args[2];
......@@ -264,7 +269,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPredecessors")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t src = args[2];
......@@ -272,7 +277,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSuccessors")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t src = args[2];
......@@ -281,7 +286,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeId")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsAll")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray src = args[2];
......@@ -290,9 +295,8 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsAll")
*rv = ConvertEdgeArrayToPackedFunc(ret);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsOne")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray src = args[2];
......@@ -301,7 +305,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsOne")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray eids = args[2];
......@@ -310,7 +314,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroFindEdges")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t vid = args[2];
......@@ -319,7 +323,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_1")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray vids = args[2];
......@@ -328,7 +332,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_2")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t vid = args[2];
......@@ -337,7 +341,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_1")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray vids = args[2];
......@@ -346,7 +350,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_2")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
std::string order = args[2];
......@@ -355,7 +359,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdges")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t vid = args[2];
......@@ -363,7 +367,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegree")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray vids = args[2];
......@@ -371,7 +375,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegrees")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
dgl_id_t vid = args[2];
......@@ -379,7 +383,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegree")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
IdArray vids = args[2];
......@@ -387,17 +391,16 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegrees")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_type_t etype = args[1];
bool transpose = args[2];
std::string fmt = args[3];
*rv = ConvertNDArrayVectorToPackedFunc(
hg->GetAdj(etype, transpose, fmt));
*rv = ConvertNDArrayVectorToPackedFunc(hg->GetAdj(etype, transpose, fmt));
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> vids = args[1];
bool relabel_nodes = args[2];
......@@ -413,7 +416,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroVertexSubgraph")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> eids = args[1];
bool preserve_nodes = args[2];
......@@ -430,13 +433,14 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeSubgraph")
///////////////////////// HeteroSubgraph members /////////////////////////
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroSubgraphRef subg = args[0];
*rv = HeteroGraphRef(subg->graph);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetInducedVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
DGL_REGISTER_GLOBAL(
"heterograph_index._CAPI_DGLHeteroSubgraphGetInducedVertices")
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroSubgraphRef subg = args[0];
List<Value> induced_verts;
for (IdArray arr : subg->induced_vertices) {
......@@ -446,7 +450,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetInducedVertices
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetInducedEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroSubgraphRef subg = args[0];
List<Value> induced_edges;
for (IdArray arr : subg->induced_edges) {
......@@ -455,10 +459,11 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetInducedEdges")
*rv = induced_edges;
});
///////////////////////// Global functions and algorithms /////////////////////////
///////////////////////// Global functions and algorithms
////////////////////////////
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAsNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
int bits = args[1];
HeteroGraphPtr bhg_ptr = hg.sptr();
......@@ -473,7 +478,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAsNumBits")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyTo")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
int device_type = args[1];
int device_id = args[2];
......@@ -485,7 +490,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyTo")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPinMemory_")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
hgindex->PinMemory_();
......@@ -493,7 +498,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPinMemory_")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroUnpinMemory_")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
hgindex->UnpinMemory_();
......@@ -501,7 +506,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroUnpinMemory_")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroRecordStream")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
DGLStreamHandle stream = args[1];
auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
......@@ -510,7 +515,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroRecordStream")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyToSharedMem")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
std::string name = args[1];
List<Value> ntypes = args[2];
......@@ -519,7 +524,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyToSharedMem")
auto ntypes_vec = ListValueToVector<std::string>(ntypes);
auto etypes_vec = ListValueToVector<std::string>(etypes);
std::set<std::string> fmts_set;
for (const auto &fmt : fmts) {
for (const auto& fmt : fmts) {
std::string fmt_data = fmt->data;
fmts_set.insert(fmt_data);
}
......@@ -529,7 +534,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyToSharedMem")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFromSharedMem")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
std::string name = args[0];
HeteroGraphPtr hg;
std::vector<std::string> ntypes;
......@@ -537,9 +542,9 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFromSharedMem")
std::tie(hg, ntypes, etypes) = HeteroGraph::CreateFromSharedMem(name);
List<Value> ntypes_list;
List<Value> etypes_list;
for (const auto &ntype : ntypes)
for (const auto& ntype : ntypes)
ntypes_list.push_back(Value(MakeValue(ntype)));
for (const auto &etype : etypes)
for (const auto& etype : etypes)
etypes_list.push_back(Value(MakeValue(etype)));
List<ObjectRef> ret;
ret.push_back(HeteroGraphRef(hg));
......@@ -549,7 +554,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFromSharedMem")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroJointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef meta_graph = args[0];
List<HeteroGraphRef> component_graphs = args[1];
CHECK(component_graphs.size() > 1)
......@@ -561,8 +566,8 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroJointUnion")
for (const auto& component : component_graphs) {
component_ptrs.push_back(component.sptr());
CHECK_EQ(component->NumBits(), bits)
<< "Expect graphs to joint union have the same index dtype(int" << bits
<< "), but got int" << component->NumBits();
<< "Expect graphs to joint union have the same index dtype(int"
<< bits << "), but got int" << component->NumBits();
CHECK_EQ(component->Context(), ctx)
<< "Expect graphs to joint union have the same context" << ctx
<< "), but got " << component->Context();
......@@ -570,10 +575,10 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroJointUnion")
auto hgptr = JointUnionHeteroGraph(meta_graph.sptr(), component_ptrs);
*rv = HeteroGraphRef(hgptr);
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointUnion_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef meta_graph = args[0];
List<HeteroGraphRef> component_graphs = args[1];
CHECK(component_graphs.size() > 0)
......@@ -594,86 +599,86 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointUnion_v2")
auto hgptr = DisjointUnionHeteroGraph2(meta_graph.sptr(), component_ptrs);
*rv = HeteroGraphRef(hgptr);
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointPartitionBySizes_v2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
DGL_REGISTER_GLOBAL(
"heterograph_index._CAPI_DGLHeteroDisjointPartitionBySizes_v2")
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
const IdArray vertex_sizes = args[1];
const IdArray edge_sizes = args[2];
std::vector<HeteroGraphPtr> ret;
ret = DisjointPartitionHeteroBySizes2(hg->meta_graph(), hg.sptr(),
vertex_sizes, edge_sizes);
ret = DisjointPartitionHeteroBySizes2(
hg->meta_graph(), hg.sptr(), vertex_sizes, edge_sizes);
List<HeteroGraphRef> ret_list;
for (HeteroGraphPtr hgptr : ret) {
ret_list.push_back(HeteroGraphRef(hgptr));
}
*rv = ret_list;
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
const IdArray vertex_sizes = args[1];
const IdArray edge_sizes = args[2];
const int64_t bits = hg->NumBits();
std::vector<HeteroGraphPtr> ret;
ATEN_ID_BITS_SWITCH(bits, IdType, {
ret = DisjointPartitionHeteroBySizes<IdType>(hg->meta_graph(), hg.sptr(),
vertex_sizes, edge_sizes);
ret = DisjointPartitionHeteroBySizes<IdType>(
hg->meta_graph(), hg.sptr(), vertex_sizes, edge_sizes);
});
List<HeteroGraphRef> ret_list;
for (HeteroGraphPtr hgptr : ret) {
ret_list.push_back(HeteroGraphRef(hgptr));
}
*rv = ret_list;
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSlice")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
const IdArray num_nodes_per_type = args[1];
const IdArray start_nid_per_type = args[2];
const IdArray num_edges_per_type = args[3];
const IdArray start_eid_per_type = args[4];
auto hgptr = SliceHeteroGraph(hg->meta_graph(), hg.sptr(), num_nodes_per_type,
start_nid_per_type, num_edges_per_type, start_eid_per_type);
auto hgptr = SliceHeteroGraph(
hg->meta_graph(), hg.sptr(), num_nodes_per_type, start_nid_per_type,
num_edges_per_type, start_eid_per_type);
*rv = HeteroGraphRef(hgptr);
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetCreatedFormats")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> format_list;
dgl_format_code_t code = hg->GetRelationGraph(0)->GetCreatedFormats();
for (auto format : CodeToSparseFormats(code)) {
format_list.push_back(
Value(MakeValue(ToStringSparseFormat(format))));
format_list.push_back(Value(MakeValue(ToStringSparseFormat(format))));
}
*rv = format_list;
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetAllowedFormats")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> format_list;
dgl_format_code_t code = hg->GetRelationGraph(0)->GetAllowedFormats();
for (auto format : CodeToSparseFormats(code)) {
format_list.push_back(
Value(MakeValue(ToStringSparseFormat(format))));
format_list.push_back(Value(MakeValue(ToStringSparseFormat(format))));
}
*rv = format_list;
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFormat")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
dgl_format_code_t code = hg->GetRelationGraph(0)->GetAllowedFormats();
auto get_format_f = [&](size_t etype_b, size_t etype_e) {
for (auto etype = etype_b; etype < etype_e; ++etype) {
auto bg = std::dynamic_pointer_cast<UnitGraph>(hg->GetRelationGraph(etype));
for (auto format : CodeToSparseFormats(code))
bg->GetFormat(format);
auto bg =
std::dynamic_pointer_cast<UnitGraph>(hg->GetRelationGraph(etype));
for (auto format : CodeToSparseFormats(code)) bg->GetFormat(format);
}
};
......@@ -682,10 +687,10 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFormat")
#else
get_format_f(0, hg->NumEdgeTypes());
#endif
});
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFormatGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
List<Value> formats = args[1];
std::vector<SparseFormat> formats_vec;
......@@ -693,13 +698,12 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFormatGraph")
std::string fmt = val->data;
formats_vec.push_back(ParseSparseFormat(fmt));
}
auto hgptr = hg->GetGraphInFormat(
SparseFormatsToCode(formats_vec));
auto hgptr = hg->GetGraphInFormat(SparseFormatsToCode(formats_vec));
*rv = HeteroGraphRef(hgptr);
});
});
DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLInSubgraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
const auto& nodes = ListValueToVector<IdArray>(args[1]);
bool relabel_nodes = args[2];
......@@ -709,7 +713,7 @@ DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLInSubgraph")
});
DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLOutSubgraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
const auto& nodes = ListValueToVector<IdArray>(args[1]);
bool relabel_nodes = args[2];
......@@ -719,27 +723,30 @@ DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLOutSubgraph")
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLAsImmutableGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
*rv = GraphRef(hg->AsImmutableGraph());
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortOutEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
NDArray tag = args[1];
int64_t num_tag = args[2];
CHECK_EQ(hg->Context().device_type, kDGLCPU) << "Only support sorting by tag on cpu";
CHECK_EQ(hg->Context().device_type, kDGLCPU)
<< "Only support sorting by tag on cpu";
CHECK(aten::IsValidIdArray(tag));
CHECK_EQ(tag->ctx.device_type, kDGLCPU) << "Only support sorting by tag on cpu";
CHECK_EQ(tag->ctx.device_type, kDGLCPU)
<< "Only support sorting by tag on cpu";
const auto csr = hg->GetCSRMatrix(0);
NDArray tag_pos = aten::NullArray();
aten::CSRMatrix output;
std::tie(output, tag_pos) = aten::CSRSortByTag(csr, tag, num_tag);
HeteroGraphPtr output_hg = CreateFromCSR(hg->NumVertexTypes(), output, ALL_CODE);
HeteroGraphPtr output_hg =
CreateFromCSR(hg->NumVertexTypes(), output, ALL_CODE);
List<ObjectRef> ret;
ret.push_back(HeteroGraphRef(output_hg));
ret.push_back(Value(MakeValue(tag_pos)));
......@@ -747,14 +754,16 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortOutEdges")
});
DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortInEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
NDArray tag = args[1];
int64_t num_tag = args[2];
CHECK_EQ(hg->Context().device_type, kDGLCPU) << "Only support sorting by tag on cpu";
CHECK_EQ(hg->Context().device_type, kDGLCPU)
<< "Only support sorting by tag on cpu";
CHECK(aten::IsValidIdArray(tag));
CHECK_EQ(tag->ctx.device_type, kDGLCPU) << "Only support sorting by tag on cpu";
CHECK_EQ(tag->ctx.device_type, kDGLCPU)
<< "Only support sorting by tag on cpu";
const auto csc = hg->GetCSCMatrix(0);
......@@ -762,7 +771,8 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortInEdges")
aten::CSRMatrix output;
std::tie(output, tag_pos) = aten::CSRSortByTag(csc, tag, num_tag);
HeteroGraphPtr output_hg = CreateFromCSC(hg->NumVertexTypes(), output, ALL_CODE);
HeteroGraphPtr output_hg =
CreateFromCSC(hg->NumVertexTypes(), output, ALL_CODE);
List<ObjectRef> ret;
ret.push_back(HeteroGraphRef(output_hg));
ret.push_back(Value(MakeValue(tag_pos)));
......@@ -770,7 +780,7 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortInEdges")
});
DGL_REGISTER_GLOBAL("heterograph._CAPI_DGLFindSrcDstNtypes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
GraphRef metagraph = args[0];
std::unordered_set<uint64_t> dst_set;
std::unordered_set<uint64_t> src_set;
......@@ -793,7 +803,8 @@ DGL_REGISTER_GLOBAL("heterograph._CAPI_DGLFindSrcDstNtypes")
else if (is_dst)
dstlist.push_back(Value(MakeValue(static_cast<int64_t>(nid))));
else
// If a node type is isolated, put it in srctype as defined in the Python docstring.
// If a node type is isolated, put it in srctype as defined in the
// Python docstring.
srclist.push_back(Value(MakeValue(static_cast<int64_t>(nid))));
}
ret_list.push_back(srclist);
......@@ -802,24 +813,24 @@ DGL_REGISTER_GLOBAL("heterograph._CAPI_DGLFindSrcDstNtypes")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroReverse")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
HeteroGraphRef hg = args[0];
CHECK_GT(hg->NumEdgeTypes(), 0);
auto g = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
std::vector<HeteroGraphPtr> rev_ugs;
const auto &ugs = g->relation_graphs();
const auto& ugs = g->relation_graphs();
rev_ugs.resize(ugs.size());
for (size_t i = 0; i < ugs.size(); ++i) {
const auto &rev_ug = ugs[i]->Reverse();
const auto& rev_ug = ugs[i]->Reverse();
rev_ugs[i] = rev_ug;
}
// node types are not changed
const auto& num_nodes = g->NumVerticesPerType();
const auto& meta_edges = hg->meta_graph()->Edges("eid");
// reverse the metagraph
const auto& rev_meta = ImmutableGraph::CreateFromCOO(hg->meta_graph()->NumVertices(),
meta_edges.dst, meta_edges.src);
const auto& rev_meta = ImmutableGraph::CreateFromCOO(
hg->meta_graph()->NumVertices(), meta_edges.dst, meta_edges.src);
*rv = CreateHeteroGraph(rev_meta, rev_ugs, num_nodes);
});
} // namespace dgl
src/graph/immutable_graph.cc
View file @ 81831111
......@@ -4,13 +4,14 @@
* \brief DGL immutable graph index implementation
*/
#include <dgl/base_heterograph.h>
#include <dgl/immutable_graph.h>
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/smart_ptr_serializer.h>
#include <dgl/base_heterograph.h>
#include <dmlc/io.h>
#include <dmlc/type_traits.h>
#include <string.h>
#include <bitset>
#include <numeric>
#include <tuple>
......@@ -23,7 +24,8 @@ using namespace dgl::runtime;
namespace dgl {
namespace {
inline std::string GetSharedMemName(const std::string &name, const std::string &edge_dir) {
inline std::string GetSharedMemName(
const std::string &name, const std::string &edge_dir) {
return name + "_" + edge_dir;
}
......@@ -39,8 +41,9 @@ struct GraphIndexMetadata {
};
/*
* Serialize the metadata of a graph index and place it in a shared-memory tensor.
* In this way, another process can reconstruct a GraphIndex from a shared-memory tensor.
* Serialize the metadata of a graph index and place it in a shared-memory
* tensor. In this way, another process can reconstruct a GraphIndex from a
* shared-memory tensor.
*/
NDArray SerializeMetadata(ImmutableGraphPtr gidx, const std::string &name) {
#ifndef _WIN32
......@@ -51,8 +54,9 @@ NDArray SerializeMetadata(ImmutableGraphPtr gidx, const std::string &name) {
meta.has_out_csr = gidx->HasOutCSR();
meta.has_coo = false;
NDArray meta_arr = NDArray::EmptyShared(name, {sizeof(meta)}, DGLDataType{kDGLInt, 8, 1},
DGLContext{kDGLCPU, 0}, true);
NDArray meta_arr = NDArray::EmptyShared(
name, {sizeof(meta)}, DGLDataType{kDGLInt, 8, 1}, DGLContext{kDGLCPU, 0},
true);
memcpy(meta_arr->data, &meta, sizeof(meta));
return meta_arr;
#else
......@@ -67,8 +71,9 @@ NDArray SerializeMetadata(ImmutableGraphPtr gidx, const std::string &name) {
GraphIndexMetadata DeserializeMetadata(const std::string &name) {
GraphIndexMetadata meta;
#ifndef _WIN32
NDArray meta_arr = NDArray::EmptyShared(name, {sizeof(meta)}, DGLDataType{kDGLInt, 8, 1},
DGLContext{kDGLCPU, 0}, false);
NDArray meta_arr = NDArray::EmptyShared(
name, {sizeof(meta)}, DGLDataType{kDGLInt, 8, 1}, DGLContext{kDGLCPU, 0},
false);
memcpy(&meta, meta_arr->data, sizeof(meta));
#else
LOG(FATAL) << "CSR graph doesn't support shared memory in Windows yet";
......@@ -77,18 +82,23 @@ GraphIndexMetadata DeserializeMetadata(const std::string &name) {
}
std::tuple<IdArray, IdArray, IdArray> MapFromSharedMemory(
const std::string &shared_mem_name, int64_t num_verts, int64_t num_edges, bool is_create) {
const std::string &shared_mem_name, int64_t num_verts, int64_t num_edges,
bool is_create) {
#ifndef _WIN32
const int64_t file_size = (num_verts + 1 + num_edges * 2) * sizeof(dgl_id_t);
IdArray sm_array = IdArray::EmptyShared(
shared_mem_name, {file_size}, DGLDataType{kDGLInt, 8, 1}, DGLContext{kDGLCPU, 0}, is_create);
shared_mem_name, {file_size}, DGLDataType{kDGLInt, 8, 1},
DGLContext{kDGLCPU, 0}, is_create);
// Create views from the shared memory array. Note that we don't need to save
// the sm_array because the refcount is maintained by the view arrays.
IdArray indptr = sm_array.CreateView({num_verts + 1}, DGLDataType{kDGLInt, 64, 1});
IdArray indices = sm_array.CreateView({num_edges}, DGLDataType{kDGLInt, 64, 1},
IdArray indptr =
sm_array.CreateView({num_verts + 1}, DGLDataType{kDGLInt, 64, 1});
IdArray indices = sm_array.CreateView(
{num_edges}, DGLDataType{kDGLInt, 64, 1},
(num_verts + 1) * sizeof(dgl_id_t));
IdArray edge_ids = sm_array.CreateView({num_edges}, DGLDataType{kDGLInt, 64, 1},
IdArray edge_ids = sm_array.CreateView(
{num_edges}, DGLDataType{kDGLInt, 64, 1},
(num_verts + 1 + num_edges) * sizeof(dgl_id_t));
return std::make_tuple(indptr, indices, edge_ids);
#else
......@@ -106,10 +116,9 @@ std::tuple<IdArray, IdArray, IdArray> MapFromSharedMemory(
CSR::CSR(int64_t num_vertices, int64_t num_edges) {
CHECK(!(num_vertices == 0 && num_edges != 0));
adj_ = aten::CSRMatrix{num_vertices, num_vertices,
aten::NewIdArray(num_vertices + 1),
aten::NewIdArray(num_edges),
aten::NewIdArray(num_edges)};
adj_ = aten::CSRMatrix{
num_vertices, num_vertices, aten::NewIdArray(num_vertices + 1),
aten::NewIdArray(num_edges), aten::NewIdArray(num_edges)};
adj_.sorted = false;
}
......@@ -123,8 +132,10 @@ CSR::CSR(IdArray indptr, IdArray indices, IdArray edge_ids) {
adj_.sorted = false;
}
CSR::CSR(IdArray indptr, IdArray indices, IdArray edge_ids,
const std::string &shared_mem_name): shared_mem_name_(shared_mem_name) {
CSR::CSR(
IdArray indptr, IdArray indices, IdArray edge_ids,
const std::string &shared_mem_name)
: shared_mem_name_(shared_mem_name) {
CHECK(aten::IsValidIdArray(indptr));
CHECK(aten::IsValidIdArray(indices));
CHECK(aten::IsValidIdArray(edge_ids));
......@@ -133,8 +144,8 @@ CSR::CSR(IdArray indptr, IdArray indices, IdArray edge_ids,
const int64_t num_edges = indices->shape[0];
adj_.num_rows = num_verts;
adj_.num_cols = num_verts;
std::tie(adj_.indptr, adj_.indices, adj_.data) = MapFromSharedMemory(
shared_mem_name, num_verts, num_edges, true);
std::tie(adj_.indptr, adj_.indices, adj_.data) =
MapFromSharedMemory(shared_mem_name, num_verts, num_edges, true);
// copy the given data into the shared memory arrays
adj_.indptr.CopyFrom(indptr);
adj_.indices.CopyFrom(indices);
......@@ -142,19 +153,18 @@ CSR::CSR(IdArray indptr, IdArray indices, IdArray edge_ids,
adj_.sorted = false;
}
CSR::CSR(const std::string &shared_mem_name,
int64_t num_verts, int64_t num_edges): shared_mem_name_(shared_mem_name) {
CSR::CSR(
const std::string &shared_mem_name, int64_t num_verts, int64_t num_edges)
: shared_mem_name_(shared_mem_name) {
CHECK(!(num_verts == 0 && num_edges != 0));
adj_.num_rows = num_verts;
adj_.num_cols = num_verts;
std::tie(adj_.indptr, adj_.indices, adj_.data) = MapFromSharedMemory(
shared_mem_name, num_verts, num_edges, false);
std::tie(adj_.indptr, adj_.indices, adj_.data) =
MapFromSharedMemory(shared_mem_name, num_verts, num_edges, false);
adj_.sorted = false;
}
bool CSR::IsMultigraph() const {
return aten::CSRHasDuplicate(adj_);
}
bool CSR::IsMultigraph() const { return aten::CSRHasDuplicate(adj_); }
EdgeArray CSR::OutEdges(dgl_id_t vid) const {
CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
......@@ -204,7 +214,7 @@ IdArray CSR::EdgeId(dgl_id_t src, dgl_id_t dst) const {
}
EdgeArray CSR::EdgeIds(IdArray src_ids, IdArray dst_ids) const {
const auto& arrs = aten::CSRGetDataAndIndices(adj_, src_ids, dst_ids);
const auto &arrs = aten::CSRGetDataAndIndices(adj_, src_ids, dst_ids);
return EdgeArray{arrs[0], arrs[1], arrs[2]};
}
......@@ -212,14 +222,15 @@ EdgeArray CSR::Edges(const std::string &order) const {
CHECK(order.empty() || order == std::string("srcdst"))
<< "CSR only support Edges of order \"srcdst\","
<< " but got \"" << order << "\".";
const auto& coo = aten::CSRToCOO(adj_, false);
const auto &coo = aten::CSRToCOO(adj_, false);
return EdgeArray{coo.row, coo.col, coo.data};
}
Subgraph CSR::VertexSubgraph(IdArray vids) const {
CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
const auto& submat = aten::CSRSliceMatrix(adj_, vids, vids);
IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
const auto &submat = aten::CSRSliceMatrix(adj_, vids, vids);
IdArray sub_eids =
aten::Range(0, submat.data->shape[0], NumBits(), Context());
CSRPtr subcsr(new CSR(submat.indptr, submat.indices, sub_eids));
subcsr->adj_.sorted = this->adj_.sorted;
Subgraph subg;
......@@ -230,21 +241,21 @@ Subgraph CSR::VertexSubgraph(IdArray vids) const {
}
CSRPtr CSR::Transpose() const {
const auto& trans = aten::CSRTranspose(adj_);
const auto &trans = aten::CSRTranspose(adj_);
return CSRPtr(new CSR(trans.indptr, trans.indices, trans.data));
}
COOPtr CSR::ToCOO() const {
const auto& coo = aten::CSRToCOO(adj_, true);
const auto &coo = aten::CSRToCOO(adj_, true);
return COOPtr(new COO(NumVertices(), coo.row, coo.col));
}
CSR CSR::CopyTo(const DGLContext& ctx) const {
CSR CSR::CopyTo(const DGLContext &ctx) const {
if (Context() == ctx) {
return *this;
} else {
CSR ret(adj_.indptr.CopyTo(ctx),
adj_.indices.CopyTo(ctx),
CSR ret(
adj_.indptr.CopyTo(ctx), adj_.indices.CopyTo(ctx),
adj_.data.CopyTo(ctx));
return ret;
}
......@@ -264,8 +275,8 @@ CSR CSR::AsNumBits(uint8_t bits) const {
if (NumBits() == bits) {
return *this;
} else {
CSR ret(aten::AsNumBits(adj_.indptr, bits),
aten::AsNumBits(adj_.indices, bits),
CSR ret(
aten::AsNumBits(adj_.indptr, bits), aten::AsNumBits(adj_.indices, bits),
aten::AsNumBits(adj_.data, bits));
return ret;
}
......@@ -274,8 +285,8 @@ CSR CSR::AsNumBits(uint8_t bits) const {
DGLIdIters CSR::SuccVec(dgl_id_t vid) const {
// TODO(minjie): This still assumes the data type and device context
// of this graph. Should fix later.
const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
const dgl_id_t* indices_data = static_cast<dgl_id_t*>(adj_.indices->data);
const dgl_id_t *indptr_data = static_cast<dgl_id_t *>(adj_.indptr->data);
const dgl_id_t *indices_data = static_cast<dgl_id_t *>(adj_.indices->data);
const dgl_id_t start = indptr_data[vid];
const dgl_id_t end = indptr_data[vid + 1];
return DGLIdIters(indices_data + start, indices_data + end);
......@@ -284,29 +295,28 @@ DGLIdIters CSR::SuccVec(dgl_id_t vid) const {
DGLIdIters CSR::OutEdgeVec(dgl_id_t vid) const {
// TODO(minjie): This still assumes the data type and device context
// of this graph. Should fix later.
const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
const dgl_id_t* eid_data = static_cast<dgl_id_t*>(adj_.data->data);
const dgl_id_t *indptr_data = static_cast<dgl_id_t *>(adj_.indptr->data);
const dgl_id_t *eid_data = static_cast<dgl_id_t *>(adj_.data->data);
const dgl_id_t start = indptr_data[vid];
const dgl_id_t end = indptr_data[vid + 1];
return DGLIdIters(eid_data + start, eid_data + end);
}
bool CSR::Load(dmlc::Stream *fs) {
fs->Read(const_cast<dgl::aten::CSRMatrix*>(&adj_));
fs->Read(const_cast<dgl::aten::CSRMatrix *>(&adj_));
return true;
}
void CSR::Save(dmlc::Stream *fs) const {
fs->Write(adj_);
}
void CSR::Save(dmlc::Stream *fs) const { fs->Write(adj_); }
//////////////////////////////////////////////////////////
//
// COO graph implementation
//
//////////////////////////////////////////////////////////
COO::COO(int64_t num_vertices, IdArray src, IdArray dst,
bool row_sorted, bool col_sorted) {
COO::COO(
int64_t num_vertices, IdArray src, IdArray dst, bool row_sorted,
bool col_sorted) {
CHECK(aten::IsValidIdArray(src));
CHECK(aten::IsValidIdArray(dst));
CHECK_EQ(src->shape[0], dst->shape[0]);
......@@ -314,9 +324,7 @@ COO::COO(int64_t num_vertices, IdArray src, IdArray dst,
aten::NullArray(), row_sorted, col_sorted};
}
bool COO::IsMultigraph() const {
return aten::COOHasDuplicate(adj_);
}
bool COO::IsMultigraph() const { return aten::COOHasDuplicate(adj_); }
std::pair<dgl_id_t, dgl_id_t> COO::FindEdge(dgl_id_t eid) const {
CHECK(eid < NumEdges()) << "Invalid edge id: " << eid;
......@@ -327,17 +335,17 @@ std::pair<dgl_id_t, dgl_id_t> COO::FindEdge(dgl_id_t eid) const {
EdgeArray COO::FindEdges(IdArray eids) const {
CHECK(aten::IsValidIdArray(eids)) << "Invalid edge id array";
BUG_IF_FAIL(aten::IsNullArray(adj_.data)) <<
"FindEdges requires the internal COO matrix not having EIDs.";
return EdgeArray{aten::IndexSelect(adj_.row, eids),
aten::IndexSelect(adj_.col, eids),
BUG_IF_FAIL(aten::IsNullArray(adj_.data))
<< "FindEdges requires the internal COO matrix not having EIDs.";
return EdgeArray{
aten::IndexSelect(adj_.row, eids), aten::IndexSelect(adj_.col, eids),
eids};
}
EdgeArray COO::Edges(const std::string &order) const {
CHECK(order.empty() || order == std::string("eid"))
<< "COO only support Edges of order \"eid\", but got \""
<< order << "\".";
<< "COO only support Edges of order \"eid\", but got \"" << order
<< "\".";
IdArray rst_eid = aten::Range(0, NumEdges(), NumBits(), Context());
return EdgeArray{adj_.row, adj_.col, rst_eid};
}
......@@ -366,17 +374,15 @@ Subgraph COO::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
}
CSRPtr COO::ToCSR() const {
const auto& csr = aten::COOToCSR(adj_);
const auto &csr = aten::COOToCSR(adj_);
return CSRPtr(new CSR(csr.indptr, csr.indices, csr.data));
}
COO COO::CopyTo(const DGLContext& ctx) const {
COO COO::CopyTo(const DGLContext &ctx) const {
if (Context() == ctx) {
return *this;
} else {
COO ret(NumVertices(),
adj_.row.CopyTo(ctx),
adj_.col.CopyTo(ctx));
COO ret(NumVertices(), adj_.row.CopyTo(ctx), adj_.col.CopyTo(ctx));
return ret;
}
}
......@@ -390,8 +396,8 @@ COO COO::AsNumBits(uint8_t bits) const {
if (NumBits() == bits) {
return *this;
} else {
COO ret(NumVertices(),
aten::AsNumBits(adj_.row, bits),
COO ret(
NumVertices(), aten::AsNumBits(adj_.row, bits),
aten::AsNumBits(adj_.col, bits));
return ret;
}
......@@ -411,13 +417,14 @@ BoolArray ImmutableGraph::HasVertices(IdArray vids) const {
CSRPtr ImmutableGraph::GetInCSR() const {
if (!in_csr_) {
if (out_csr_) {
const_cast<ImmutableGraph*>(this)->in_csr_ = out_csr_->Transpose();
const_cast<ImmutableGraph *>(this)->in_csr_ = out_csr_->Transpose();
if (out_csr_->IsSharedMem())
LOG(WARNING) << "We just construct an in-CSR from a shared-memory out CSR. "
LOG(WARNING)
<< "We just construct an in-CSR from a shared-memory out CSR. "
<< "It may dramatically increase memory consumption.";
} else {
CHECK(coo_) << "None of CSR, COO exist";
const_cast<ImmutableGraph*>(this)->in_csr_ = coo_->Transpose()->ToCSR();
const_cast<ImmutableGraph *>(this)->in_csr_ = coo_->Transpose()->ToCSR();
}
}
return in_csr_;
......@@ -427,13 +434,14 @@ CSRPtr ImmutableGraph::GetInCSR() const {
CSRPtr ImmutableGraph::GetOutCSR() const {
if (!out_csr_) {
if (in_csr_) {
const_cast<ImmutableGraph*>(this)->out_csr_ = in_csr_->Transpose();
const_cast<ImmutableGraph *>(this)->out_csr_ = in_csr_->Transpose();
if (in_csr_->IsSharedMem())
LOG(WARNING) << "We just construct an out-CSR from a shared-memory in CSR. "
LOG(WARNING)
<< "We just construct an out-CSR from a shared-memory in CSR. "
<< "It may dramatically increase memory consumption.";
} else {
CHECK(coo_) << "None of CSR, COO exist";
const_cast<ImmutableGraph*>(this)->out_csr_ = coo_->ToCSR();
const_cast<ImmutableGraph *>(this)->out_csr_ = coo_->ToCSR();
}
}
return out_csr_;
......@@ -443,10 +451,10 @@ CSRPtr ImmutableGraph::GetOutCSR() const {
COOPtr ImmutableGraph::GetCOO() const {
if (!coo_) {
if (in_csr_) {
const_cast<ImmutableGraph*>(this)->coo_ = in_csr_->ToCOO()->Transpose();
const_cast<ImmutableGraph *>(this)->coo_ = in_csr_->ToCOO()->Transpose();
} else {
CHECK(out_csr_) << "Both CSR are missing.";
const_cast<ImmutableGraph*>(this)->coo_ = out_csr_->ToCOO();
const_cast<ImmutableGraph *>(this)->coo_ = out_csr_->ToCOO();
}
}
return coo_;
......@@ -457,7 +465,7 @@ EdgeArray ImmutableGraph::Edges(const std::string &order) const {
// arbitrary order
if (in_csr_) {
// transpose
const auto& edges = in_csr_->Edges(order);
const auto &edges = in_csr_->Edges(order);
return EdgeArray{edges.dst, edges.src, edges.id};
} else {
return AnyGraph()->Edges(order);
......@@ -489,16 +497,19 @@ Subgraph ImmutableGraph::EdgeSubgraph(IdArray eids, bool preserve_nodes) const {
return sg;
}
std::vector<IdArray> ImmutableGraph::GetAdj(bool transpose, const std::string &fmt) const {
// TODO(minjie): Our current semantics of adjacency matrix is row for dst nodes and col for
// src nodes. Therefore, we need to flip the transpose flag. For example, transpose=False
// is equal to in edge CSR.
// We have this behavior because previously we use framework's SPMM and we don't cache
// reverse adj. This is not intuitive and also not consistent with networkx's
// to_scipy_sparse_matrix. With the upcoming custom kernel change, we should change the
// behavior and make row for src and col for dst.
std::vector<IdArray> ImmutableGraph::GetAdj(
bool transpose, const std::string &fmt) const {
// TODO(minjie): Our current semantics of adjacency matrix is row for dst
// nodes and col for
// src nodes. Therefore, we need to flip the transpose flag. For example,
// transpose=False is equal to in edge CSR. We have this behavior because
// previously we use framework's SPMM and we don't cache reverse adj. This
// is not intuitive and also not consistent with networkx's
// to_scipy_sparse_matrix. With the upcoming custom kernel change, we should
// change the behavior and make row for src and col for dst.
if (fmt == std::string("csr")) {
return transpose? GetOutCSR()->GetAdj(false, "csr") : GetInCSR()->GetAdj(false, "csr");
return transpose ? GetOutCSR()->GetAdj(false, "csr")
: GetInCSR()->GetAdj(false, "csr");
} else if (fmt == std::string("coo")) {
return GetCOO()->GetAdj(!transpose, fmt);
} else {
......@@ -508,7 +519,8 @@ std::vector<IdArray> ImmutableGraph::GetAdj(bool transpose, const std::string &f
}
ImmutableGraphPtr ImmutableGraph::CreateFromCSR(
IdArray indptr, IdArray indices, IdArray edge_ids, const std::string &edge_dir) {
IdArray indptr, IdArray indices, IdArray edge_ids,
const std::string &edge_dir) {
CSRPtr csr(new CSR(indptr, indices, edge_ids));
if (edge_dir == "in") {
return ImmutableGraphPtr(new ImmutableGraph(csr, nullptr));
......@@ -530,17 +542,19 @@ ImmutableGraphPtr ImmutableGraph::CreateFromCSR(const std::string &name) {
GraphIndexMetadata meta = DeserializeMetadata(GetSharedMemName(name, "meta"));
CSRPtr in_csr, out_csr;
if (meta.has_in_csr) {
in_csr = CSRPtr(new CSR(GetSharedMemName(name, "in"), meta.num_nodes, meta.num_edges));
in_csr = CSRPtr(
new CSR(GetSharedMemName(name, "in"), meta.num_nodes, meta.num_edges));
}
if (meta.has_out_csr) {
out_csr = CSRPtr(new CSR(GetSharedMemName(name, "out"), meta.num_nodes, meta.num_edges));
out_csr = CSRPtr(
new CSR(GetSharedMemName(name, "out"), meta.num_nodes, meta.num_edges));
}
return ImmutableGraphPtr(new ImmutableGraph(in_csr, out_csr, name));
}
ImmutableGraphPtr ImmutableGraph::CreateFromCOO(
int64_t num_vertices, IdArray src, IdArray dst,
bool row_sorted, bool col_sorted) {
int64_t num_vertices, IdArray src, IdArray dst, bool row_sorted,
bool col_sorted) {
COOPtr coo(new COO(num_vertices, src, dst, row_sorted, col_sorted));
return std::make_shared<ImmutableGraph>(coo);
}
......@@ -550,13 +564,14 @@ ImmutableGraphPtr ImmutableGraph::ToImmutable(GraphPtr graph) {
if (ig) {
return ig;
} else {
const auto& adj = graph->GetAdj(true, "csr");
const auto &adj = graph->GetAdj(true, "csr");
CSRPtr csr(new CSR(adj[0], adj[1], adj[2]));
return ImmutableGraph::CreateFromCSR(adj[0], adj[1], adj[2], "out");
}
}
ImmutableGraphPtr ImmutableGraph::CopyTo(ImmutableGraphPtr g, const DGLContext& ctx) {
ImmutableGraphPtr ImmutableGraph::CopyTo(
ImmutableGraphPtr g, const DGLContext &ctx) {
if (ctx == g->Context()) {
return g;
}
......@@ -569,16 +584,20 @@ ImmutableGraphPtr ImmutableGraph::CopyTo(ImmutableGraphPtr g, const DGLContext&
return ImmutableGraphPtr(new ImmutableGraph(new_incsr, new_outcsr));
}
ImmutableGraphPtr ImmutableGraph::CopyToSharedMem(ImmutableGraphPtr g, const std::string &name) {
ImmutableGraphPtr ImmutableGraph::CopyToSharedMem(
ImmutableGraphPtr g, const std::string &name) {
CSRPtr new_incsr, new_outcsr;
std::string shared_mem_name = GetSharedMemName(name, "in");
new_incsr = CSRPtr(new CSR(g->GetInCSR()->CopyToSharedMem(shared_mem_name)));
shared_mem_name = GetSharedMemName(name, "out");
new_outcsr = CSRPtr(new CSR(g->GetOutCSR()->CopyToSharedMem(shared_mem_name)));
new_outcsr =
CSRPtr(new CSR(g->GetOutCSR()->CopyToSharedMem(shared_mem_name)));
auto new_g = ImmutableGraphPtr(new ImmutableGraph(new_incsr, new_outcsr, name));
new_g->serialized_shared_meta_ = SerializeMetadata(new_g, GetSharedMemName(name, "meta"));
auto new_g =
ImmutableGraphPtr(new ImmutableGraph(new_incsr, new_outcsr, name));
new_g->serialized_shared_meta_ =
SerializeMetadata(new_g, GetSharedMemName(name, "meta"));
return new_g;
}
......@@ -598,8 +617,8 @@ ImmutableGraphPtr ImmutableGraph::AsNumBits(ImmutableGraphPtr g, uint8_t bits) {
ImmutableGraphPtr ImmutableGraph::Reverse() const {
if (coo_) {
return ImmutableGraphPtr(new ImmutableGraph(
out_csr_, in_csr_, coo_->Transpose()));
return ImmutableGraphPtr(
new ImmutableGraph(out_csr_, in_csr_, coo_->Transpose()));
} else {
return ImmutableGraphPtr(new ImmutableGraph(out_csr_, in_csr_));
}
......@@ -628,54 +647,53 @@ HeteroGraphPtr ImmutableGraph::AsHeteroGraph() const {
aten::CSRMatrix in_csr, out_csr;
aten::COOMatrix coo;
if (in_csr_)
in_csr = GetInCSR()->ToCSRMatrix();
if (out_csr_)
out_csr = GetOutCSR()->ToCSRMatrix();
if (coo_)
coo = GetCOO()->ToCOOMatrix();
if (in_csr_) in_csr = GetInCSR()->ToCSRMatrix();
if (out_csr_) out_csr = GetOutCSR()->ToCSRMatrix();
if (coo_) coo = GetCOO()->ToCOOMatrix();
auto g = UnitGraph::CreateUnitGraphFrom(
1, in_csr, out_csr, coo,
in_csr_ != nullptr,
out_csr_ != nullptr,
1, in_csr, out_csr, coo, in_csr_ != nullptr, out_csr_ != nullptr,
coo_ != nullptr);
return HeteroGraphPtr(new HeteroGraph(g->meta_graph(), {g}));
}
DGL_REGISTER_GLOBAL("transform._CAPI_DGLAsHeteroGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
ImmutableGraphPtr ig = std::dynamic_pointer_cast<ImmutableGraph>(g.sptr());
ImmutableGraphPtr ig =
std::dynamic_pointer_cast<ImmutableGraph>(g.sptr());
CHECK(ig) << "graph is not readonly";
*rv = HeteroGraphRef(ig->AsHeteroGraph());
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLImmutableGraphCopyTo")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
const int device_type = args[1];
const int device_id = args[2];
DGLContext ctx;
ctx.device_type = static_cast<DGLDeviceType>(device_type);
ctx.device_id = device_id;
ImmutableGraphPtr ig = CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
ImmutableGraphPtr ig =
CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
*rv = ImmutableGraph::CopyTo(ig, ctx);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLImmutableGraphCopyToSharedMem")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
std::string name = args[1];
ImmutableGraphPtr ig = CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
ImmutableGraphPtr ig =
CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
*rv = ImmutableGraph::CopyToSharedMem(ig, name);
});
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLImmutableGraphAsNumBits")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
int bits = args[1];
ImmutableGraphPtr ig = CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
ImmutableGraphPtr ig =
CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
*rv = ImmutableGraph::AsNumBits(ig, bits);
});
......
src/graph/metis_partition.cc
View file @ 81831111
......@@ -4,9 +4,10 @@
* \brief Call Metis partitioning
*/
#include <metis.h>
#include <dgl/graph_op.h>
#include <dgl/packed_func_ext.h>
#include <metis.h>
#include "../c_api_common.h"
using namespace dgl::runtime;
......@@ -25,12 +26,12 @@ IdArray MetisPartition(GraphPtr g, int k, NDArray vwgt_arr, bool obj_cut) {
idx_t nvtxs = g->NumVertices();
idx_t ncon = 1; // # balacing constraints.
idx_t *xadj = static_cast<idx_t*>(mat.indptr->data);
idx_t *adjncy = static_cast<idx_t*>(mat.indices->data);
idx_t *xadj = static_cast<idx_t *>(mat.indptr->data);
idx_t *adjncy = static_cast<idx_t *>(mat.indices->data);
idx_t nparts = k;
IdArray part_arr = aten::NewIdArray(nvtxs);
idx_t objval = 0;
idx_t *part = static_cast<idx_t*>(part_arr->data);
idx_t *part = static_cast<idx_t *>(part_arr->data);
int64_t vwgt_len = vwgt_arr->shape[0];
CHECK_EQ(sizeof(idx_t), vwgt_arr->dtype.bits / 8)
......@@ -40,7 +41,7 @@ IdArray MetisPartition(GraphPtr g, int k, NDArray vwgt_arr, bool obj_cut) {
idx_t *vwgt = NULL;
if (vwgt_len > 0) {
ncon = vwgt_len / g->NumVertices();
vwgt = static_cast<idx_t*>(vwgt_arr->data);
vwgt = static_cast<idx_t *>(vwgt_arr->data);
}
idx_t options[METIS_NOPTIONS];
......@@ -56,7 +57,8 @@ IdArray MetisPartition(GraphPtr g, int k, NDArray vwgt_arr, bool obj_cut) {
options[METIS_OPTION_OBJTYPE] = METIS_OBJTYPE_VOL;
}
int ret = METIS_PartGraphKway(&nvtxs, // The number of vertices
int ret = METIS_PartGraphKway(
&nvtxs, // The number of vertices
&ncon, // The number of balancing constraints.
xadj, // indptr
adjncy, // indices
......@@ -99,7 +101,7 @@ IdArray MetisPartition(GraphPtr g, int k, NDArray vwgt_arr, bool obj_cut) {
#endif // !defined(_WIN32)
DGL_REGISTER_GLOBAL("transform._CAPI_DGLMetisPartition")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
int k = args[1];
NDArray vwgt = args[2];
......
src/graph/network.cc
View file @ 81831111
......@@ -5,21 +5,20 @@
*/
#include "./network.h"
#include <stdlib.h>
#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
#include <dgl/runtime/ndarray.h>
#include <dgl/runtime/parallel_for.h>
#include <dgl/packed_func_ext.h>
#include <dgl/immutable_graph.h>
#include <dgl/nodeflow.h>
#include <stdlib.h>
#include <unordered_map>
#include "../rpc/network/common.h"
#include "../rpc/network/communicator.h"
#include "../rpc/network/socket_communicator.h"
#include "../rpc/network/msg_queue.h"
#include "../rpc/network/common.h"
#include "../rpc/network/socket_communicator.h"
using dgl::network::StringPrintf;
using namespace dgl::runtime;
......@@ -29,11 +28,8 @@ const bool AUTO_FREE = true;
namespace dgl {
namespace network {
NDArray CreateNDArrayFromRaw(std::vector<int64_t> shape,
DGLDataType dtype,
DGLContext ctx,
void* raw,
NDArray CreateNDArrayFromRaw(
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx, void* raw,
bool auto_free) {
return NDArray::CreateFromRaw(shape, dtype, ctx, raw, auto_free);
}
......@@ -74,16 +70,16 @@ char* ArrayMeta::Serialize(int64_t* size) {
*(reinterpret_cast<int*>(pointer)) = ndarray_count_;
pointer += sizeof(ndarray_count_);
// Write data type
memcpy(pointer,
reinterpret_cast<DGLDataType*>(data_type_.data()),
memcpy(
pointer, reinterpret_cast<DGLDataType*>(data_type_.data()),
sizeof(DGLDataType) * data_type_.size());
pointer += (sizeof(DGLDataType) * data_type_.size());
// Write size of data_shape_
*(reinterpret_cast<size_t*>(pointer)) = data_shape_.size();
pointer += sizeof(data_shape_.size());
// Write data of data_shape_
memcpy(pointer,
reinterpret_cast<char*>(data_shape_.data()),
memcpy(
pointer, reinterpret_cast<char*>(data_shape_.data()),
sizeof(int64_t) * data_shape_.size());
}
*size = buffer_size;
......@@ -103,8 +99,7 @@ void ArrayMeta::Deserialize(char* buffer, int64_t size) {
data_size += sizeof(int);
// Read data type
data_type_.resize(ndarray_count_);
memcpy(data_type_.data(), buffer,
ndarray_count_ * sizeof(DGLDataType));
memcpy(data_type_.data(), buffer, ndarray_count_ * sizeof(DGLDataType));
buffer += ndarray_count_ * sizeof(DGLDataType);
data_size += ndarray_count_ * sizeof(DGLDataType);
// Read size of data_shape_
......@@ -113,8 +108,7 @@ void ArrayMeta::Deserialize(char* buffer, int64_t size) {
data_size += sizeof(size_t);
data_shape_.resize(count);
// Read data of data_shape_
memcpy(data_shape_.data(), buffer,
count * sizeof(int64_t));
memcpy(data_shape_.data(), buffer, count * sizeof(int64_t));
data_size += count * sizeof(int64_t);
}
CHECK_EQ(data_size, size);
......@@ -170,11 +164,11 @@ void KVStoreMsg::Deserialize(char* buffer, int64_t size) {
CHECK_EQ(data_size, size);
}
////////////////////////////////// Basic Networking Components ////////////////////////////////
////////////////////////////////// Basic Networking Components
///////////////////////////////////
DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
std::string type = args[0];
int64_t msg_queue_size = args[1];
network::Sender* sender = nullptr;
......@@ -188,7 +182,7 @@ DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderCreate")
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
std::string type = args[0];
int64_t msg_queue_size = args[1];
network::Receiver* receiver = nullptr;
......@@ -202,21 +196,22 @@ DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverCreate")
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLFinalizeSender")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
network::Sender* sender = static_cast<network::Sender*>(chandle);
sender->Finalize();
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLFinalizeReceiver")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
network::Receiver* receiver =
static_cast<network::SocketReceiver*>(chandle);
receiver->Finalize();
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderAddReceiver")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
std::string ip = args[1];
int port = args[2];
......@@ -232,7 +227,7 @@ DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderAddReceiver")
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderConnect")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
network::Sender* sender = static_cast<network::Sender*>(chandle);
const int max_try_times = 1024;
......@@ -242,12 +237,13 @@ DGL_REGISTER_GLOBAL("network._CAPI_DGLSenderConnect")
});
DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverWait")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
std::string ip = args[1];
int port = args[2];
int num_sender = args[3];
network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
network::Receiver* receiver =
static_cast<network::SocketReceiver*>(chandle);
std::string addr;
if (receiver->NetType() == "socket") {
addr = StringPrintf("socket://%s:%d", ip.c_str(), port);
......@@ -259,12 +255,11 @@ DGL_REGISTER_GLOBAL("network._CAPI_DGLReceiverWait")
}
});
////////////////////////// Distributed Sampler Components ////////////////////////////////
////////////////////////// Distributed Sampler Components
///////////////////////////////////
DGL_REGISTER_GLOBAL("network._CAPI_SenderSendNodeFlow")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
int recv_id = args[1];
GraphRef g = args[2];
......@@ -348,7 +343,7 @@ DGL_REGISTER_GLOBAL("network._CAPI_SenderSendNodeFlow")
});
DGL_REGISTER_GLOBAL("network._CAPI_SenderSendSamplerEndSignal")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
int recv_id = args[1];
ArrayMeta meta(kFinalMsg);
......@@ -361,9 +356,10 @@ DGL_REGISTER_GLOBAL("network._CAPI_SenderSendSamplerEndSignal")
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
network::Receiver* receiver =
static_cast<network::SocketReceiver*>(chandle);
int send_id = 0;
Message recv_msg;
CHECK_EQ(receiver->Recv(&recv_msg, &send_id), REMOVE_SUCCESS);
......@@ -377,71 +373,50 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
CHECK_EQ(receiver->RecvFrom(&array_0, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[0], 1);
nf->node_mapping = CreateNDArrayFromRaw(
{meta.data_shape_[1]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_0.data,
AUTO_FREE);
{meta.data_shape_[1]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_0.data, AUTO_FREE);
// edge_mapping
Message array_1;
CHECK_EQ(receiver->RecvFrom(&array_1, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[2], 1);
nf->edge_mapping = CreateNDArrayFromRaw(
{meta.data_shape_[3]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_1.data,
AUTO_FREE);
{meta.data_shape_[3]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_1.data, AUTO_FREE);
// layer_offset
Message array_2;
CHECK_EQ(receiver->RecvFrom(&array_2, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[4], 1);
nf->layer_offsets = CreateNDArrayFromRaw(
{meta.data_shape_[5]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_2.data,
AUTO_FREE);
{meta.data_shape_[5]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_2.data, AUTO_FREE);
// flow_offset
Message array_3;
CHECK_EQ(receiver->RecvFrom(&array_3, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[6], 1);
nf->flow_offsets = CreateNDArrayFromRaw(
{meta.data_shape_[7]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_3.data,
AUTO_FREE);
{meta.data_shape_[7]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_3.data, AUTO_FREE);
// CSR indptr
Message array_4;
CHECK_EQ(receiver->RecvFrom(&array_4, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[8], 1);
NDArray indptr = CreateNDArrayFromRaw(
{meta.data_shape_[9]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_4.data,
AUTO_FREE);
{meta.data_shape_[9]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_4.data, AUTO_FREE);
// CSR indice
Message array_5;
CHECK_EQ(receiver->RecvFrom(&array_5, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[10], 1);
NDArray indice = CreateNDArrayFromRaw(
{meta.data_shape_[11]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_5.data,
AUTO_FREE);
{meta.data_shape_[11]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_5.data, AUTO_FREE);
// CSR edge_ids
Message array_6;
CHECK_EQ(receiver->RecvFrom(&array_6, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[12], 1);
NDArray edge_ids = CreateNDArrayFromRaw(
{meta.data_shape_[13]},
DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0},
array_6.data,
AUTO_FREE);
{meta.data_shape_[13]}, DGLDataType{kDGLInt, 64, 1},
DGLContext{kDGLCPU, 0}, array_6.data, AUTO_FREE);
// Create CSR
CSRPtr csr(new CSR(indptr, indice, edge_ids));
nf->graph = GraphPtr(new ImmutableGraph(csr, nullptr));
......@@ -453,14 +428,11 @@ DGL_REGISTER_GLOBAL("network._CAPI_ReceiverRecvNodeFlow")
}
});
////////////////////////// Distributed KVStore Components
///////////////////////////////////
////////////////////////// Distributed KVStore Components ////////////////////////////////
static void send_kv_message(network::Sender* sender,
KVStoreMsg* kv_msg,
int recv_id,
bool auto_free) {
static void send_kv_message(
network::Sender* sender, KVStoreMsg* kv_msg, int recv_id, bool auto_free) {
int64_t kv_size = 0;
char* kv_data = kv_msg->Serialize(&kv_size);
// Send kv_data
......@@ -471,23 +443,18 @@ static void send_kv_message(network::Sender* sender,
send_kv_msg.deallocator = DefaultMessageDeleter;
}
CHECK_EQ(sender->Send(send_kv_msg, recv_id), ADD_SUCCESS);
if (kv_msg->msg_type != kFinalMsg &&
kv_msg->msg_type != kBarrierMsg &&
kv_msg->msg_type != kIPIDMsg &&
kv_msg->msg_type != kGetShapeMsg) {
if (kv_msg->msg_type != kFinalMsg && kv_msg->msg_type != kBarrierMsg &&
kv_msg->msg_type != kIPIDMsg && kv_msg->msg_type != kGetShapeMsg) {
// Send ArrayMeta
ArrayMeta meta(kv_msg->msg_type);
if (kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
if (kv_msg->msg_type != kInitMsg && kv_msg->msg_type != kGetShapeBackMsg) {
meta.AddArray(kv_msg->id);
}
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kInitMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
meta.AddArray(kv_msg->data);
}
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kPushMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kPushMsg &&
kv_msg->msg_type != kPullBackMsg) {
meta.AddArray(kv_msg->shape);
}
......@@ -501,8 +468,7 @@ static void send_kv_message(network::Sender* sender,
}
CHECK_EQ(sender->Send(send_meta_msg, recv_id), ADD_SUCCESS);
// Send ID NDArray
if (kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeMsg &&
if (kv_msg->msg_type != kInitMsg && kv_msg->msg_type != kGetShapeMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
Message send_id_msg;
send_id_msg.data = static_cast<char*>(kv_msg->id->data);
......@@ -514,8 +480,7 @@ static void send_kv_message(network::Sender* sender,
CHECK_EQ(sender->Send(send_id_msg, recv_id), ADD_SUCCESS);
}
// Send data NDArray
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kInitMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
Message send_data_msg;
......@@ -528,8 +493,7 @@ static void send_kv_message(network::Sender* sender,
CHECK_EQ(sender->Send(send_data_msg, recv_id), ADD_SUCCESS);
}
// Send shape NDArray
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kPushMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kPushMsg &&
kv_msg->msg_type != kPullBackMsg) {
Message send_shape_msg;
send_shape_msg.data = static_cast<char*>(kv_msg->shape->data);
......@@ -544,17 +508,15 @@ static void send_kv_message(network::Sender* sender,
}
static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
KVStoreMsg *kv_msg = new KVStoreMsg();
KVStoreMsg* kv_msg = new KVStoreMsg();
// Recv kv_Msg
Message recv_kv_msg;
int send_id;
CHECK_EQ(receiver->Recv(&recv_kv_msg, &send_id), REMOVE_SUCCESS);
kv_msg->Deserialize(recv_kv_msg.data, recv_kv_msg.size);
recv_kv_msg.deallocator(&recv_kv_msg);
if (kv_msg->msg_type == kFinalMsg ||
kv_msg->msg_type == kBarrierMsg ||
kv_msg->msg_type == kIPIDMsg ||
kv_msg->msg_type == kGetShapeMsg) {
if (kv_msg->msg_type == kFinalMsg || kv_msg->msg_type == kBarrierMsg ||
kv_msg->msg_type == kIPIDMsg || kv_msg->msg_type == kGetShapeMsg) {
return kv_msg;
}
// Recv ArrayMeta
......@@ -563,21 +525,16 @@ static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
ArrayMeta meta(recv_meta_msg.data, recv_meta_msg.size);
recv_meta_msg.deallocator(&recv_meta_msg);
// Recv ID NDArray
if (kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
if (kv_msg->msg_type != kInitMsg && kv_msg->msg_type != kGetShapeBackMsg) {
Message recv_id_msg;
CHECK_EQ(receiver->RecvFrom(&recv_id_msg, send_id), REMOVE_SUCCESS);
CHECK_EQ(meta.data_shape_[0], 1);
kv_msg->id = CreateNDArrayFromRaw(
{meta.data_shape_[1]},
meta.data_type_[0],
DGLContext{kDGLCPU, 0},
recv_id_msg.data,
AUTO_FREE);
{meta.data_shape_[1]}, meta.data_type_[0], DGLContext{kDGLCPU, 0},
recv_id_msg.data, AUTO_FREE);
}
// Recv Data NDArray
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kInitMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kInitMsg &&
kv_msg->msg_type != kGetShapeBackMsg) {
Message recv_data_msg;
CHECK_EQ(receiver->RecvFrom(&recv_data_msg, send_id), REMOVE_SUCCESS);
......@@ -585,18 +542,14 @@ static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
CHECK_GE(ndim, 1);
std::vector<int64_t> vec_shape;
for (int i = 0; i < ndim; ++i) {
vec_shape.push_back(meta.data_shape_[3+i]);
vec_shape.push_back(meta.data_shape_[3 + i]);
}
kv_msg->data = CreateNDArrayFromRaw(
vec_shape,
meta.data_type_[1],
DGLContext{kDGLCPU, 0},
recv_data_msg.data,
AUTO_FREE);
vec_shape, meta.data_type_[1], DGLContext{kDGLCPU, 0},
recv_data_msg.data, AUTO_FREE);
}
// Recv Shape
if (kv_msg->msg_type != kPullMsg &&
kv_msg->msg_type != kPushMsg &&
if (kv_msg->msg_type != kPullMsg && kv_msg->msg_type != kPushMsg &&
kv_msg->msg_type != kPullBackMsg) {
Message recv_shape_msg;
CHECK_EQ(receiver->RecvFrom(&recv_shape_msg, send_id), REMOVE_SUCCESS);
......@@ -604,20 +557,17 @@ static KVStoreMsg* recv_kv_message(network::Receiver* receiver) {
CHECK_GE(ndim, 1);
std::vector<int64_t> vec_shape;
for (int i = 0; i < ndim; ++i) {
vec_shape.push_back(meta.data_shape_[1+i]);
vec_shape.push_back(meta.data_shape_[1 + i]);
}
kv_msg->shape = CreateNDArrayFromRaw(
vec_shape,
meta.data_type_[0],
DGLContext{kDGLCPU, 0},
recv_shape_msg.data,
AUTO_FREE);
vec_shape, meta.data_type_[0], DGLContext{kDGLCPU, 0},
recv_shape_msg.data, AUTO_FREE);
}
return kv_msg;
}
DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
int args_count = 0;
CommunicatorHandle chandle = args[args_count++];
int recv_id = args[args_count++];
......@@ -628,23 +578,18 @@ DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
if (kv_msg.msg_type != kFinalMsg && kv_msg.msg_type != kBarrierMsg) {
std::string name = args[args_count++];
kv_msg.name = name;
if (kv_msg.msg_type != kIPIDMsg &&
kv_msg.msg_type != kInitMsg &&
if (kv_msg.msg_type != kIPIDMsg && kv_msg.msg_type != kInitMsg &&
kv_msg.msg_type != kGetShapeMsg &&
kv_msg.msg_type != kGetShapeBackMsg) {
kv_msg.id = args[args_count++];
}
if (kv_msg.msg_type != kPullMsg &&
kv_msg.msg_type != kIPIDMsg &&
kv_msg.msg_type != kInitMsg &&
kv_msg.msg_type != kGetShapeMsg &&
if (kv_msg.msg_type != kPullMsg && kv_msg.msg_type != kIPIDMsg &&
kv_msg.msg_type != kInitMsg && kv_msg.msg_type != kGetShapeMsg &&
kv_msg.msg_type != kGetShapeBackMsg) {
kv_msg.data = args[args_count++];
}
if (kv_msg.msg_type != kIPIDMsg &&
kv_msg.msg_type != kPullMsg &&
kv_msg.msg_type != kPushMsg &&
kv_msg.msg_type != kPullBackMsg &&
if (kv_msg.msg_type != kIPIDMsg && kv_msg.msg_type != kPullMsg &&
kv_msg.msg_type != kPushMsg && kv_msg.msg_type != kPullBackMsg &&
kv_msg.msg_type != kGetShapeMsg) {
kv_msg.shape = args[args_count++];
}
......@@ -653,63 +598,64 @@ DGL_REGISTER_GLOBAL("network._CAPI_SenderSendKVMsg")
});
DGL_REGISTER_GLOBAL("network.CAPI_ReceiverRecvKVMsg")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
CommunicatorHandle chandle = args[0];
network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle);
network::Receiver* receiver =
static_cast<network::SocketReceiver*>(chandle);
*rv = recv_kv_message(receiver);
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgType")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->msg_type;
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgRank")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->rank;
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgName")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->name;
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->id;
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgData")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->data;
});
DGL_REGISTER_GLOBAL("network._CAPI_ReceiverGetKVMsgShape")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
*rv = msg->shape;
});
DGL_REGISTER_GLOBAL("network._CAPI_DeleteKVMsg")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
KVMsgHandle chandle = args[0];
network::KVStoreMsg* msg = static_cast<KVStoreMsg*>(chandle);
delete msg;
});
DGL_REGISTER_GLOBAL("network._CAPI_FastPull")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue* rv) {
std::string name = args[0];
int local_machine_id = args[1];
int machine_count = args[2];
......@@ -722,7 +668,8 @@ DGL_REGISTER_GLOBAL("network._CAPI_FastPull")
CommunicatorHandle chandle_receiver = args[9];
std::string str_flag = args[10];
network::Sender* sender = static_cast<network::Sender*>(chandle_sender);
network::Receiver* receiver = static_cast<network::SocketReceiver*>(chandle_receiver);
network::Receiver* receiver =
static_cast<network::SocketReceiver*>(chandle_receiver);
int64_t ID_size = ID.GetSize() / sizeof(int64_t);
int64_t* ID_data = static_cast<int64_t*>(ID->data);
int64_t* pb_data = static_cast<int64_t*>(pb->data);
......@@ -785,15 +732,13 @@ DGL_REGISTER_GLOBAL("network._CAPI_FastPull")
kv_msg.msg_type = MessageType::kPullMsg;
kv_msg.rank = client_id;
kv_msg.name = name;
kv_msg.id = CreateNDArrayFromRaw({static_cast<int64_t>(remote_ids[i].size())},
ID->dtype,
DGLContext{kDGLCPU, 0},
remote_ids[i].data(),
!AUTO_FREE);
int lower = i*group_count;
int higher = (i+1)*group_count-1;
kv_msg.id = CreateNDArrayFromRaw(
{static_cast<int64_t>(remote_ids[i].size())}, ID->dtype,
DGLContext{kDGLCPU, 0}, remote_ids[i].data(), !AUTO_FREE);
int lower = i * group_count;
int higher = (i + 1) * group_count - 1;
#ifndef _WIN32 // windows does not support rand_r()
int s_id = (rand_r(&seed) % (higher-lower+1))+lower;
int s_id = (rand_r(&seed) % (higher - lower + 1)) + lower;
send_kv_message(sender, &kv_msg, s_id, !AUTO_FREE);
#else
LOG(FATAL) << "KVStore does not support Windows yet.";
......@@ -801,40 +746,38 @@ DGL_REGISTER_GLOBAL("network._CAPI_FastPull")
msg_count++;
}
}
char *return_data = new char[ID_size*row_size];
char* return_data = new char[ID_size * row_size];
const int64_t local_ids_size = local_ids.size();
// Copy local data
runtime::parallel_for(0, local_ids_size, [&](size_t b, size_t e) {
for (auto i = b; i < e; ++i) {
CHECK_GE(ID_size*row_size, local_ids_orginal[i] * row_size + row_size);
CHECK_GE(
ID_size * row_size, local_ids_orginal[i] * row_size + row_size);
CHECK_GE(data_size, local_ids[i] * row_size + row_size);
CHECK_GE(local_ids[i], 0);
memcpy(return_data + local_ids_orginal[i] * row_size,
local_data_char + local_ids[i] * row_size,
row_size);
memcpy(
return_data + local_ids_orginal[i] * row_size,
local_data_char + local_ids[i] * row_size, row_size);
}
});
// Recv remote message
for (int i = 0; i < msg_count; ++i) {
KVStoreMsg *kv_msg = recv_kv_message(receiver);
KVStoreMsg* kv_msg = recv_kv_message(receiver);
int64_t id_size = kv_msg->id.GetSize() / sizeof(int64_t);
int part_id = kv_msg->rank / group_count;
char* data_char = static_cast<char*>(kv_msg->data->data);
for (int64_t n = 0; n < id_size; ++n) {
memcpy(return_data + remote_ids_original[part_id][n] * row_size,
data_char + n * row_size,
row_size);
memcpy(
return_data + remote_ids_original[part_id][n] * row_size,
data_char + n * row_size, row_size);
}
delete kv_msg;
}
// Get final tensor
local_data_shape[0] = ID_size;
NDArray res_tensor = CreateNDArrayFromRaw(
local_data_shape,
local_data->dtype,
DGLContext{kDGLCPU, 0},
return_data,
AUTO_FREE);
local_data_shape, local_data->dtype, DGLContext{kDGLCPU, 0},
return_data, AUTO_FREE);
*rv = res_tensor;
});
......
src/graph/network.h
View file @ 81831111
......@@ -6,12 +6,12 @@
#ifndef DGL_GRAPH_NETWORK_H_
#define DGL_GRAPH_NETWORK_H_
#include <dmlc/logging.h>
#include <dgl/runtime/ndarray.h>
#include <dmlc/logging.h>
#include <string.h>
#include <vector>
#include <string>
#include <vector>
#include "../c_api_common.h"
#include "../rpc/network/msg_queue.h"
......@@ -24,10 +24,8 @@ namespace network {
/*!
* \brief Create NDArray from raw data
*/
NDArray CreateNDArrayFromRaw(std::vector<int64_t> shape,
DGLDataType dtype,
DGLContext ctx,
void* raw);
NDArray CreateNDArrayFromRaw(
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx, void* raw);
/*!
* \brief Message type for DGL distributed training
......@@ -75,7 +73,6 @@ enum MessageType {
kGetShapeBackMsg = 9
};
/*!
* \brief Meta data for NDArray message
*/
......@@ -85,8 +82,7 @@ class ArrayMeta {
* \brief ArrayMeta constructor.
* \param msg_type type of message
*/
explicit ArrayMeta(int msg_type)
: msg_type_(msg_type), ndarray_count_(0) {}
explicit ArrayMeta(int msg_type) : msg_type_(msg_type), ndarray_count_(0) {}
/*!
* \brief Construct ArrayMeta from binary data buffer.
......@@ -101,16 +97,12 @@ class ArrayMeta {
/*!
* \return message type
*/
inline int msg_type() const {
return msg_type_;
}
inline int msg_type() const { return msg_type_; }
/*!
* \return count of ndarray
*/
inline int ndarray_count() const {
return ndarray_count_;
}
inline int ndarray_count() const { return ndarray_count_; }
/*!
* \brief Add NDArray meta data to ArrayMeta
......
src/graph/nodeflow.cc
View file @ 81831111
......@@ -5,8 +5,8 @@
*/
#include <dgl/immutable_graph.h>
#include <dgl/packed_func_ext.h>
#include <dgl/nodeflow.h>
#include <dgl/packed_func_ext.h>
#include <string>
......@@ -19,15 +19,16 @@ using dgl::runtime::PackedFunc;
namespace dgl {
std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::string &fmt,
size_t layer0_size, size_t layer1_start,
size_t layer1_end, bool remap) {
std::vector<IdArray> GetNodeFlowSlice(
const ImmutableGraph &graph, const std::string &fmt, size_t layer0_size,
size_t layer1_start, size_t layer1_end, bool remap) {
CHECK_GE(layer1_start, layer0_size);
if (fmt == std::string("csr")) {
dgl_id_t first_vid = layer1_start - layer0_size;
auto csr = aten::CSRSliceRows(graph.GetInCSR()->ToCSRMatrix(), layer1_start, layer1_end);
auto csr = aten::CSRSliceRows(
graph.GetInCSR()->ToCSRMatrix(), layer1_start, layer1_end);
if (remap) {
dgl_id_t *eid_data = static_cast<dgl_id_t*>(csr.data->data);
dgl_id_t *eid_data = static_cast<dgl_id_t *>(csr.data->data);
const dgl_id_t first_eid = eid_data[0];
IdArray new_indices = aten::Sub(csr.indices, first_vid);
IdArray new_data = aten::Sub(csr.data, first_eid);
......@@ -37,14 +38,14 @@ std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::st
}
} else if (fmt == std::string("coo")) {
auto csr = graph.GetInCSR()->ToCSRMatrix();
const dgl_id_t* indptr = static_cast<dgl_id_t*>(csr.indptr->data);
const dgl_id_t* indices = static_cast<dgl_id_t*>(csr.indices->data);
const dgl_id_t* edge_ids = static_cast<dgl_id_t*>(csr.data->data);
const dgl_id_t *indptr = static_cast<dgl_id_t *>(csr.indptr->data);
const dgl_id_t *indices = static_cast<dgl_id_t *>(csr.indices->data);
const dgl_id_t *edge_ids = static_cast<dgl_id_t *>(csr.data->data);
int64_t nnz = indptr[layer1_end] - indptr[layer1_start];
IdArray idx = aten::NewIdArray(2 * nnz);
IdArray eid = aten::NewIdArray(nnz);
int64_t *idx_data = static_cast<int64_t*>(idx->data);
dgl_id_t *eid_data = static_cast<dgl_id_t*>(eid->data);
int64_t *idx_data = static_cast<int64_t *>(idx->data);
dgl_id_t *eid_data = static_cast<dgl_id_t *>(eid->data);
size_t num_edges = 0;
for (size_t i = layer1_start; i < layer1_end; i++) {
for (dgl_id_t j = indptr[i]; j < indptr[i + 1]; j++) {
......@@ -65,10 +66,12 @@ std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::st
eid_data[i] = edge_ids[edge_start + i] - first_eid;
}
} else {
std::copy(indices + indptr[layer1_start],
indices + indptr[layer1_end], idx_data + nnz);
std::copy(edge_ids + indptr[layer1_start],
edge_ids + indptr[layer1_end], eid_data);
std::copy(
indices + indptr[layer1_start], indices + indptr[layer1_end],
idx_data + nnz);
std::copy(
edge_ids + indptr[layer1_start], edge_ids + indptr[layer1_end],
eid_data);
}
return std::vector<IdArray>{idx, eid};
} else {
......@@ -78,14 +81,15 @@ std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::st
}
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetBlockAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef g = args[0];
std::string format = args[1];
int64_t layer0_size = args[2];
int64_t start = args[3];
int64_t end = args[4];
const bool remap = args[5];
auto ig = CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
auto ig =
CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
auto res = GetNodeFlowSlice(*ig, format, layer0_size, start, end, remap);
*rv = ConvertNDArrayVectorToPackedFunc(res);
});
......
src/graph/pickle.cc
View file @ 81831111
......@@ -3,13 +3,14 @@
* \file graph/pickle.cc
* \brief Functions for pickle and unpickle a graph
*/
#include <dgl/graph_serializer.h>
#include <dgl/immutable_graph.h>
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
#include <dgl/immutable_graph.h>
#include <dgl/graph_serializer.h>
#include <dmlc/memory_io.h>
#include "./heterograph.h"
#include "../c_api_common.h"
#include "./heterograph.h"
#include "unit_graph.h"
using namespace dgl::runtime;
......@@ -91,7 +92,7 @@ HeteroPickleStates HeteroForkingPickle(HeteroGraphPtr graph) {
return states;
}
HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates &states) {
char *buf = const_cast<char *>(states.meta.c_str()); // a readonly stream?
dmlc::MemoryFixedSizeStream ifs(buf, states.meta.size());
dmlc::Stream *strm = &ifs;
......@@ -108,10 +109,10 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
auto array_itr = states.arrays.begin();
for (dgl_type_t etype = 0; etype < metagraph->NumEdges(); ++etype) {
const auto& pair = metagraph->FindEdge(etype);
const auto &pair = metagraph->FindEdge(etype);
const dgl_type_t srctype = pair.first;
const dgl_type_t dsttype = pair.second;
const int64_t num_vtypes = (srctype == dsttype)? 1 : 2;
const int64_t num_vtypes = (srctype == dsttype) ? 1 : 2;
int64_t num_src = num_nodes_per_type[srctype];
int64_t num_dst = num_nodes_per_type[dsttype];
SparseFormat fmt;
......@@ -126,7 +127,8 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
bool csorted;
CHECK(strm->Read(&rsorted)) << "Invalid flag 'rsorted'";
CHECK(strm->Read(&csorted)) << "Invalid flag 'csorted'";
auto coo = aten::COOMatrix(num_src, num_dst, row, col, aten::NullArray(), rsorted, csorted);
auto coo = aten::COOMatrix(
num_src, num_dst, row, col, aten::NullArray(), rsorted, csorted);
// TODO(zihao) fix
relgraph = CreateFromCOO(num_vtypes, coo, ALL_CODE);
break;
......@@ -138,7 +140,8 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
const auto &edge_id = *(array_itr++);
bool sorted;
CHECK(strm->Read(&sorted)) << "Invalid flag 'sorted'";
auto csr = aten::CSRMatrix(num_src, num_dst, indptr, indices, edge_id, sorted);
auto csr =
aten::CSRMatrix(num_src, num_dst, indptr, indices, edge_id, sorted);
// TODO(zihao) fix
relgraph = CreateFromCSR(num_vtypes, csr, ALL_CODE);
break;
......@@ -157,17 +160,18 @@ HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states) {
}
// For backward compatibility
HeteroGraphPtr HeteroUnpickleOld(const HeteroPickleStates& states) {
HeteroGraphPtr HeteroUnpickleOld(const HeteroPickleStates &states) {
const auto metagraph = states.metagraph;
const auto &num_nodes_per_type = states.num_nodes_per_type;
CHECK_EQ(states.adjs.size(), metagraph->NumEdges());
std::vector<HeteroGraphPtr> relgraphs(metagraph->NumEdges());
for (dgl_type_t etype = 0; etype < metagraph->NumEdges(); ++etype) {
const auto& pair = metagraph->FindEdge(etype);
const auto &pair = metagraph->FindEdge(etype);
const dgl_type_t srctype = pair.first;
const dgl_type_t dsttype = pair.second;
const int64_t num_vtypes = (srctype == dsttype)? 1 : 2;
const SparseFormat fmt = static_cast<SparseFormat>(states.adjs[etype]->format);
const int64_t num_vtypes = (srctype == dsttype) ? 1 : 2;
const SparseFormat fmt =
static_cast<SparseFormat>(states.adjs[etype]->format);
switch (fmt) {
case SparseFormat::kCOO:
relgraphs[etype] = UnitGraph::CreateFromCOO(
......@@ -202,7 +206,7 @@ HeteroGraphPtr HeteroForkingUnpickle(const HeteroPickleStates &states) {
auto array_itr = states.arrays.begin();
for (dgl_type_t etype = 0; etype < metagraph->NumEdges(); ++etype) {
const auto& pair = metagraph->FindEdge(etype);
const auto &pair = metagraph->FindEdge(etype);
const dgl_type_t srctype = pair.first;
const dgl_type_t dsttype = pair.second;
const int64_t num_vtypes = (srctype == dsttype) ? 1 : 2;
......@@ -227,7 +231,8 @@ HeteroGraphPtr HeteroForkingUnpickle(const HeteroPickleStates &states) {
bool csorted;
CHECK(strm->Read(&rsorted)) << "Invalid flag 'rsorted'";
CHECK(strm->Read(&csorted)) << "Invalid flag 'csorted'";
coo = aten::COOMatrix(num_src, num_dst, row, col, aten::NullArray(), rsorted, csorted);
coo = aten::COOMatrix(
num_src, num_dst, row, col, aten::NullArray(), rsorted, csorted);
}
if (created_formats & CSR_CODE) {
CHECK_GE(states.arrays.end() - array_itr, 3);
......@@ -258,13 +263,13 @@ HeteroGraphPtr HeteroForkingUnpickle(const HeteroPickleStates &states) {
}
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickleStatesGetVersion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef st = args[0];
*rv = st->version;
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickleStatesGetMeta")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef st = args[0];
DGLByteArray buf;
buf.data = st->meta.c_str();
......@@ -273,50 +278,50 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickleStatesGetMeta")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickleStatesGetArrays")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef st = args[0];
*rv = ConvertNDArrayVectorToPackedFunc(st->arrays);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickleStatesGetArraysNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef st = args[0];
*rv = static_cast<int64_t>(st->arrays.size());
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLCreateHeteroPickleStates")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
const int version = args[0];
std::string meta = args[1];
const List<Value> arrays = args[2];
std::shared_ptr<HeteroPickleStates> st( new HeteroPickleStates );
std::shared_ptr<HeteroPickleStates> st(new HeteroPickleStates);
st->version = version == 0 ? 1 : version;
st->meta = meta;
st->arrays.reserve(arrays.size());
for (const auto& ref : arrays) {
for (const auto &ref : arrays) {
st->arrays.push_back(ref->data);
}
*rv = HeteroPickleStatesRef(st);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPickle")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroGraphRef ref = args[0];
std::shared_ptr<HeteroPickleStates> st( new HeteroPickleStates );
std::shared_ptr<HeteroPickleStates> st(new HeteroPickleStates);
*st = HeteroPickle(ref.sptr());
*rv = HeteroPickleStatesRef(st);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroForkingPickle")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroGraphRef ref = args[0];
std::shared_ptr<HeteroPickleStates> st( new HeteroPickleStates );
std::shared_ptr<HeteroPickleStates> st(new HeteroPickleStates);
*st = HeteroForkingPickle(ref.sptr());
*rv = HeteroPickleStatesRef(st);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroUnpickle")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef ref = args[0];
HeteroGraphPtr graph;
switch (ref->version) {
......@@ -334,24 +339,23 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroUnpickle")
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroForkingUnpickle")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
HeteroPickleStatesRef ref = args[0];
HeteroGraphPtr graph = HeteroForkingUnpickle(*ref.sptr());
*rv = HeteroGraphRef(graph);
});
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLCreateHeteroPickleStatesOld")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
GraphRef metagraph = args[0];
IdArray num_nodes_per_type = args[1];
List<SparseMatrixRef> adjs = args[2];
std::shared_ptr<HeteroPickleStates> st( new HeteroPickleStates );
std::shared_ptr<HeteroPickleStates> st(new HeteroPickleStates);
st->version = 0;
st->metagraph = metagraph.sptr();
st->num_nodes_per_type = num_nodes_per_type.ToVector<int64_t>();
st->adjs.reserve(adjs.size());
for (const auto& ref : adjs)
st->adjs.push_back(ref.sptr());
for (const auto &ref : adjs) st->adjs.push_back(ref.sptr());
*rv = HeteroPickleStatesRef(st);
});
} // namespace dgl
src/graph/sampler.cc
View file @ 81831111
......@@ -3,18 +3,20 @@
* \file graph/sampler.cc
* \brief DGL sampler implementation
*/
#include <dgl/sampler.h>
#include <dgl/array.h>
#include <dgl/immutable_graph.h>
#include <dgl/runtime/container.h>
#include <dgl/packed_func_ext.h>
#include <dgl/random.h>
#include <dgl/runtime/container.h>
#include <dgl/runtime/parallel_for.h>
#include <dgl/sampler.h>
#include <dmlc/omp.h>
#include <algorithm>
#include <cstdlib>
#include <cmath>
#include <cstdlib>
#include <numeric>
#include "../c_api_common.h"
using namespace dgl::runtime;
......@@ -25,21 +27,21 @@ namespace {
/*
* ArrayHeap is used to sample elements from vector
*/
template<typename ValueType>
template <typename ValueType>
class ArrayHeap {
public:
explicit ArrayHeap(const std::vector<ValueType>& prob) {
explicit ArrayHeap(const std::vector<ValueType> &prob) {
vec_size_ = prob.size();
bit_len_ = ceil(log2(vec_size_));
limit_ = 1UL << bit_len_;
// allocate twice the size
heap_.resize(limit_ << 1, 0);
// allocate the leaves
for (size_t i = limit_; i < vec_size_+limit_; ++i) {
heap_[i] = prob[i-limit_];
for (size_t i = limit_; i < vec_size_ + limit_; ++i) {
heap_[i] = prob[i - limit_];
}
// iterate up the tree (this is O(m))
for (int i = bit_len_-1; i >= 0; --i) {
for (int i = bit_len_ - 1; i >= 0; --i) {
for (size_t j = (1UL << i); j < (1UL << (i + 1)); ++j) {
heap_[j] = heap_[j << 1] + heap_[(j << 1) + 1];
}
......@@ -54,7 +56,7 @@ class ArrayHeap {
size_t i = index + limit_;
heap_[i] = 0;
i /= 2;
for (int j = bit_len_-1; j >= 0; --j) {
for (int j = bit_len_ - 1; j >= 0; --j) {
// Using heap_[i] = heap_[i] - w will loss some precision in float.
// Using addition to re-calculate the weight layer by layer.
heap_[i] = heap_[i << 1] + heap_[(i << 1) + 1];
......@@ -93,7 +95,7 @@ class ArrayHeap {
/*
* Sample a vector by given the size n
*/
size_t SampleWithoutReplacement(size_t n, std::vector<size_t>* samples) {
size_t SampleWithoutReplacement(size_t n, std::vector<size_t> *samples) {
// sample n elements
size_t i = 0;
for (; i < n; ++i) {
......@@ -116,20 +118,14 @@ class ArrayHeap {
};
///////////////////////// Samplers //////////////////////////
class EdgeSamplerObject: public Object {
class EdgeSamplerObject : public Object {
public:
EdgeSamplerObject(const GraphPtr gptr,
IdArray seed_edges,
const int64_t batch_size,
const int64_t num_workers,
const bool replacement,
const bool reset,
const std::string neg_mode,
const int64_t neg_sample_size,
const int64_t chunk_size,
const bool exclude_positive,
const bool check_false_neg,
IdArray relations) {
EdgeSamplerObject(
const GraphPtr gptr, IdArray seed_edges, const int64_t batch_size,
const int64_t num_workers, const bool replacement, const bool reset,
const std::string neg_mode, const int64_t neg_sample_size,
const int64_t chunk_size, const bool exclude_positive,
const bool check_false_neg, IdArray relations) {
gptr_ = gptr;
seed_edges_ = seed_edges;
relations_ = relations;
......@@ -147,24 +143,22 @@ class EdgeSamplerObject: public Object {
~EdgeSamplerObject() {}
virtual void Fetch(DGLRetValue* rv) = 0;
virtual void Fetch(DGLRetValue *rv) = 0;
virtual void Reset() = 0;
protected:
virtual void randomSample(size_t set_size, size_t num, std::vector<size_t>* out) = 0;
virtual void randomSample(size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t>* out) = 0;
NegSubgraph genNegEdgeSubgraph(const Subgraph &pos_subg,
const std::string &neg_mode,
int64_t neg_sample_size,
bool exclude_positive,
bool check_false_neg);
NegSubgraph genChunkedNegEdgeSubgraph(const Subgraph &pos_subg,
const std::string &neg_mode,
int64_t neg_sample_size,
bool exclude_positive,
bool check_false_neg);
virtual void randomSample(
size_t set_size, size_t num, std::vector<size_t> *out) = 0;
virtual void randomSample(
size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t> *out) = 0;
NegSubgraph genNegEdgeSubgraph(
const Subgraph &pos_subg, const std::string &neg_mode,
int64_t neg_sample_size, bool exclude_positive, bool check_false_neg);
NegSubgraph genChunkedNegEdgeSubgraph(
const Subgraph &pos_subg, const std::string &neg_mode,
int64_t neg_sample_size, bool exclude_positive, bool check_false_neg);
GraphPtr gptr_;
IdArray seed_edges_;
......@@ -184,7 +178,7 @@ class EdgeSamplerObject: public Object {
/*
* Uniformly sample integers from [0, set_size) without replacement.
*/
void RandomSample(size_t set_size, size_t num, std::vector<size_t>* out) {
void RandomSample(size_t set_size, size_t num, std::vector<size_t> *out) {
if (num < set_size) {
std::unordered_set<size_t> sampled_idxs;
while (sampled_idxs.size() < num) {
......@@ -194,13 +188,13 @@ void RandomSample(size_t set_size, size_t num, std::vector<size_t>* out) {
} else {
// If we need to sample all elements in the set, we don't need to
// generate random numbers.
for (size_t i = 0; i < set_size; i++)
out->push_back(i);
for (size_t i = 0; i < set_size; i++) out->push_back(i);
}
}
void RandomSample(size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t>* out) {
void RandomSample(
size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t> *out) {
std::unordered_map<size_t, int> sampled_idxs;
for (auto v : exclude) {
sampled_idxs.insert(std::pair<size_t, int>(v, 0));
......@@ -231,9 +225,9 @@ void RandomSample(size_t set_size, size_t num, const std::vector<size_t> &exclud
* For a sparse array whose non-zeros are represented by nz_idxs,
* negate the sparse array and outputs the non-zeros in the negated array.
*/
void NegateArray(const std::vector<size_t> &nz_idxs,
size_t arr_size,
std::vector<size_t>* out) {
void NegateArray(
const std::vector<size_t> &nz_idxs, size_t arr_size,
std::vector<size_t> *out) {
// nz_idxs must have been sorted.
auto it = nz_idxs.begin();
size_t i = 0;
......@@ -253,12 +247,10 @@ void NegateArray(const std::vector<size_t> &nz_idxs,
/*
* Uniform sample vertices from a list of vertices.
*/
void GetUniformSample(const dgl_id_t* edge_id_list,
const dgl_id_t* vid_list,
const size_t ver_len,
const size_t max_num_neighbor,
std::vector<dgl_id_t>* out_ver,
std::vector<dgl_id_t>* out_edge) {
void GetUniformSample(
const dgl_id_t *edge_id_list, const dgl_id_t *vid_list,
const size_t ver_len, const size_t max_num_neighbor,
std::vector<dgl_id_t> *out_ver, std::vector<dgl_id_t> *out_edge) {
// Copy vid_list to output
if (ver_len <= max_num_neighbor) {
out_ver->insert(out_ver->end(), vid_list, vid_list + ver_len);
......@@ -292,16 +284,15 @@ void GetUniformSample(const dgl_id_t* edge_id_list,
/*
* Non-uniform sample via ArrayHeap
*
* \param probability Transition probability on the entire graph, indexed by edge ID
* \param probability Transition probability on the entire graph, indexed by
* edge ID
*/
template<typename ValueType>
void GetNonUniformSample(const ValueType* probability,
const dgl_id_t* edge_id_list,
const dgl_id_t* vid_list,
const size_t ver_len,
const size_t max_num_neighbor,
std::vector<dgl_id_t>* out_ver,
std::vector<dgl_id_t>* out_edge) {
template <typename ValueType>
void GetNonUniformSample(
const ValueType *probability, const dgl_id_t *edge_id_list,
const dgl_id_t *vid_list, const size_t ver_len,
const size_t max_num_neighbor, std::vector<dgl_id_t> *out_ver,
std::vector<dgl_id_t> *out_edge) {
// Copy vid_list to output
if (ver_len <= max_num_neighbor) {
out_ver->insert(out_ver->end(), vid_list, vid_list + ver_len);
......@@ -333,8 +324,8 @@ void GetNonUniformSample(const ValueType* probability,
struct neigh_list {
std::vector<dgl_id_t> neighs;
std::vector<dgl_id_t> edges;
neigh_list(const std::vector<dgl_id_t> &_neighs,
const std::vector<dgl_id_t> &_edges)
neigh_list(
const std::vector<dgl_id_t> &_neighs, const std::vector<dgl_id_t> &_edges)
: neighs(_neighs), edges(_edges) {}
};
......@@ -350,12 +341,11 @@ struct neighbor_info {
}
};
NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
std::vector<dgl_id_t> edge_list,
NodeFlow ConstructNodeFlow(
std::vector<dgl_id_t> neighbor_list, std::vector<dgl_id_t> edge_list,
std::vector<size_t> layer_offsets,
std::vector<std::pair<dgl_id_t, int> > *sub_vers,
std::vector<neighbor_info> *neigh_pos,
const std::string &edge_type,
std::vector<std::pair<dgl_id_t, int>> *sub_vers,
std::vector<neighbor_info> *neigh_pos, const std::string &edge_type,
int64_t num_edges, int num_hops) {
NodeFlow nf = NodeFlow::Create();
uint64_t num_vertices = sub_vers->size();
......@@ -371,9 +361,9 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
// Construct sub_csr_graph, we treat nodeflow as multigraph by default
auto subg_csr = CSRPtr(new CSR(num_vertices, num_edges));
dgl_id_t* indptr_out = static_cast<dgl_id_t*>(subg_csr->indptr()->data);
dgl_id_t* col_list_out = static_cast<dgl_id_t*>(subg_csr->indices()->data);
dgl_id_t* eid_out = static_cast<dgl_id_t*>(subg_csr->edge_ids()->data);
dgl_id_t *indptr_out = static_cast<dgl_id_t *>(subg_csr->indptr()->data);
dgl_id_t *col_list_out = static_cast<dgl_id_t *>(subg_csr->indices()->data);
dgl_id_t *eid_out = static_cast<dgl_id_t *>(subg_csr->edge_ids()->data);
size_t collected_nedges = 0;
// The data from the previous steps:
......@@ -385,12 +375,13 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
layer_ver_maps.resize(num_hops);
size_t out_node_idx = 0;
for (int layer_id = num_hops - 1; layer_id >= 0; layer_id--) {
// We sort the vertices in a layer so that we don't need to sort the neighbor Ids
// after remap to a subgraph. However, we don't need to sort the first layer
// because we want the order of the nodes in the first layer is the same as
// the input seed nodes.
// We sort the vertices in a layer so that we don't need to sort the
// neighbor Ids after remap to a subgraph. However, we don't need to sort
// the first layer because we want the order of the nodes in the first layer
// is the same as the input seed nodes.
if (layer_id > 0) {
std::sort(sub_vers->begin() + layer_offsets[layer_id],
std::sort(
sub_vers->begin() + layer_offsets[layer_id],
sub_vers->begin() + layer_offsets[layer_id + 1],
[](const std::pair<dgl_id_t, dgl_id_t> &a1,
const std::pair<dgl_id_t, dgl_id_t> &a2) {
......@@ -399,37 +390,40 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
}
// Save the sampled vertices and its layer Id.
for (size_t i = layer_offsets[layer_id]; i < layer_offsets[layer_id + 1]; i++) {
for (size_t i = layer_offsets[layer_id]; i < layer_offsets[layer_id + 1];
i++) {
node_map_data[out_node_idx++] = sub_vers->at(i).first;
layer_ver_maps[layer_id].insert(std::pair<dgl_id_t, dgl_id_t>(sub_vers->at(i).first,
ver_id++));
layer_ver_maps[layer_id].insert(
std::pair<dgl_id_t, dgl_id_t>(sub_vers->at(i).first, ver_id++));
CHECK_EQ(sub_vers->at(i).second, layer_id);
}
}
CHECK(out_node_idx == num_vertices);
// sampling algorithms have to start from the seed nodes, so the seed nodes are
// in the first layer and the input nodes are in the last layer.
// When we expose the sampled graph to a Python user, we say the input nodes
// are in the first layer and the seed nodes are in the last layer.
// Thus, when we copy sampled results to a CSR, we need to reverse the order of layers.
// sampling algorithms have to start from the seed nodes, so the seed nodes
// are in the first layer and the input nodes are in the last layer. When we
// expose the sampled graph to a Python user, we say the input nodes are in
// the first layer and the seed nodes are in the last layer. Thus, when we
// copy sampled results to a CSR, we need to reverse the order of layers.
std::fill(indptr_out, indptr_out + num_vertices + 1, 0);
size_t row_idx = layer_offsets[num_hops] - layer_offsets[num_hops - 1];
layer_off_data[0] = 0;
layer_off_data[1] = layer_offsets[num_hops] - layer_offsets[num_hops - 1];
int out_layer_idx = 1;
for (int layer_id = num_hops - 2; layer_id >= 0; layer_id--) {
// Because we don't sort the vertices in the first layer above, we can't sort
// the neighbor positions of the vertices in the first layer either.
// Because we don't sort the vertices in the first layer above, we can't
// sort the neighbor positions of the vertices in the first layer either.
if (layer_id > 0) {
std::sort(neigh_pos->begin() + layer_offsets[layer_id],
std::sort(
neigh_pos->begin() + layer_offsets[layer_id],
neigh_pos->begin() + layer_offsets[layer_id + 1],
[](const neighbor_info &a1, const neighbor_info &a2) {
return a1.id < a2.id;
});
}
for (size_t i = layer_offsets[layer_id]; i < layer_offsets[layer_id + 1]; i++) {
for (size_t i = layer_offsets[layer_id]; i < layer_offsets[layer_id + 1];
i++) {
dgl_id_t dst_id = sub_vers->at(i).first;
CHECK_EQ(dst_id, neigh_pos->at(i).id);
size_t pos = neigh_pos->at(i).pos;
......@@ -441,18 +435,22 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
auto neigh_it = neighbor_list.begin() + pos;
for (size_t i = 0; i < nedges; i++) {
dgl_id_t neigh = *(neigh_it + i);
CHECK(layer_ver_maps[layer_id + 1].find(neigh) != layer_ver_maps[layer_id + 1].end());
col_list_out[collected_nedges + i] = layer_ver_maps[layer_id + 1][neigh];
CHECK(
layer_ver_maps[layer_id + 1].find(neigh) !=
layer_ver_maps[layer_id + 1].end());
col_list_out[collected_nedges + i] =
layer_ver_maps[layer_id + 1][neigh];
}
// We can simply copy the edge Ids.
std::copy_n(edge_list.begin() + pos,
nedges, edge_map_data + collected_nedges);
std::copy_n(
edge_list.begin() + pos, nedges, edge_map_data + collected_nedges);
collected_nedges += nedges;
indptr_out[row_idx+1] = indptr_out[row_idx] + nedges;
indptr_out[row_idx + 1] = indptr_out[row_idx] + nedges;
row_idx++;
}
layer_off_data[out_layer_idx + 1] = layer_off_data[out_layer_idx]
+ layer_offsets[layer_id + 1] - layer_offsets[layer_id];
layer_off_data[out_layer_idx + 1] = layer_off_data[out_layer_idx] +
layer_offsets[layer_id + 1] -
layer_offsets[layer_id];
out_layer_idx++;
}
CHECK_EQ(row_idx, num_vertices);
......@@ -464,7 +462,8 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
flow_off_data[0] = 0;
int out_flow_idx = 0;
for (size_t i = 0; i < layer_offsets.size() - 2; i++) {
size_t num_edges = indptr_out[layer_off_data[i + 2]] - indptr_out[layer_off_data[i + 1]];
size_t num_edges =
indptr_out[layer_off_data[i + 2]] - indptr_out[layer_off_data[i + 1]];
flow_off_data[out_flow_idx + 1] = flow_off_data[out_flow_idx] + num_edges;
out_flow_idx++;
}
......@@ -482,23 +481,21 @@ NodeFlow ConstructNodeFlow(std::vector<dgl_id_t> neighbor_list,
return nf;
}
template<typename ValueType>
NodeFlow SampleSubgraph(const ImmutableGraph *graph,
const std::vector<dgl_id_t>& seeds,
const ValueType* probability,
const std::string &edge_type,
int num_hops,
size_t num_neighbor,
const bool add_self_loop) {
template <typename ValueType>
NodeFlow SampleSubgraph(
const ImmutableGraph *graph, const std::vector<dgl_id_t> &seeds,
const ValueType *probability, const std::string &edge_type, int num_hops,
size_t num_neighbor, const bool add_self_loop) {
CHECK_EQ(graph->NumBits(), 64) << "32 bit graph is not supported yet";
const size_t num_seeds = seeds.size();
auto orig_csr = edge_type == "in" ? graph->GetInCSR() : graph->GetOutCSR();
const dgl_id_t* val_list = static_cast<dgl_id_t*>(orig_csr->edge_ids()->data);
const dgl_id_t* col_list = static_cast<dgl_id_t*>(orig_csr->indices()->data);
const dgl_id_t* indptr = static_cast<dgl_id_t*>(orig_csr->indptr()->data);
const dgl_id_t *val_list =
static_cast<dgl_id_t *>(orig_csr->edge_ids()->data);
const dgl_id_t *col_list = static_cast<dgl_id_t *>(orig_csr->indices()->data);
const dgl_id_t *indptr = static_cast<dgl_id_t *>(orig_csr->indptr()->data);
std::unordered_set<dgl_id_t> sub_ver_map; // The vertex Ids in a layer.
std::vector<std::pair<dgl_id_t, int> > sub_vers;
std::vector<std::pair<dgl_id_t, int>> sub_vers;
sub_vers.reserve(num_seeds * 10);
// add seed vertices
for (size_t i = 0; i < num_seeds; ++i) {
......@@ -526,37 +523,38 @@ NodeFlow SampleSubgraph(const ImmutableGraph *graph,
// sampled nodes in a layer, and clear it when entering a new layer.
sub_ver_map.clear();
// Previous iteration collects all nodes in sub_vers, which are collected
// in the previous layer. sub_vers is used both as a node collection and a queue.
for (size_t idx = layer_offsets[layer_id - 1]; idx < layer_offsets[layer_id]; idx++) {
// in the previous layer. sub_vers is used both as a node collection and a
// queue.
for (size_t idx = layer_offsets[layer_id - 1];
idx < layer_offsets[layer_id]; idx++) {
dgl_id_t dst_id = sub_vers[idx].first;
const int cur_node_level = sub_vers[idx].second;
tmp_sampled_src_list.clear();
tmp_sampled_edge_list.clear();
dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id);
dgl_id_t ver_len = *(indptr + dst_id + 1) - *(indptr + dst_id);
if (probability == nullptr) { // uniform-sample
GetUniformSample(val_list + *(indptr + dst_id),
col_list + *(indptr + dst_id),
ver_len,
num_neighbor,
&tmp_sampled_src_list,
GetUniformSample(
val_list + *(indptr + dst_id), col_list + *(indptr + dst_id),
ver_len, num_neighbor, &tmp_sampled_src_list,
&tmp_sampled_edge_list);
} else { // non-uniform-sample
GetNonUniformSample(probability,
val_list + *(indptr + dst_id),
col_list + *(indptr + dst_id),
ver_len,
num_neighbor,
&tmp_sampled_src_list,
&tmp_sampled_edge_list);
}
// If we need to add self loop and it doesn't exist in the sampled neighbor list.
if (add_self_loop && std::find(tmp_sampled_src_list.begin(), tmp_sampled_src_list.end(),
GetNonUniformSample(
probability, val_list + *(indptr + dst_id),
col_list + *(indptr + dst_id), ver_len, num_neighbor,
&tmp_sampled_src_list, &tmp_sampled_edge_list);
}
// If we need to add self loop and it doesn't exist in the sampled
// neighbor list.
if (add_self_loop &&
std::find(
tmp_sampled_src_list.begin(), tmp_sampled_src_list.end(),
dst_id) == tmp_sampled_src_list.end()) {
tmp_sampled_src_list.push_back(dst_id);
const dgl_id_t *src_list = col_list + *(indptr + dst_id);
const dgl_id_t *eid_list = val_list + *(indptr + dst_id);
// TODO(zhengda) this operation has O(N) complexity. It can be pretty slow.
// TODO(zhengda) this operation has O(N) complexity. It can be pretty
// slow.
const dgl_id_t *src = std::find(src_list, src_list + ver_len, dst_id);
// If there doesn't exist a self loop in the graph.
// we have to add -1 as the edge id for the self-loop edge.
......@@ -566,7 +564,8 @@ NodeFlow SampleSubgraph(const ImmutableGraph *graph,
tmp_sampled_edge_list.push_back(eid_list[src - src_list]);
}
CHECK_EQ(tmp_sampled_src_list.size(), tmp_sampled_edge_list.size());
neigh_pos.emplace_back(dst_id, neighbor_list.size(), tmp_sampled_src_list.size());
neigh_pos.emplace_back(
dst_id, neighbor_list.size(), tmp_sampled_src_list.size());
// Then push the vertices
for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) {
neighbor_list.push_back(tmp_sampled_src_list[i]);
......@@ -578,8 +577,8 @@ NodeFlow SampleSubgraph(const ImmutableGraph *graph,
num_edges += tmp_sampled_src_list.size();
for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) {
// We need to add the neighbor in the hashtable here. This ensures that
// the vertex in the queue is unique. If we see a vertex before, we don't
// need to add it to the queue again.
// the vertex in the queue is unique. If we see a vertex before, we
// don't need to add it to the queue again.
auto ret = sub_ver_map.insert(tmp_sampled_src_list[i]);
// If the sampled neighbor is inserted to the map successfully.
if (ret.second) {
......@@ -591,76 +590,69 @@ NodeFlow SampleSubgraph(const ImmutableGraph *graph,
CHECK_EQ(layer_offsets[layer_id + 1], sub_vers.size());
}
return ConstructNodeFlow(neighbor_list, edge_list, layer_offsets, &sub_vers, &neigh_pos,
edge_type, num_edges, num_hops);
return ConstructNodeFlow(
neighbor_list, edge_list, layer_offsets, &sub_vers, &neigh_pos, edge_type,
num_edges, num_hops);
}
} // namespace
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
NodeFlow nflow = args[0];
*rv = nflow->graph;
});
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetNodeMapping")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
NodeFlow nflow = args[0];
*rv = nflow->node_mapping;
});
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetEdgeMapping")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
NodeFlow nflow = args[0];
*rv = nflow->edge_mapping;
});
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetLayerOffsets")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
NodeFlow nflow = args[0];
*rv = nflow->layer_offsets;
});
DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetBlockOffsets")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
NodeFlow nflow = args[0];
*rv = nflow->flow_offsets;
});
template<typename ValueType>
NodeFlow SamplerOp::NeighborSample(const ImmutableGraph *graph,
const std::vector<dgl_id_t>& seeds,
const std::string &edge_type,
int num_hops, int expand_factor,
const bool add_self_loop,
const ValueType *probability) {
return SampleSubgraph(graph,
seeds,
probability,
edge_type,
num_hops + 1,
expand_factor,
template <typename ValueType>
NodeFlow SamplerOp::NeighborSample(
const ImmutableGraph *graph, const std::vector<dgl_id_t> &seeds,
const std::string &edge_type, int num_hops, int expand_factor,
const bool add_self_loop, const ValueType *probability) {
return SampleSubgraph(
graph, seeds, probability, edge_type, num_hops + 1, expand_factor,
add_self_loop);
}
namespace {
void ConstructLayers(const dgl_id_t *indptr,
const dgl_id_t *indices,
const std::vector<dgl_id_t>& seed_array,
IdArray layer_sizes,
std::vector<dgl_id_t> *layer_offsets,
std::vector<dgl_id_t> *node_mapping,
std::vector<int64_t> *actl_layer_sizes,
std::vector<float> *probabilities) {
void ConstructLayers(
const dgl_id_t *indptr, const dgl_id_t *indices,
const std::vector<dgl_id_t> &seed_array, IdArray layer_sizes,
std::vector<dgl_id_t> *layer_offsets, std::vector<dgl_id_t> *node_mapping,
std::vector<int64_t> *actl_layer_sizes, std::vector<float> *probabilities) {
/*
* Given a graph and a collection of seed nodes, this function constructs NodeFlow
* layers via uniform layer-wise sampling, and return the resultant layers and their
* corresponding probabilities.
* Given a graph and a collection of seed nodes, this function constructs
* NodeFlow layers via uniform layer-wise sampling, and return the resultant
* layers and their corresponding probabilities.
*/
std::copy(seed_array.begin(), seed_array.end(), std::back_inserter(*node_mapping));
std::copy(
seed_array.begin(), seed_array.end(), std::back_inserter(*node_mapping));
actl_layer_sizes->push_back(node_mapping->size());
probabilities->insert(probabilities->end(), node_mapping->size(), 1);
const int64_t* layer_sizes_data = static_cast<int64_t*>(layer_sizes->data);
const int64_t *layer_sizes_data = static_cast<int64_t *>(layer_sizes->data);
const int64_t num_layers = layer_sizes->shape[0];
size_t curr = 0;
......@@ -674,14 +666,15 @@ namespace {
}
std::vector<dgl_id_t> candidate_vector;
std::copy(candidate_set.begin(), candidate_set.end(),
std::copy(
candidate_set.begin(), candidate_set.end(),
std::back_inserter(candidate_vector));
std::unordered_map<dgl_id_t, size_t> n_occurrences;
auto n_candidates = candidate_vector.size();
for (int64_t j = 0; j != layer_size; ++j) {
auto dst = candidate_vector[
RandomEngine::ThreadLocal()->RandInt(n_candidates)];
auto dst =
candidate_vector[RandomEngine::ThreadLocal()->RandInt(n_candidates)];
if (!n_occurrences.insert(std::make_pair(dst, 1)).second) {
++n_occurrences[dst];
}
......@@ -703,21 +696,18 @@ namespace {
for (const auto &size : *actl_layer_sizes) {
layer_offsets->push_back(size + layer_offsets->back());
}
}
}
void ConstructFlows(const dgl_id_t *indptr,
const dgl_id_t *indices,
const dgl_id_t *eids,
void ConstructFlows(
const dgl_id_t *indptr, const dgl_id_t *indices, const dgl_id_t *eids,
const std::vector<dgl_id_t> &node_mapping,
const std::vector<int64_t> &actl_layer_sizes,
std::vector<dgl_id_t> *sub_indptr,
std::vector<dgl_id_t> *sub_indices,
std::vector<dgl_id_t> *sub_eids,
std::vector<dgl_id_t> *flow_offsets,
std::vector<dgl_id_t> *sub_indptr, std::vector<dgl_id_t> *sub_indices,
std::vector<dgl_id_t> *sub_eids, std::vector<dgl_id_t> *flow_offsets,
std::vector<dgl_id_t> *edge_mapping) {
/*
* Given a graph and a sequence of NodeFlow layers, this function constructs dense
* subgraphs (flows) between consecutive layers.
* Given a graph and a sequence of NodeFlow layers, this function constructs
* dense subgraphs (flows) between consecutive layers.
*/
auto n_flows = actl_layer_sizes.size() - 1;
for (int64_t i = 0; i < actl_layer_sizes.front() + 1; i++)
......@@ -742,7 +732,9 @@ namespace {
neighbor_indices.push_back(std::make_pair(ret->second, eids[k]));
}
}
auto cmp = [](const id_pair p, const id_pair q)->bool { return p.first < q.first; };
auto cmp = [](const id_pair p, const id_pair q) -> bool {
return p.first < q.first;
};
std::sort(neighbor_indices.begin(), neighbor_indices.end(), cmp);
for (const auto &pair : neighbor_indices) {
sub_indices->push_back(pair.first);
......@@ -755,44 +747,32 @@ namespace {
}
sub_eids->resize(sub_indices->size());
std::iota(sub_eids->begin(), sub_eids->end(), 0);
}
}
} // namespace
NodeFlow SamplerOp::LayerUniformSample(const ImmutableGraph *graph,
const std::vector<dgl_id_t>& seeds,
const std::string &neighbor_type,
IdArray layer_sizes) {
const auto g_csr = neighbor_type == "in" ? graph->GetInCSR() : graph->GetOutCSR();
const dgl_id_t *indptr = static_cast<dgl_id_t*>(g_csr->indptr()->data);
const dgl_id_t *indices = static_cast<dgl_id_t*>(g_csr->indices()->data);
const dgl_id_t *eids = static_cast<dgl_id_t*>(g_csr->edge_ids()->data);
NodeFlow SamplerOp::LayerUniformSample(
const ImmutableGraph *graph, const std::vector<dgl_id_t> &seeds,
const std::string &neighbor_type, IdArray layer_sizes) {
const auto g_csr =
neighbor_type == "in" ? graph->GetInCSR() : graph->GetOutCSR();
const dgl_id_t *indptr = static_cast<dgl_id_t *>(g_csr->indptr()->data);
const dgl_id_t *indices = static_cast<dgl_id_t *>(g_csr->indices()->data);
const dgl_id_t *eids = static_cast<dgl_id_t *>(g_csr->edge_ids()->data);
std::vector<dgl_id_t> layer_offsets;
std::vector<dgl_id_t> node_mapping;
std::vector<int64_t> actl_layer_sizes;
std::vector<float> probabilities;
ConstructLayers(indptr,
indices,
seeds,
layer_sizes,
&layer_offsets,
&node_mapping,
&actl_layer_sizes,
&probabilities);
ConstructLayers(
indptr, indices, seeds, layer_sizes, &layer_offsets, &node_mapping,
&actl_layer_sizes, &probabilities);
std::vector<dgl_id_t> sub_indptr, sub_indices, sub_edge_ids;
std::vector<dgl_id_t> flow_offsets;
std::vector<dgl_id_t> edge_mapping;
ConstructFlows(indptr,
indices,
eids,
node_mapping,
actl_layer_sizes,
&sub_indptr,
&sub_indices,
&sub_edge_ids,
&flow_offsets,
&edge_mapping);
ConstructFlows(
indptr, indices, eids, node_mapping, actl_layer_sizes, &sub_indptr,
&sub_indices, &sub_edge_ids, &flow_offsets, &edge_mapping);
// sanity check
CHECK_GT(sub_indptr.size(), 0);
CHECK_EQ(sub_indptr[0], 0);
......@@ -800,8 +780,8 @@ NodeFlow SamplerOp::LayerUniformSample(const ImmutableGraph *graph,
CHECK_EQ(sub_indices.size(), sub_edge_ids.size());
NodeFlow nf = NodeFlow::Create();
auto sub_csr = CSRPtr(new CSR(aten::VecToIdArray(sub_indptr),
aten::VecToIdArray(sub_indices),
auto sub_csr = CSRPtr(new CSR(
aten::VecToIdArray(sub_indptr), aten::VecToIdArray(sub_indices),
aten::VecToIdArray(sub_edge_ids)));
if (neighbor_type == std::string("in")) {
......@@ -830,24 +810,22 @@ void BuildCsr(const ImmutableGraph &g, const std::string neigh_type) {
}
}
template<typename ValueType>
std::vector<NodeFlow> NeighborSamplingImpl(const ImmutableGraphPtr gptr,
const IdArray seed_nodes,
const int64_t batch_start_id,
const int64_t batch_size,
const int64_t max_num_workers,
const int64_t expand_factor,
const int64_t num_hops,
const std::string neigh_type,
const bool add_self_loop,
const ValueType *probability) {
template <typename ValueType>
std::vector<NodeFlow> NeighborSamplingImpl(
const ImmutableGraphPtr gptr, const IdArray seed_nodes,
const int64_t batch_start_id, const int64_t batch_size,
const int64_t max_num_workers, const int64_t expand_factor,
const int64_t num_hops, const std::string neigh_type,
const bool add_self_loop, const ValueType *probability) {
// process args
CHECK(aten::IsValidIdArray(seed_nodes));
const dgl_id_t* seed_nodes_data = static_cast<dgl_id_t*>(seed_nodes->data);
const dgl_id_t *seed_nodes_data = static_cast<dgl_id_t *>(seed_nodes->data);
const int64_t num_seeds = seed_nodes->shape[0];
const int64_t num_workers = std::min(max_num_workers,
const int64_t num_workers = std::min(
max_num_workers,
(num_seeds + batch_size - 1) / batch_size - batch_start_id);
// We need to make sure we have the right CSR before we enter parallel sampling.
// We need to make sure we have the right CSR before we enter parallel
// sampling.
BuildCsr(*gptr, neigh_type);
// generate node flows
std::vector<NodeFlow> nflows(num_workers);
......@@ -858,8 +836,8 @@ std::vector<NodeFlow> NeighborSamplingImpl(const ImmutableGraphPtr gptr,
const int64_t end = std::min(start + batch_size, num_seeds);
// TODO(minjie): the vector allocation/copy is unnecessary
std::vector<dgl_id_t> worker_seeds(end - start);
std::copy(seed_nodes_data + start, seed_nodes_data + end,
worker_seeds.begin());
std::copy(
seed_nodes_data + start, seed_nodes_data + end, worker_seeds.begin());
nflows[i] = SamplerOp::NeighborSample(
gptr.get(), worker_seeds, neigh_type, num_hops, expand_factor,
add_self_loop, probability);
......@@ -869,7 +847,7 @@ std::vector<NodeFlow> NeighborSamplingImpl(const ImmutableGraphPtr gptr,
}
DGL_REGISTER_GLOBAL("sampling._CAPI_UniformSampling")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
// arguments
const GraphRef g = args[0];
const IdArray seed_nodes = args[1];
......@@ -896,7 +874,7 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_UniformSampling")
});
DGL_REGISTER_GLOBAL("sampling._CAPI_NeighborSampling")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
// arguments
const GraphRef g = args[0];
const IdArray seed_nodes = args[1];
......@@ -926,17 +904,17 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_NeighborSampling")
<< "NeighborSampling only support CPU sampling";
ATEN_FLOAT_TYPE_SWITCH(
probability->dtype,
FloatType,
"transition probability",
{
probability->dtype, FloatType, "transition probability", {
const FloatType *prob;
if (aten::IsNullArray(probability)) {
prob = nullptr;
} else {
CHECK(probability->shape[0] == static_cast<int64_t>(gptr->NumEdges()))
<< "transition probability must have same number of elements as edges";
CHECK(
probability->shape[0] ==
static_cast<int64_t>(gptr->NumEdges()))
<< "transition probability must have same number of elements "
"as edges";
CHECK(probability.IsContiguous())
<< "transition probability must be contiguous tensor";
prob = static_cast<const FloatType *>(probability->data);
......@@ -951,7 +929,7 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_NeighborSampling")
});
DGL_REGISTER_GLOBAL("sampling._CAPI_LayerSampling")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
// arguments
GraphRef g = args[0];
const IdArray seed_nodes = args[1];
......@@ -971,11 +949,14 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_LayerSampling")
CHECK_EQ(layer_sizes->ctx.device_type, kDGLCPU)
<< "LayerSampler only support CPU sampling";
const dgl_id_t* seed_nodes_data = static_cast<dgl_id_t*>(seed_nodes->data);
const dgl_id_t *seed_nodes_data =
static_cast<dgl_id_t *>(seed_nodes->data);
const int64_t num_seeds = seed_nodes->shape[0];
const int64_t num_workers = std::min(max_num_workers,
const int64_t num_workers = std::min(
max_num_workers,
(num_seeds + batch_size - 1) / batch_size - batch_start_id);
// We need to make sure we have the right CSR before we enter parallel sampling.
// We need to make sure we have the right CSR before we enter parallel
// sampling.
BuildCsr(*gptr, neigh_type);
// generate node flows
std::vector<NodeFlow> nflows(num_workers);
......@@ -986,7 +967,8 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_LayerSampling")
const int64_t end = std::min(start + batch_size, num_seeds);
// TODO(minjie): the vector allocation/copy is unnecessary
std::vector<dgl_id_t> worker_seeds(end - start);
std::copy(seed_nodes_data + start, seed_nodes_data + end,
std::copy(
seed_nodes_data + start, seed_nodes_data + end,
worker_seeds.begin());
nflows[i] = SamplerOp::LayerUniformSample(
gptr.get(), worker_seeds, neigh_type, layer_sizes);
......@@ -1002,9 +984,8 @@ void BuildCoo(const ImmutableGraph &g) {
assert(coo);
}
dgl_id_t global2local_map(dgl_id_t global_id,
std::unordered_map<dgl_id_t, dgl_id_t> *map) {
dgl_id_t global2local_map(
dgl_id_t global_id, std::unordered_map<dgl_id_t, dgl_id_t> *map) {
auto it = map->find(global_id);
if (it == map->end()) {
dgl_id_t local_id = map->size();
......@@ -1020,7 +1001,8 @@ inline bool IsNegativeHeadMode(const std::string &mode) {
}
IdArray GetGlobalVid(IdArray induced_nid, IdArray subg_nid) {
IdArray gnid = IdArray::Empty({subg_nid->shape[0]}, subg_nid->dtype, subg_nid->ctx);
IdArray gnid =
IdArray::Empty({subg_nid->shape[0]}, subg_nid->dtype, subg_nid->ctx);
const dgl_id_t *induced_nid_data = static_cast<dgl_id_t *>(induced_nid->data);
const dgl_id_t *subg_nid_data = static_cast<dgl_id_t *>(subg_nid->data);
dgl_id_t *gnid_data = static_cast<dgl_id_t *>(gnid->data);
......@@ -1030,14 +1012,14 @@ IdArray GetGlobalVid(IdArray induced_nid, IdArray subg_nid) {
return gnid;
}
IdArray CheckExistence(GraphPtr gptr, IdArray neg_src, IdArray neg_dst,
IdArray induced_nid) {
return gptr->HasEdgesBetween(GetGlobalVid(induced_nid, neg_src),
GetGlobalVid(induced_nid, neg_dst));
IdArray CheckExistence(
GraphPtr gptr, IdArray neg_src, IdArray neg_dst, IdArray induced_nid) {
return gptr->HasEdgesBetween(
GetGlobalVid(induced_nid, neg_src), GetGlobalVid(induced_nid, neg_dst));
}
IdArray CheckExistence(GraphPtr gptr, IdArray relations,
IdArray neg_src, IdArray neg_dst,
IdArray CheckExistence(
GraphPtr gptr, IdArray relations, IdArray neg_src, IdArray neg_dst,
IdArray induced_nid, IdArray neg_eid) {
neg_src = GetGlobalVid(induced_nid, neg_src);
neg_dst = GetGlobalVid(induced_nid, neg_dst);
......@@ -1051,8 +1033,7 @@ IdArray CheckExistence(GraphPtr gptr, IdArray relations,
int64_t num_neg_edges = neg_src->shape[0];
for (int64_t i = 0; i < num_neg_edges; i++) {
// If the edge doesn't exist, we don't need to do anything.
if (!exist_data[i])
continue;
if (!exist_data[i]) continue;
// If the edge exists, we need to double check if the relations match.
// If they match, this negative edge isn't really a negative edge.
dgl_id_t eid1 = neg_eid_data[i];
......@@ -1073,7 +1054,8 @@ IdArray CheckExistence(GraphPtr gptr, IdArray relations,
return exist;
}
std::vector<dgl_id_t> Global2Local(const std::vector<size_t> &ids,
std::vector<dgl_id_t> Global2Local(
const std::vector<size_t> &ids,
const std::unordered_map<dgl_id_t, dgl_id_t> &map) {
std::vector<dgl_id_t> local_ids(ids.size());
for (size_t i = 0; i < ids.size(); i++) {
......@@ -1084,35 +1066,36 @@ std::vector<dgl_id_t> Global2Local(const std::vector<size_t> &ids,
return local_ids;
}
NegSubgraph EdgeSamplerObject::genNegEdgeSubgraph(const Subgraph &pos_subg,
const std::string &neg_mode,
int64_t neg_sample_size,
bool exclude_positive,
bool check_false_neg) {
NegSubgraph EdgeSamplerObject::genNegEdgeSubgraph(
const Subgraph &pos_subg, const std::string &neg_mode,
int64_t neg_sample_size, bool exclude_positive, bool check_false_neg) {
int64_t num_tot_nodes = gptr_->NumVertices();
if (neg_sample_size > num_tot_nodes)
neg_sample_size = num_tot_nodes;
if (neg_sample_size > num_tot_nodes) neg_sample_size = num_tot_nodes;
std::vector<IdArray> adj = pos_subg.graph->GetAdj(false, "coo");
IdArray coo = adj[0];
int64_t num_pos_edges = coo->shape[0] / 2;
int64_t num_neg_edges = num_pos_edges * neg_sample_size;
IdArray neg_dst = IdArray::Empty({num_neg_edges}, coo->dtype, coo->ctx);
IdArray neg_src = IdArray::Empty({num_neg_edges}, coo->dtype, coo->ctx);
IdArray induced_neg_eid = IdArray::Empty({num_neg_edges}, coo->dtype, coo->ctx);
IdArray induced_neg_eid =
IdArray::Empty({num_neg_edges}, coo->dtype, coo->ctx);
// These are vids in the positive subgraph.
const dgl_id_t *dst_data = static_cast<const dgl_id_t *>(coo->data);
const dgl_id_t *src_data = static_cast<const dgl_id_t *>(coo->data) + num_pos_edges;
const dgl_id_t *src_data =
static_cast<const dgl_id_t *>(coo->data) + num_pos_edges;
const dgl_id_t *induced_vid_data =
static_cast<const dgl_id_t *>(pos_subg.induced_vertices->data);
const dgl_id_t *induced_eid_data =
static_cast<const dgl_id_t *>(pos_subg.induced_edges->data);
size_t num_pos_nodes = pos_subg.graph->NumVertices();
std::vector<size_t> pos_nodes(induced_vid_data, induced_vid_data + num_pos_nodes);
std::vector<size_t> pos_nodes(
induced_vid_data, induced_vid_data + num_pos_nodes);
dgl_id_t *neg_dst_data = static_cast<dgl_id_t *>(neg_dst->data);
dgl_id_t *neg_src_data = static_cast<dgl_id_t *>(neg_src->data);
dgl_id_t *induced_neg_eid_data = static_cast<dgl_id_t *>(induced_neg_eid->data);
dgl_id_t *induced_neg_eid_data =
static_cast<dgl_id_t *>(induced_neg_eid->data);
const dgl_id_t *unchanged;
dgl_id_t *neg_unchanged;
......@@ -1206,7 +1189,8 @@ NegSubgraph EdgeSamplerObject::genNegEdgeSubgraph(const Subgraph &pos_subg,
for (int64_t j = 0; j < neg_sample_size; j++) {
neg_unchanged[neg_idx + j] = local_unchanged;
dgl_id_t local_changed = global2local_map(neg_vids[j + prev_neg_offset], &neg_map);
dgl_id_t local_changed =
global2local_map(neg_vids[j + prev_neg_offset], &neg_map);
neg_changed[neg_idx + j] = local_changed;
// induced negative eid references to the positive one.
induced_neg_eid_data[neg_idx + j] = induced_eid_data[i];
......@@ -1215,8 +1199,10 @@ NegSubgraph EdgeSamplerObject::genNegEdgeSubgraph(const Subgraph &pos_subg,
// Now we know the number of vertices in the negative graph.
int64_t num_neg_nodes = neg_map.size();
IdArray induced_neg_vid = IdArray::Empty({num_neg_nodes}, coo->dtype, coo->ctx);
dgl_id_t *induced_neg_vid_data = static_cast<dgl_id_t *>(induced_neg_vid->data);
IdArray induced_neg_vid =
IdArray::Empty({num_neg_nodes}, coo->dtype, coo->ctx);
dgl_id_t *induced_neg_vid_data =
static_cast<dgl_id_t *>(induced_neg_vid->data);
for (auto it = neg_map.begin(); it != neg_map.end(); it++) {
induced_neg_vid_data[it->second] = it->first;
}
......@@ -1241,24 +1227,22 @@ NegSubgraph EdgeSamplerObject::genNegEdgeSubgraph(const Subgraph &pos_subg,
if (aten::IsNullArray(relations_)) {
neg_subg.exist = CheckExistence(gptr_, neg_src, neg_dst, induced_neg_vid);
} else {
neg_subg.exist = CheckExistence(gptr_, relations_, neg_src, neg_dst,
induced_neg_vid, induced_neg_eid);
neg_subg.exist = CheckExistence(
gptr_, relations_, neg_src, neg_dst, induced_neg_vid,
induced_neg_eid);
}
}
return neg_subg;
}
NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_subg,
const std::string &neg_mode,
int64_t neg_sample_size,
bool exclude_positive,
bool check_false_neg) {
NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(
const Subgraph &pos_subg, const std::string &neg_mode,
int64_t neg_sample_size, bool exclude_positive, bool check_false_neg) {
int64_t num_tot_nodes = gptr_->NumVertices();
std::vector<IdArray> adj = pos_subg.graph->GetAdj(false, "coo");
IdArray coo = adj[0];
int64_t num_pos_edges = coo->shape[0] / 2;
if (neg_sample_size > num_tot_nodes)
neg_sample_size = num_tot_nodes;
if (neg_sample_size > num_tot_nodes) neg_sample_size = num_tot_nodes;
int64_t chunk_size = chunk_size_;
CHECK_GT(chunk_size, 0) << "chunk size has to be positive";
......@@ -1275,26 +1259,29 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
int64_t num_all_neg_edges = num_neg_edges + num_neg_edges_last_chunk;
// We should include the last chunk.
if (last_chunk_size > 0)
num_chunks++;
if (last_chunk_size > 0) num_chunks++;
IdArray neg_dst = IdArray::Empty({num_all_neg_edges}, coo->dtype, coo->ctx);
IdArray neg_src = IdArray::Empty({num_all_neg_edges}, coo->dtype, coo->ctx);
IdArray induced_neg_eid = IdArray::Empty({num_all_neg_edges}, coo->dtype, coo->ctx);
IdArray induced_neg_eid =
IdArray::Empty({num_all_neg_edges}, coo->dtype, coo->ctx);
// These are vids in the positive subgraph.
const dgl_id_t *dst_data = static_cast<const dgl_id_t *>(coo->data);
const dgl_id_t *src_data = static_cast<const dgl_id_t *>(coo->data) + num_pos_edges;
const dgl_id_t *src_data =
static_cast<const dgl_id_t *>(coo->data) + num_pos_edges;
const dgl_id_t *induced_vid_data =
static_cast<const dgl_id_t *>(pos_subg.induced_vertices->data);
const dgl_id_t *induced_eid_data =
static_cast<const dgl_id_t *>(pos_subg.induced_edges->data);
int64_t num_pos_nodes = pos_subg.graph->NumVertices();
std::vector<dgl_id_t> pos_nodes(induced_vid_data, induced_vid_data + num_pos_nodes);
std::vector<dgl_id_t> pos_nodes(
induced_vid_data, induced_vid_data + num_pos_nodes);
dgl_id_t *neg_dst_data = static_cast<dgl_id_t *>(neg_dst->data);
dgl_id_t *neg_src_data = static_cast<dgl_id_t *>(neg_src->data);
dgl_id_t *induced_neg_eid_data = static_cast<dgl_id_t *>(induced_neg_eid->data);
dgl_id_t *induced_neg_eid_data =
static_cast<dgl_id_t *>(induced_neg_eid->data);
const dgl_id_t *unchanged;
dgl_id_t *neg_unchanged;
......@@ -1312,9 +1299,7 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
// We first sample all negative edges.
std::vector<size_t> global_neg_vids;
std::vector<size_t> local_neg_vids;
randomSample(num_tot_nodes,
num_chunks * neg_sample_size,
&global_neg_vids);
randomSample(num_tot_nodes, num_chunks * neg_sample_size, &global_neg_vids);
CHECK_EQ(num_chunks * neg_sample_size, global_neg_vids.size());
std::unordered_map<dgl_id_t, dgl_id_t> neg_map;
......@@ -1336,7 +1321,7 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
// to reduce computation overhead.
local_neg_vids.resize(global_neg_vids.size());
for (size_t i = 0; i < global_neg_vids.size(); i++) {
local_neg_vids[i] = global2local_map(global_neg_vids[i], &neg_map);;
local_neg_vids[i] = global2local_map(global_neg_vids[i], &neg_map);
}
for (int64_t i_chunk = 0; i_chunk < num_chunks; i_chunk++) {
......@@ -1353,12 +1338,14 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
for (int64_t in_chunk = 0; in_chunk != chunk_size1; ++in_chunk) {
// For each positive node in a chunk.
dgl_id_t global_unchanged = induced_vid_data[unchanged[pos_edge_idx + in_chunk]];
dgl_id_t global_unchanged =
induced_vid_data[unchanged[pos_edge_idx + in_chunk]];
dgl_id_t local_unchanged = global2local_map(global_unchanged, &neg_map);
for (int64_t j = 0; j < neg_sample_size; ++j) {
neg_unchanged[neg_idx] = local_unchanged;
neg_changed[neg_idx] = local_neg_vids[neg_node_idx + j];
induced_neg_eid_data[neg_idx] = induced_eid_data[pos_edge_idx + in_chunk];
induced_neg_eid_data[neg_idx] =
induced_eid_data[pos_edge_idx + in_chunk];
neg_idx++;
}
}
......@@ -1366,8 +1353,10 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
// Now we know the number of vertices in the negative graph.
int64_t num_neg_nodes = neg_map.size();
IdArray induced_neg_vid = IdArray::Empty({num_neg_nodes}, coo->dtype, coo->ctx);
dgl_id_t *induced_neg_vid_data = static_cast<dgl_id_t *>(induced_neg_vid->data);
IdArray induced_neg_vid =
IdArray::Empty({num_neg_nodes}, coo->dtype, coo->ctx);
dgl_id_t *induced_neg_vid_data =
static_cast<dgl_id_t *>(induced_neg_vid->data);
for (auto it = neg_map.begin(); it != neg_map.end(); it++) {
induced_neg_vid_data[it->second] = it->first;
}
......@@ -1380,18 +1369,21 @@ NegSubgraph EdgeSamplerObject::genChunkedNegEdgeSubgraph(const Subgraph &pos_sub
neg_subg.induced_vertices = induced_neg_vid;
neg_subg.induced_edges = induced_neg_eid;
if (IsNegativeHeadMode(neg_mode)) {
neg_subg.head_nid = aten::VecToIdArray(Global2Local(global_neg_vids, neg_map));
neg_subg.head_nid =
aten::VecToIdArray(Global2Local(global_neg_vids, neg_map));
neg_subg.tail_nid = aten::VecToIdArray(local_pos_vids);
} else {
neg_subg.head_nid = aten::VecToIdArray(local_pos_vids);
neg_subg.tail_nid = aten::VecToIdArray(Global2Local(global_neg_vids, neg_map));
neg_subg.tail_nid =
aten::VecToIdArray(Global2Local(global_neg_vids, neg_map));
}
if (check_false_neg) {
if (aten::IsNullArray(relations_)) {
neg_subg.exist = CheckExistence(gptr_, neg_src, neg_dst, induced_neg_vid);
} else {
neg_subg.exist = CheckExistence(gptr_, relations_, neg_src, neg_dst,
induced_neg_vid, induced_neg_eid);
neg_subg.exist = CheckExistence(
gptr_, relations_, neg_src, neg_dst, induced_neg_vid,
induced_neg_eid);
}
}
return neg_subg;
......@@ -1408,52 +1400,38 @@ inline SubgraphRef ConvertRef(const NegSubgraph &subg) {
} // namespace
DGL_REGISTER_GLOBAL("sampling._CAPI_GetNegEdgeExistence")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
SubgraphRef g = args[0];
auto gptr = std::dynamic_pointer_cast<NegSubgraph>(g.sptr());
*rv = gptr->exist;
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_GetEdgeSubgraphHead")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
SubgraphRef g = args[0];
auto gptr = std::dynamic_pointer_cast<NegSubgraph>(g.sptr());
*rv = gptr->head_nid;
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_GetEdgeSubgraphTail")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
SubgraphRef g = args[0];
auto gptr = std::dynamic_pointer_cast<NegSubgraph>(g.sptr());
*rv = gptr->tail_nid;
});
class UniformEdgeSamplerObject: public EdgeSamplerObject {
public:
explicit UniformEdgeSamplerObject(const GraphPtr gptr,
IdArray seed_edges,
const int64_t batch_size,
const int64_t num_workers,
const bool replacement,
const bool reset,
const std::string neg_mode,
const int64_t neg_sample_size,
const int64_t chunk_size,
const bool exclude_positive,
const bool check_false_neg,
IdArray relations)
: EdgeSamplerObject(gptr,
seed_edges,
batch_size,
num_workers,
replacement,
reset,
neg_mode,
neg_sample_size,
chunk_size,
exclude_positive,
check_false_neg,
relations) {
});
class UniformEdgeSamplerObject : public EdgeSamplerObject {
public:
explicit UniformEdgeSamplerObject(
const GraphPtr gptr, IdArray seed_edges, const int64_t batch_size,
const int64_t num_workers, const bool replacement, const bool reset,
const std::string neg_mode, const int64_t neg_sample_size,
const int64_t chunk_size, const bool exclude_positive,
const bool check_false_neg, IdArray relations)
: EdgeSamplerObject(
gptr, seed_edges, batch_size, num_workers, replacement, reset,
neg_mode, neg_sample_size, chunk_size, exclude_positive,
check_false_neg, relations) {
batch_curr_id_ = 0;
num_seeds_ = seed_edges->shape[0];
max_batch_id_ = (num_seeds_ + batch_size - 1) / batch_size;
......@@ -1463,8 +1441,9 @@ public:
}
~UniformEdgeSamplerObject() {}
void Fetch(DGLRetValue* rv) {
const int64_t num_workers = std::min(num_workers_, max_batch_id_ - batch_curr_id_);
void Fetch(DGLRetValue *rv) {
const int64_t num_workers =
std::min(num_workers_, max_batch_id_ - batch_curr_id_);
// generate subgraphs.
std::vector<SubgraphRef> positive_subgs(num_workers);
std::vector<SubgraphRef> negative_subgs(num_workers);
......@@ -1477,11 +1456,13 @@ public:
IdArray worker_seeds;
if (replacement_ == false) {
worker_seeds = seed_edges_.CreateView({num_edges}, DGLDataType{kDGLInt, 64, 1},
worker_seeds = seed_edges_.CreateView(
{num_edges}, DGLDataType{kDGLInt, 64, 1},
sizeof(dgl_id_t) * start);
} else {
std::vector<dgl_id_t> seeds;
const dgl_id_t *seed_edge_ids = static_cast<const dgl_id_t *>(seed_edges_->data);
const dgl_id_t *seed_edge_ids =
static_cast<const dgl_id_t *>(seed_edges_->data);
// sampling of each edge is a standalone event
for (int64_t i = 0; i < num_edges; ++i) {
int64_t seed = static_cast<const int64_t>(
......@@ -1497,29 +1478,29 @@ public:
const dgl_id_t *dst_ids = static_cast<const dgl_id_t *>(arr.dst->data);
std::vector<dgl_id_t> src_vec(src_ids, src_ids + num_edges);
std::vector<dgl_id_t> dst_vec(dst_ids, dst_ids + num_edges);
// TODO(zhengda) what if there are duplicates in the src and dst vectors.
// TODO(zhengda) what if there are duplicates in the src and dst
// vectors.
Subgraph subg = gptr_->EdgeSubgraph(worker_seeds, false);
positive_subgs[i] = ConvertRef(subg);
// For chunked negative sampling, we accept "chunk-head" for corrupting head
// nodes and "chunk-tail" for corrupting tail nodes.
// For chunked negative sampling, we accept "chunk-head" for corrupting
// head nodes and "chunk-tail" for corrupting tail nodes.
if (neg_mode_.substr(0, 5) == "chunk") {
NegSubgraph neg_subg = genChunkedNegEdgeSubgraph(subg, neg_mode_.substr(6),
neg_sample_size_,
exclude_positive_,
NegSubgraph neg_subg = genChunkedNegEdgeSubgraph(
subg, neg_mode_.substr(6), neg_sample_size_, exclude_positive_,
check_false_neg_);
negative_subgs[i] = ConvertRef(neg_subg);
} else if (neg_mode_ == "head" || neg_mode_ == "tail") {
NegSubgraph neg_subg = genNegEdgeSubgraph(subg, neg_mode_,
neg_sample_size_,
exclude_positive_,
NegSubgraph neg_subg = genNegEdgeSubgraph(
subg, neg_mode_, neg_sample_size_, exclude_positive_,
check_false_neg_);
negative_subgs[i] = ConvertRef(neg_subg);
}
}
});
if (neg_mode_.size() > 0) {
positive_subgs.insert(positive_subgs.end(), negative_subgs.begin(), negative_subgs.end());
positive_subgs.insert(
positive_subgs.end(), negative_subgs.begin(), negative_subgs.end());
}
batch_curr_id_ += num_workers;
......@@ -1535,20 +1516,22 @@ public:
if (replacement_ == false) {
// Now we should shuffle the data and reset the sampler.
dgl_id_t *seed_ids = static_cast<dgl_id_t *>(seed_edges_->data);
std::shuffle(seed_ids, seed_ids + seed_edges_->shape[0],
std::shuffle(
seed_ids, seed_ids + seed_edges_->shape[0],
std::default_random_engine());
}
}
DGL_DECLARE_OBJECT_TYPE_INFO(UniformEdgeSamplerObject, Object);
private:
void randomSample(size_t set_size, size_t num, std::vector<size_t>* out) {
private:
void randomSample(size_t set_size, size_t num, std::vector<size_t> *out) {
RandomSample(set_size, num, out);
}
void randomSample(size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t>* out) {
void randomSample(
size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t> *out) {
RandomSample(set_size, num, exclude, out);
}
......@@ -1557,17 +1540,19 @@ private:
int64_t num_seeds_;
};
class UniformEdgeSampler: public ObjectRef {
class UniformEdgeSampler : public ObjectRef {
public:
UniformEdgeSampler() {}
explicit UniformEdgeSampler(std::shared_ptr<runtime::Object> obj): ObjectRef(obj) {}
explicit UniformEdgeSampler(std::shared_ptr<runtime::Object> obj)
: ObjectRef(obj) {}
UniformEdgeSamplerObject* operator->() const {
return static_cast<UniformEdgeSamplerObject*>(obj_.get());
UniformEdgeSamplerObject *operator->() const {
return static_cast<UniformEdgeSamplerObject *>(obj_.get());
}
std::shared_ptr<UniformEdgeSamplerObject> sptr() const {
return CHECK_NOTNULL(std::dynamic_pointer_cast<UniformEdgeSamplerObject>(obj_));
return CHECK_NOTNULL(
std::dynamic_pointer_cast<UniformEdgeSamplerObject>(obj_));
}
operator bool() const { return this->defined(); }
......@@ -1575,7 +1560,7 @@ class UniformEdgeSampler: public ObjectRef {
};
DGL_REGISTER_GLOBAL("sampling._CAPI_CreateUniformEdgeSampler")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
// arguments
GraphRef g = args[0];
IdArray seed_edges = args[1];
......@@ -1603,64 +1588,42 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_CreateUniformEdgeSampler")
}
BuildCoo(*gptr);
auto o = std::make_shared<UniformEdgeSamplerObject>(gptr,
seed_edges,
batch_size,
max_num_workers,
replacement,
reset,
neg_mode,
neg_sample_size,
chunk_size,
exclude_positive,
check_false_neg,
relations);
auto o = std::make_shared<UniformEdgeSamplerObject>(
gptr, seed_edges, batch_size, max_num_workers, replacement, reset,
neg_mode, neg_sample_size, chunk_size, exclude_positive,
check_false_neg, relations);
*rv = o;
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_FetchUniformEdgeSample")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
UniformEdgeSampler sampler = args[0];
sampler->Fetch(rv);
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_ResetUniformEdgeSample")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
UniformEdgeSampler sampler = args[0];
sampler->Reset();
});
});
template<typename ValueType>
class WeightedEdgeSamplerObject: public EdgeSamplerObject {
template <typename ValueType>
class WeightedEdgeSamplerObject : public EdgeSamplerObject {
public:
explicit WeightedEdgeSamplerObject(const GraphPtr gptr,
IdArray seed_edges,
NDArray edge_weight,
NDArray node_weight,
const int64_t batch_size,
const int64_t num_workers,
const bool replacement,
const bool reset,
const std::string neg_mode,
const int64_t neg_sample_size,
const int64_t chunk_size,
const bool exclude_positive,
const bool check_false_neg,
explicit WeightedEdgeSamplerObject(
const GraphPtr gptr, IdArray seed_edges, NDArray edge_weight,
NDArray node_weight, const int64_t batch_size, const int64_t num_workers,
const bool replacement, const bool reset, const std::string neg_mode,
const int64_t neg_sample_size, const int64_t chunk_size,
const bool exclude_positive, const bool check_false_neg,
IdArray relations)
: EdgeSamplerObject(gptr,
seed_edges,
batch_size,
num_workers,
replacement,
reset,
neg_mode,
neg_sample_size,
chunk_size,
exclude_positive,
check_false_neg,
relations) {
: EdgeSamplerObject(
gptr, seed_edges, batch_size, num_workers, replacement, reset,
neg_mode, neg_sample_size, chunk_size, exclude_positive,
check_false_neg, relations) {
const int64_t num_edges = edge_weight->shape[0];
const ValueType *edge_prob = static_cast<const ValueType*>(edge_weight->data);
const ValueType *edge_prob =
static_cast<const ValueType *>(edge_weight->data);
std::vector<ValueType> eprob(num_edges);
for (int64_t i = 0; i < num_edges; ++i) {
eprob[i] = edge_prob[i];
......@@ -1672,7 +1635,8 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
if (num_nodes == 0) {
node_selector_ = nullptr;
} else {
const ValueType *node_prob = static_cast<const ValueType*>(node_weight->data);
const ValueType *node_prob =
static_cast<const ValueType *>(node_weight->data);
std::vector<ValueType> nprob(num_nodes);
for (size_t i = 0; i < num_nodes; ++i) {
nprob[i] = node_prob[i];
......@@ -1687,27 +1651,26 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
gptr_->FindEdge(0);
}
~WeightedEdgeSamplerObject() {
}
~WeightedEdgeSamplerObject() {}
void Fetch(DGLRetValue* rv) {
const int64_t num_workers = std::min(num_workers_, max_batch_id_ - curr_batch_id_);
void Fetch(DGLRetValue *rv) {
const int64_t num_workers =
std::min(num_workers_, max_batch_id_ - curr_batch_id_);
// generate subgraphs.
std::vector<SubgraphRef> positive_subgs(num_workers);
std::vector<SubgraphRef> negative_subgs(num_workers);
#pragma omp parallel for
for (int i = 0; i < num_workers; i++) {
const dgl_id_t *seed_edge_ids = static_cast<const dgl_id_t *>(seed_edges_->data);
const dgl_id_t *seed_edge_ids =
static_cast<const dgl_id_t *>(seed_edges_->data);
std::vector<size_t> edge_ids(batch_size_);
if (replacement_ == false) {
size_t n = batch_size_;
size_t num_ids = 0;
#pragma omp critical
{
num_ids = edge_selector_->SampleWithoutReplacement(n, &edge_ids);
}
{ num_ids = edge_selector_->SampleWithoutReplacement(n, &edge_ids); }
edge_ids.resize(num_ids);
for (size_t i = 0; i < num_ids; ++i) {
edge_ids[i] = seed_edge_ids[edge_ids[i]];
......@@ -1730,18 +1693,16 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
// TODO(zhengda) what if there are duplicates in the src and dst vectors.
Subgraph subg = gptr_->EdgeSubgraph(worker_seeds, false);
positive_subgs[i] = ConvertRef(subg);
// For chunked negative sampling, we accept "chunk-head" for corrupting head
// nodes and "chunk-tail" for corrupting tail nodes.
// For chunked negative sampling, we accept "chunk-head" for corrupting
// head nodes and "chunk-tail" for corrupting tail nodes.
if (neg_mode_.substr(0, 5) == "chunk") {
NegSubgraph neg_subg = genChunkedNegEdgeSubgraph(subg, neg_mode_.substr(6),
neg_sample_size_,
exclude_positive_,
NegSubgraph neg_subg = genChunkedNegEdgeSubgraph(
subg, neg_mode_.substr(6), neg_sample_size_, exclude_positive_,
check_false_neg_);
negative_subgs[i] = ConvertRef(neg_subg);
} else if (neg_mode_ == "head" || neg_mode_ == "tail") {
NegSubgraph neg_subg = genNegEdgeSubgraph(subg, neg_mode_,
neg_sample_size_,
exclude_positive_,
NegSubgraph neg_subg = genNegEdgeSubgraph(
subg, neg_mode_, neg_sample_size_, exclude_positive_,
check_false_neg_);
negative_subgs[i] = ConvertRef(neg_subg);
}
......@@ -1753,7 +1714,8 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
}
if (neg_mode_.size() > 0) {
positive_subgs.insert(positive_subgs.end(), negative_subgs.begin(), negative_subgs.end());
positive_subgs.insert(
positive_subgs.end(), negative_subgs.begin(), negative_subgs.end());
}
*rv = List<SubgraphRef>(positive_subgs);
}
......@@ -1762,7 +1724,8 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
curr_batch_id_ = 0;
if (replacement_ == false) {
const int64_t num_edges = edge_weight_->shape[0];
const ValueType *edge_prob = static_cast<const ValueType*>(edge_weight_->data);
const ValueType *edge_prob =
static_cast<const ValueType *>(edge_weight_->data);
std::vector<ValueType> eprob(num_edges);
for (int64_t i = 0; i < num_edges; ++i) {
eprob[i] = edge_prob[i];
......@@ -1775,8 +1738,8 @@ class WeightedEdgeSamplerObject: public EdgeSamplerObject {
DGL_DECLARE_OBJECT_TYPE_INFO(WeightedEdgeSamplerObject<ValueType>, Object);
private:
void randomSample(size_t set_size, size_t num, std::vector<size_t>* out) {
private:
void randomSample(size_t set_size, size_t num, std::vector<size_t> *out) {
if (num < set_size) {
std::unordered_set<size_t> sampled_idxs;
while (sampled_idxs.size() < num) {
......@@ -1792,13 +1755,13 @@ private:
} else {
// If we need to sample all elements in the set, we don't need to
// generate random numbers.
for (size_t i = 0; i < set_size; i++)
out->push_back(i);
for (size_t i = 0; i < set_size; i++) out->push_back(i);
}
}
void randomSample(size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t>* out) {
void randomSample(
size_t set_size, size_t num, const std::vector<size_t> &exclude,
std::vector<size_t> *out) {
std::unordered_map<size_t, int> sampled_idxs;
for (auto v : exclude) {
sampled_idxs.insert(std::pair<size_t, int>(v, 0));
......@@ -1830,7 +1793,7 @@ private:
}
}
private:
private:
std::shared_ptr<ArrayHeap<ValueType>> edge_selector_;
std::shared_ptr<ArrayHeap<ValueType>> node_selector_;
......@@ -1841,17 +1804,19 @@ private:
template class WeightedEdgeSamplerObject<float>;
class FloatWeightedEdgeSampler: public ObjectRef {
class FloatWeightedEdgeSampler : public ObjectRef {
public:
FloatWeightedEdgeSampler() {}
explicit FloatWeightedEdgeSampler(std::shared_ptr<runtime::Object> obj): ObjectRef(obj) {}
explicit FloatWeightedEdgeSampler(std::shared_ptr<runtime::Object> obj)
: ObjectRef(obj) {}
WeightedEdgeSamplerObject<float>* operator->() const {
return static_cast<WeightedEdgeSamplerObject<float>*>(obj_.get());
WeightedEdgeSamplerObject<float> *operator->() const {
return static_cast<WeightedEdgeSamplerObject<float> *>(obj_.get());
}
std::shared_ptr<WeightedEdgeSamplerObject<float>> sptr() const {
return CHECK_NOTNULL(std::dynamic_pointer_cast<WeightedEdgeSamplerObject<float>>(obj_));
return CHECK_NOTNULL(
std::dynamic_pointer_cast<WeightedEdgeSamplerObject<float>>(obj_));
}
operator bool() const { return this->defined(); }
......@@ -1859,7 +1824,7 @@ class FloatWeightedEdgeSampler: public ObjectRef {
};
DGL_REGISTER_GLOBAL("sampling._CAPI_CreateWeightedEdgeSampler")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
// arguments
GraphRef g = args[0];
IdArray seed_edges = args[1];
......@@ -1881,13 +1846,17 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_CreateWeightedEdgeSampler")
CHECK(aten::IsValidIdArray(seed_edges));
CHECK_EQ(seed_edges->ctx.device_type, kDGLCPU)
<< "WeightedEdgeSampler only support CPU sampling";
CHECK(edge_weight->dtype.code == kDGLFloat) << "edge_weight should be FloatType";
CHECK(edge_weight->dtype.bits == 32) << "WeightedEdgeSampler only support float weight";
CHECK(edge_weight->dtype.code == kDGLFloat)
<< "edge_weight should be FloatType";
CHECK(edge_weight->dtype.bits == 32)
<< "WeightedEdgeSampler only support float weight";
CHECK_EQ(edge_weight->ctx.device_type, kDGLCPU)
<< "WeightedEdgeSampler only support CPU sampling";
if (node_weight->shape[0] > 0) {
CHECK(node_weight->dtype.code == kDGLFloat) << "node_weight should be FloatType";
CHECK(node_weight->dtype.bits == 32) << "WeightedEdgeSampler only support float weight";
CHECK(node_weight->dtype.code == kDGLFloat)
<< "node_weight should be FloatType";
CHECK(node_weight->dtype.bits == 32)
<< "WeightedEdgeSampler only support float weight";
CHECK_EQ(node_weight->ctx.device_type, kDGLCPU)
<< "WeightedEdgeSampler only support CPU sampling";
}
......@@ -1899,36 +1868,26 @@ DGL_REGISTER_GLOBAL("sampling._CAPI_CreateWeightedEdgeSampler")
BuildCoo(*gptr);
const int64_t num_seeds = seed_edges->shape[0];
const int64_t num_workers = std::min(max_num_workers,
(num_seeds + batch_size - 1) / batch_size);
auto o = std::make_shared<WeightedEdgeSamplerObject<float>>(gptr,
seed_edges,
edge_weight,
node_weight,
batch_size,
num_workers,
replacement,
reset,
neg_mode,
neg_sample_size,
chunk_size,
exclude_positive,
check_false_neg,
relations);
const int64_t num_workers =
std::min(max_num_workers, (num_seeds + batch_size - 1) / batch_size);
auto o = std::make_shared<WeightedEdgeSamplerObject<float>>(
gptr, seed_edges, edge_weight, node_weight, batch_size, num_workers,
replacement, reset, neg_mode, neg_sample_size, chunk_size,
exclude_positive, check_false_neg, relations);
*rv = o;
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_FetchWeightedEdgeSample")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
FloatWeightedEdgeSampler sampler = args[0];
sampler->Fetch(rv);
});
});
DGL_REGISTER_GLOBAL("sampling._CAPI_ResetWeightedEdgeSample")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
.set_body([](DGLArgs args, DGLRetValue *rv) {
FloatWeightedEdgeSampler sampler = args[0];
sampler->Reset();
});
});
} // namespace dgl
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment