unit_graph.cc 54 KB
Newer Older
1
2
/*!
 *  Copyright (c) 2019 by Contributors
Minjie Wang's avatar
Minjie Wang committed
3
4
 * \file graph/unit_graph.cc
 * \brief UnitGraph graph implementation
5
6
 */
#include <dgl/array.h>
7
#include <dgl/base_heterograph.h>
8
9
#include <dgl/immutable_graph.h>
#include <dgl/lazy.h>
10
11

#include "../c_api_common.h"
12
#include "./unit_graph.h"
13
14

namespace dgl {
15

16
namespace {
17
18
19

using namespace dgl::aten;

Minjie Wang's avatar
Minjie Wang committed
20
21
22
23
24
25
26
27
28
29
// create metagraph of one node type
inline GraphPtr CreateUnitGraphMetaGraph1() {
  // a self-loop edge 0->0
  std::vector<int64_t> row_vec(1, 0);
  std::vector<int64_t> col_vec(1, 0);
  IdArray row = aten::VecToIdArray(row_vec);
  IdArray col = aten::VecToIdArray(col_vec);
  GraphPtr g = ImmutableGraph::CreateFromCOO(1, row, col);
  return g;
}
30

Minjie Wang's avatar
Minjie Wang committed
31
32
33
34
35
// create metagraph of two node types
inline GraphPtr CreateUnitGraphMetaGraph2() {
  // an edge 0->1
  std::vector<int64_t> row_vec(1, 0);
  std::vector<int64_t> col_vec(1, 1);
36
37
38
39
40
  IdArray row = aten::VecToIdArray(row_vec);
  IdArray col = aten::VecToIdArray(col_vec);
  GraphPtr g = ImmutableGraph::CreateFromCOO(2, row, col);
  return g;
}
Minjie Wang's avatar
Minjie Wang committed
41
42
43
44
45
46
47
48
49
50
51
52

inline GraphPtr CreateUnitGraphMetaGraph(int num_vtypes) {
  static GraphPtr mg1 = CreateUnitGraphMetaGraph1();
  static GraphPtr mg2 = CreateUnitGraphMetaGraph2();
  if (num_vtypes == 1)
    return mg1;
  else if (num_vtypes == 2)
    return mg2;
  else
    LOG(FATAL) << "Invalid number of vertex types. Must be 1 or 2.";
  return {};
}
53
54

};  // namespace
55
56
57
58
59
60
61

//////////////////////////////////////////////////////////
//
// COO graph implementation
//
//////////////////////////////////////////////////////////

Minjie Wang's avatar
Minjie Wang committed
62
class UnitGraph::COO : public BaseHeteroGraph {
63
 public:
64
65
  COO(GraphPtr metagraph, int64_t num_src, int64_t num_dst, IdArray src,
      IdArray dst, bool row_sorted = false, bool col_sorted = false)
Minjie Wang's avatar
Minjie Wang committed
66
    : BaseHeteroGraph(metagraph) {
67
68
69
    CHECK(aten::IsValidIdArray(src));
    CHECK(aten::IsValidIdArray(dst));
    CHECK_EQ(src->shape[0], dst->shape[0]) << "Input arrays should have the same length.";
70
71
72
    adj_ = aten::COOMatrix{num_src, num_dst, src, dst,
        NullArray(),
        row_sorted, col_sorted};
73
  }
74

75
76
77
78
  COO(GraphPtr metagraph, const aten::COOMatrix& coo)
    : BaseHeteroGraph(metagraph), adj_(coo) {
    // Data index should not be inherited. Edges in COO format are always
    // assigned ids from 0 to num_edges - 1.
79
    CHECK(!COOHasData(coo)) << "[BUG] COO should not contain data.";
80
    adj_.data = aten::NullArray();
81
  }
82

83
84
85
86
87
88
89
90
91
92
93
  COO() {
    // set magic num_rows/num_cols to mark it as undefined
    // adj_.num_rows == 0 and adj_.num_cols == 0 means empty UnitGraph which is supported
    adj_.num_rows = -1;
    adj_.num_cols = -1;
  };

  bool defined() const {
    return (adj_.num_rows >= 0) && (adj_.num_cols >= 0);
  }

Minjie Wang's avatar
Minjie Wang committed
94
95
  inline dgl_type_t SrcType() const {
    return 0;
96
  }
Minjie Wang's avatar
Minjie Wang committed
97
98
99
100
101
102
103

  inline dgl_type_t DstType() const {
    return NumVertexTypes() == 1? 0 : 1;
  }

  inline dgl_type_t EdgeType() const {
    return 0;
104
105
106
  }

  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
Minjie Wang's avatar
Minjie Wang committed
107
    LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
108
109
110
111
112
      << "The relation graph is simply this graph itself.";
    return {};
  }

  void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
Minjie Wang's avatar
Minjie Wang committed
113
    LOG(FATAL) << "UnitGraph graph is not mutable.";
114
115
116
  }

  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
Minjie Wang's avatar
Minjie Wang committed
117
    LOG(FATAL) << "UnitGraph graph is not mutable.";
118
119
120
  }

  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
Minjie Wang's avatar
Minjie Wang committed
121
    LOG(FATAL) << "UnitGraph graph is not mutable.";
122
123
124
  }

  void Clear() override {
Minjie Wang's avatar
Minjie Wang committed
125
    LOG(FATAL) << "UnitGraph graph is not mutable.";
126
127
  }

128
129
130
131
  DLDataType DataType() const override {
    return adj_.row->dtype;
  }

132
133
134
135
136
137
138
139
  DLContext Context() const override {
    return adj_.row->ctx;
  }

  uint8_t NumBits() const override {
    return adj_.row->dtype.bits;
  }

140
141
142
143
144
145
146
147
148
149
150
151
  COO AsNumBits(uint8_t bits) const {
    if (NumBits() == bits)
      return *this;

    COO ret(
        meta_graph_,
        adj_.num_rows, adj_.num_cols,
        aten::AsNumBits(adj_.row, bits),
        aten::AsNumBits(adj_.col, bits));
    return ret;
  }

152
153
  COO CopyTo(const DLContext &ctx,
             const DGLStreamHandle &stream = nullptr) const {
154
155
    if (Context() == ctx)
      return *this;
156
    return COO(meta_graph_, adj_.CopyTo(ctx, stream));
157
158
  }

159
  bool IsMultigraph() const override {
160
    return aten::COOHasDuplicate(adj_);
161
162
163
164
165
166
167
  }

  bool IsReadonly() const override {
    return true;
  }

  uint64_t NumVertices(dgl_type_t vtype) const override {
Minjie Wang's avatar
Minjie Wang committed
168
    if (vtype == SrcType()) {
169
      return adj_.num_rows;
Minjie Wang's avatar
Minjie Wang committed
170
    } else if (vtype == DstType()) {
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
      return adj_.num_cols;
    } else {
      LOG(FATAL) << "Invalid vertex type: " << vtype;
      return 0;
    }
  }

  uint64_t NumEdges(dgl_type_t etype) const override {
    return adj_.row->shape[0];
  }

  bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const override {
    return vid < NumVertices(vtype);
  }

  BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return {};
  }

  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
192
193
194
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
    return aten::COOIsNonZero(adj_, src, dst);
195
196
197
  }

  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
198
199
200
    CHECK(aten::IsValidIdArray(src_ids)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst_ids)) << "Invalid vertex id array.";
    return aten::COOIsNonZero(adj_, src_ids, dst_ids);
201
202
203
  }

  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
204
205
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
    return aten::COOGetRowDataAndIndices(aten::COOTranspose(adj_), dst).second;
206
207
208
  }

  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
209
210
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    return aten::COOGetRowDataAndIndices(adj_, src).second;
211
212
213
  }

  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
214
215
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
216
    return aten::COOGetAllData(adj_, src, dst);
217
218
  }

219
  EdgeArray EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const override {
220
221
222
223
    CHECK(aten::IsValidIdArray(src)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst)) << "Invalid vertex id array.";
    const auto& arrs = aten::COOGetDataAndIndices(adj_, src, dst);
    return EdgeArray{arrs[0], arrs[1], arrs[2]};
224
225
  }

226
227
228
229
  IdArray EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const override {
    return aten::COOGetData(adj_, src, dst);
  }

230
231
  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
    CHECK(eid < NumEdges(etype)) << "Invalid edge id: " << eid;
232
233
    const dgl_id_t src = aten::IndexSelect<int64_t>(adj_.row, eid);
    const dgl_id_t dst = aten::IndexSelect<int64_t>(adj_.col, eid);
234
235
236
237
    return std::pair<dgl_id_t, dgl_id_t>(src, dst);
  }

  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
238
    CHECK(aten::IsValidIdArray(eids)) << "Invalid edge id array";
239
    BUG_IF_FAIL(aten::IsNullArray(adj_.data)) <<
