unit_graph.cc 48.3 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:
Minjie Wang's avatar
Minjie Wang committed
64
65
  COO(GraphPtr metagraph, int64_t num_src, int64_t num_dst, IdArray src, IdArray dst)
    : BaseHeteroGraph(metagraph) {
66
67
68
    CHECK(aten::IsValidIdArray(src));
    CHECK(aten::IsValidIdArray(dst));
    CHECK_EQ(src->shape[0], dst->shape[0]) << "Input arrays should have the same length.";
69
70
    adj_ = aten::COOMatrix{num_src, num_dst, src, dst};
  }
71

72
73
74
75
  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.
76
    CHECK(!COOHasData(coo)) << "[BUG] COO should not contain data.";
77
    adj_.data = aten::NullArray();
78
  }
79

Minjie Wang's avatar
Minjie Wang committed
80
81
  inline dgl_type_t SrcType() const {
    return 0;
82
  }
Minjie Wang's avatar
Minjie Wang committed
83
84
85
86
87
88
89

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

  inline dgl_type_t EdgeType() const {
    return 0;
90
91
92
  }

  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
Minjie Wang's avatar
Minjie Wang committed
93
    LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
94
95
96
97
98
      << "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
99
    LOG(FATAL) << "UnitGraph graph is not mutable.";
100
101
102
  }

  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
Minjie Wang's avatar
Minjie Wang committed
103
    LOG(FATAL) << "UnitGraph graph is not mutable.";
104
105
106
  }

  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
Minjie Wang's avatar
Minjie Wang committed
107
    LOG(FATAL) << "UnitGraph graph is not mutable.";
108
109
110
  }

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

114
115
116
117
  DLDataType DataType() const override {
    return adj_.row->dtype;
  }

118
119
120
121
122
123
124
125
  DLContext Context() const override {
    return adj_.row->ctx;
  }

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

126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
  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;
  }

  COO CopyTo(const DLContext& ctx) const {
    if (Context() == ctx)
      return *this;
141
    return COO(meta_graph_, adj_.CopyTo(ctx));
142
143
  }

144
  bool IsMultigraph() const override {
145
    return aten::COOHasDuplicate(adj_);
146
147
148
149
150
151
152
  }

  bool IsReadonly() const override {
    return true;
  }

  uint64_t NumVertices(dgl_type_t vtype) const override {
Minjie Wang's avatar
Minjie Wang committed
153
    if (vtype == SrcType()) {
154
      return adj_.num_rows;
Minjie Wang's avatar
Minjie Wang committed
155
    } else if (vtype == DstType()) {
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
      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 {
177
178
179
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
    return aten::COOIsNonZero(adj_, src, dst);
180
181
182
  }

  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
183
184
185
    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);
186
187
188
  }

  IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const override {
189
190
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
    return aten::COOGetRowDataAndIndices(aten::COOTranspose(adj_), dst).second;
191
192
193
  }

  IdArray Successors(dgl_type_t etype, dgl_id_t src) const override {
194
195
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    return aten::COOGetRowDataAndIndices(adj_, src).second;
196
197
198
  }

  IdArray EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const override {
199
200
201
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
    return aten::COOGetData(adj_, src, dst);
202
203
204
  }

  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
205
206
207
208
    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]};
209
210
211
212
  }

  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;
213
214
    const dgl_id_t src = aten::IndexSelect<int64_t>(adj_.row, eid);
    const dgl_id_t dst = aten::IndexSelect<int64_t>(adj_.col, eid);
215
216
217
218
    return std::pair<dgl_id_t, dgl_id_t>(src, dst);
  }

  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
219
    CHECK(aten::IsValidIdArray(eids)) << "Invalid edge id array";
220
221
222
223
224
225
    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 {
226
227
228
229
230
    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};
231
232
233
  }

  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
234
235
236
237
    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};
238
239
240
  }

  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
241
242
243
244
    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};
245
246
247
  }

  EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const override {
248
249
250
251
    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};
252
253
254
255
256
257
258
259
260
261
262
  }

  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 {
263
264
    CHECK(HasVertex(DstType(), vid)) << "Invalid dst vertex id: " << vid;
    return aten::COOGetRowNNZ(aten::COOTranspose(adj_), vid);
265
266
267
  }

  DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const override {
268
269
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    return aten::COOGetRowNNZ(aten::COOTranspose(adj_), vids);
270
271
272
  }

  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
273
274
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
    return aten::COOGetRowNNZ(adj_, vid);
275
276
277
  }

  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
278
279
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
    return aten::COOGetRowNNZ(adj_, vids);
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
  }

  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)};
    }
  }

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
  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();
  }

  SparseFormat SelectFormat(dgl_type_t etype, SparseFormat preferred_format) const override {
    LOG(FATAL) << "Not enabled for COO graph";
328
    return SparseFormat::kAny;
329
330
  }

331
332
333
334
335
336
337
338
339
340
  std::string GetRestrictFormat() const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return std::string("");
  }

  dgl_format_code_t GetFormatInUse() const override {
    LOG(FATAL) << "Not enabled for COO graph";
    return 0;
  }

341
  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
342
343
344
345
346
347
348
349
350
351
352
353
    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;
354
355
356
357
358
359
360
361
362
363
364
365
366
367
  }

  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
368
          meta_graph(), new_nsrc, new_ndst, new_src, new_dst);
