scheduler.cc

/*!
 *  Copyright (c) 2018 by Contributors
 * \file scheduler/scheduler.cc
 * \brief DGL Scheduler implementation
 */
#include <dgl/scheduler.h>

#include <unordered_map>
#include <vector>

namespace dgl {
namespace sched {

template <class IdType>
std::vector<IdArray> DegreeBucketing(
    const IdArray& msg_ids, const IdArray& vids, const IdArray& recv_ids) {
  auto n_msgs = msg_ids->shape[0];

  const IdType* vid_data = static_cast<IdType*>(vids->data);
  const IdType* msg_id_data = static_cast<IdType*>(msg_ids->data);
  const IdType* recv_id_data = static_cast<IdType*>(recv_ids->data);

  // in edge: dst->msgs
  std::unordered_map<IdType, std::vector<IdType>> in_edges;
  for (IdType i = 0; i < n_msgs; ++i) {
    in_edges[vid_data[i]].push_back(msg_id_data[i]);
  }

  // bkt: deg->dsts
  std::unordered_map<IdType, std::vector<IdType>> bkt;
  for (const auto& it : in_edges) {
    bkt[it.second.size()].push_back(it.first);
  }

  std::unordered_set<IdType> zero_deg_nodes;
  for (IdType i = 0; i < recv_ids->shape[0]; ++i) {
    if (in_edges.find(recv_id_data[i]) == in_edges.end()) {
      zero_deg_nodes.insert(recv_id_data[i]);
    }
  }
  auto n_zero_deg = zero_deg_nodes.size();

  // calc output size
  IdType n_deg = bkt.size();
  IdType n_dst = in_edges.size();
  IdType n_mid_sec = bkt.size();  // zero deg won't affect message size
  if (n_zero_deg > 0) {
    n_deg += 1;
    n_dst += n_zero_deg;
  }

  // initialize output
  IdArray degs = IdArray::Empty({n_deg}, vids->dtype, vids->ctx);
  IdArray nids = IdArray::Empty({n_dst}, vids->dtype, vids->ctx);
  IdArray nid_section = IdArray::Empty({n_deg}, vids->dtype, vids->ctx);
  IdArray mids = IdArray::Empty({n_msgs}, vids->dtype, vids->ctx);
  IdArray mid_section = IdArray::Empty({n_mid_sec}, vids->dtype, vids->ctx);
  IdType* deg_ptr = static_cast<IdType*>(degs->data);
  IdType* nid_ptr = static_cast<IdType*>(nids->data);
  IdType* nsec_ptr = static_cast<IdType*>(nid_section->data);
  IdType* mid_ptr = static_cast<IdType*>(mids->data);
  IdType* msec_ptr = static_cast<IdType*>(mid_section->data);

  // fill in bucketing ordering
  for (const auto& it : bkt) {  // for each bucket
    const IdType deg = it.first;
    const IdType bucket_size = it.second.size();
    *deg_ptr++ = deg;
    *nsec_ptr++ = bucket_size;
    *msec_ptr++ = deg * bucket_size;
    for (const auto dst : it.second) {  // for each dst in this bucket
      *nid_ptr++ = dst;
      for (const auto mid : in_edges[dst]) {  // for each in edge of dst
        *mid_ptr++ = mid;
      }
    }
  }

  if (n_zero_deg > 0) {
    *deg_ptr = 0;
    *nsec_ptr = n_zero_deg;
    for (const auto dst : zero_deg_nodes) {
      *nid_ptr++ = dst;
    }
  }

  std::vector<IdArray> ret;
  ret.push_back(std::move(degs));
  ret.push_back(std::move(nids));
  ret.push_back(std::move(nid_section));
  ret.push_back(std::move(mids));
  ret.push_back(std::move(mid_section));

  return ret;
}

template std::vector<IdArray> DegreeBucketing<int32_t>(
    const IdArray& msg_ids, const IdArray& vids, const IdArray& recv_ids);

template std::vector<IdArray> DegreeBucketing<int64_t>(
    const IdArray& msg_ids, const IdArray& vids, const IdArray& recv_ids);

template <class IdType>
std::vector<IdArray> GroupEdgeByNodeDegree(
    const IdArray& uids, const IdArray& vids, const IdArray& eids) {
  auto n_edge = eids->shape[0];
  const IdType* eid_data = static_cast<IdType*>(eids->data);
  const IdType* uid_data = static_cast<IdType*>(uids->data);
  const IdType* vid_data = static_cast<IdType*>(vids->data);

  // node2edge: group_by nodes uid -> (eid, the other end vid)
  std::unordered_map<IdType, std::vector<std::pair<IdType, IdType>>> node2edge;
  for (IdType i = 0; i < n_edge; ++i) {
    node2edge[uid_data[i]].emplace_back(eid_data[i], vid_data[i]);
  }

  // bkt: deg -> group_by node uid
  std::unordered_map<IdType, std::vector<IdType>> bkt;
  for (const auto& it : node2edge) {
    bkt[it.second.size()].push_back(it.first);
  }

  // number of unique degree
  IdType n_deg = bkt.size();

  // initialize output
  IdArray degs = IdArray::Empty({n_deg}, eids->dtype, eids->ctx);
  IdArray new_uids = IdArray::Empty({n_edge}, uids->dtype, uids->ctx);
  IdArray new_vids = IdArray::Empty({n_edge}, vids->dtype, vids->ctx);
  IdArray new_eids = IdArray::Empty({n_edge}, eids->dtype, eids->ctx);
  IdArray sections = IdArray::Empty({n_deg}, eids->dtype, eids->ctx);
  IdType* deg_ptr = static_cast<IdType*>(degs->data);
  IdType* uid_ptr = static_cast<IdType*>(new_uids->data);
  IdType* vid_ptr = static_cast<IdType*>(new_vids->data);
  IdType* eid_ptr = static_cast<IdType*>(new_eids->data);
  IdType* sec_ptr = static_cast<IdType*>(sections->data);

  // fill in bucketing ordering
  for (const auto& it : bkt) {  // for each bucket
    // degree of this bucket
    const IdType deg = it.first;
    // number of edges in this bucket
    const IdType bucket_size = it.second.size();
    *deg_ptr++ = deg;
    *sec_ptr++ = deg * bucket_size;
    for (const auto u : it.second) {           // for uid in this bucket
      for (const auto& pair : node2edge[u]) {  // for each edge of uid
        *uid_ptr++ = u;
        *vid_ptr++ = pair.second;
        *eid_ptr++ = pair.first;
      }
    }
  }

  std::vector<IdArray> ret;
  ret.push_back(std::move(degs));
  ret.push_back(std::move(new_uids));
  ret.push_back(std::move(new_vids));
  ret.push_back(std::move(new_eids));
  ret.push_back(std::move(sections));

  return ret;
}

template std::vector<IdArray> GroupEdgeByNodeDegree<int32_t>(
    const IdArray& uids, const IdArray& vids, const IdArray& eids);

template std::vector<IdArray> GroupEdgeByNodeDegree<int64_t>(
    const IdArray& uids, const IdArray& vids, const IdArray& eids);

}  // namespace sched

}  // namespace dgl