240
      "FindEdges requires the internal COO matrix not having EIDs.";
241
242
243
244
245
246
    return EdgeArray{aten::IndexSelect(adj_.row, eids),
                     aten::IndexSelect(adj_.col, eids),
                     eids};
  }

  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
247
248
249
250
251
    IdArray ret_src, ret_eid;
    std::tie(ret_eid, ret_src) = aten::COOGetRowDataAndIndices(
        aten::COOTranspose(adj_), vid);
    IdArray ret_dst = aten::Full(vid, ret_src->shape[0], NumBits(), ret_src->ctx);
    return EdgeArray{ret_src, ret_dst, ret_eid};
252
253
254
  }

  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
255
256
257
258
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    auto coosubmat = aten::COOSliceRows(aten::COOTranspose(adj_), vids);
    auto row = aten::IndexSelect(vids, coosubmat.row);
    return EdgeArray{coosubmat.col, row, coosubmat.data};
259
260
261
  }

  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
262
263
264
265
    IdArray ret_dst, ret_eid;
    std::tie(ret_eid, ret_dst) = aten::COOGetRowDataAndIndices(adj_, vid);
    IdArray ret_src = aten::Full(vid, ret_dst->shape[0], NumBits(), ret_dst->ctx);
    return EdgeArray{ret_src, ret_dst, ret_eid};
266
267
268
  }

  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
269
270
271
272
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    auto coosubmat = aten::COOSliceRows(adj_, vids);
    auto row = aten::IndexSelect(vids, coosubmat.row);
    return EdgeArray{row, coosubmat.col, coosubmat.data};
273
274
275
276
277
278
279
280
281
282
283
  }

  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override {
    CHECK(order.empty() || order == std::string("eid"))
      << "COO only support Edges of order \"eid\", but got \""
      << order << "\".";
    IdArray rst_eid = aten::Range(0, NumEdges(etype), NumBits(), Context());
    return EdgeArray{adj_.row, adj_.col, rst_eid};
  }

  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override {
284
285
    CHECK(HasVertex(DstType(), vid)) << "Invalid dst vertex id: " << vid;
    return aten::COOGetRowNNZ(aten::COOTranspose(adj_), vid);
286
287
288
  }

  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
289
290
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    return aten::COOGetRowNNZ(aten::COOTranspose(adj_), vids);
291
292
293
  }

  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
294
295
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
    return aten::COOGetRowNNZ(adj_, vid);
296
297
298
  }

  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
299
300
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    return aten::COOGetRowNNZ(adj_, vids);
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
  }

  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override {
    LOG(INFO) << "Not enabled for COO graph.";
    return {};
  }

  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
    LOG(INFO) << "Not enabled for COO graph.";
    return {};
  }

  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override {
    LOG(INFO) << "Not enabled for COO graph.";
    return {};
  }

  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
    LOG(INFO) << "Not enabled for COO graph.";
    return {};
  }

  std::vector<IdArray> GetAdj(
      dgl_type_t etype, bool transpose, const std::string &fmt) const override {
    CHECK(fmt == "coo") << "Not valid adj format request.";
    if (transpose) {
      return {aten::HStack(adj_.col, adj_.row)};
    } else {
      return {aten::HStack(adj_.row, adj_.col)};
    }
  }

333
334
335
336
337
338
339
340
341
342
343
344
345
346
  aten::COOMatrix GetCOOMatrix(dgl_type_t etype) const override {
    return adj_;
  }

  aten::CSRMatrix GetCSCMatrix(dgl_type_t etype) const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return aten::CSRMatrix();
  }

  aten::CSRMatrix GetCSRMatrix(dgl_type_t etype) const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return aten::CSRMatrix();
  }

347
  SparseFormat SelectFormat(dgl_type_t etype, dgl_format_code_t preferred_formats) const override {
348
    LOG(FATAL) << "Not enabled for COO graph";
349
    return SparseFormat::kCOO;
350
351
  }

352
  dgl_format_code_t GetAllowedFormats() const override {
353
    LOG(FATAL) << "Not enabled for COO graph";
354
    return 0;
355
356
  }

357
  dgl_format_code_t GetCreatedFormats() const override {
358
359
360
361
    LOG(FATAL) << "Not enabled for COO graph";
    return 0;
  }

362
  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
363
364
365
366
367
368
369
370
371
372
373
374
    CHECK_EQ(vids.size(), NumVertexTypes()) << "Number of vertex types mismatch";
    auto srcvids = vids[SrcType()], dstvids = vids[DstType()];
    CHECK(aten::IsValidIdArray(srcvids)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dstvids)) << "Invalid vertex id array.";
    HeteroSubgraph subg;
    const auto& submat = aten::COOSliceMatrix(adj_, srcvids, dstvids);
    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
    subg.graph = std::make_shared<COO>(meta_graph(), submat.num_rows, submat.num_cols,
        submat.row, submat.col);
    subg.induced_vertices = vids;
    subg.induced_edges.emplace_back(submat.data);
    return subg;
375
376
377
378
379
380
381
382
383
384
385
386
387
388
  }

  HeteroSubgraph EdgeSubgraph(
      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override {
    CHECK_EQ(eids.size(), 1) << "Edge type number mismatch.";
    HeteroSubgraph subg;
    if (!preserve_nodes) {
      IdArray new_src = aten::IndexSelect(adj_.row, eids[0]);
      IdArray new_dst = aten::IndexSelect(adj_.col, eids[0]);
      subg.induced_vertices.emplace_back(aten::Relabel_({new_src}));
      subg.induced_vertices.emplace_back(aten::Relabel_({new_dst}));
      const auto new_nsrc = subg.induced_vertices[0]->shape[0];
      const auto new_ndst = subg.induced_vertices[1]->shape[0];
      subg.graph = std::make_shared<COO>(
Minjie Wang's avatar
Minjie Wang committed
389
          meta_graph(), new_nsrc, new_ndst, new_src, new_dst);
390
391
392
393
      subg.induced_edges = eids;
    } else {
      IdArray new_src = aten::IndexSelect(adj_.row, eids[0]);
      IdArray new_dst = aten::IndexSelect(adj_.col, eids[0]);
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
394
395
396
397
      subg.induced_vertices.emplace_back(
          aten::Range(0, NumVertices(SrcType()), NumBits(), Context()));
      subg.induced_vertices.emplace_back(
          aten::Range(0, NumVertices(DstType()), NumBits(), Context()));
398
      subg.graph = std::make_shared<COO>(
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
399
          meta_graph(), NumVertices(SrcType()), NumVertices(DstType()), new_src, new_dst);
400
401
402
403
404
      subg.induced_edges = eids;
    }
    return subg;
  }

405
  HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const override {
406
407
408
409
    LOG(FATAL) << "Not enabled for COO graph.";
    return nullptr;
  }

410
411
412
413
  aten::COOMatrix adj() const {
    return adj_;
  }

414
415
416
417
418
419
420
421
422
  /*!
   * \brief Determines whether the graph is "hypersparse", i.e. having significantly more
   * nodes than edges.
   */
  bool IsHypersparse() const {
    return (NumVertices(SrcType()) / 8 > NumEdges(EdgeType())) &&
           (NumVertices(SrcType()) > 1000000);
  }

423
424
425
426
427
428
429
430
431
432
433
434
435
  bool Load(dmlc::Stream* fs) {
    auto meta_imgraph = Serializer::make_shared<ImmutableGraph>();
    CHECK(fs->Read(&meta_imgraph)) << "Invalid meta graph";
    meta_graph_ = meta_imgraph;
    CHECK(fs->Read(&adj_)) << "Invalid adj matrix";
    return true;
  }
  void Save(dmlc::Stream* fs) const {
    auto meta_graph_ptr = ImmutableGraph::ToImmutable(meta_graph());
    fs->Write(meta_graph_ptr);
    fs->Write(adj_);
  }

436
 private:
437
438
  friend class Serializer;

439
440
441
442
443
444
445
446
447
448
449
  /*! \brief internal adjacency matrix. Data array is empty */
  aten::COOMatrix adj_;
};

//////////////////////////////////////////////////////////
//
// CSR graph implementation
//
//////////////////////////////////////////////////////////

