array_sort.cc

/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/cpu/array_sort.cc
 * \brief Array sort CPU implementation
 */
#include <dgl/array.h>
#ifdef PARALLEL_ALGORITHMS
#include <parallel/algorithm>
#endif
#include <algorithm>
#include <iterator>

namespace {

template <typename V1, typename V2>
struct PairRef {
  PairRef() = delete;
  PairRef(const PairRef& other) = default;
  PairRef(PairRef&& other) = default;
  PairRef(V1 *const r, V2 *const c)
    : row(r), col(c) {}

  PairRef& operator=(const PairRef& other) {
    *row = *other.row;
    *col = *other.col;
    return *this;
  }
  PairRef& operator=(const std::pair<V1, V2>& val) {
    *row = std::get<0>(val);
    *col = std::get<1>(val);
    return *this;
  }

  operator std::pair<V1, V2>() const {
    return std::make_pair(*row, *col);
  }

  void Swap(const PairRef& other) const {
    std::swap(*row, *other.row);
    std::swap(*col, *other.col);
  }

  V1 *row;
  V2 *col;
};

using std::swap;
template <typename V1, typename V2>
void swap(const PairRef<V1, V2>& r1, const PairRef<V1, V2>& r2) {
  r1.Swap(r2);
}

template <typename V1, typename V2>
struct PairIterator : public std::iterator<std::random_access_iterator_tag,
                                           std::pair<V1, V2>,
                                           std::ptrdiff_t,
                                           std::pair<V1*, V2*>,
                                           PairRef<V1, V2>> {
  PairIterator() = default;
  PairIterator(const PairIterator& other) = default;
  PairIterator(PairIterator&& other) = default;
  PairIterator(V1 *r, V2 *c): row(r), col(c) {}

  PairIterator& operator=(const PairIterator& other) = default;
  PairIterator& operator=(PairIterator&& other) = default;
  ~PairIterator() = default;

  bool operator==(const PairIterator& other) const {
    return row == other.row;
  }

  bool operator!=(const PairIterator& other) const {
    return row != other.row;
  }

  bool operator<(const PairIterator& other) const {
    return row < other.row;
  }

  bool operator>(const PairIterator& other) const {
    return row > other.row;
  }

  bool operator<=(const PairIterator& other) const {
    return row <= other.row;
  }

  bool operator>=(const PairIterator& other) const {
    return row >= other.row;
  }

  PairIterator& operator+=(const std::ptrdiff_t& movement) {
    row += movement;
    col += movement;
    return *this;
  }

  PairIterator& operator-=(const std::ptrdiff_t& movement) {
    row -= movement;
    col -= movement;
    return *this;
  }

  PairIterator& operator++() {
    return operator+=(1);
  }

  PairIterator& operator--() {
    return operator-=(1);
  }

  PairIterator operator++(int) {
    PairIterator ret(*this);
    operator++();
    return ret;
  }

  PairIterator operator--(int) {
    PairIterator ret(*this);
    operator--();
    return ret;
  }

  PairIterator operator+(const std::ptrdiff_t& movement) const {
    PairIterator ret(*this);
    ret += movement;
    return ret;
  }

  PairIterator operator-(const std::ptrdiff_t& movement) const {
    PairIterator ret(*this);
    ret -= movement;
    return ret;
  }

  std::ptrdiff_t operator-(const PairIterator& other) const {
    return row - other.row;
  }

  PairRef<V1, V2> operator*() const {
    return PairRef<V1, V2>(row, col);
  }
  PairRef<V1, V2> operator*() {
    return PairRef<V1, V2>(row, col);
  }

  V1 *row;
  V2 *col;
};

}  // namespace

namespace dgl {
using runtime::NDArray;
namespace aten {
namespace impl {

template <DLDeviceType XPU, typename IdType>
std::pair<IdArray, IdArray> Sort(IdArray array) {
  const int64_t nitem = array->shape[0];
  IdArray val = array.Clone();
  IdArray idx = aten::Range(0, nitem, 64, array->ctx);
  IdType* val_data = val.Ptr<IdType>();
  int64_t* idx_data = idx.Ptr<int64_t>();
  typedef std::pair<IdType, int64_t> Pair;
#ifdef PARALLEL_ALGORITHMS
  __gnu_parallel::sort(
#else
  std::sort(
#endif
      PairIterator<IdType, int64_t>(val_data, idx_data),
      PairIterator<IdType, int64_t>(val_data, idx_data) + nitem,
      [] (const Pair& a, const Pair& b) {
        return std::get<0>(a) < std::get<0>(b);
      });
  return std::make_pair(val, idx);
}

template std::pair<IdArray, IdArray> Sort<kDLCPU, int32_t>(IdArray);
template std::pair<IdArray, IdArray> Sort<kDLCPU, int64_t>(IdArray);

}  // namespace impl
}  // namespace aten
}  // namespace dgl