[Kernel] Line graph with new Kernel (#1744)

* line graph

* update docstring

* Add test and compile OK

* Add python API

* Fix linty

* Fix

* Fix

* upd

* upd

* Fix

* merge

* Fix

* fix

* upd

* lint

* Fix
Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal>
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

include/dgl/aten/coo.h

@@ -514,6 +514,38 @@ std::vector<COOMatrix> DisjointPartitionCooBySizes(
   const std::vector<uint64_t> &src_vertex_cumsum,
   const std::vector<uint64_t> &dst_vertex_cumsum);
 
+/*!
+ * \brief Create a LineGraph of input coo
+ * 
+ * A = [[0, 0, 1],
+ *      [1, 0, 1],
+ *      [1, 1, 0]]
+ * A.row = [0, 1, 1, 2, 2]
+ * A.col = [2, 0, 2, 0, 1]
+ * A.eid = [0, 1, 2, 3, 4]
+ *
+ * B = COOLineGraph(A, backtracking=False)
+ *
+ * B = [[0, 0, 0, 0, 1],
+ *      [1, 0, 0, 0, 0],
+ *      [0, 0, 0, 1, 0],
+ *      [0, 0, 0, 0, 0],
+ *      [0, 1, 0, 0, 0]]
+ *
+ * C = COOLineGraph(A, backtracking=True)
+ *
+ * C = [[0, 0, 0, 1, 1],
+ *      [1, 0, 0, 0, 0],
+ *      [0, 0, 0, 1, 1],
+ *      [1, 0, 0, 0, 0],
+ *      [0, 1, 1, 0, 0]]
+ *
+ * \param coo COOMatrix to create the LineGraph
+ * \param backtracking whether the pair of (v, u) (u, v) edges are treated as linked
+ * \return LineGraph in COO format
+ */
+COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking);
+
 }  // namespace aten
 }  // namespace dgl
 

include/dgl/base_heterograph.h