/*! \brief CSR graph */
Minjie Wang's avatar
Minjie Wang committed
450
class UnitGraph::CSR : public BaseHeteroGraph {
451
 public:
Minjie Wang's avatar
Minjie Wang committed
452
  CSR(GraphPtr metagraph, int64_t num_src, int64_t num_dst,
453
      IdArray indptr, IdArray indices, IdArray edge_ids)
Minjie Wang's avatar
Minjie Wang committed
454
    : BaseHeteroGraph(metagraph) {
455
456
    CHECK(aten::IsValidIdArray(indptr));
    CHECK(aten::IsValidIdArray(indices));
457
458
459
    if (aten::IsValidIdArray(edge_ids))
      CHECK((indices->shape[0] == edge_ids->shape[0]) || aten::IsNullArray(edge_ids))
        << "edge id arrays should have the same length as indices if not empty";
460
461
    CHECK_EQ(num_src, indptr->shape[0] - 1)
      << "number of nodes do not match the length of indptr minus 1.";
462

463
464
465
    adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
  }

466
  CSR(GraphPtr metagraph, const aten::CSRMatrix& csr)
Da Zheng's avatar
Da Zheng committed
467
468
    : BaseHeteroGraph(metagraph), adj_(csr) {
  }
469

470
471
472
473
474
475
476
477
478
479
480
  CSR() {
    // set magic num_rows/num_cols to mark it as undefined
    // adj_.num_rows == 0 and adj_.num_cols == 0 means empty UnitGraph which is supported
    adj_.num_rows = -1;
    adj_.num_cols = -1;
  };

  bool defined() const {
    return (adj_.num_rows >= 0) || (adj_.num_cols >= 0);
  }

Minjie Wang's avatar
Minjie Wang committed
481
482
  inline dgl_type_t SrcType() const {
    return 0;
483
  }
Minjie Wang's avatar
Minjie Wang committed
484
485
486
487
488
489
490

  inline dgl_type_t DstType() const {
    return NumVertexTypes() == 1? 0 : 1;
  }

  inline dgl_type_t EdgeType() const {
    return 0;
491
492
493
  }

  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
Minjie Wang's avatar
Minjie Wang committed
494
    LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
495
496
497
498
499
      << "The relation graph is simply this graph itself.";
    return {};
  }

  void AddVertices(dgl_type_t vtype, uint64_t num_vertices) override {
Minjie Wang's avatar
Minjie Wang committed
500
    LOG(FATAL) << "UnitGraph graph is not mutable.";
501
502
503
  }

  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
Minjie Wang's avatar
Minjie Wang committed
504
    LOG(FATAL) << "UnitGraph graph is not mutable.";
505
506
507
  }

  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
Minjie Wang's avatar
Minjie Wang committed
508
    LOG(FATAL) << "UnitGraph graph is not mutable.";
509
510
511
  }

  void Clear() override {
Minjie Wang's avatar
Minjie Wang committed
512
    LOG(FATAL) << "UnitGraph graph is not mutable.";
513
514
  }

515
516
517
518
  DLDataType DataType() const override {
    return adj_.indices->dtype;
  }

519
520
521
522
523
524
525
526
  DLContext Context() const override {
    return adj_.indices->ctx;
  }

  uint8_t NumBits() const override {
    return adj_.indices->dtype.bits;
  }

527
528
529
530
531
  CSR AsNumBits(uint8_t bits) const {
    if (NumBits() == bits) {
      return *this;
    } else {
      CSR ret(
Minjie Wang's avatar
Minjie Wang committed
532
          meta_graph_,
533
534
535
536
537
538
539
540
          adj_.num_rows, adj_.num_cols,
          aten::AsNumBits(adj_.indptr, bits),
          aten::AsNumBits(adj_.indices, bits),
          aten::AsNumBits(adj_.data, bits));
      return ret;
    }
  }

541
542
  CSR CopyTo(const DLContext &ctx,
             const DGLStreamHandle &stream = nullptr) const {
543
544
545
    if (Context() == ctx) {
      return *this;
    } else {
546
      return CSR(meta_graph_, adj_.CopyTo(ctx, stream));
547
548
549
    }
  }

550
  bool IsMultigraph() const override {
551
    return aten::CSRHasDuplicate(adj_);
552
553
554
555
556
557
558
  }

  bool IsReadonly() const override {
    return true;
  }

  uint64_t NumVertices(dgl_type_t vtype) const override {
Minjie Wang's avatar
Minjie Wang committed
559
    if (vtype == SrcType()) {
560
      return adj_.num_rows;
Minjie Wang's avatar
Minjie Wang committed
561
    } else if (vtype == DstType()) {
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
      return adj_.num_cols;
    } else {
      LOG(FATAL) << "Invalid vertex type: " << vtype;
      return 0;
    }
  }

  uint64_t NumEdges(dgl_type_t etype) const override {
    return adj_.indices->shape[0];
  }

  bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const override {
    return vid < NumVertices(vtype);
  }

  BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return {};
  }

  bool HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
Minjie Wang's avatar
Minjie Wang committed
583
584
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
585
586
587
588
    return aten::CSRIsNonZero(adj_, src, dst);
  }

  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
589
590
    CHECK(aten::IsValidIdArray(src_ids)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst_ids)) << "Invalid vertex id array.";
591
592
593
594
595
596
597
598
599
    return aten::CSRIsNonZero(adj_, src_ids, dst_ids);
  }

  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
    LOG(INFO) << "Not enabled for CSR graph.";
    return {};
  }

  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
Minjie Wang's avatar
Minjie Wang committed
600
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
601
602
603
604
    return aten::CSRGetRowColumnIndices(adj_, src);
  }

  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
Minjie Wang's avatar
Minjie Wang committed
605
606
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
607
    return aten::CSRGetAllData(adj_, src, dst);
608
609
  }

610
  EdgeArray EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const override {
611
612
    CHECK(aten::IsValidIdArray(src)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst)) << "Invalid vertex id array.";
613
614
615
616
    const auto& arrs = aten::CSRGetDataAndIndices(adj_, src, dst);
    return EdgeArray{arrs[0], arrs[1], arrs[2]};
  }

617
618
619
620
  IdArray EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const override {
    return aten::CSRGetData(adj_, src, dst);
  }

621
  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
622
    LOG(FATAL) << "Not enabled for CSR graph.";
623
624
625
626
    return {};
  }

  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
627
    LOG(FATAL) << "Not enabled for CSR graph.";
628
629
630
631
    return {};
  }

  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
632
    LOG(FATAL) << "Not enabled for CSR graph.";
633
634
635
636
    return {};
  }

  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
637
    LOG(FATAL) << "Not enabled for CSR graph.";
638
639
640
641
    return {};
  }

  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
Minjie Wang's avatar
Minjie Wang committed
642
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
643
644
645
646
647
648
649
    IdArray ret_dst = aten::CSRGetRowColumnIndices(adj_, vid);
    IdArray ret_eid = aten::CSRGetRowData(adj_, vid);
    IdArray ret_src = aten::Full(vid, ret_dst->shape[0], NumBits(), ret_dst->ctx);
    return EdgeArray{ret_src, ret_dst, ret_eid};
  }

  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
650
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
651
652
653
654
655
656
657
658
659
660
661
662
    auto csrsubmat = aten::CSRSliceRows(adj_, vids);
    auto coosubmat = aten::CSRToCOO(csrsubmat, false);
    // Note that the row id in the csr submat is relabled, so
    // we need to recover it using an index select.
    auto row = aten::IndexSelect(vids, coosubmat.row);
    return EdgeArray{row, coosubmat.col, coosubmat.data};
  }

  EdgeArray Edges(dgl_type_t etype, const std::string &order = "") const override {
    CHECK(order.empty() || order == std::string("srcdst"))
      << "CSR only support Edges of order \"srcdst\","
      << " but got \"" << order << "\".";
663
664
665
666
667
    auto coo = aten::CSRToCOO(adj_, false);
    if (order == std::string("srcdst")) {
      // make sure the coo is sorted if an order is requested
      coo = aten::COOSort(coo, true);
    }
668
669
670
671
    return EdgeArray{coo.row, coo.col, coo.data};
  }

  uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const override {
672
    LOG(FATAL) << "Not enabled for CSR graph.";
673
674
675
676
    return {};
  }

  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
677
    LOG(FATAL) << "Not enabled for CSR graph.";
678
679
680
681
    return {};
  }

  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
Minjie Wang's avatar
Minjie Wang committed
682
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
683
684
685
686
    return aten::CSRGetRowNNZ(adj_, vid);
  }

  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
687
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
688
689
690
691
692
693
    return aten::CSRGetRowNNZ(adj_, vids);
  }

  DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const override {
    // TODO(minjie): This still assumes the data type and device context
    //   of this graph. Should fix later.
694
    CHECK_EQ(NumBits(), 64);
695
696
697
698
699
700
701
    const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
    const dgl_id_t* indices_data = static_cast<dgl_id_t*>(adj_.indices->data);
    const dgl_id_t start = indptr_data[vid];
    const dgl_id_t end = indptr_data[vid + 1];
    return DGLIdIters(indices_data + start, indices_data + end);
  }

