[Kernel] Matrix toSimple (#1756)

* Matrix to simple

* Pass test

* new impl

* Fix test

* Fix lint

* trigger

* upd

* Fix comments
Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal>

include/dgl/aten/coo.h

@@ -11,6 +11,7 @@
 #include <dmlc/serializer.h>
 #include <vector>
 #include <utility>
+#include <tuple>
 #include "./types.h"
 #include "./array_ops.h"
 #include "./spmat.h"
@@ -406,6 +407,30 @@ COOMatrix COORowWiseTopk(
 COOMatrix DisjointUnionCoo(
   const std::vector<COOMatrix>& coos);
 
+/*!
+ * \brief COOMatrix toSimple.
+ *
+ * A = [[0, 0, 0],
+ *      [3, 0, 2],
+ *      [1, 1, 0],
+ *      [0, 0, 4]]
+ * 
+ * B, cnt, edge_map = COOToSimple(A)
+ *
+ * B = [[0, 0, 0],
+ *      [1, 0, 1],
+ *      [1, 1, 0],
+ *      [0, 0, 1]]
+ * cnt = [3, 2, 1, 1, 4]
+ * edge_map = [0, 0, 0, 1, 1, 2, 3, 4, 4, 4, 4]
+ *
+ * \return The simplified COOMatrix
+ *         The count recording the number of duplicated edges from the original graph.
+ *         The edge mapping from the edge IDs of original graph to those of the
+ *         returned graph.
+ */
+std::tuple<COOMatrix, IdArray, IdArray> COOToSimple(const COOMatrix& coo);
+
 /*!
  * \brief Split a COOMatrix into multiple disjoin components.
  *

include/dgl/aten/csr.h

@@ -9,6 +9,7 @@
 #include <dmlc/io.h>
 #include <dmlc/serializer.h>
 #include <vector>
+#include <tuple>
 #include "./types.h"
 #include "./array_ops.h"
 #include "./spmat.h"
@@ -401,6 +402,30 @@ COOMatrix CSRRowWiseTopk(
 CSRMatrix DisjointUnionCsr(
   const std::vector<CSRMatrix>& csrs);
 
+/*!
+ * \brief CSRMatrix toSimple.
+ *
+ * A = [[0, 0, 0],
+ *      [3, 0, 2],
+ *      [1, 1, 0],
+ *      [0, 0, 4]]
+ * 
+ * B, cnt, edge_map = CSRToSimple(A)
+ *
+ * B = [[0, 0, 0],
+ *      [1, 0, 1],
+ *      [1, 1, 0],
+ *      [0, 0, 1]]
+ * cnt = [3, 2, 1, 1, 4]
+ * edge_map = [0, 0, 0, 1, 1, 2, 3, 4, 4, 4, 4]
+ *
+ * \return The simplified CSRMatrix
+ *         The count recording the number of duplicated edges from the original graph.
+ *         The edge mapping from the edge IDs of original graph to those of the
+ *         returned graph.
+ */
+std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(const CSRMatrix& csr);
+
 /*!
  * \brief Split a CSRMatrix into multiple disjoin components.
  *

python/dgl/transform.py

@@ -1214,6 +1214,8 @@ def to_simple(g, return_counts='count', writeback_mapping=None):
 
     This function does not preserve node and edge features.
 
+    TODO(xiangsx): Don't save writeback_mapping into g, but put it into return value.
+
     Parameters
     ----------
     g : DGLHeteroGraph
@@ -1234,7 +1236,7 @@ def to_simple(g, return_counts='count', writeback_mapping=None):
     Examples
     --------
     Consider the following graph
-    >>> g = dgl.graph([(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)])
+    >>> g = dgl.graph(([0, 1, 2, 1, 1, 1], [1, 3, 2, 3, 4, 4]))
     >>> sg = dgl.to_simple(g, return_counts='weights', writeback_mapping='new_eid')
 
     The returned graph would have duplicate edges connecting (1, 3) and (1, 4) removed:

src/array/array.cc

@@ -509,6 +509,17 @@ COOMatrix CSRRowWiseTopk(
   return ret;
 }
 
+std::tuple<CSRMatrix, IdArray, IdArray>
+CSRToSimple(const CSRMatrix& csr) {
+  std::tuple<CSRMatrix, IdArray, IdArray> ret;
+
+  CSRMatrix sorted_csr = (CSRIsSorted(csr)) ? csr : CSRSort(csr);
+  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRToSimple", {
+    ret = impl::CSRToSimple<XPU, IdType>(sorted_csr);
+  });
+  return ret;
+}
+
 ///////////////////////// COO routines //////////////////////////
 
 bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
@@ -678,6 +689,52 @@ std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
   return ret;
 }
 
+std::tuple<COOMatrix, IdArray, IdArray>
+COOToSimple(const COOMatrix& coo) {
+  // coo column sorted
+  const COOMatrix sorted_coo = COOSort(coo, true);
+  const IdArray eids_shuffled = COOHasData(sorted_coo) ?
+    sorted_coo.data :
+    Range(0, sorted_coo.row->shape[0], sorted_coo.row->dtype.bits, sorted_coo.row->ctx);
+  const auto &coalesced_result = COOCoalesce(sorted_coo);
+  const COOMatrix &coalesced_adj = coalesced_result.first;
+  const IdArray &count = coalesced_result.second;
+
+  /*
+   * eids_shuffled actually already contains the mapping from old edge space to the
+   * new one:
+   *
+   * * eids_shuffled[0:count[0]] indicates the original edge IDs that coalesced into new
+   *   edge #0.
+   * * eids_shuffled[count[0]:count[0] + count[1]] indicates those that coalesced into
+   *   new edge #1.
+   * * eids_shuffled[count[0] + count[1]:count[0] + count[1] + count[2]] indicates those
+   *   that coalesced into new edge #2.
+   * * etc.
+   *
+   * Here, we need to translate eids_shuffled to an array "eids_remapped" such that
+   * eids_remapped[i] indicates the new edge ID the old edge #i is mapped to.  The
+   * translation can simply be achieved by (in numpy code):
+   *
+   *     new_eid_for_eids_shuffled = np.range(len(count)).repeat(count)
+   *     eids_remapped = np.zeros_like(new_eid_for_eids_shuffled)
+   *     eids_remapped[eids_shuffled] = new_eid_for_eids_shuffled
+   */
+  const IdArray new_eids = Range(
+    0, coalesced_adj.row->shape[0], coalesced_adj.row->dtype.bits, coalesced_adj.row->ctx);
+  const IdArray eids_remapped = Scatter(Repeat(new_eids, count), eids_shuffled);
+
+  COOMatrix ret = COOMatrix(
+    coalesced_adj.num_rows,
+    coalesced_adj.num_cols,
+    coalesced_adj.row,
+    coalesced_adj.col,
+    NullArray(),
+    true,
+    true);
+  return std::make_tuple(ret, count, eids_remapped);
+}
+
 ///////////////////////// Graph Traverse routines //////////////////////////
 Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source) {
   Frontiers ret;

src/array/array_op.h

@@ -152,6 +152,8 @@ template <DLDeviceType XPU, typename IdType, typename DType>
 COOMatrix CSRRowWiseTopk(
     CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending);
 
+template <DLDeviceType XPU, typename IdType>
+std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(CSRMatrix csr);
 
 ///////////////////////////////////////////////////////////////////////////////////////////
 

src/array/cpu/csr_to_simple.cc

+/*!
+ *  Copyright (c) 2020 by Contributors
+ * \file array/cpu/csr_to_simple.cc
+ * \brief CSR sorting
+ */
+#include <dgl/array.h>
+#include <numeric>
+#include <algorithm>
+#include <vector>
+
+namespace dgl {
+namespace aten {
+namespace impl {
+
+template <DLDeviceType XPU, typename IdType>
+std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(CSRMatrix csr) {
+  if (!csr.sorted)
+    csr = CSRSort(csr);
+
+  const IdType *indptr_data = static_cast<IdType*>(csr.indptr->data);
+  const IdType *indices_data = static_cast<IdType*>(csr.indices->data);
+
+  std::vector<IdType> indptr;
+  std::vector<IdType> indices;
+  std::vector<IdType> count;
+  indptr.resize(csr.indptr->shape[0]);
+  indptr[0] = 0;
+
+  for (int64_t i = 1; i < csr.indptr->shape[0]; ++i) {
+    if (indptr_data[i-1] == indptr_data[i]) {
+      indptr[i] = indptr[i-1];
+      continue;
+    }
+
+    int64_t cnt = 1;
+    int64_t dup_cnt = 1;
+    indices.push_back(indices_data[indptr_data[i-1]]);
+    for (int64_t j = indptr_data[i-1]+1; j < indptr_data[i]; ++j) {
+      if (indices_data[j-1] == indices_data[j]) {
+        ++dup_cnt;
+        continue;
+      }
+      count.push_back(dup_cnt);
+      dup_cnt = 1;
+      indices.push_back(indices_data[j]);
+      ++cnt;
+    }
+    count.push_back(dup_cnt);
+    indptr[i] = indptr[i-1] + cnt;
+  }
+
+  CSRMatrix res_csr = CSRMatrix(
+    csr.num_rows,
+    csr.num_cols,
+    IdArray::FromVector(indptr),
+    IdArray::FromVector(indices),
+    NullArray(),
+    true);
+
+  const IdArray &edge_count = IdArray::FromVector(count);
+  const IdArray new_eids = Range(
+      0, res_csr.indices->shape[0], sizeof(IdType) * 8, csr.indptr->ctx);
+  const IdArray eids_remapped = CSRHasData(csr) ?
+    Scatter(Repeat(new_eids, edge_count), csr.data) :
+    Repeat(new_eids, edge_count);
+
+  return std::make_tuple(res_csr, edge_count, eids_remapped);
+}
+
+template std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple<kDLCPU, int32_t>(CSRMatrix);
+template std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple<kDLCPU, int64_t>(CSRMatrix);
+
+}  // namespace impl
+}  // namespace aten
+}  // namespace dgl

src/graph/transform/to_simple.cc

@@ -10,6 +10,7 @@
 #include <dgl/packed_func_ext.h>
 #include <vector>
 #include <utility>
+#include "../heterograph.h"
 #include "../unit_graph.h"
 #include "../../c_api_common.h"
 
@@ -24,51 +25,14 @@ std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>>
 ToSimpleGraph(const HeteroGraphPtr graph) {
   const int64_t num_etypes = graph->NumEdgeTypes();
   const auto metagraph = graph->meta_graph();
+  const auto &ugs = std::dynamic_pointer_cast<HeteroGraph>(graph)->relation_graphs();
 
   std::vector<IdArray> counts(num_etypes), edge_maps(num_etypes);
   std::vector<HeteroGraphPtr> rel_graphs(num_etypes);
 
   for (int64_t etype = 0; etype < num_etypes; ++etype) {
-    const auto vtypes = graph->GetEndpointTypes(etype);
-    const COOMatrix adj = graph->GetCOOMatrix(etype);
-    const COOMatrix sorted_adj = COOSort(adj, true);
-    const IdArray eids_shuffled = sorted_adj.data;
-    const auto &coalesced_result = COOCoalesce(sorted_adj);
-    const COOMatrix &coalesced_adj = coalesced_result.first;
-    const IdArray &count = coalesced_result.second;
-
-    /*
-     * eids_shuffled actually already contains the mapping from old edge space to the
-     * new one:
-     *
-     * * eids_shuffled[0:count[0]] indicates the original edge IDs that coalesced into new
-     *   edge #0.
-     * * eids_shuffled[count[0]:count[0] + count[1]] indicates those that coalesced into
-     *   new edge #1.
-     * * eids_shuffled[count[0] + count[1]:count[0] + count[1] + count[2]] indicates those
-     *   that coalesced into new edge #2.
-     * * etc.
-     *
-     * Here, we need to translate eids_shuffled to an array "eids_remapped" such that
-     * eids_remapped[i] indicates the new edge ID the old edge #i is mapped to.  The
-     * translation can simply be achieved by (in numpy code):
-     *
-     *     new_eid_for_eids_shuffled = np.range(len(count)).repeat(count)
-     *     eids_remapped = np.zeros_like(new_eid_for_eids_shuffled)
-     *     eids_remapped[eids_shuffled] = new_eid_for_eids_shuffled
-     */
-    const IdArray new_eids = Range(
-        0, coalesced_adj.row->shape[0], coalesced_adj.row->dtype.bits, coalesced_adj.row->ctx);
-    const IdArray eids_remapped = Scatter(Repeat(new_eids, count), eids_shuffled);
-
-    edge_maps[etype] = eids_remapped;
-    counts[etype] = count;
-    rel_graphs[etype] = UnitGraph::CreateFromCOO(
-        vtypes.first == vtypes.second ? 1 : 2,
-        coalesced_adj.num_rows,
-        coalesced_adj.num_cols,
-        coalesced_adj.row,
-        coalesced_adj.col);
+    const auto result = ugs[etype]->ToSimple();
+    std::tie(rel_graphs[etype], counts[etype], edge_maps[etype]) = result;
   }
 
   const HeteroGraphPtr result = CreateHeteroGraph(

src/graph/unit_graph.cc

@@ -1571,4 +1571,44 @@ UnitGraphPtr UnitGraph::Reverse() const {
   return UnitGraphPtr(new UnitGraph(meta_graph(), new_incsr, new_outcsr, new_coo));
 }
 
+std::tuple<UnitGraphPtr, IdArray, IdArray>
+UnitGraph::ToSimple() const {
+  CSRPtr new_incsr = nullptr, new_outcsr = nullptr;
+  COOPtr new_coo = nullptr;
+  IdArray count;
+  IdArray edge_map;
+
+  auto avail_fmt = SelectFormat(SparseFormat::kAny);
+  switch (avail_fmt) {
+    case SparseFormat::kCOO: {
+      auto ret = aten::COOToSimple(coo_->adj());
+      count = std::get<1>(ret);
+      edge_map = std::get<2>(ret);
+      new_coo = COOPtr(new COO(coo_->meta_graph(), std::get<0>(ret)));
+      break;
+    }
+    case SparseFormat::kCSR: {
+      auto ret = aten::CSRToSimple(in_csr_->adj());
+      count = std::get<1>(ret);
+      edge_map = std::get<2>(ret);
+      new_incsr = CSRPtr(new CSR(in_csr_->meta_graph(), std::get<0>(ret)));
+      break;
+    }
+    case SparseFormat::kCSC: {
+      auto ret = aten::CSRToSimple(out_csr_->adj());
+      count = std::get<1>(ret);
+      edge_map = std::get<2>(ret);
+      new_outcsr = CSRPtr(new CSR(out_csr_->meta_graph(), std::get<0>(ret)));
+      break;
+    }
+    default:
+      LOG(FATAL) << "At lease one of COO, CSR or CSC adj should exist.";
+      break;
+  }
+
+  return std::make_tuple(UnitGraphPtr(new UnitGraph(meta_graph(), new_incsr, new_outcsr, new_coo)),
+                         count,
+                         edge_map);
+}
+
 }  // namespace dgl

src/graph/unit_graph.h

@@ -16,6 +16,7 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <tuple>
 
 #include "../c_api_common.h"
 
@@ -269,6 +270,13 @@ class UnitGraph : public BaseHeteroGraph {
   /*! \return the reversed graph */
   UnitGraphPtr Reverse() const;
 
+  /*! \return the simpled (no-multi-edge) graph
+   *          the count recording the number of duplicated edges from the original graph.
+   *          the edge mapping from the edge IDs of original graph to those of the
+   *          returned graph.
+   */
+  std::tuple<UnitGraphPtr, IdArray, IdArray>ToSimple() const;
+
  private:
   friend class Serializer;
   friend class HeteroGraph;

tests/compute/test_transform.py

@@ -721,10 +721,10 @@ if __name__ == '__main__':
     # test_remove_self_loop()
     # test_add_self_loop()
     # test_partition_with_halo()
-    test_metis_partition()
+    # test_metis_partition()
     # test_compact()
-    # test_to_simple()
+    test_to_simple("int32")
     # test_in_subgraph("int32")
     # test_out_subgraph()
-    test_to_block("int32")
+    # test_to_block("int32")
     # test_remove_edges()

tests/cpp/test_aten.cc

@@ -257,17 +257,206 @@ void _TestConcat(DLContext ctx) {
   ASSERT_TRUE(ArrayEQ<IDX>(d, td));
 }
 
-TEST(ArrayTest, TestConcat) {
-  _TestConcat<int32_t>(CPU);
-  _TestConcat<int64_t>(CPU);
-  _TestConcat<float>(CPU);
-  _TestConcat<double>(CPU);
-#ifdef DGL_USE_CUDA
-  _TestConcat<int32_t>(GPU);
-  _TestConcat<int64_t>(GPU);
-  _TestConcat<float>(GPU);
-  _TestConcat<double>(GPU);
-#endif
+template <typename IdType>
+void _TestToSimpleCsr(DLContext ctx) {
+ /* 
+  * A = [[0, 0, 0, 0],
+  *      [1, 0, 0, 1],
+  *      [1, 1, 1, 1],
+  *      [3, 2, 2, 3],
+  *      [2, 0, 0, 2]]
+  *
+  * B = CSRToSimple(A)
+  * B = [[0, 0, 0, 0],
+  *      [1, 0, 0, 1],
+  *      [1, 1, 1, 1],
+  *      [1, 1, 1, 1],
+  *      [1, 0, 0, 1]]
+  */
+  IdArray a_indptr =
+    aten::VecToIdArray(std::vector<IdType>({0, 0, 2, 6, 16, 20}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray a_indices =
+    aten::VecToIdArray(std::vector<IdType>({0, 3, 0, 1, 2, 3, 0, 0, 0,
+                                            1, 1, 2, 2, 3, 3, 3, 0, 0,
+                                            3, 3}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray b_indptr =
+    aten::VecToIdArray(std::vector<IdType>({0, 0, 2, 6, 10, 12}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray b_indices =
+    aten::VecToIdArray(std::vector<IdType>({0, 3, 0, 1, 2, 3, 0, 1, 2,
+                                            3, 0, 3}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray cnt =
+    aten::VecToIdArray(std::vector<IdType>({1, 1, 1, 1, 1, 1, 3, 2, 2, 3, 2, 2}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray map =
+    aten::VecToIdArray(std::vector<IdType>({0, 1, 2, 3, 4, 5, 6, 6, 6,
+                                            7, 7, 8, 8, 9, 9, 9, 10, 10,
+                                            11, 11}),
+                                           sizeof(IdType)*8, CTX);
+  const aten::CSRMatrix &csr_a = aten::CSRMatrix(
+    5,
+    4,
+    a_indptr,
+    a_indices,
+    aten::NullArray(),
+    true);
+  auto ret = CSRToSimple(csr_a);
+  aten::CSRMatrix csr_b = std::get<0>(ret);
+  IdArray ecnt = std::get<1>(ret);
+  IdArray emap = std::get<2>(ret);
+  ASSERT_EQ(csr_b.num_rows, 5);
+  ASSERT_EQ(csr_b.num_cols, 4);
+  ASSERT_TRUE(ArrayEQ<IdType>(csr_b.indptr, b_indptr));
+  ASSERT_TRUE(ArrayEQ<IdType>(csr_b.indices, b_indices));
+  ASSERT_TRUE(ArrayEQ<IdType>(ecnt, cnt));
+  ASSERT_TRUE(ArrayEQ<IdType>(emap, map));
+  ASSERT_TRUE(csr_b.sorted);
+
+  // a not sorted
+  a_indices =
+    aten::VecToIdArray(std::vector<IdType>({0, 3, 0, 1, 2, 3, 3, 0, 0,
+                                            1, 1, 2, 2, 3, 3, 0, 0, 3,
+                                            0, 3}),
+                                           sizeof(IdType)*8, CTX);
+  map =
+    aten::VecToIdArray(std::vector<IdType>({0, 1, 2, 3, 4, 5, 9, 6, 6,
+                                            7, 7, 8, 8, 9, 9, 6, 10, 11,
+                                            10, 11}),
+                                           sizeof(IdType)*8, CTX);
+  const aten::CSRMatrix &csr_a2 = aten::CSRMatrix(
+    5,
+    4,
+    a_indptr,
+    a_indices,
+    aten::NullArray(),
+    false);
+  ret = CSRToSimple(csr_a2);
+  csr_b = std::get<0>(ret);
+  ecnt = std::get<1>(ret);
+  emap = std::get<2>(ret);
+  ASSERT_EQ(csr_b.num_rows, 5);
+  ASSERT_EQ(csr_b.num_cols, 4);
+  ASSERT_TRUE(ArrayEQ<IdType>(csr_b.indptr, b_indptr));
+  ASSERT_TRUE(ArrayEQ<IdType>(csr_b.indices, b_indices));
+  ASSERT_TRUE(ArrayEQ<IdType>(ecnt, cnt));
+  ASSERT_TRUE(ArrayEQ<IdType>(emap, map));
+  ASSERT_TRUE(csr_b.sorted);
+}
+
+TEST(MatrixTest, TestToSimpleCsr) {
+  _TestToSimpleCsr<int32_t>(CPU);
+  _TestToSimpleCsr<int64_t>(CPU);
+}
+
+template <typename IdType>
+void _TestToSimpleCoo(DLContext ctx) {
+ /* 
+  * A = [[0, 0, 0, 0],
+  *      [1, 0, 0, 1],
+  *      [1, 1, 1, 1],
+  *      [3, 2, 2, 3],
+  *      [2, 0, 0, 2]]
+  *
+  * B = CSRToSimple(A)
+  * B = [[0, 0, 0, 0],
+  *      [1, 0, 0, 1],
+  *      [1, 1, 1, 1],
+  *      [1, 1, 1, 1],
+  *      [1, 0, 0, 1]]
+  */
+  IdArray a_row =
+    aten::VecToIdArray(std::vector<IdType>({1, 1, 2, 2, 2, 2, 3, 3, 3,
+                                            3, 3, 3, 3, 3, 3, 3, 4, 4,
+                                            4, 4}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray a_col = 
+    aten::VecToIdArray(std::vector<IdType>({0, 3, 0, 1, 2, 3, 0, 0, 0,
+                                            1, 1, 2, 2, 3, 3, 3, 0, 0,
+                                            3, 3}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray b_row =
+    aten::VecToIdArray(std::vector<IdType>({1, 1, 2, 2, 2, 2, 3, 3, 3,
+                                            3, 4, 4}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray b_col = 
+    aten::VecToIdArray(std::vector<IdType>({0, 3, 0, 1, 2, 3, 0, 1, 2,
+                                            3, 0, 3}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray cnt =
+    aten::VecToIdArray(std::vector<IdType>({1, 1, 1, 1, 1, 1, 3, 2, 2, 3, 2, 2}),
+                                           sizeof(IdType)*8, CTX);
+  IdArray map =
+    aten::VecToIdArray(std::vector<IdType>({0, 1, 2, 3, 4, 5, 6, 6, 6,
+                                            7, 7, 8, 8, 9, 9, 9, 10, 10,
+                                            11, 11}),
+                                           sizeof(IdType)*8, CTX);
+  const aten::COOMatrix &coo_a = aten::COOMatrix(
+    5,
+    4,
+    a_row,
+    a_col,
+    aten::NullArray(),
+    true,
+    true);
+  auto ret = COOToSimple(coo_a);
+  aten::COOMatrix coo_b = std::get<0>(ret);
+  IdArray ecnt = std::get<1>(ret);
+  IdArray emap = std::get<2>(ret);
+  ASSERT_EQ(coo_b.num_rows, 5);
+  ASSERT_EQ(coo_b.num_cols, 4);
+  ASSERT_TRUE(ArrayEQ<IdType>(coo_b.row, b_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(coo_b.col, b_col));
+  ASSERT_TRUE(ArrayEQ<IdType>(ecnt, cnt));
+  ASSERT_TRUE(ArrayEQ<IdType>(emap, map));
+  ASSERT_FALSE(COOHasData(coo_b));
+  ASSERT_TRUE(coo_b.row_sorted);
+  ASSERT_TRUE(coo_b.col_sorted);
+
+  // a not sorted
+  a_row =
+    aten::VecToIdArray(std::vector<IdType>({1, 2, 1, 2, 2, 2, 3, 3, 3,
+                                            3, 3, 3, 3, 3, 3, 3, 4, 4,
+                                            4, 4}),
+                                           sizeof(IdType)*8, CTX);
+  a_col = 
+    aten::VecToIdArray(std::vector<IdType>({0, 0, 3, 1, 2, 3, 0, 0, 0,
+                                            1, 1, 2, 2, 3, 3, 3, 0, 3,
+                                            0, 3}),
+                                           sizeof(IdType)*8, CTX);
+  map =
+    aten::VecToIdArray(std::vector<IdType>({0, 2, 1, 3, 4, 5, 6, 6, 6,
+                                            7, 7, 8, 8, 9, 9, 9, 10, 11,
+                                            10, 11}),
+                                           sizeof(IdType)*8, CTX);
+  const aten::COOMatrix &coo_a2 = aten::COOMatrix(
+    5,
+    4,
+    a_row,
+    a_col,
+    aten::NullArray(),
+    false,
+    false);
+  ret = COOToSimple(coo_a2);
+  coo_b = std::get<0>(ret);
+  ecnt = std::get<1>(ret);
+  emap = std::get<2>(ret);
+  ASSERT_EQ(coo_b.num_rows, 5);
+  ASSERT_EQ(coo_b.num_cols, 4);
+  ASSERT_TRUE(ArrayEQ<IdType>(coo_b.row, b_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(coo_b.col, b_col));
+  ASSERT_TRUE(ArrayEQ<IdType>(ecnt, cnt));
+  ASSERT_TRUE(ArrayEQ<IdType>(emap, map));
+  ASSERT_FALSE(COOHasData(coo_b));
+  ASSERT_TRUE(coo_b.row_sorted);
+  ASSERT_TRUE(coo_b.col_sorted);
+}
+
+TEST(MatrixTest, TestToSimpleCoo) {
+  _TestToSimpleCoo<int32_t>(CPU);
+  _TestToSimpleCoo<int64_t>(CPU);
 }