369
370
371
372
      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
373
374
375
376
      subg.induced_vertices.emplace_back(
          aten::Range(0, NumVertices(SrcType()), NumBits(), Context()));
      subg.induced_vertices.emplace_back(
          aten::Range(0, NumVertices(DstType()), NumBits(), Context()));
377
      subg.graph = std::make_shared<COO>(
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
378
          meta_graph(), NumVertices(SrcType()), NumVertices(DstType()), new_src, new_dst);
379
380
381
382
383
      subg.induced_edges = eids;
    }
    return subg;
  }

384
385
386
387
388
  HeteroGraphPtr GetGraphInFormat(SparseFormat restrict_format) const override {
    LOG(FATAL) << "Not enabled for COO graph.";
    return nullptr;
  }

389
390
391
392
  aten::COOMatrix adj() const {
    return adj_;
  }

393
394
395
396
397
398
399
400
401
  /*!
   * \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);
  }

402
403
404
405
406
407
408
409
410
411
412
413
414
  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_);
  }

415
 private:
416
417
418
419
  friend class Serializer;

  COO() {}

420
421
422
423
424
425
426
427
428
429
430
  /*! \brief internal adjacency matrix. Data array is empty */
  aten::COOMatrix adj_;
};

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

/*! \brief CSR graph */
Minjie Wang's avatar
Minjie Wang committed
431
class UnitGraph::CSR : public BaseHeteroGraph {
432
 public:
Minjie Wang's avatar
Minjie Wang committed
433
  CSR(GraphPtr metagraph, int64_t num_src, int64_t num_dst,
434
      IdArray indptr, IdArray indices, IdArray edge_ids)
Minjie Wang's avatar
Minjie Wang committed
435
    : BaseHeteroGraph(metagraph) {
436
437
438
439
440
    CHECK(aten::IsValidIdArray(indptr));
    CHECK(aten::IsValidIdArray(indices));
    CHECK(aten::IsValidIdArray(edge_ids));
    CHECK_EQ(indices->shape[0], edge_ids->shape[0])
      << "indices and edge id arrays should have the same length";
441

442
443
444
    adj_ = aten::CSRMatrix{num_src, num_dst, indptr, indices, edge_ids};
  }

445
  CSR(GraphPtr metagraph, const aten::CSRMatrix& csr)
Da Zheng's avatar
Da Zheng committed
446
447
    : BaseHeteroGraph(metagraph), adj_(csr) {
  }
448

Minjie Wang's avatar
Minjie Wang committed
449
450
  inline dgl_type_t SrcType() const {
    return 0;
451
  }
Minjie Wang's avatar
Minjie Wang committed
452
453
454
455
456
457
458

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

  inline dgl_type_t EdgeType() const {
    return 0;
459
460
461
  }

  HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const override {
Minjie Wang's avatar
Minjie Wang committed
462
    LOG(FATAL) << "The method shouldn't be called for UnitGraph graph. "
463
464
465
466
467
      << "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
468
    LOG(FATAL) << "UnitGraph graph is not mutable.";
469
470
471
  }

  void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) override {
Minjie Wang's avatar
Minjie Wang committed
472
    LOG(FATAL) << "UnitGraph graph is not mutable.";
473
474
475
  }

  void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) override {
Minjie Wang's avatar
Minjie Wang committed
476
    LOG(FATAL) << "UnitGraph graph is not mutable.";
477
478
479
  }

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

483
484
485
486
  DLDataType DataType() const override {
    return adj_.indices->dtype;
  }

487
488
489
490
491
492
493
494
  DLContext Context() const override {
    return adj_.indices->ctx;
  }

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

495
496
497
498
499
  CSR AsNumBits(uint8_t bits) const {
    if (NumBits() == bits) {
      return *this;
    } else {
      CSR ret(
Minjie Wang's avatar
Minjie Wang committed
500
          meta_graph_,
501
502
503
504
505
506
507
508
509
510
511
512
          adj_.num_rows, adj_.num_cols,
          aten::AsNumBits(adj_.indptr, bits),
          aten::AsNumBits(adj_.indices, bits),
          aten::AsNumBits(adj_.data, bits));
      return ret;
    }
  }

  CSR CopyTo(const DLContext& ctx) const {
    if (Context() == ctx) {
      return *this;
    } else {
513
      return CSR(meta_graph_, adj_.CopyTo(ctx));
514
515
516
    }
  }