702
703
704
705
706
707
708
709
710
711
  DGLIdIters32 SuccVec32(dgl_type_t etype, dgl_id_t vid) {
    // TODO(minjie): This still assumes the data type and device context
    //   of this graph. Should fix later.
    const int32_t* indptr_data = static_cast<int32_t*>(adj_.indptr->data);
    const int32_t* indices_data = static_cast<int32_t*>(adj_.indices->data);
    const int32_t start = indptr_data[vid];
    const int32_t end = indptr_data[vid + 1];
    return DGLIdIters32(indices_data + start, indices_data + end);
  }

712
713
714
  DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
    // TODO(minjie): This still assumes the data type and device context
    //   of this graph. Should fix later.
715
    CHECK_EQ(NumBits(), 64);
716
717
718
719
720
721
722
723
    const dgl_id_t* indptr_data = static_cast<dgl_id_t*>(adj_.indptr->data);
    const dgl_id_t* eid_data = static_cast<dgl_id_t*>(adj_.data->data);
    const dgl_id_t start = indptr_data[vid];
    const dgl_id_t end = indptr_data[vid + 1];
    return DGLIdIters(eid_data + start, eid_data + end);
  }

  DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const override {
724
    LOG(FATAL) << "Not enabled for CSR graph.";
725
726
727
728
    return {};
  }

  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
729
    LOG(FATAL) << "Not enabled for CSR graph.";
730
731
732
733
734
735
736
737
738
    return {};
  }

  std::vector<IdArray> GetAdj(
      dgl_type_t etype, bool transpose, const std::string &fmt) const override {
    CHECK(!transpose && fmt == "csr") << "Not valid adj format request.";
    return {adj_.indptr, adj_.indices, adj_.data};
  }

739
740
741
742
743
744
745
746
747
748
749
750
751
752
  aten::COOMatrix GetCOOMatrix(dgl_type_t etype) const override {
    LOG(FATAL) << "Not enabled for CSR graph";
    return aten::COOMatrix();
  }

  aten::CSRMatrix GetCSCMatrix(dgl_type_t etype) const override {
    LOG(FATAL) << "Not enabled for CSR graph";
    return aten::CSRMatrix();
  }

  aten::CSRMatrix GetCSRMatrix(dgl_type_t etype) const override {
    return adj_;
  }

753
  SparseFormat SelectFormat(dgl_type_t etype, dgl_format_code_t preferred_formats) const override {
754
    LOG(FATAL) << "Not enabled for CSR graph";
755
    return SparseFormat::kCSR;
756
757
  }

758
759
760
  dgl_format_code_t GetAllowedFormats() const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return 0;
761
762
  }

763
  dgl_format_code_t GetCreatedFormats() const override {
764
765
766
767
    LOG(FATAL) << "Not enabled for CSR graph";
    return 0;
  }

768
  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
Minjie Wang's avatar
Minjie Wang committed
769
770
771
772
    CHECK_EQ(vids.size(), NumVertexTypes()) << "Number of vertex types mismatch";
    auto srcvids = vids[SrcType()], dstvids = vids[DstType()];
    CHECK(aten::IsValidIdArray(srcvids)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dstvids)) << "Invalid vertex id array.";
773
    HeteroSubgraph subg;
Minjie Wang's avatar
Minjie Wang committed
774
    const auto& submat = aten::CSRSliceMatrix(adj_, srcvids, dstvids);
775
    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
Minjie Wang's avatar
Minjie Wang committed
776
    subg.graph = std::make_shared<CSR>(meta_graph(), submat.num_rows, submat.num_cols,
777
778
779
780
781
782
783
784
        submat.indptr, submat.indices, sub_eids);
    subg.induced_vertices = vids;
    subg.induced_edges.emplace_back(submat.data);
    return subg;
  }

  HeteroSubgraph EdgeSubgraph(
      const std::vector<IdArray>& eids, bool preserve_nodes = false) const override {
785
    LOG(FATAL) << "Not enabled for CSR graph.";
786
787
788
    return {};
  }

789
  HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const override {
790
791
792
793
    LOG(FATAL) << "Not enabled for CSR graph.";
    return nullptr;
  }

794
795
796
797
  aten::CSRMatrix adj() const {
    return adj_;
  }

798
799
800
801
802
803
804
805
806
807
808
809
810
  bool Load(dmlc::Stream* fs) {
    auto meta_imgraph = Serializer::make_shared<ImmutableGraph>();
    CHECK(fs->Read(&meta_imgraph)) << "Invalid meta graph";
    meta_graph_ = meta_imgraph;
    CHECK(fs->Read(&adj_)) << "Invalid adj matrix";
    return true;
  }
  void Save(dmlc::Stream* fs) const {
    auto meta_graph_ptr = ImmutableGraph::ToImmutable(meta_graph());
    fs->Write(meta_graph_ptr);
    fs->Write(adj_);
  }

811
 private:
812
813
  friend class Serializer;

814
815
816
817
818
819
  /*! \brief internal adjacency matrix. Data array stores edge ids */
  aten::CSRMatrix adj_;
};

//////////////////////////////////////////////////////////
//
Minjie Wang's avatar
Minjie Wang committed
820
// unit graph implementation
821
822
823
//
//////////////////////////////////////////////////////////

824
825
826
827
DLDataType UnitGraph::DataType() const {
  return GetAny()->DataType();
}

Minjie Wang's avatar
Minjie Wang committed
828
DLContext UnitGraph::Context() const {
829
830
831
  return GetAny()->Context();
}

Minjie Wang's avatar
Minjie Wang committed
832
uint8_t UnitGraph::NumBits() const {
833
834
835
  return GetAny()->NumBits();
}

Minjie Wang's avatar
Minjie Wang committed
836
bool UnitGraph::IsMultigraph() const {
837
  const SparseFormat fmt = SelectFormat(CSC_CODE);
838
839
  const auto ptr = GetFormat(fmt);
  return ptr->IsMultigraph();
840
841
}

Minjie Wang's avatar
Minjie Wang committed
842
uint64_t UnitGraph::NumVertices(dgl_type_t vtype) const {
843
  const SparseFormat fmt = SelectFormat(ALL_CODE);
844
845
846
  const auto ptr = GetFormat(fmt);
  // TODO(BarclayII): we have a lot of special handling for CSC.
  // Need to have a UnitGraph::CSC backend instead.
847
  if (fmt == SparseFormat::kCSC)
Minjie Wang's avatar
Minjie Wang committed
848
    vtype = (vtype == SrcType()) ? DstType() : SrcType();
849
  return ptr->NumVertices(vtype);
850
851
}

Minjie Wang's avatar
Minjie Wang committed
852
uint64_t UnitGraph::NumEdges(dgl_type_t etype) const {
853
854
855
  return GetAny()->NumEdges(etype);
}

Minjie Wang's avatar
Minjie Wang committed
856
bool UnitGraph::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
857
  const SparseFormat fmt = SelectFormat(ALL_CODE);
858
  const auto ptr = GetFormat(fmt);
859
  if (fmt == SparseFormat::kCSC)
Minjie Wang's avatar
Minjie Wang committed
860
    vtype = (vtype == SrcType()) ? DstType() : SrcType();
861
  return ptr->HasVertex(vtype, vid);
862
863
}

Minjie Wang's avatar
Minjie Wang committed
864
BoolArray UnitGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
865
  CHECK(aten::IsValidIdArray(vids)) << "Invalid id array input";
866
867
868
  return aten::LT(vids, NumVertices(vtype));
}

Minjie Wang's avatar
Minjie Wang committed
869
bool UnitGraph::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
870
  const SparseFormat fmt = SelectFormat(CSC_CODE);
871
  const auto ptr = GetFormat(fmt);
872
  if (fmt == SparseFormat::kCSC)
873
874
875
    return ptr->HasEdgeBetween(etype, dst, src);
  else
    return ptr->HasEdgeBetween(etype, src, dst);
876
877
}

Minjie Wang's avatar
Minjie Wang committed
878
BoolArray UnitGraph::HasEdgesBetween(
879
    dgl_type_t etype, IdArray src, IdArray dst) const {
880
  const SparseFormat fmt = SelectFormat(CSC_CODE);
881
  const auto ptr = GetFormat(fmt);
882
  if (fmt == SparseFormat::kCSC)
883
884
885
    return ptr->HasEdgesBetween(etype, dst, src);
  else
    return ptr->HasEdgesBetween(etype, src, dst);
886
887
}

Minjie Wang's avatar
Minjie Wang committed
888
IdArray UnitGraph::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
889
  const SparseFormat fmt = SelectFormat(CSC_CODE);
890
  const auto ptr = GetFormat(fmt);
891
  if (fmt == SparseFormat::kCSC)
892
893
894
    return ptr->Successors(etype, dst);
  else
    return ptr->Predecessors(etype, dst);
895
896
}

Minjie Wang's avatar
Minjie Wang committed
897
IdArray UnitGraph::Successors(dgl_type_t etype, dgl_id_t src) const {
898
  const SparseFormat fmt = SelectFormat(CSR_CODE);
899
900
  const auto ptr = GetFormat(fmt);
  return ptr->Successors(etype, src);
901
902
}