@@ -682,7 +682,6 @@ HeteroSubgraph InEdgeGraph(const HeteroGraphPtr graph, const std::vector<IdArray
  */
 HeteroSubgraph OutEdgeGraph(const HeteroGraphPtr graph, const std::vector<IdArray>& nodes);
 
-
 /*!
  * \brief Union multiple graphs into one with each input graph as one disjoint component.
  *

python/dgl/transform.py

@@ -20,6 +20,7 @@ from . import ndarray as nd
 
 __all__ = [
     'line_graph',
+    'line_heterograph',
     'khop_adj',
     'khop_graph',
     'reverse',
@@ -162,6 +163,53 @@ def line_graph(g, backtracking=True, shared=False):
     node_frame = g._edge_frame if shared else None
     return DGLGraph(graph_data, node_frame)
 
+def line_heterograph(g, backtracking=True):
+    """Return the line graph of this graph.
+
+    The graph should be an directed homogeneous graph. Aother type of graphs
+    are not supported right now.
+
+    All node features and edge features are not copied to the output
+
+    Parameters
+    ----------
+    backtracking : bool
+        Whether the pair of (v, u) (u, v) edges are treated as linked. Default True.
+
+    Returns
+    -------
+    G : DGLHeteroGraph
+        The line graph of this graph.
+
+    Examples:
+    A = [[0, 0, 1],
+            [1, 0, 1],
+            [1, 1, 0]]
+    >>> g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]), 'user', 'follows')
+    >>> lg = g.line_graph()
+    >>> lg
+    ... Graph(num_nodes=5, num_edges=8,
+    ... ndata_schemes={}
+    ... edata_schemes={})
+    >>> lg.edges()
+    ... (tensor([0, 0, 1, 2, 2, 3, 4, 4]), tensor([3, 4, 0, 3, 4, 0, 1, 2]))
+    >>>
+    >>> lg = g.line_graph(backtracking=False)
+    >>> lg
+    ... Graph(num_nodes=5, num_edges=4,
+    ... ndata_schemes={}
+    ... edata_schemes={})
+    >>> lg.edges()
+    ... (tensor([0, 1, 2, 4]), tensor([4, 0, 3, 1]))
+
+    """
+    assert g.is_homograph(), \
+        'line_heterograph only support directed homogeneous graph right now'
+
+    hgidx = _CAPI_DGLHeteroLineGraph(g._graph, backtracking)
+    hg = DGLHeteroGraph(hgidx, g._etypes, g._ntypes)
+    return hg
+
 def khop_adj(g, k):
     """Return the matrix of :math:`A^k` where :math:`A` is the adjacency matrix of :math:`g`,
     where a row represents the destination and a column represents the source.

src/array/array.cc

@@ -710,6 +710,13 @@ std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
   return ret;
 }
 
+COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking) {
+  COOMatrix ret;
+  ATEN_COO_SWITCH(coo, XPU, IdType, "COOLineGraph", {
+    ret = impl::COOLineGraph<XPU, IdType>(coo, backtracking);
+  });
+  return ret;
+}
 
 COOMatrix UnionCoo(const std::vector<COOMatrix>& coos) {
   COOMatrix ret;

src/array/array_op.h

@@ -249,7 +249,8 @@ Frontiers DGLDFSLabeledEdges(const CSRMatrix& csr,
                              const bool has_nontree_edge,
                              const bool return_labels);
 
-
+template <DLDeviceType XPU, typename IdType>
+COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking);
 
 }  // namespace impl
 }  // namespace aten

src/array/cpu/coo_linegraph.cc

+/*!
+ *  Copyright (c) 2020 by Contributors
+ * \file array/cpu/coo_line_graph.cc
+ * \brief COO LineGraph
+ */
+
+#include <dgl/array.h>
+#include <numeric>
+#include <algorithm>
+#include <vector>
+#include <iterator>
+
+namespace dgl {
+namespace aten {
+namespace impl {
+
+template <DLDeviceType XPU, typename IdType>
+COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking) {
+  const int64_t nnz = coo.row->shape[0];
+  IdType* coo_row = coo.row.Ptr<IdType>();
+  IdType* coo_col = coo.col.Ptr<IdType>();
+  IdArray data = COOHasData(coo) ? coo.data : Range(0,
+                                                    nnz,
+                                                    coo.row->dtype.bits,
+                                                    coo.row->ctx);
+  IdType* data_data = data.Ptr<IdType>();
+  std::vector<IdType> new_row;
+  std::vector<IdType> new_col;
+
+  for (int64_t i = 0; i < nnz; ++i) {
+    IdType u = coo_row[i];
+    IdType v = coo_col[i];
+    for (int64_t j = 0; j < nnz; ++j) {
+      // no self-loop
+      if (i == j)
+        continue;
+
+      // succ_u == v
+      // if not backtracking succ_u != u
+      if (v == coo_row[j] && (backtracking || u != coo_col[j])) {
+        new_row.push_back(data_data[i]);
+        new_col.push_back(data_data[j]);
+      }
+    }
+  }
+
+  COOMatrix res = COOMatrix(nnz, nnz, NDArray::FromVector(new_row), NDArray::FromVector(new_col),
+    NullArray(), false, false);
+  return res;
+}
+
+
+template COOMatrix COOLineGraph<kDLCPU, int32_t>(const COOMatrix &coo, bool backtracking);
+template COOMatrix COOLineGraph<kDLCPU, int64_t>(const COOMatrix &coo, bool backtracking);
+
+}  // namespace impl
+}  // namespace aten
+}  // namespace dgl

src/graph/heterograph.cc

@@ -409,4 +409,16 @@ GraphPtr HeteroGraph::AsImmutableGraph() const {
   return unit_graph->AsImmutableGraph();
 }
 
+HeteroGraphPtr HeteroGraph::LineGraph(bool backtracking) const {
+  CHECK_EQ(1, meta_graph_->NumEdges()) << "Only support Homogeneous graph now (one edge type)";
+  CHECK_EQ(1, meta_graph_->NumVertices()) << "Only support Homogeneous graph now (one node type)";
+  CHECK_EQ(1, relation_graphs_.size()) << "Only support Homogeneous graph now";
+  UnitGraphPtr ug = relation_graphs_[0];
+
+  const auto &ulg = ug->LineGraph(backtracking);
+  std::vector<HeteroGraphPtr> rel_graph = {ulg};
+  std::vector<int64_t> num_nodes_per_type = {static_cast<int64_t>(ulg->NumVertices(0))};
+  return HeteroGraphPtr(new HeteroGraph(meta_graph_, rel_graph, std::move(num_nodes_per_type)));
+}
+
 }  // namespace dgl

src/graph/heterograph.h