517
  bool IsMultigraph() const override {
518
    return aten::CSRHasDuplicate(adj_);
519
520
521
522
523
524
525
  }

  bool IsReadonly() const override {
    return true;
  }

  uint64_t NumVertices(dgl_type_t vtype) const override {
Minjie Wang's avatar
Minjie Wang committed
526
    if (vtype == SrcType()) {
527
      return adj_.num_rows;
Minjie Wang's avatar
Minjie Wang committed
528
    } else if (vtype == DstType()) {
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
      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
550
551
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
552
553
554
555
    return aten::CSRIsNonZero(adj_, src, dst);
  }

  BoolArray HasEdgesBetween(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const override {
556
557
    CHECK(aten::IsValidIdArray(src_ids)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst_ids)) << "Invalid vertex id array.";
558
559
560
561
562
563
564
565
566
    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
567
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
568
569
570
571
    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
572
573
    CHECK(HasVertex(SrcType(), src)) << "Invalid src vertex id: " << src;
    CHECK(HasVertex(DstType(), dst)) << "Invalid dst vertex id: " << dst;
574
575
576
577
    return aten::CSRGetData(adj_, src, dst);
  }

  EdgeArray EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const override {
578
579
    CHECK(aten::IsValidIdArray(src)) << "Invalid vertex id array.";
    CHECK(aten::IsValidIdArray(dst)) << "Invalid vertex id array.";
580
581
582
583
584
    const auto& arrs = aten::CSRGetDataAndIndices(adj_, src, dst);
    return EdgeArray{arrs[0], arrs[1], arrs[2]};
  }

  std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_type_t etype, dgl_id_t eid) const override {
585
    LOG(FATAL) << "Not enabled for CSR graph.";
586
587
588
589
    return {};
  }

  EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const override {
590
    LOG(FATAL) << "Not enabled for CSR graph.";
591
592
593
594
    return {};
  }

  EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const override {
595
    LOG(FATAL) << "Not enabled for CSR graph.";
596
597
598
599
    return {};
  }

  EdgeArray InEdges(dgl_type_t etype, IdArray vids) const override {
600
    LOG(FATAL) << "Not enabled for CSR graph.";
601
602
603
604
    return {};
  }

  EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const override {
Minjie Wang's avatar
Minjie Wang committed
605
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
606
607
608
609
610
611
612
    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 {
613
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    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 << "\".";
    const auto& coo = aten::CSRToCOO(adj_, false);
    return EdgeArray{coo.row, coo.col, coo.data};
  }

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

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

  uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const override {
Minjie Wang's avatar
Minjie Wang committed
641
    CHECK(HasVertex(SrcType(), vid)) << "Invalid src vertex id: " << vid;
642
643
644
645
    return aten::CSRGetRowNNZ(adj_, vid);
  }

  DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const override {
646
    CHECK(aten::IsValidIdArray(vids)) << "Invalid vertex id array.";
647
648
649
650
651
652
    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.
653
    CHECK_EQ(NumBits(), 64);
654
655
656
657
658
659
660
    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);
  }

661
662
663
664
665
666
667
668
669
670
  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);
  }

671
672
673
  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.
674
    CHECK_EQ(NumBits(), 64);
675
676
677
678
679
680
681
682
    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 {
683
    LOG(FATAL) << "Not enabled for CSR graph.";
684
685
686
687
    return {};
  }

  DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const override {
688
    LOG(FATAL) << "Not enabled for CSR graph.";
689
690
691
692
693
694
695
696
697
    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};
  }

698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
  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_;
  }

  SparseFormat SelectFormat(dgl_type_t etype, SparseFormat preferred_format) const override {
    LOG(FATAL) << "Not enabled for CSR graph";
714
    return SparseFormat::kAny;
715
716
  }

717
718
719
720
721
722
723
724
725
726
  std::string GetRestrictFormat() const override {
    LOG(FATAL) << "Not enabled for CSR graph";
    return std::string("");
  }

  dgl_format_code_t GetFormatInUse() const override {
    LOG(FATAL) << "Not enabled for CSR graph";
    return 0;
  }

727
  HeteroSubgraph VertexSubgraph(const std::vector<IdArray>& vids) const override {
Minjie Wang's avatar
Minjie Wang committed
728
729
730
731
    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.";
732
    HeteroSubgraph subg;
Minjie Wang's avatar
Minjie Wang committed
733
    const auto& submat = aten::CSRSliceMatrix(adj_, srcvids, dstvids);
734
    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
Minjie Wang's avatar
Minjie Wang committed
735
    subg.graph = std::make_shared<CSR>(meta_graph(), submat.num_rows, submat.num_cols,
736
737
738
739
740
741
742
743
        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 {
744
    LOG(FATAL) << "Not enabled for CSR graph.";
745
746
747
    return {};
  }

748
749
750
751
752
  HeteroGraphPtr GetGraphInFormat(SparseFormat restrict_format) const override {
    LOG(FATAL) << "Not enabled for CSR graph.";
    return nullptr;
  }

753
754
755
756
  aten::CSRMatrix adj() const {
    return adj_;
  }

757
758
759
760
761
762
763
764
765
766
767
768
769
770
  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_);
  }


771
 private:
772
773
774
775
  friend class Serializer;

  CSR() {};

776
777
778
779
780
781
  /*! \brief internal adjacency matrix. Data array stores edge ids */
  aten::CSRMatrix adj_;
};

//////////////////////////////////////////////////////////
//
Minjie Wang's avatar
Minjie Wang committed
782
// unit graph implementation
783
784
785
//
//////////////////////////////////////////////////////////

786
787
788
789
DLDataType UnitGraph::DataType() const {
  return GetAny()->DataType();
}

Minjie Wang's avatar
Minjie Wang committed
790
DLContext UnitGraph::Context() const {
791
792
793
  return GetAny()->Context();
}

Minjie Wang's avatar
Minjie Wang committed
794
uint8_t UnitGraph::NumBits() const {
795
796
797
  return GetAny()->NumBits();
}

Minjie Wang's avatar
Minjie Wang committed
798
bool UnitGraph::IsMultigraph() const {
799
800
801
  return GetAny()->IsMultigraph();
}