Minjie Wang's avatar
Minjie Wang committed
903
IdArray UnitGraph::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
904
  const SparseFormat fmt = SelectFormat(CSR_CODE);
905
  const auto ptr = GetFormat(fmt);
906
  if (fmt == SparseFormat::kCSC)
907
908
909
    return ptr->EdgeId(etype, dst, src);
  else
    return ptr->EdgeId(etype, src, dst);
910
911
}

912
EdgeArray UnitGraph::EdgeIdsAll(dgl_type_t etype, IdArray src, IdArray dst) const {
913
  const SparseFormat fmt = SelectFormat(CSR_CODE);
914
  const auto ptr = GetFormat(fmt);
915
  if (fmt == SparseFormat::kCSC) {
916
    EdgeArray edges = ptr->EdgeIdsAll(etype, dst, src);
917
918
    return EdgeArray{edges.dst, edges.src, edges.id};
  } else {
919
920
921
922
923
    return ptr->EdgeIdsAll(etype, src, dst);
  }
}

IdArray UnitGraph::EdgeIdsOne(dgl_type_t etype, IdArray src, IdArray dst) const {
924
  const SparseFormat fmt = SelectFormat(CSR_CODE);
925
926
927
928
929
  const auto ptr = GetFormat(fmt);
  if (fmt == SparseFormat::kCSC) {
    return ptr->EdgeIdsOne(etype, dst, src);
  } else {
    return ptr->EdgeIdsOne(etype, src, dst);
930
931
932
  }
}

Minjie Wang's avatar
Minjie Wang committed
933
std::pair<dgl_id_t, dgl_id_t> UnitGraph::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
934
  const SparseFormat fmt = SelectFormat(COO_CODE);
935
936
  const auto ptr = GetFormat(fmt);
  return ptr->FindEdge(etype, eid);
937
938
}

Minjie Wang's avatar
Minjie Wang committed
939
EdgeArray UnitGraph::FindEdges(dgl_type_t etype, IdArray eids) const {
940
  const SparseFormat fmt = SelectFormat(COO_CODE);
941
942
  const auto ptr = GetFormat(fmt);
  return ptr->FindEdges(etype, eids);
943
944
}

Minjie Wang's avatar
Minjie Wang committed
945
EdgeArray UnitGraph::InEdges(dgl_type_t etype, dgl_id_t vid) const {
946
  const SparseFormat fmt = SelectFormat(CSC_CODE);
947
  const auto ptr = GetFormat(fmt);
948
  if (fmt == SparseFormat::kCSC) {
949
950
951
952
953
    const EdgeArray& ret = ptr->OutEdges(etype, vid);
    return {ret.dst, ret.src, ret.id};
  } else {
    return ptr->InEdges(etype, vid);
  }
954
955
}

Minjie Wang's avatar
Minjie Wang committed
956
EdgeArray UnitGraph::InEdges(dgl_type_t etype, IdArray vids) const {
957
  const SparseFormat fmt = SelectFormat(CSC_CODE);
958
  const auto ptr = GetFormat(fmt);
959
  if (fmt == SparseFormat::kCSC) {
960
961
962
963
964
    const EdgeArray& ret = ptr->OutEdges(etype, vids);
    return {ret.dst, ret.src, ret.id};
  } else {
    return ptr->InEdges(etype, vids);
  }
965
966
}

Minjie Wang's avatar
Minjie Wang committed
967
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
968
  const SparseFormat fmt = SelectFormat(CSR_CODE);
969
970
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdges(etype, vid);
971
972
}

Minjie Wang's avatar
Minjie Wang committed
973
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, IdArray vids) const {
974
  const SparseFormat fmt = SelectFormat(CSR_CODE);
975
976
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdges(etype, vids);
977
978
}

Minjie Wang's avatar
Minjie Wang committed
979
EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
980
981
  SparseFormat fmt;
  if (order == std::string("eid")) {
982
    fmt = SelectFormat(COO_CODE);
983
  } else if (order.empty()) {
984
    // arbitrary order
985
    fmt = SelectFormat(ALL_CODE);
986
  } else if (order == std::string("srcdst")) {
987
    fmt = SelectFormat(CSR_CODE);
988
989
  } else {
    LOG(FATAL) << "Unsupported order request: " << order;
990
    return {};
991
  }
992
993

  const auto& edges = GetFormat(fmt)->Edges(etype, order);
994
  if (fmt == SparseFormat::kCSC)
995
996
997
    return EdgeArray{edges.dst, edges.src, edges.id};
  else
    return edges;
998
999
}

Minjie Wang's avatar
Minjie Wang committed
1000
uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
1001
  SparseFormat fmt = SelectFormat(CSC_CODE);
1002
  const auto ptr = GetFormat(fmt);
1003
  if (fmt == SparseFormat::kCSC)
1004
1005
1006
    return ptr->OutDegree(etype, vid);
  else
    return ptr->InDegree(etype, vid);
1007
1008
}

Minjie Wang's avatar
Minjie Wang committed
1009
DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
1010
  SparseFormat fmt = SelectFormat(CSC_CODE);
1011
  const auto ptr = GetFormat(fmt);
1012
  if (fmt == SparseFormat::kCSC)
1013
1014
1015
    return ptr->OutDegrees(etype, vids);
  else
    return ptr->InDegrees(etype, vids);
1016
1017
}

Minjie Wang's avatar
Minjie Wang committed
1018
uint64_t UnitGraph::OutDegree(dgl_type_t etype, dgl_id_t vid) const {
1019
  SparseFormat fmt = SelectFormat(CSR_CODE);
1020
1021
  const auto ptr = GetFormat(fmt);
  return ptr->OutDegree(etype, vid);
1022
1023
}

Minjie Wang's avatar
Minjie Wang committed
1024
DegreeArray UnitGraph::OutDegrees(dgl_type_t etype, IdArray vids) const {
1025
  SparseFormat fmt = SelectFormat(CSR_CODE);
1026
1027
  const auto ptr = GetFormat(fmt);
  return ptr->OutDegrees(etype, vids);
1028
1029
}

Minjie Wang's avatar
Minjie Wang committed
1030
DGLIdIters UnitGraph::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
1031
  SparseFormat fmt = SelectFormat(CSR_CODE);
1032
1033
  const auto ptr = GetFormat(fmt);
  return ptr->SuccVec(etype, vid);
1034
1035
}

1036
DGLIdIters32 UnitGraph::SuccVec32(dgl_type_t etype, dgl_id_t vid) const {
1037
  SparseFormat fmt = SelectFormat(CSR_CODE);
1038
1039
1040
1041
1042
  const auto ptr = std::dynamic_pointer_cast<CSR>(GetFormat(fmt));
  CHECK_NOTNULL(ptr);
  return ptr->SuccVec32(etype, vid);
}

Minjie Wang's avatar
Minjie Wang committed
1043
DGLIdIters UnitGraph::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
1044
  SparseFormat fmt = SelectFormat(CSR_CODE);
1045
1046
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdgeVec(etype, vid);
1047
1048
}

Minjie Wang's avatar
Minjie Wang committed
1049
DGLIdIters UnitGraph::PredVec(dgl_type_t etype, dgl_id_t vid) const {
1050
  SparseFormat fmt = SelectFormat(CSC_CODE);
1051
  const auto ptr = GetFormat(fmt);
1052
  if (fmt == SparseFormat::kCSC)
1053
1054
1055
    return ptr->SuccVec(etype, vid);
  else
    return ptr->PredVec(etype, vid);
1056
1057
}

Minjie Wang's avatar
Minjie Wang committed
1058
DGLIdIters UnitGraph::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
1059
  SparseFormat fmt = SelectFormat(CSC_CODE);
1060
  const auto ptr = GetFormat(fmt);
1061
  if (fmt == SparseFormat::kCSC)
1062
1063
1064
    return ptr->OutEdgeVec(etype, vid);
  else
    return ptr->InEdgeVec(etype, vid);
1065
1066
}

Minjie Wang's avatar
Minjie Wang committed
1067
std::vector<IdArray> UnitGraph::GetAdj(
1068
1069
1070
1071
1072
1073
1074
1075
1076
    dgl_type_t etype, bool transpose, const std::string &fmt) const {
  // TODO(minjie): Our current semantics of adjacency matrix is row for dst nodes and col for
  //   src nodes. Therefore, we need to flip the transpose flag. For example, transpose=False
  //   is equal to in edge CSR.
  //   We have this behavior because previously we use framework's SPMM and we don't cache
  //   reverse adj. This is not intuitive and also not consistent with networkx's
  //   to_scipy_sparse_matrix. With the upcoming custom kernel change, we should change the
  //   behavior and make row for src and col for dst.
  if (fmt == std::string("csr")) {
1077
    return !transpose ? GetOutCSR()->GetAdj(etype, false, "csr")
1078
1079
      : GetInCSR()->GetAdj(etype, false, "csr");
  } else if (fmt == std::string("coo")) {
1080
    return GetCOO()->GetAdj(etype, transpose, fmt);
1081
1082
1083
1084
1085
1086
  } else {
    LOG(FATAL) << "unsupported adjacency matrix format: " << fmt;
    return {};
  }
}