@@ -220,6 +220,10 @@ class HeteroGraph : public BaseHeteroGraph {
   /*! \brief Copy the data to another context */
   static HeteroGraphPtr CopyTo(HeteroGraphPtr g, const DLContext& ctx);
 
+
+  /*! \brief Creat a LineGraph of self */
+  HeteroGraphPtr LineGraph(bool backtracking) const;
+
   const std::vector<UnitGraphPtr>& relation_graphs() const {
     return relation_graphs_;
   }

src/graph/heterograph_capi.cc

@@ -597,6 +597,7 @@ DGL_REGISTER_GLOBAL("heterograph._CAPI_DGLFindSrcDstNtypes")
     *rv = ret_list;
   });
 
+
 DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroReverse")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
     HeteroGraphRef hg = args[0];

src/graph/transform/line_graph.cc

+/*!
+ *  Copyright (c) 2020 by Contributors
+ * \file graph/transform/line_graph.cc
+ * \brief Line graph implementation
+ */
+
+#include <dgl/base_heterograph.h>
+#include <dgl/transform.h>
+#include <dgl/array.h>
+#include <dgl/packed_func_ext.h>
+#include <vector>
+#include <utility>
+#include "../../c_api_common.h"
+#include "../heterograph.h"
+
+namespace dgl {
+
+using namespace dgl::runtime;
+using namespace dgl::aten;
+
+namespace transform {
+
+/*!
+ * \brief Create Line Graph
+ * \param hg Graph
+ * \param backtracking whether the pair of (v, u) (u, v) edges are treated as linked
+ * \return The Line Graph
+ */
+HeteroGraphPtr CreateLineGraph(
+    HeteroGraphPtr hg,
+    bool backtracking) {
+  const auto hgp = std::dynamic_pointer_cast<HeteroGraph>(hg);
+  return hgp->LineGraph(backtracking);
+}
+
+DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroLineGraph")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+  HeteroGraphRef hg = args[0];
+  bool backtracking = args[1];
+
+  auto hgptr = CreateLineGraph(hg.sptr(), backtracking);
+  *rv = HeteroGraphRef(hgptr);
+});
+
+};  // namespace transform
+};  // namespace dgl

src/graph/unit_graph.cc

@@ -1494,6 +1494,31 @@ GraphPtr UnitGraph::AsImmutableGraph() const {
   return GraphPtr(new dgl::ImmutableGraph(in_csr_ptr, out_csr_ptr, coo_ptr));
 }
 