Minjie Wang's avatar
Minjie Wang committed
802
uint64_t UnitGraph::NumVertices(dgl_type_t vtype) const {
803
  const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
804
805
806
  const auto ptr = GetFormat(fmt);
  // TODO(BarclayII): we have a lot of special handling for CSC.
  // Need to have a UnitGraph::CSC backend instead.
807
  if (fmt == SparseFormat::kCSC)
Minjie Wang's avatar
Minjie Wang committed
808
    vtype = (vtype == SrcType()) ? DstType() : SrcType();
809
  return ptr->NumVertices(vtype);
810
811
}

Minjie Wang's avatar
Minjie Wang committed
812
uint64_t UnitGraph::NumEdges(dgl_type_t etype) const {
813
814
815
  return GetAny()->NumEdges(etype);
}

Minjie Wang's avatar
Minjie Wang committed
816
bool UnitGraph::HasVertex(dgl_type_t vtype, dgl_id_t vid) const {
817
  const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
818
  const auto ptr = GetFormat(fmt);
819
  if (fmt == SparseFormat::kCSC)
Minjie Wang's avatar
Minjie Wang committed
820
    vtype = (vtype == SrcType()) ? DstType() : SrcType();
821
  return ptr->HasVertex(vtype, vid);
822
823
}

Minjie Wang's avatar
Minjie Wang committed
824
BoolArray UnitGraph::HasVertices(dgl_type_t vtype, IdArray vids) const {
825
  CHECK(aten::IsValidIdArray(vids)) << "Invalid id array input";
826
827
828
  return aten::LT(vids, NumVertices(vtype));
}

Minjie Wang's avatar
Minjie Wang committed
829
bool UnitGraph::HasEdgeBetween(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
830
  const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
831
  const auto ptr = GetFormat(fmt);
832
  if (fmt == SparseFormat::kCSC)
833
834
835
    return ptr->HasEdgeBetween(etype, dst, src);
  else
    return ptr->HasEdgeBetween(etype, src, dst);
836
837
}

Minjie Wang's avatar
Minjie Wang committed
838
BoolArray UnitGraph::HasEdgesBetween(
839
    dgl_type_t etype, IdArray src, IdArray dst) const {
840
  const SparseFormat fmt = SelectFormat(SparseFormat::kAny);
841
  const auto ptr = GetFormat(fmt);
842
  if (fmt == SparseFormat::kCSC)
843
844
845
    return ptr->HasEdgesBetween(etype, dst, src);
  else
    return ptr->HasEdgesBetween(etype, src, dst);
846
847
}

Minjie Wang's avatar
Minjie Wang committed
848
IdArray UnitGraph::Predecessors(dgl_type_t etype, dgl_id_t dst) const {
849
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
850
  const auto ptr = GetFormat(fmt);
851
  if (fmt == SparseFormat::kCSC)
852
853
854
    return ptr->Successors(etype, dst);
  else
    return ptr->Predecessors(etype, dst);
855
856
}

Minjie Wang's avatar
Minjie Wang committed
857
IdArray UnitGraph::Successors(dgl_type_t etype, dgl_id_t src) const {
858
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
859
860
  const auto ptr = GetFormat(fmt);
  return ptr->Successors(etype, src);
861
862
}

Minjie Wang's avatar
Minjie Wang committed
863
IdArray UnitGraph::EdgeId(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const {
864
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
865
  const auto ptr = GetFormat(fmt);
866
  if (fmt == SparseFormat::kCSC)
867
868
869
    return ptr->EdgeId(etype, dst, src);
  else
    return ptr->EdgeId(etype, src, dst);
870
871
}

Minjie Wang's avatar
Minjie Wang committed
872
EdgeArray UnitGraph::EdgeIds(dgl_type_t etype, IdArray src, IdArray dst) const {
873
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
874
  const auto ptr = GetFormat(fmt);
875
  if (fmt == SparseFormat::kCSC) {
876
    EdgeArray edges = ptr->EdgeIds(etype, dst, src);
877
878
    return EdgeArray{edges.dst, edges.src, edges.id};
  } else {
879
    return ptr->EdgeIds(etype, src, dst);
880
881
882
  }
}

Minjie Wang's avatar
Minjie Wang committed
883
std::pair<dgl_id_t, dgl_id_t> UnitGraph::FindEdge(dgl_type_t etype, dgl_id_t eid) const {
884
  const SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
885
886
  const auto ptr = GetFormat(fmt);
  return ptr->FindEdge(etype, eid);
887
888
}

Minjie Wang's avatar
Minjie Wang committed
889
EdgeArray UnitGraph::FindEdges(dgl_type_t etype, IdArray eids) const {
890
  const SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
891
892
  const auto ptr = GetFormat(fmt);
  return ptr->FindEdges(etype, eids);
893
894
}

Minjie Wang's avatar
Minjie Wang committed
895
EdgeArray UnitGraph::InEdges(dgl_type_t etype, dgl_id_t vid) const {
896
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
897
  const auto ptr = GetFormat(fmt);
898
  if (fmt == SparseFormat::kCSC) {
899
900
901
902
903
    const EdgeArray& ret = ptr->OutEdges(etype, vid);
    return {ret.dst, ret.src, ret.id};
  } else {
    return ptr->InEdges(etype, vid);
  }
904
905
}

Minjie Wang's avatar
Minjie Wang committed
906
EdgeArray UnitGraph::InEdges(dgl_type_t etype, IdArray vids) const {
907
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
908
  const auto ptr = GetFormat(fmt);
909
  if (fmt == SparseFormat::kCSC) {
910
911
912
913
914
    const EdgeArray& ret = ptr->OutEdges(etype, vids);
    return {ret.dst, ret.src, ret.id};
  } else {
    return ptr->InEdges(etype, vids);
  }
915
916
}

Minjie Wang's avatar
Minjie Wang committed
917
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, dgl_id_t vid) const {
918
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
919
920
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdges(etype, vid);
921
922
}