Minjie Wang's avatar
Minjie Wang committed
1087
HeteroSubgraph UnitGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
1088
  // We prefer to generate a subgraph from out-csr.
1089
  SparseFormat fmt = SelectFormat(CSR_CODE);
1090
  HeteroSubgraph sg = GetFormat(fmt)->VertexSubgraph(vids);
1091
  HeteroSubgraph ret;
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111

  CSRPtr subcsr = nullptr;
  CSRPtr subcsc = nullptr;
  COOPtr subcoo = nullptr;
  switch (fmt) {
    case SparseFormat::kCSR:
      subcsr = std::dynamic_pointer_cast<CSR>(sg.graph);
      break;
    case SparseFormat::kCSC:
      subcsc = std::dynamic_pointer_cast<CSR>(sg.graph);
      break;
    case SparseFormat::kCOO:
      subcoo = std::dynamic_pointer_cast<COO>(sg.graph);
      break;
    default:
      LOG(FATAL) << "[BUG] unsupported format " << static_cast<int>(fmt);
      return ret;
  }

  ret.graph = HeteroGraphPtr(new UnitGraph(meta_graph(), subcsc, subcsr, subcoo));
1112
1113
1114
1115
1116
  ret.induced_vertices = std::move(sg.induced_vertices);
  ret.induced_edges = std::move(sg.induced_edges);
  return ret;
}

Minjie Wang's avatar
Minjie Wang committed
1117
HeteroSubgraph UnitGraph::EdgeSubgraph(
1118
    const std::vector<IdArray>& eids, bool preserve_nodes) const {
1119
  SparseFormat fmt = SelectFormat(COO_CODE);
1120
  auto sg = GetFormat(fmt)->EdgeSubgraph(eids, preserve_nodes);
1121
  HeteroSubgraph ret;
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141

  CSRPtr subcsr = nullptr;
  CSRPtr subcsc = nullptr;
  COOPtr subcoo = nullptr;
  switch (fmt) {
    case SparseFormat::kCSR:
      subcsr = std::dynamic_pointer_cast<CSR>(sg.graph);
      break;
    case SparseFormat::kCSC:
      subcsc = std::dynamic_pointer_cast<CSR>(sg.graph);
      break;
    case SparseFormat::kCOO:
      subcoo = std::dynamic_pointer_cast<COO>(sg.graph);
      break;
    default:
      LOG(FATAL) << "[BUG] unsupported format " << static_cast<int>(fmt);
      return ret;
  }

  ret.graph = HeteroGraphPtr(new UnitGraph(meta_graph(), subcsc, subcsr, subcoo));
1142
1143
1144
1145
1146
  ret.induced_vertices = std::move(sg.induced_vertices);
  ret.induced_edges = std::move(sg.induced_edges);
  return ret;
}

Minjie Wang's avatar
Minjie Wang committed
1147
HeteroGraphPtr UnitGraph::CreateFromCOO(
1148
1149
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
    IdArray row, IdArray col,
1150
    bool row_sorted, bool col_sorted,
1151
    dgl_format_code_t formats) {
Minjie Wang's avatar
Minjie Wang committed
1152
1153
1154
1155
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(num_src, num_dst);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
1156
1157
  COOPtr coo(new COO(mg, num_src, num_dst, row, col,
      row_sorted, col_sorted));
1158
1159

  return HeteroGraphPtr(
1160
      new UnitGraph(mg, nullptr, nullptr, coo, formats));
1161
1162
}

1163
1164
HeteroGraphPtr UnitGraph::CreateFromCOO(
    int64_t num_vtypes, const aten::COOMatrix& mat,
1165
    dgl_format_code_t formats) {
1166
1167
1168
1169
1170
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(mat.num_rows, mat.num_cols);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  COOPtr coo(new COO(mg, mat));
1171

1172
  return HeteroGraphPtr(
1173
      new UnitGraph(mg, nullptr, nullptr, coo, formats));
1174
1175
}

Minjie Wang's avatar
Minjie Wang committed
1176
1177
HeteroGraphPtr UnitGraph::CreateFromCSR(
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
1178
    IdArray indptr, IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
Minjie Wang's avatar
Minjie Wang committed
1179
1180
1181
1182
1183
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(num_src, num_dst);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  CSRPtr csr(new CSR(mg, num_src, num_dst, indptr, indices, edge_ids));
1184
  return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, formats));
1185
1186
}

1187
1188
HeteroGraphPtr UnitGraph::CreateFromCSR(
    int64_t num_vtypes, const aten::CSRMatrix& mat,
1189
    dgl_format_code_t formats) {
1190
1191
1192
1193
1194
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(mat.num_rows, mat.num_cols);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  CSRPtr csr(new CSR(mg, mat));
1195
  return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, formats));
1196
1197
}

1198
1199
HeteroGraphPtr UnitGraph::CreateFromCSC(
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
1200
    IdArray indptr, IdArray indices, IdArray edge_ids, dgl_format_code_t formats) {
1201
1202
1203
1204
1205
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(num_src, num_dst);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  CSRPtr csc(new CSR(mg, num_src, num_dst, indptr, indices, edge_ids));
1206
  return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, formats));
1207
1208
1209
1210
}

HeteroGraphPtr UnitGraph::CreateFromCSC(
    int64_t num_vtypes, const aten::CSRMatrix& mat,
1211
    dgl_format_code_t formats) {
1212
1213
1214
1215
1216
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(mat.num_rows, mat.num_cols);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  CSRPtr csc(new CSR(mg, mat));
1217
  return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, formats));
1218
1219
}

Minjie Wang's avatar
Minjie Wang committed
1220
HeteroGraphPtr UnitGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
1221
1222
1223
  if (g->NumBits() == bits) {
    return g;
  } else {
Minjie Wang's avatar
Minjie Wang committed
1224
    auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
1225
    CHECK_NOTNULL(bg);
1226
1227
1228
1229
1230
1231
    CSRPtr new_incsr =
      (bg->in_csr_->defined())? CSRPtr(new CSR(bg->in_csr_->AsNumBits(bits))) : nullptr;
    CSRPtr new_outcsr =
      (bg->out_csr_->defined())? CSRPtr(new CSR(bg->out_csr_->AsNumBits(bits))) : nullptr;
    COOPtr new_coo =
      (bg->coo_->defined())? COOPtr(new COO(bg->coo_->AsNumBits(bits))) : nullptr;
1232
    return HeteroGraphPtr(
1233
        new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->formats_));
1234
1235
1236
  }
}

1237
1238
HeteroGraphPtr UnitGraph::CopyTo(HeteroGraphPtr g, const DLContext &ctx,
                                 const DGLStreamHandle &stream) {
1239
1240
  if (ctx == g->Context()) {
    return g;
1241
1242
1243
  } else {
    auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
    CHECK_NOTNULL(bg);
1244
1245
1246
1247
1248
1249
1250
1251
1252
    CSRPtr new_incsr = (bg->in_csr_->defined())
                           ? CSRPtr(new CSR(bg->in_csr_->CopyTo(ctx, stream)))
                           : nullptr;
    CSRPtr new_outcsr = (bg->out_csr_->defined())
                            ? CSRPtr(new CSR(bg->out_csr_->CopyTo(ctx, stream)))
                            : nullptr;
    COOPtr new_coo = (bg->coo_->defined())
                         ? COOPtr(new COO(bg->coo_->CopyTo(ctx, stream)))
                         : nullptr;
1253
    return HeteroGraphPtr(
1254
        new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->formats_));
1255
1256
1257
  }
}

1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
void UnitGraph::InvalidateCSR() {
  this->out_csr_ = CSRPtr(new CSR());
}

void UnitGraph::InvalidateCSC() {
  this->in_csr_ = CSRPtr(new CSR());
}

void UnitGraph::InvalidateCOO() {
  this->coo_ = COOPtr(new COO());
}

1270
UnitGraph::UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo,
1271
                     dgl_format_code_t formats)