+HeteroGraphPtr UnitGraph::LineGraph(bool backtracking) const {
+  // TODO(xiangsx) currently we only support homogeneous graph
+  auto fmt = SelectFormat(SparseFormat::kAny);
+  switch (fmt) {
+    case SparseFormat::kCOO: {
+      return CreateFromCOO(1, aten::COOLineGraph(coo_->adj(), backtracking), SparseFormat::kAny);
+    }
+    case SparseFormat::kCSR: {
+      const aten::CSRMatrix csr = GetCSRMatrix(0);
+      const aten::COOMatrix coo = aten::COOLineGraph(aten::CSRToCOO(csr, true), backtracking);
+      return CreateFromCOO(1, coo, SparseFormat::kAny);
+    }
+    case SparseFormat::kCSC: {
+      const aten::CSRMatrix csc = GetCSCMatrix(0);
+      const aten::CSRMatrix csr = aten::CSRTranspose(csc);
+      const aten::COOMatrix coo = aten::COOLineGraph(aten::CSRToCOO(csr, true), backtracking);
+      return CreateFromCOO(1, coo, SparseFormat::kAny);
+    }
+    default:
+      LOG(FATAL) << "None of CSC, CSR, COO exist";
+      break;
+  }
+  return nullptr;
+}
+
 constexpr uint64_t kDGLSerialize_UnitGraphMagic = 0xDD2E60F0F6B4A127;
 
 bool UnitGraph::Load(dmlc::Stream* fs) {

src/graph/unit_graph.h

@@ -267,6 +267,9 @@ class UnitGraph : public BaseHeteroGraph {
   /*! \return Save UnitGraph to stream, using CSRMatrix */
   void Save(dmlc::Stream* fs) const;
 
+  /*! \brief Creat a LineGraph of self */
+  HeteroGraphPtr LineGraph(bool backtracking) const;
+
   /*! \return the reversed graph */
   UnitGraphPtr Reverse() const;
 

tests/compute/test_heterograph.py

@@ -2085,5 +2085,4 @@ if __name__ == '__main__':
     # test_dtype_cast()
     # test_reverse("int32")
     test_format()
-
     pass

tests/compute/test_transform.py

@@ -35,6 +35,52 @@ def test_line_graph():
     L.ndata['w'] = data
     assert F.allclose(G.edata['w'], data)
 
+@parametrize_dtype
+def test_hetero_linegraph(index_dtype):
+    g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]),
+        'user', 'follows', index_dtype=index_dtype)
+    lg = dgl.line_heterograph(g)
+    assert lg.number_of_nodes() == 5
+    assert lg.number_of_edges() == 8
+    row, col = lg.edges()
+    assert np.array_equal(F.asnumpy(row),
+                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
+    assert np.array_equal(F.asnumpy(col),
+                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))
+
+    lg = dgl.line_heterograph(g, backtracking=False)
+    assert lg.number_of_nodes() == 5
+    assert lg.number_of_edges() == 4
+    row, col = lg.edges()
+    assert np.array_equal(F.asnumpy(row),
+                          np.array([0, 1, 2, 4]))
+    assert np.array_equal(F.asnumpy(col),
+                          np.array([4, 0, 3, 1]))
+    g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]), 
+        'user', 'follows', restrict_format='csr', index_dtype=index_dtype)
+    lg = dgl.line_heterograph(g)
+    assert lg.number_of_nodes() == 5
+    assert lg.number_of_edges() == 8
+    row, col = lg.edges()
+    assert np.array_equal(F.asnumpy(row),
+                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
+    assert np.array_equal(F.asnumpy(col),
+                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))
+
+    g = dgl.graph(([0, 1, 1, 2, 2],[2, 0, 2, 0, 1]), 
+        'user', 'follows', restrict_format='csc', index_dtype=index_dtype)
+    lg = dgl.line_heterograph(g)
+    assert lg.number_of_nodes() == 5
+    assert lg.number_of_edges() == 8
+    row, col, eid = lg.edges('all')
+    row = F.asnumpy(row)
+    col = F.asnumpy(col)
+    eid = F.asnumpy(eid).astype(int)
+    order = np.argsort(eid)
+    assert np.array_equal(row[order],
+                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
+    assert np.array_equal(col[order],
+                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))
 
 def test_no_backtracking():
     N = 5
@@ -722,6 +768,7 @@ if __name__ == '__main__':
     # test_add_self_loop()
     # test_partition_with_halo()
     # test_metis_partition()
+    test_hetero_linegraph('int32')
     # test_compact()
     test_to_simple("int32")
     # test_in_subgraph("int32")

tests/cpp/test_aten.cc

@@ -1238,6 +1238,113 @@ TEST(ArrayTest, CumSum) {
 #endif
 }
 