Minjie Wang's avatar
Minjie Wang committed
923
EdgeArray UnitGraph::OutEdges(dgl_type_t etype, IdArray vids) const {
924
  const SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
925
926
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdges(etype, vids);
927
928
}

Minjie Wang's avatar
Minjie Wang committed
929
EdgeArray UnitGraph::Edges(dgl_type_t etype, const std::string &order) const {
930
931
  SparseFormat fmt;
  if (order == std::string("eid")) {
932
    fmt = SelectFormat(SparseFormat::kCOO);
933
  } else if (order.empty()) {
934
    // arbitrary order
935
    fmt = SelectFormat(SparseFormat::kAny);
936
  } else if (order == std::string("srcdst")) {
937
    fmt = SelectFormat(SparseFormat::kCSR);
938
939
  } else {
    LOG(FATAL) << "Unsupported order request: " << order;
940
    return {};
941
  }
942
943

  const auto& edges = GetFormat(fmt)->Edges(etype, order);
944
  if (fmt == SparseFormat::kCSC)
945
946
947
    return EdgeArray{edges.dst, edges.src, edges.id};
  else
    return edges;
948
949
}

Minjie Wang's avatar
Minjie Wang committed
950
uint64_t UnitGraph::InDegree(dgl_type_t etype, dgl_id_t vid) const {
951
  SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
952
  const auto ptr = GetFormat(fmt);
953
  if (fmt == SparseFormat::kCSC)
954
955
956
    return ptr->OutDegree(etype, vid);
  else
    return ptr->InDegree(etype, vid);
957
958
}

Minjie Wang's avatar
Minjie Wang committed
959
DegreeArray UnitGraph::InDegrees(dgl_type_t etype, IdArray vids) const {
960
  SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
961
  const auto ptr = GetFormat(fmt);
962
  if (fmt == SparseFormat::kCSC)
963
964
965
    return ptr->OutDegrees(etype, vids);
  else
    return ptr->InDegrees(etype, vids);
966
967
}

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

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

Minjie Wang's avatar
Minjie Wang committed
980
DGLIdIters UnitGraph::SuccVec(dgl_type_t etype, dgl_id_t vid) const {
981
  SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
982
983
  const auto ptr = GetFormat(fmt);
  return ptr->SuccVec(etype, vid);
984
985
}

986
987
988
989
990
991
992
DGLIdIters32 UnitGraph::SuccVec32(dgl_type_t etype, dgl_id_t vid) const {
  SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
  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
993
DGLIdIters UnitGraph::OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
994
  SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
995
996
  const auto ptr = GetFormat(fmt);
  return ptr->OutEdgeVec(etype, vid);
997
998
}

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

Minjie Wang's avatar
Minjie Wang committed
1008
DGLIdIters UnitGraph::InEdgeVec(dgl_type_t etype, dgl_id_t vid) const {
1009
  SparseFormat fmt = SelectFormat(SparseFormat::kCSC);
1010
  const auto ptr = GetFormat(fmt);
1011
  if (fmt == SparseFormat::kCSC)
1012
1013
1014
    return ptr->OutEdgeVec(etype, vid);
  else
    return ptr->InEdgeVec(etype, vid);
1015
1016
}

Minjie Wang's avatar
Minjie Wang committed
1017
std::vector<IdArray> UnitGraph::GetAdj(
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
    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")) {
    return transpose? GetOutCSR()->GetAdj(etype, false, "csr")
      : GetInCSR()->GetAdj(etype, false, "csr");
  } else if (fmt == std::string("coo")) {
    return GetCOO()->GetAdj(etype, !transpose, fmt);
  } else {
    LOG(FATAL) << "unsupported adjacency matrix format: " << fmt;
    return {};
  }
}

Minjie Wang's avatar
Minjie Wang committed
1037
HeteroSubgraph UnitGraph::VertexSubgraph(const std::vector<IdArray>& vids) const {
1038
  // We prefer to generate a subgraph from out-csr.
1039
  SparseFormat fmt = SelectFormat(SparseFormat::kCSR);
1040
  HeteroSubgraph sg = GetFormat(fmt)->VertexSubgraph(vids);
1041
  HeteroSubgraph ret;
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061

  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));
1062
1063
1064
1065
1066
  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
1067
HeteroSubgraph UnitGraph::EdgeSubgraph(
1068
    const std::vector<IdArray>& eids, bool preserve_nodes) const {
1069
  SparseFormat fmt = SelectFormat(SparseFormat::kCOO);
1070
  auto sg = GetFormat(fmt)->EdgeSubgraph(eids, preserve_nodes);
1071
  HeteroSubgraph ret;
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091

  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));
1092
1093
1094
1095
1096
  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
