union_partition.cc

/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/cpu/coo_union_partition.cc
 * \brief COO union and partition
 */
#include <dgl/array.h>
#include <vector>

namespace dgl {
namespace aten {
///////////////////////// COO Based Operations/////////////////////////
std::vector<COOMatrix> DisjointPartitionCooBySizes(
  const COOMatrix &coo,
  const uint64_t batch_size,
  const std::vector<uint64_t> &edge_cumsum,
  const std::vector<uint64_t> &src_vertex_cumsum,
  const std::vector<uint64_t> &dst_vertex_cumsum) {
  CHECK_EQ(edge_cumsum.size(), batch_size + 1);
  CHECK_EQ(src_vertex_cumsum.size(), batch_size + 1);
  CHECK_EQ(dst_vertex_cumsum.size(), batch_size + 1);
  std::vector<COOMatrix> ret;
  ret.resize(batch_size);

  for (size_t g = 0; g < batch_size; ++g) {
    IdArray result_src = IndexSelect(coo.row,
                                     edge_cumsum[g],
                                     edge_cumsum[g + 1]) - src_vertex_cumsum[g];
    IdArray result_dst = IndexSelect(coo.col,
                                     edge_cumsum[g],
                                     edge_cumsum[g + 1]) - dst_vertex_cumsum[g];
    IdArray result_data = NullArray();
    // has data index array
    if (COOHasData(coo)) {
      result_data = IndexSelect(coo.data,
                                edge_cumsum[g],
                                edge_cumsum[g + 1]) - edge_cumsum[g];
    }

    COOMatrix sub_coo = COOMatrix(
        src_vertex_cumsum[g+1]-src_vertex_cumsum[g],
        dst_vertex_cumsum[g+1]-dst_vertex_cumsum[g],
        result_src,
        result_dst,
        result_data,
        coo.row_sorted,
        coo.col_sorted);
    ret[g] = sub_coo;
  }

  return ret;
}

COOMatrix COOSliceContiguousChunk(
  const COOMatrix &coo,
  const std::vector<uint64_t> &edge_range,
  const std::vector<uint64_t> &src_vertex_range,
  const std::vector<uint64_t> &dst_vertex_range) {
  IdArray result_src = NullArray(coo.row->dtype, coo.row->ctx);
  IdArray result_dst = NullArray(coo.row->dtype, coo.row->ctx);
  if (edge_range[1] != edge_range[0]) {
    // The chunk has edges
    result_src = IndexSelect(coo.row,
                             edge_range[0],
                             edge_range[1]) - src_vertex_range[0];
    result_dst = IndexSelect(coo.col,
                             edge_range[0],
                             edge_range[1]) - dst_vertex_range[0];
  }

  IdArray result_data = NullArray();
  // has data index array
  if (COOHasData(coo)) {
    result_data = IndexSelect(coo.data,
                              edge_range[0],
                              edge_range[1]) - edge_range[0];
  }

  COOMatrix sub_coo = COOMatrix(
    src_vertex_range[1]-src_vertex_range[0],
    dst_vertex_range[1]-dst_vertex_range[0],
    result_src,
    result_dst,
    result_data,
    coo.row_sorted,
    coo.col_sorted);

  return sub_coo;
}

///////////////////////// CSR Based Operations/////////////////////////
CSRMatrix DisjointUnionCsr(const std::vector<CSRMatrix>& csrs) {
  uint64_t src_offset = 0, dst_offset = 0;
  int64_t indices_offset = 0;
  bool has_data = false;
  bool sorted = true;

  // check if data index array
  for (size_t i = 0; i < csrs.size(); ++i) {
    CHECK_SAME_DTYPE(csrs[0].indptr, csrs[i].indptr);
    CHECK_SAME_CONTEXT(csrs[0].indices, csrs[i].indices);
    has_data |= CSRHasData(csrs[i]);
  }

  std::vector<IdArray> res_indptr;
  std::vector<IdArray> res_indices;
  std::vector<IdArray> res_data;
  res_indptr.resize(csrs.size());
  res_indices.resize(csrs.size());

  for (size_t i = 0; i < csrs.size(); ++i) {
    const aten::CSRMatrix &csr = csrs[i];
    sorted &= csr.sorted;
    IdArray indptr = csr.indptr + indices_offset;
    IdArray indices = csr.indices + dst_offset;
    if (i > 0)
      indptr = IndexSelect(indptr,
                           1,
                           indptr->shape[0]);
    res_indptr[i] = indptr;
    res_indices[i] = indices;
    src_offset += csr.num_rows;
    dst_offset += csr.num_cols;

    // any one of input csr has data index array
    if (has_data) {
      IdArray edges_data;
      if (CSRHasData(csr) == false) {
        edges_data = Range(indices_offset,
                           indices_offset + csr.indices->shape[0],
                           csr.indices->dtype.bits,
                           csr.indices->ctx);
      } else {
        edges_data = csr.data + indices_offset;
      }
      res_data.push_back(edges_data);
      indices_offset += csr.indices->shape[0];
    }
  }

  IdArray result_indptr = Concat(res_indptr);
  IdArray result_indices = Concat(res_indices);
  IdArray result_data = has_data ? Concat(res_data) : NullArray();

  return CSRMatrix(
    src_offset, dst_offset,
    result_indptr,
    result_indices,
    result_data,
    sorted);
}

std::vector<CSRMatrix> DisjointPartitionCsrBySizes(
  const CSRMatrix &csr,
  const uint64_t batch_size,
  const std::vector<uint64_t> &edge_cumsum,
  const std::vector<uint64_t> &src_vertex_cumsum,
  const std::vector<uint64_t> &dst_vertex_cumsum) {
  CHECK_EQ(edge_cumsum.size(), batch_size + 1);
  CHECK_EQ(src_vertex_cumsum.size(), batch_size + 1);
  CHECK_EQ(dst_vertex_cumsum.size(), batch_size + 1);
  std::vector<CSRMatrix> ret;
  ret.resize(batch_size);

  for (size_t g = 0; g < batch_size; ++g) {
    uint64_t num_src = src_vertex_cumsum[g+1]-src_vertex_cumsum[g];
    IdArray result_indptr;
    if (g == 0) {
      result_indptr = IndexSelect(csr.indptr,
                                  0,
                                  src_vertex_cumsum[1] + 1) - edge_cumsum[0];
    } else {
      result_indptr = IndexSelect(csr.indptr,
                                  src_vertex_cumsum[g],
                                  src_vertex_cumsum[g+1] + 1) - edge_cumsum[g];
    }

    IdArray result_indices = IndexSelect(csr.indices,
                                         edge_cumsum[g],
                                         edge_cumsum[g+1]) - dst_vertex_cumsum[g];

    IdArray result_data = NullArray();
    // has data index array
    if (CSRHasData(csr)) {
      result_data = IndexSelect(csr.data,
                                edge_cumsum[g],
                                edge_cumsum[g+1]) - edge_cumsum[g];
    }

    CSRMatrix sub_csr = CSRMatrix(
      num_src,
      dst_vertex_cumsum[g+1]-dst_vertex_cumsum[g],
      result_indptr,
      result_indices,
      result_data,
      csr.sorted);
    ret[g] = sub_csr;
  }

  return ret;
}

CSRMatrix CSRSliceContiguousChunk(
  const CSRMatrix &csr,
  const std::vector<uint64_t> &edge_range,
  const std::vector<uint64_t> &src_vertex_range,
  const std::vector<uint64_t> &dst_vertex_range) {
  int64_t indptr_len = src_vertex_range[1] - src_vertex_range[0] + 1;
  IdArray result_indptr = Full(0, indptr_len, csr.indptr->dtype.bits, csr.indptr->ctx);
  IdArray result_indices = NullArray(csr.indptr->dtype, csr.indptr->ctx);
  IdArray result_data = NullArray();
  if (edge_range[1] != edge_range[0]) {
    // The chunk has edges
    result_indptr = IndexSelect(csr.indptr,
                                src_vertex_range[0],
                                src_vertex_range[1] + 1) - edge_range[0];
    result_indices = IndexSelect(csr.indices,
                                 edge_range[0],
                                 edge_range[1]) - dst_vertex_range[0];
    if (CSRHasData(csr)) {
      result_data = IndexSelect(csr.data,
                                edge_range[0],
                                edge_range[1]) - edge_range[0];
    }
  }

  CSRMatrix sub_csr = CSRMatrix(
    src_vertex_range[1]-src_vertex_range[0],
    dst_vertex_range[1]-dst_vertex_range[0],
    result_indptr,
    result_indices,
    result_data,
    csr.sorted);

  return sub_csr;
}

}  // namespace aten
}  // namespace dgl