+template <typename IdType>
+void _TestLineGraphCOO(DLContext ctx) {
+  /*
+   * A = [[0, 0, 1, 0],
+   *      [1, 0, 1, 0],
+   *      [1, 1, 0, 0],
+   *      [0, 0, 0, 1]]
+   * row: 0 1 1 2 2 3
+   * col: 2 0 2 0 1 3
+   * ID:  0 1 2 3 4 5
+   *
+   * B = COOLineGraph(A, backtracking=False)
+   *
+   * B = [[0, 0, 0, 0, 1, 0],
+   *      [1, 0, 0, 0, 0, 0],
+   *      [0, 0, 0, 1, 0, 0],
+   *      [0, 0, 0, 0, 0, 0],
+   *      [0, 1, 0, 0, 0, 0],
+   *      [0, 0, 0, 0, 0, 0]]
+   *
+   * C = COOLineGraph(A, backtracking=True)
+   *
+   * C = [[0, 0, 0, 1, 1, 0],
+   *      [1, 0, 0, 0, 0, 0],
+   *      [0, 0, 0, 1, 1, 0],
+   *      [1, 0, 0, 0, 0, 0],
+   *      [0, 1, 1, 0, 0, 0],
+   *      [0, 0, 0, 0, 0, 0]]
+   */
+  IdArray a_row =
+    aten::VecToIdArray(std::vector<IdType>({0, 1, 1, 2, 2, 3}), sizeof(IdType)*8, CTX);
+  IdArray a_col =
+    aten::VecToIdArray(std::vector<IdType>({2, 0, 2, 0, 1, 3}), sizeof(IdType)*8, CTX);
+  IdArray b_row =
+    aten::VecToIdArray(std::vector<IdType>({0, 1, 2, 4}), sizeof(IdType)*8, CTX);
+  IdArray b_col =
+    aten::VecToIdArray(std::vector<IdType>({4, 0, 3, 1}), sizeof(IdType)*8, CTX);
+  IdArray c_row =
+    aten::VecToIdArray(std::vector<IdType>({0, 0, 1, 2, 2, 3, 4, 4}), sizeof(IdType)*8, CTX);
+  IdArray c_col =
+    aten::VecToIdArray(std::vector<IdType>({3, 4, 0, 3, 4, 0, 1, 2}), sizeof(IdType)*8, CTX);
+
+  const aten::COOMatrix &coo_a = aten::COOMatrix(
+    4,
+    4,
+    a_row,
+    a_col,
+    aten::NullArray(),
+    true,
+    false);
+
+  const aten::COOMatrix &l_coo = COOLineGraph(coo_a, false);
+  ASSERT_EQ(l_coo.num_rows, 6);
+  ASSERT_EQ(l_coo.num_cols, 6);
+  ASSERT_TRUE(ArrayEQ<IdType>(l_coo.row, b_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(l_coo.col, b_col));
+  ASSERT_FALSE(l_coo.row_sorted);
+  ASSERT_FALSE(l_coo.col_sorted);
+
+  const aten::COOMatrix &l_coo2 = COOLineGraph(coo_a, true);
+  ASSERT_EQ(l_coo2.num_rows, 6);
+  ASSERT_EQ(l_coo2.num_cols, 6);
+  ASSERT_TRUE(ArrayEQ<IdType>(l_coo2.row, c_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(l_coo2.col, c_col));
+  ASSERT_FALSE(l_coo2.row_sorted);
+  ASSERT_FALSE(l_coo2.col_sorted);
+
+  IdArray a_data =
+    aten::VecToIdArray(std::vector<IdType>({4, 5, 0, 1, 2, 3}), sizeof(IdType)*8, CTX);
+  b_row =
+    aten::VecToIdArray(std::vector<IdType>({4, 5, 0, 2}), sizeof(IdType)*8, CTX);
+  b_col =
+    aten::VecToIdArray(std::vector<IdType>({2, 4, 1, 5}), sizeof(IdType)*8, CTX);
+  c_row =
+    aten::VecToIdArray(std::vector<IdType>({4, 4, 5, 0, 0, 1, 2, 2}), sizeof(IdType)*8, CTX);
+  c_col =
+    aten::VecToIdArray(std::vector<IdType>({1, 2, 4, 1, 2, 4, 5, 0}), sizeof(IdType)*8, CTX);
+  const aten::COOMatrix &coo_ad = aten::COOMatrix(
+    4,
+    4,
+    a_row,
+    a_col,
+    a_data,
+    true,
+    false);
+  const aten::COOMatrix &ld_coo = COOLineGraph(coo_ad, false);
+  ASSERT_EQ(ld_coo.num_rows, 6);
+  ASSERT_EQ(ld_coo.num_cols, 6);
+  ASSERT_TRUE(ArrayEQ<IdType>(ld_coo.row, b_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(ld_coo.col, b_col));
+  ASSERT_FALSE(ld_coo.row_sorted);
+  ASSERT_FALSE(ld_coo.col_sorted);
+
+  const aten::COOMatrix &ld_coo2 = COOLineGraph(coo_ad, true);
+  ASSERT_EQ(ld_coo2.num_rows, 6);
+  ASSERT_EQ(ld_coo2.num_cols, 6);
+  ASSERT_TRUE(ArrayEQ<IdType>(ld_coo2.row, c_row));
+  ASSERT_TRUE(ArrayEQ<IdType>(ld_coo2.col, c_col));
+  ASSERT_FALSE(ld_coo2.row_sorted);
+  ASSERT_FALSE(ld_coo2.col_sorted);
+}
+
+TEST(LineGraphTest, LineGraphCOO) {
+  _TestLineGraphCOO<int32_t>(CPU);
+  _TestLineGraphCOO<int64_t>(CPU);
+}
+
 template <typename IDX>
 void _TestNonZero() {
   BoolArray a = aten::VecToIdArray(std::vector<IDX>({1, 0, 1, 1, 0, 0, 1}));