1097
HeteroGraphPtr UnitGraph::CreateFromCOO(
1098
1099
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
    IdArray row, IdArray col,
1100
    SparseFormat restrict_format) {
Minjie Wang's avatar
Minjie Wang committed
1101
1102
1103
1104
1105
  CHECK(num_vtypes == 1 || num_vtypes == 2);
  if (num_vtypes == 1)
    CHECK_EQ(num_src, num_dst);
  auto mg = CreateUnitGraphMetaGraph(num_vtypes);
  COOPtr coo(new COO(mg, num_src, num_dst, row, col));
1106
1107
1108

  return HeteroGraphPtr(
      new UnitGraph(mg, nullptr, nullptr, coo, restrict_format));
1109
1110
}

1111
1112
1113
1114
1115
1116
1117
1118
HeteroGraphPtr UnitGraph::CreateFromCOO(
    int64_t num_vtypes, const aten::COOMatrix& mat,
    SparseFormat restrict_format) {
  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));
1119

1120
1121
1122
1123
  return HeteroGraphPtr(
      new UnitGraph(mg, nullptr, nullptr, coo, restrict_format));
}

Minjie Wang's avatar
Minjie Wang committed
1124
1125
HeteroGraphPtr UnitGraph::CreateFromCSR(
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
1126
    IdArray indptr, IdArray indices, IdArray edge_ids, SparseFormat restrict_format) {
Minjie Wang's avatar
Minjie Wang committed
1127
1128
1129
1130
1131
  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));
1132
  return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, restrict_format));
1133
1134
}

1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
HeteroGraphPtr UnitGraph::CreateFromCSR(
    int64_t num_vtypes, const aten::CSRMatrix& mat,
    SparseFormat restrict_format) {
  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));
  return HeteroGraphPtr(new UnitGraph(mg, nullptr, csr, nullptr, restrict_format));
}

1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
HeteroGraphPtr UnitGraph::CreateFromCSC(
    int64_t num_vtypes, int64_t num_src, int64_t num_dst,
    IdArray indptr, IdArray indices, IdArray edge_ids, SparseFormat restrict_format) {
  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));
  return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, restrict_format));
}

HeteroGraphPtr UnitGraph::CreateFromCSC(
    int64_t num_vtypes, const aten::CSRMatrix& mat,
    SparseFormat restrict_format) {
  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));
  return HeteroGraphPtr(new UnitGraph(mg, csc, nullptr, nullptr, restrict_format));
}

Minjie Wang's avatar
Minjie Wang committed
1168
HeteroGraphPtr UnitGraph::AsNumBits(HeteroGraphPtr g, uint8_t bits) {
1169
1170
1171
  if (g->NumBits() == bits) {
    return g;
  } else {
Minjie Wang's avatar
Minjie Wang committed
1172
    auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
1173
    CHECK_NOTNULL(bg);
1174
1175
1176
    CSRPtr new_incsr = (bg->in_csr_)? CSRPtr(new CSR(bg->in_csr_->AsNumBits(bits))) : nullptr;
    CSRPtr new_outcsr = (bg->out_csr_)? CSRPtr(new CSR(bg->out_csr_->AsNumBits(bits))) : nullptr;
    COOPtr new_coo = (bg->coo_)? COOPtr(new COO(bg->coo_->AsNumBits(bits))) : nullptr;
1177
    return HeteroGraphPtr(
1178
        new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->restrict_format_));
1179
1180
1181
  }
}

Minjie Wang's avatar
Minjie Wang committed
1182
HeteroGraphPtr UnitGraph::CopyTo(HeteroGraphPtr g, const DLContext& ctx) {
1183
1184
  if (ctx == g->Context()) {
    return g;
1185
1186
1187
1188
1189
1190
1191
1192
  } else {
    auto bg = std::dynamic_pointer_cast<UnitGraph>(g);
    CHECK_NOTNULL(bg);
    CSRPtr new_incsr = (bg->in_csr_)? CSRPtr(new CSR(bg->in_csr_->CopyTo(ctx))) : nullptr;
    CSRPtr new_outcsr = (bg->out_csr_)? CSRPtr(new CSR(bg->out_csr_->CopyTo(ctx))) : nullptr;
    COOPtr new_coo = (bg->coo_)? COOPtr(new COO(bg->coo_->CopyTo(ctx))) : nullptr;
    return HeteroGraphPtr(
        new UnitGraph(g->meta_graph(), new_incsr, new_outcsr, new_coo, bg->restrict_format_));
1193
1194
1195
  }
}

1196
1197
UnitGraph::UnitGraph(GraphPtr metagraph, CSRPtr in_csr, CSRPtr out_csr, COOPtr coo,
                     SparseFormat restrict_format)
Minjie Wang's avatar
Minjie Wang committed
1198
  : BaseHeteroGraph(metagraph), in_csr_(in_csr), out_csr_(out_csr), coo_(coo) {
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
  restrict_format_ = AutoDetectFormat(in_csr, out_csr, coo, restrict_format);
  switch (restrict_format) {
  case SparseFormat::kCSC:
    in_csr_ = GetInCSR();
    coo_ = nullptr;
    out_csr_ = nullptr;
    break;
  case SparseFormat::kCSR:
    out_csr_ = GetOutCSR();
    coo_ = nullptr;
    in_csr_ = nullptr;
    break;
  case SparseFormat::kCOO:
    coo_ = GetCOO();
    in_csr_ = nullptr;
    out_csr_ = nullptr;
    break;
  default:
    break;
1218
1219
  }

1220
1221
1222
  CHECK(GetAny()) << "At least one graph structure should exist.";
}

