wehub-resource-sync
65 lines
2.3 KiB
C++
65 lines
2.3 KiB
C++
/**
|
|
* Copyright (c) 2018 by Contributors
|
|
* @file dgl/scheduler.h
|
|
* @brief Operations on graph index.
|
|
*/
|
|
#ifndef DGL_SCHEDULER_H_
|
|
#define DGL_SCHEDULER_H_
|
|
|
|
#include <vector>
|
|
|
|
#include "runtime/ndarray.h"
|
|
|
|
namespace dgl {
|
|
|
|
typedef dgl::runtime::NDArray IdArray;
|
|
|
|
namespace sched {
|
|
|
|
/**
|
|
* @brief Generate degree bucketing schedule
|
|
* @tparam IdType Graph's index data type, can be int32_t or int64_t
|
|
* @param msg_ids The edge id for each message
|
|
* @param vids The destination vertex for each message
|
|
* @param recv_ids The recv nodes (for checking zero degree nodes)
|
|
* @note If there are multiple messages going into the same destination vertex,
|
|
* then there will be multiple copies of the destination vertex in vids.
|
|
* @return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
|
|
* degrees: degrees for each bucket
|
|
* nids: destination node ids
|
|
* nid_section: number of nodes in each bucket (used to split nids)
|
|
* mids: message ids
|
|
* mid_section: number of messages in each bucket (used to split mids)
|
|
*/
|
|
template <class IdType>
|
|
std::vector<IdArray> DegreeBucketing(
|
|
const IdArray& msg_ids, const IdArray& vids, const IdArray& recv_ids);
|
|
|
|
/**
|
|
* @brief Generate degree bucketing schedule for group_apply edge
|
|
* @tparam IdType Graph's index data type, can be int32_t or int64_t
|
|
* @param uids One end vertex of edge by which edges are grouped
|
|
* @param vids The other end vertex of edge
|
|
* @param eids Edge ids
|
|
* @note This function always generate group_apply schedule based on degrees of
|
|
* nodes in uids. Therefore, if group_apply by source nodes, then uids
|
|
* should be source. If group_apply by destination nodes, then uids
|
|
* should be destination.
|
|
* @return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
|
|
* degrees: degrees for each bucket
|
|
* new_uids: uids reordered by degree bucket
|
|
* new_vids: vids reordered by degree bucket
|
|
* new_edis: eids reordered by degree bucket
|
|
* sections: number of edges in each degree bucket (used to partition
|
|
* new_uids, new_vids, and new_eids)
|
|
*/
|
|
template <class IdType>
|
|
std::vector<IdArray> GroupEdgeByNodeDegree(
|
|
const IdArray& uids, const IdArray& vids, const IdArray& eids);
|
|
|
|
} // namespace sched
|
|
|
|
} // namespace dgl
|
|
|
|
#endif // DGL_SCHEDULER_H_
|