Minjie Wang's avatar
Minjie Wang committed
1272
  : BaseHeteroGraph(metagraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
1273
1274
1275
1276
1277
1278
1279
1280
1281
  if (!in_csr_) {
    in_csr_ = CSRPtr(new CSR());
  }
  if (!out_csr_) {
    out_csr_ = CSRPtr(new CSR());
  }
  if (!coo_) {
    coo_ = COOPtr(new COO());
  }
1282
1283
1284
1285
1286
  formats_ = formats;
  dgl_format_code_t created = GetCreatedFormats();
  if ((formats | created) != formats)
    LOG(FATAL) << "Graph created from formats: " << CodeToStr(created) <<
      ", which is not compatible with available formats: " << CodeToStr(formats);
1287
1288
1289
  CHECK(GetAny()) << "At least one graph structure should exist.";
}

1290
1291
1292
1293
1294
1295
1296
HeteroGraphPtr UnitGraph::CreateHomographFrom(
    const aten::CSRMatrix &in_csr,
    const aten::CSRMatrix &out_csr,
    const aten::COOMatrix &coo,
    bool has_in_csr,
    bool has_out_csr,
    bool has_coo,
1297
    dgl_format_code_t formats) {
1298
1299
1300
1301
1302
1303
1304
1305
  auto mg = CreateUnitGraphMetaGraph1();

  CSRPtr in_csr_ptr = nullptr;
  CSRPtr out_csr_ptr = nullptr;
  COOPtr coo_ptr = nullptr;

  if (has_in_csr)
    in_csr_ptr = CSRPtr(new CSR(mg, in_csr));
1306
1307
  else
    in_csr_ptr = CSRPtr(new CSR());
1308
1309
  if (has_out_csr)
    out_csr_ptr = CSRPtr(new CSR(mg, out_csr));
1310
1311
  else
    out_csr_ptr = CSRPtr(new CSR());
1312
1313
  if (has_coo)
    coo_ptr = COOPtr(new COO(mg, coo));
1314
1315
  else
    coo_ptr = COOPtr(new COO());
1316

1317
  return HeteroGraphPtr(new UnitGraph(mg, in_csr_ptr, out_csr_ptr, coo_ptr, formats));
1318
1319
}

1320
1321
UnitGraph::CSRPtr UnitGraph::GetInCSR(bool inplace) const {
  if (inplace)
1322
    if (!(formats_ & CSC_CODE))
1323
1324
      LOG(FATAL) << "The graph have restricted sparse format " <<
        CodeToStr(formats_) << ", cannot create CSC matrix.";
1325
  CSRPtr ret = in_csr_;
1326
1327
  // Prefers converting from COO since it is parallelized.
  // TODO(BarclayII): need benchmarking.
1328
  if (!in_csr_->defined()) {
1329
1330
1331
    if (coo_->defined()) {
      const auto& newadj = aten::COOToCSR(
            aten::COOTranspose(coo_->adj()));
1332

1333
      if (inplace)
1334
1335
1336
        *(const_cast<UnitGraph*>(this)->in_csr_) = CSR(meta_graph(), newadj);
      else
        ret = std::make_shared<CSR>(meta_graph(), newadj);
1337
    } else {
1338
1339
      CHECK(out_csr_->defined()) << "None of CSR, COO exist";
      const auto& newadj = aten::CSRTranspose(out_csr_->adj());
1340

1341
      if (inplace)
1342
1343
1344
        *(const_cast<UnitGraph*>(this)->in_csr_) = CSR(meta_graph(), newadj);
      else
        ret = std::make_shared<CSR>(meta_graph(), newadj);
1345
1346
    }
  }
1347
  return ret;
1348
1349
1350
}

/* !\brief Return out csr. If not exist, transpose the other one.*/
1351
1352
UnitGraph::CSRPtr UnitGraph::GetOutCSR(bool inplace) const {
  if (inplace)
1353
    if (!(formats_ & CSR_CODE))
1354
1355
      LOG(FATAL) << "The graph have restricted sparse format " <<
        CodeToStr(formats_) << ", cannot create CSR matrix.";
1356
  CSRPtr ret = out_csr_;
1357
1358
  // Prefers converting from COO since it is parallelized.
  // TODO(BarclayII): need benchmarking.
1359
  if (!out_csr_->defined()) {
1360
1361
    if (coo_->defined()) {
      const auto& newadj = aten::COOToCSR(coo_->adj());
1362

1363
      if (inplace)
1364
1365
1366
        *(const_cast<UnitGraph*>(this)->out_csr_) = CSR(meta_graph(), newadj);
      else
        ret = std::make_shared<CSR>(meta_graph(), newadj);
1367
    } else {
1368
1369
      CHECK(in_csr_->defined()) << "None of CSR, COO exist";
      const auto& newadj = aten::CSRTranspose(in_csr_->adj());
1370

1371
      if (inplace)
1372
1373
1374
        *(const_cast<UnitGraph*>(this)->out_csr_) = CSR(meta_graph(), newadj);
      else
        ret = std::make_shared<CSR>(meta_graph(), newadj);
1375
1376
    }
  }
1377
  return ret;
1378
1379
1380
}

/* !\brief Return coo. If not exist, create from csr.*/
1381
1382
UnitGraph::COOPtr UnitGraph::GetCOO(bool inplace) const {
  if (inplace)
1383
    if (!(formats_ & COO_CODE))
1384
1385
      LOG(FATAL) << "The graph have restricted sparse format " <<
        CodeToStr(formats_) << ", cannot create COO matrix.";
1386
  COOPtr ret = coo_;
1387
1388
  if (!coo_->defined()) {
    if (in_csr_->defined()) {
1389
      const auto& newadj = aten::COOTranspose(aten::CSRToCOO(in_csr_->adj(), true));
1390

1391
      if (inplace)
1392
1393
1394
        *(const_cast<UnitGraph*>(this)->coo_) = COO(meta_graph(), newadj);
      else
        ret = std::make_shared<COO>(meta_graph(), newadj);
1395
    } else {
1396
      CHECK(out_csr_->defined()) << "Both CSR are missing.";
1397
      const auto& newadj = aten::CSRToCOO(out_csr_->adj(), true);
1398

1399
      if (inplace)
1400
1401
1402
        *(const_cast<UnitGraph*>(this)->coo_) = COO(meta_graph(), newadj);
      else
        ret = std::make_shared<COO>(meta_graph(), newadj);
1403
1404
    }
  }
1405
  return ret;
1406
1407
}

1408
aten::CSRMatrix UnitGraph::GetCSCMatrix(dgl_type_t etype) const {
1409
1410
1411
  return GetInCSR()->adj();
}

1412
aten::CSRMatrix UnitGraph::GetCSRMatrix(dgl_type_t etype) const {
1413
1414
1415
  return GetOutCSR()->adj();
}

1416
aten::COOMatrix UnitGraph::GetCOOMatrix(dgl_type_t etype) const {
1417
1418
1419
  return GetCOO()->adj();
}

Minjie Wang's avatar
Minjie Wang committed
1420
HeteroGraphPtr UnitGraph::GetAny() const {
1421
  if (in_csr_->defined()) {
1422
    return in_csr_;
1423
  } else if (out_csr_->defined()) {
1424
1425
1426
1427
1428
1429
    return out_csr_;
  } else {
    return coo_;
  }
}

1430
dgl_format_code_t UnitGraph::GetCreatedFormats() const {
1431
  dgl_format_code_t ret = 0;
1432
  if (in_csr_->defined())
1433
    ret |= CSC_CODE;
1434
  if (out_csr_->defined())
1435
    ret |= CSR_CODE;
1436
  if (coo_->defined())
1437
    ret |= COO_CODE;
1438
1439
1440
  return ret;
}

1441
1442
1443
1444
dgl_format_code_t UnitGraph::GetAllowedFormats() const {
  return formats_;
}

1445
1446
HeteroGraphPtr UnitGraph::GetFormat(SparseFormat format) const {
  switch (format) {
1447
1448
1449
1450
  case SparseFormat::kCSR:
    return GetOutCSR();
  case SparseFormat::kCSC:
    return GetInCSR();
1451
  default:
1452
    return GetCOO();
1453
1454
1455
  }
}

1456
HeteroGraphPtr UnitGraph::GetGraphInFormat(dgl_format_code_t formats) const {
1457
  if (formats == ALL_CODE)
1458
    return HeteroGraphPtr(
1459
1460
        // TODO(xiangsx) Make it as graph storage.Clone()
        new UnitGraph(meta_graph_,
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
                      (in_csr_->defined())
                          ? CSRPtr(new CSR(*in_csr_))
                          : nullptr,
                      (out_csr_->defined())
                          ? CSRPtr(new CSR(*out_csr_))
                          : nullptr,
                      (coo_->defined())
                          ? COOPtr(new COO(*coo_))
                          : nullptr,
                      formats));
  int64_t num_vtypes = NumVertexTypes();
1472
  if (formats & COO_CODE)
1473
    return CreateFromCOO(num_vtypes, GetCOO(false)->adj(), formats);
1474
  if (formats & CSR_CODE)
1475
1476
1477
1478
1479
1480
1481
1482
1483
    return CreateFromCSR(num_vtypes, GetOutCSR(false)->adj(), formats);
  return CreateFromCSC(num_vtypes, GetInCSR(false)->adj(), formats);
}