1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
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,
    SparseFormat restrict_format) {
  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));
  if (has_out_csr)
    out_csr_ptr = CSRPtr(new CSR(mg, out_csr));
  if (has_coo)
    coo_ptr = COOPtr(new COO(mg, coo));

  return HeteroGraphPtr(new UnitGraph(mg, in_csr_ptr, out_csr_ptr, coo_ptr, restrict_format));
}

1247
1248
1249
1250
1251
1252
1253
UnitGraph::CSRPtr UnitGraph::GetInCSR(bool inplace) const {
  if (inplace)
    if (restrict_format_ != SparseFormat::kAny &&
        restrict_format_ != SparseFormat::kCSC)
      LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
        ", cannot create CSC matrix.";
  CSRPtr ret = in_csr_;
1254
1255
1256
  if (!in_csr_) {
    if (out_csr_) {
      const auto& newadj = aten::CSRTranspose(out_csr_->adj());
1257
1258
1259
      ret = std::make_shared<CSR>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->in_csr_ = ret;
1260
1261
    } else {
      CHECK(coo_) << "None of CSR, COO exist";
1262
1263
      const auto& newadj = aten::CSRSort(aten::COOToCSR(
            aten::COOTranspose(coo_->adj())));
1264
1265
1266
      ret = std::make_shared<CSR>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->in_csr_ = ret;
1267
1268
    }
  }
1269
  return ret;
1270
1271
1272
}

/* !\brief Return out csr. If not exist, transpose the other one.*/
1273
1274
1275
1276
1277
1278
1279
UnitGraph::CSRPtr UnitGraph::GetOutCSR(bool inplace) const {
  if (inplace)
    if (restrict_format_ != SparseFormat::kAny &&
        restrict_format_ != SparseFormat::kCSR)
      LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
        ", cannot create CSR matrix.";
  CSRPtr ret = out_csr_;
1280
1281
  if (!out_csr_) {
    if (in_csr_) {
1282
      const auto& newadj = aten::CSRSort(aten::CSRTranspose(in_csr_->adj()));
1283
1284
1285
      ret = std::make_shared<CSR>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->out_csr_ = ret;
1286
1287
    } else {
      CHECK(coo_) << "None of CSR, COO exist";
1288
      const auto& newadj = aten::CSRSort(aten::COOToCSR(coo_->adj()));
1289
1290
1291
      ret = std::make_shared<CSR>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->out_csr_ = ret;
1292
1293
    }
  }
1294
  return ret;
1295
1296
1297
}

/* !\brief Return coo. If not exist, create from csr.*/
1298
1299
1300
1301
1302
1303
1304
UnitGraph::COOPtr UnitGraph::GetCOO(bool inplace) const {
  if (inplace)
    if (restrict_format_ != SparseFormat::kAny &&
        restrict_format_ != SparseFormat::kCOO)
      LOG(FATAL) << "The graph have restricted sparse format " << GetRestrictFormat() <<
        ", cannot create COO matrix.";
  COOPtr ret = coo_;
1305
1306
  if (!coo_) {
    if (in_csr_) {
1307
      const auto& newadj = aten::COOTranspose(aten::CSRToCOO(in_csr_->adj(), true));
1308
1309
1310
      ret = std::make_shared<COO>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->coo_ = ret;
1311
1312
1313
    } else {
      CHECK(out_csr_) << "Both CSR are missing.";
      const auto& newadj = aten::CSRToCOO(out_csr_->adj(), true);
1314
1315
1316
      ret = std::make_shared<COO>(meta_graph(), newadj);
      if (inplace)
        const_cast<UnitGraph*>(this)->coo_ = ret;
1317
1318
    }
  }
1319
  return ret;
1320
1321
}

1322
aten::CSRMatrix UnitGraph::GetCSCMatrix(dgl_type_t etype) const {
1323
1324
1325
  return GetInCSR()->adj();
}

1326
aten::CSRMatrix UnitGraph::GetCSRMatrix(dgl_type_t etype) const {
1327
1328
1329
  return GetOutCSR()->adj();
}

1330
aten::COOMatrix UnitGraph::GetCOOMatrix(dgl_type_t etype) const {
1331
1332
1333
  return GetCOO()->adj();
}

Minjie Wang's avatar
Minjie Wang committed
1334
HeteroGraphPtr UnitGraph::GetAny() const {
1335
1336
1337
1338
1339
1340
1341
1342
1343
  if (in_csr_) {
    return in_csr_;
  } else if (out_csr_) {
    return out_csr_;
  } else {
    return coo_;
  }
}

1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
dgl_format_code_t UnitGraph::GetFormatInUse() const {
  dgl_format_code_t ret = 0;
  if (in_csr_) ret = ret | 1;
  ret = ret << 1;
  if (out_csr_) ret = ret | 1;
  ret = ret << 1;
  if (coo_) ret = ret | 1;
  return ret;
}