SparseFormat UnitGraph::SelectFormat(dgl_format_code_t preferred_formats) const {
  dgl_format_code_t common = preferred_formats & formats_;
  dgl_format_code_t created = GetCreatedFormats();
  if (common & created)
    return DecodeFormat(common & created);
1484
1485
1486
1487
1488

  // NOTE(zihao): hypersparse is currently disabled since many CUDA operators on COO have
  // not been implmented yet.
  // if (coo_->defined() && coo_->IsHypersparse())  // only allow coo for hypersparse graph.
  //   return SparseFormat::kCOO;
1489
1490
1491
  if (common)
    return DecodeFormat(common);
  return DecodeFormat(created);
1492
1493
}

1494
1495
1496
1497
GraphPtr UnitGraph::AsImmutableGraph() const {
  CHECK(NumVertexTypes() == 1) << "not a homogeneous graph";
  dgl::CSRPtr in_csr_ptr = nullptr, out_csr_ptr = nullptr;
  dgl::COOPtr coo_ptr = nullptr;
1498
  if (in_csr_->defined()) {
1499
    aten::CSRMatrix csc = GetCSCMatrix(0);
1500
    in_csr_ptr = dgl::CSRPtr(new dgl::CSR(csc.indptr, csc.indices, csc.data));
1501
  }
1502
  if (out_csr_->defined()) {
1503
    aten::CSRMatrix csr = GetCSRMatrix(0);
1504
    out_csr_ptr = dgl::CSRPtr(new dgl::CSR(csr.indptr, csr.indices, csr.data));
1505
  }
1506
  if (coo_->defined()) {
1507
1508
    aten::COOMatrix coo = GetCOOMatrix(0);
    if (!COOHasData(coo)) {
1509
      coo_ptr = dgl::COOPtr(new dgl::COO(NumVertices(0), coo.row, coo.col));
1510
1511
1512
    } else {
      IdArray new_src = Scatter(coo.row, coo.data);
      IdArray new_dst = Scatter(coo.col, coo.data);
1513
      coo_ptr = dgl::COOPtr(new dgl::COO(NumVertices(0), new_src, new_dst));
1514
1515
1516
1517
1518
    }
  }
  return GraphPtr(new dgl::ImmutableGraph(in_csr_ptr, out_csr_ptr, coo_ptr));
}

1519
1520
HeteroGraphPtr UnitGraph::LineGraph(bool backtracking) const {
  // TODO(xiangsx) currently we only support homogeneous graph
1521
  auto fmt = SelectFormat(ALL_CODE);
1522
1523
  switch (fmt) {
    case SparseFormat::kCOO: {
1524
      return CreateFromCOO(1, aten::COOLineGraph(coo_->adj(), backtracking));
1525
1526
1527
1528
    }
    case SparseFormat::kCSR: {
      const aten::CSRMatrix csr = GetCSRMatrix(0);
      const aten::COOMatrix coo = aten::COOLineGraph(aten::CSRToCOO(csr, true), backtracking);
1529
      return CreateFromCOO(1, coo);
1530
1531
1532
1533
1534
    }
    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);
1535
      return CreateFromCOO(1, coo);
1536
1537
1538
1539
1540
1541
1542
1543
    }
    default:
      LOG(FATAL) << "None of CSC, CSR, COO exist";
      break;
  }
  return nullptr;
}

1544
1545
1546
1547
1548
1549
constexpr uint64_t kDGLSerialize_UnitGraphMagic = 0xDD2E60F0F6B4A127;

bool UnitGraph::Load(dmlc::Stream* fs) {
  uint64_t magicNum;
  CHECK(fs->Read(&magicNum)) << "Invalid Magic Number";
  CHECK_EQ(magicNum, kDGLSerialize_UnitGraphMagic) << "Invalid UnitGraph Data";
1550

1551
  int64_t save_format_code, formats_code;
1552
  CHECK(fs->Read(&save_format_code)) << "Invalid format";
1553
  CHECK(fs->Read(&formats_code)) << "Invalid format";
1554
  auto save_format = static_cast<SparseFormat>(save_format_code);
1555
1556
1557
1558
1559
1560
  if (formats_code >> 32) {
    formats_ = static_cast<dgl_format_code_t>(0xffffffff & formats_code);
  } else {
    // NOTE(zihao): to be compatible with old formats.
    switch (formats_code & 0xffffffff) {
    case 0:
1561
      formats_ = ALL_CODE;
1562
1563
      break;
    case 1:
1564
      formats_ = COO_CODE;
1565
1566
      break;
    case 2:
1567
      formats_ = CSR_CODE;
1568
1569
      break;
    case 3:
1570
      formats_ = CSC_CODE;
1571
1572
1573
1574
1575
1576
      break;
    default:
      LOG(FATAL) << "Load graph failed, formats code " << formats_code <<
        "not recognized.";
    }
  }
1577

1578
  switch (save_format) {
1579
    case SparseFormat::kCOO:
1580
1581
      fs->Read(&coo_);
      break;
1582
    case SparseFormat::kCSR:
1583
1584
      fs->Read(&out_csr_);
      break;
1585
    case SparseFormat::kCSC:
1586
1587
1588
1589
1590
1591
1592
      fs->Read(&in_csr_);
      break;
    default:
      LOG(FATAL) << "unsupported format code";
      break;
  }

1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
  if (!in_csr_) {
    in_csr_ = CSRPtr(new CSR());
  }
  if (!out_csr_) {
    out_csr_ = CSRPtr(new CSR());
  }
  if (!coo_) {
    coo_ = COOPtr(new COO());
  }

1603
1604
  meta_graph_ = GetAny()->meta_graph();

1605
1606
1607
  return true;
}

1608

1609
1610
void UnitGraph::Save(dmlc::Stream* fs) const {
  fs->Write(kDGLSerialize_UnitGraphMagic);
1611
1612
  // Didn't write UnitGraph::meta_graph_, since it's included in the underlying
  // sparse matrix
1613
  auto avail_fmt = SelectFormat(ALL_CODE);
1614
  fs->Write(static_cast<int64_t>(avail_fmt));
1615
  fs->Write(static_cast<int64_t>(formats_ | 0x100000000));
1616
  switch (avail_fmt) {
1617
    case SparseFormat::kCOO:
1618
1619
      fs->Write(GetCOO());
      break;
1620
    case SparseFormat::kCSR:
1621
1622
      fs->Write(GetOutCSR());
      break;
1623
    case SparseFormat::kCSC:
1624
1625
1626
1627
1628
1629
      fs->Write(GetInCSR());
      break;
    default:
      LOG(FATAL) << "unsupported format code";
      break;
  }
1630
1631
}

1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
UnitGraphPtr UnitGraph::Reverse() const {
  CSRPtr new_incsr = out_csr_, new_outcsr = in_csr_;
  COOPtr new_coo = nullptr;
  if (coo_->defined()) {
    new_coo = COOPtr(new COO(coo_->meta_graph(), aten::COOTranspose(coo_->adj())));
  }

  return UnitGraphPtr(new UnitGraph(meta_graph(), new_incsr, new_outcsr, new_coo));
}

1642
1643
1644
1645
1646
1647
1648
std::tuple<UnitGraphPtr, IdArray, IdArray>
UnitGraph::ToSimple() const {
  CSRPtr new_incsr = nullptr, new_outcsr = nullptr;
  COOPtr new_coo = nullptr;
  IdArray count;
  IdArray edge_map;

1649
  auto avail_fmt = SelectFormat(ALL_CODE);
1650
1651
  switch (avail_fmt) {
    case SparseFormat::kCOO: {
1652
      auto ret = aten::COOToSimple(GetCOO()->adj());
1653
1654
      count = std::get<1>(ret);
      edge_map = std::get<2>(ret);
1655
      new_coo = COOPtr(new COO(meta_graph(), std::get<0>(ret)));
1656
1657
1658
      break;
    }
    case SparseFormat::kCSR: {
1659
      auto ret = aten::CSRToSimple(GetOutCSR()->adj());
1660
1661
      count = std::get<1>(ret);
      edge_map = std::get<2>(ret);
1662
      new_outcsr = CSRPtr(new CSR(meta_graph(), std::get<0>(ret)));
1663
1664
1665
      break;
    }
    case SparseFormat::kCSC: {
1666
      auto ret = aten::CSRToSimple(GetInCSR()->adj());
1667
1668
      count = std::get<1>(ret);
      edge_map = std::get<2>(ret);
1669
      new_incsr = CSRPtr(new CSR(meta_graph(), std::get<0>(ret)));
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
      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);
}

1682
}  // namespace dgl