1354
1355
HeteroGraphPtr UnitGraph::GetFormat(SparseFormat format) const {
  switch (format) {
1356
1357
1358
1359
1360
1361
1362
1363
  case SparseFormat::kCSR:
    return GetOutCSR();
  case SparseFormat::kCSC:
    return GetInCSR();
  case SparseFormat::kCOO:
    return GetCOO();
  case SparseFormat::kAny:
    return GetAny();
1364
1365
  default:  // SparseFormat::kAuto
    LOG(FATAL) << "Must specify a restrict format.";
1366
    return nullptr;
1367
1368
1369
  }
}

1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
HeteroGraphPtr UnitGraph::GetGraphInFormat(SparseFormat restrict_format) const {
    int64_t num_vtypes = NumVertexTypes();
    switch (restrict_format) {
    case SparseFormat::kCOO:
      return CreateFromCOO(
        num_vtypes, GetCOO(false)->adj(), restrict_format);
    case SparseFormat::kCSC:
      return CreateFromCSC(
        num_vtypes, GetInCSR(false)->adj(), restrict_format);
    case SparseFormat::kCSR:
      return CreateFromCSR(
        num_vtypes, GetOutCSR(false)->adj(), restrict_format);
    case SparseFormat::kAny:
      return HeteroGraphPtr(
        new UnitGraph(meta_graph_, in_csr_, out_csr_, coo_, restrict_format));
    default:  // SparseFormat::kAuto
      LOG(FATAL) << "Must specify a restrict format.";
      return nullptr;
  }
}

SparseFormat UnitGraph::AutoDetectFormat(
    CSRPtr in_csr, CSRPtr out_csr, COOPtr coo, SparseFormat restrict_format) const {
  if (restrict_format != SparseFormat::kAuto)
    return restrict_format;
  if (coo && coo->IsHypersparse())
    return SparseFormat::kCOO;
  return SparseFormat::kAny;
}

1400
SparseFormat UnitGraph::SelectFormat(SparseFormat preferred_format) const {
1401
  if (restrict_format_ != SparseFormat::kAny)
1402
    return restrict_format_;  // force to select the restricted format
1403
  else if (preferred_format != SparseFormat::kAny)
1404
1405
    return preferred_format;
  else if (in_csr_)
1406
    return SparseFormat::kCSC;
1407
  else if (out_csr_)
1408
    return SparseFormat::kCSR;
1409
  else
1410
    return SparseFormat::kCOO;
1411
1412
}

1413
1414
1415
1416
1417
1418
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;
  if (in_csr_) {
    aten::CSRMatrix csc = GetCSCMatrix(0);
1419
    in_csr_ptr = dgl::CSRPtr(new dgl::CSR(csc.indptr, csc.indices, csc.data));
1420
1421
1422
  }
  if (out_csr_) {
    aten::CSRMatrix csr = GetCSRMatrix(0);
1423
    out_csr_ptr = dgl::CSRPtr(new dgl::CSR(csr.indptr, csr.indices, csr.data));
1424
1425
1426
1427
  }
  if (coo_) {
    aten::COOMatrix coo = GetCOOMatrix(0);
    if (!COOHasData(coo)) {
1428
      coo_ptr = dgl::COOPtr(new dgl::COO(NumVertices(0), coo.row, coo.col));
1429
1430
1431
    } else {
      IdArray new_src = Scatter(coo.row, coo.data);
      IdArray new_dst = Scatter(coo.col, coo.data);
1432
      coo_ptr = dgl::COOPtr(new dgl::COO(NumVertices(0), new_src, new_dst));
1433
1434
1435
1436
1437
    }
  }
  return GraphPtr(new dgl::ImmutableGraph(in_csr_ptr, out_csr_ptr, coo_ptr));
}

1438
1439
1440
1441
1442
1443
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";
1444
1445
1446
1447
1448
1449

  int64_t format_code;
  CHECK(fs->Read(&format_code)) << "Invalid format";
  restrict_format_ = static_cast<SparseFormat>(format_code);

  switch (restrict_format_) {
1450
    case SparseFormat::kCOO:
1451
1452
      fs->Read(&coo_);
      break;
1453
    case SparseFormat::kCSR:
1454
1455
      fs->Read(&out_csr_);
      break;
1456
    case SparseFormat::kCSC:
1457
1458
1459
1460
1461
1462
1463
1464
1465
      fs->Read(&in_csr_);
      break;
    default:
      LOG(FATAL) << "unsupported format code";
      break;
  }

  meta_graph_ = GetAny()->meta_graph();

1466
1467
1468
  return true;
}

1469

1470
1471
void UnitGraph::Save(dmlc::Stream* fs) const {
  fs->Write(kDGLSerialize_UnitGraphMagic);
1472
1473
  // Didn't write UnitGraph::meta_graph_, since it's included in the underlying
  // sparse matrix
1474
  auto avail_fmt = SelectFormat(SparseFormat::kAny);
1475
1476
  fs->Write(static_cast<int64_t>(avail_fmt));
  switch (avail_fmt) {
1477
    case SparseFormat::kCOO:
1478
1479
      fs->Write(GetCOO());
      break;
1480
    case SparseFormat::kCSR:
1481
1482
      fs->Write(GetOutCSR());
      break;
1483
    case SparseFormat::kCSC:
1484
1485
1486
1487
1488
1489
      fs->Write(GetInCSR());
      break;
    default:
      LOG(FATAL) << "unsupported format code";
      break;
  }
1490
1491
}

1492
}  // namespace dgl