heterograph_capi.cc 29.8 KB
Newer Older
1
/**
2
 *  Copyright (c) 2020 by Contributors
3
4
 * @file graph/heterograph_capi.cc
 * @brief Heterograph CAPI bindings.
5
 */
6
#include <dgl/array.h>
7
#include <dgl/aten/coo.h>
8
#include <dgl/immutable_graph.h>
9
10
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/c_runtime_api.h>
11
#include <dgl/runtime/container.h>
12
#include <dgl/runtime/parallel_for.h>
13

14
#include <set>
15

16
#include "../c_api_common.h"
17
#include "./heterograph.h"
18
#include "unit_graph.h"
19
20
21
22
23
24
25
26
27
28

using namespace dgl::runtime;

namespace dgl {

///////////////////////// Unitgraph functions /////////////////////////

// XXX(minjie): Ideally, Unitgraph should be invisible to python side

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCOO")
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      int64_t nvtypes = args[0];
      int64_t num_src = args[1];
      int64_t num_dst = args[2];
      IdArray row = args[3];
      IdArray col = args[4];
      List<Value> formats = args[5];
      bool row_sorted = args[6];
      bool col_sorted = args[7];
      std::vector<SparseFormat> formats_vec;
      for (Value val : formats) {
        std::string fmt = val->data;
        formats_vec.push_back(ParseSparseFormat(fmt));
      }
      const auto code = SparseFormatsToCode(formats_vec);
      auto hgptr = CreateFromCOO(
          nvtypes, num_src, num_dst, row, col, row_sorted, col_sorted, code);
      *rv = HeteroGraphRef(hgptr);
    });
48
49

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateUnitGraphFromCSR")
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      int64_t nvtypes = args[0];
      int64_t num_src = args[1];
      int64_t num_dst = args[2];
      IdArray indptr = args[3];
      IdArray indices = args[4];
      IdArray edge_ids = args[5];
      List<Value> formats = args[6];
      bool transpose = args[7];
      std::vector<SparseFormat> formats_vec;
      for (Value val : formats) {
        std::string fmt = val->data;
        formats_vec.push_back(ParseSparseFormat(fmt));
      }
      const auto code = SparseFormatsToCode(formats_vec);
      if (!transpose) {
        auto hgptr = CreateFromCSR(
            nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
        *rv = HeteroGraphRef(hgptr);
      } else {
        auto hgptr = CreateFromCSC(
            nvtypes, num_src, num_dst, indptr, indices, edge_ids, code);
        *rv = HeteroGraphRef(hgptr);
      }
    });
75
76

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateHeteroGraph")
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      GraphRef meta_graph = args[0];
      List<HeteroGraphRef> rel_graphs = args[1];
      std::vector<HeteroGraphPtr> rel_ptrs;
      rel_ptrs.reserve(rel_graphs.size());
      for (const auto& ref : rel_graphs) {
        rel_ptrs.push_back(ref.sptr());
      }
      auto hgptr = CreateHeteroGraph(meta_graph.sptr(), rel_ptrs);
      *rv = HeteroGraphRef(hgptr);
    });

DGL_REGISTER_GLOBAL(
    "heterograph_index._CAPI_DGLHeteroCreateHeteroGraphWithNumNodes")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      GraphRef meta_graph = args[0];
      List<HeteroGraphRef> rel_graphs = args[1];
      IdArray num_nodes_per_type = args[2];
      std::vector<HeteroGraphPtr> rel_ptrs;
      rel_ptrs.reserve(rel_graphs.size());
      for (const auto& ref : rel_graphs) {
        rel_ptrs.push_back(ref.sptr());
      }
      auto hgptr = CreateHeteroGraph(
          meta_graph.sptr(), rel_ptrs, num_nodes_per_type.ToVector<int64_t>());
      *rv = HeteroGraphRef(hgptr);
    });
104

105
106
107
///////////////////////// HeteroGraph member functions /////////////////////////

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetMetaGraph")
108
109
110
111
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->meta_graph();
    });
112
113

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsMetaGraphUniBipartite")
114
115
116
117
118
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      GraphPtr mg = hg->meta_graph();
      *rv = mg->IsUniBipartite();
    });
119
120

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetRelationGraph")
121
122
123
124
125
126
127
128
129
130
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      CHECK_LE(etype, hg->NumEdgeTypes()) << "invalid edge type " << etype;
      auto unit_graph = hg->GetRelationGraph(etype);
      auto meta_graph = unit_graph->meta_graph();
      auto hgptr = CreateHeteroGraph(
          meta_graph, {unit_graph}, unit_graph->NumVerticesPerType());
      *rv = HeteroGraphRef(hgptr);
    });
131
132

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFlattenedGraph")
133
134
135
136
137
138
139
140
141
142
143
144
145
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> etypes = args[1];
      std::vector<dgl_id_t> etypes_vec;
      for (Value val : etypes) {
        // (gq) have to decompose it into two statements because of a weird MSVC
        // internal error
        dgl_id_t id = val->data;
        etypes_vec.push_back(id);
      }

      *rv = FlattenedHeteroGraphRef(hg->Flatten(etypes_vec));
    });
146
147

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddVertices")
148
149
150
151
152
153
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t vtype = args[1];
      int64_t num = args[2];
      hg->AddVertices(vtype, num);
    });
154
155

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdge")
156
157
158
159
160
161
162
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t src = args[2];
      dgl_id_t dst = args[3];
      hg->AddEdge(etype, src, dst);
    });
163
164

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAddEdges")
165
166
167
168
169
170
171
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray src = args[2];
      IdArray dst = args[3];
      hg->AddEdges(etype, src, dst);
    });
172
173

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroClear")
174
175
176
177
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      hg->Clear();
    });
178
179

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDataType")
180
181
182
183
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->DataType();
    });
184
185

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroContext")
186
187
188
189
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->Context();
    });
190
191

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsPinned")
192
193
194
195
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->IsPinned();
    });
196
197

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumBits")
198
199
200
201
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->NumBits();
    });
202
203

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsMultigraph")
204
205
206
207
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->IsMultigraph();
    });
208
209

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroIsReadonly")
210
211
212
213
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = hg->IsReadonly();
    });
214
215

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumVertices")
216
217
218
219
220
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t vtype = args[1];
      *rv = static_cast<int64_t>(hg->NumVertices(vtype));
    });
221
222

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroNumEdges")
223
224
225
226
227
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      *rv = static_cast<int64_t>(hg->NumEdges(etype));
    });
228
229

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertex")
230
231
232
233
234
235
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t vtype = args[1];
      dgl_id_t vid = args[2];
      *rv = hg->HasVertex(vtype, vid);
    });
236
237

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasVertices")
238
239
240
241
242
243
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t vtype = args[1];
      IdArray vids = args[2];
      *rv = hg->HasVertices(vtype, vids);
    });
244
245

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgeBetween")
246
247
248
249
250
251
252
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t src = args[2];
      dgl_id_t dst = args[3];
      *rv = hg->HasEdgeBetween(etype, src, dst);
    });
253
254

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroHasEdgesBetween")
255
256
257
258
259
260
261
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray src = args[2];
      IdArray dst = args[3];
      *rv = hg->HasEdgesBetween(etype, src, dst);
    });
262
263

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPredecessors")
264
265
266
267
268
269
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t dst = args[2];
      *rv = hg->Predecessors(etype, dst);
    });
270
271

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSuccessors")
272
273
274
275
276
277
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t src = args[2];
      *rv = hg->Successors(etype, src);
    });
278
279

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeId")
280
281
282
283
284
285
286
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t src = args[2];
      dgl_id_t dst = args[3];
      *rv = hg->EdgeId(etype, src, dst);
    });
287

288
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsAll")
289
290
291
292
293
294
295
296
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray src = args[2];
      IdArray dst = args[3];
      const auto& ret = hg->EdgeIdsAll(etype, src, dst);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
297
298

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeIdsOne")
299
300
301
302
303
304
305
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray src = args[2];
      IdArray dst = args[3];
      *rv = hg->EdgeIdsOne(etype, src, dst);
    });
306

307
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroFindEdges")
308
309
310
311
312
313
314
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray eids = args[2];
      const auto& ret = hg->FindEdges(etype, eids);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
315
316

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_1")
317
318
319
320
321
322
323
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t vid = args[2];
      const auto& ret = hg->InEdges(etype, vid);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
324
325

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInEdges_2")
326
327
328
329
330
331
332
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray vids = args[2];
      const auto& ret = hg->InEdges(etype, vids);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
333
334

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_1")
335
336
337
338
339
340
341
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t vid = args[2];
      const auto& ret = hg->OutEdges(etype, vid);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
342
343

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutEdges_2")
344
345
346
347
348
349
350
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray vids = args[2];
      const auto& ret = hg->OutEdges(etype, vids);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
351
352

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdges")
353
354
355
356
357
358
359
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      std::string order = args[2];
      const auto& ret = hg->Edges(etype, order);
      *rv = ConvertEdgeArrayToPackedFunc(ret);
    });
360
361

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegree")
362
363
364
365
366
367
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t vid = args[2];
      *rv = static_cast<int64_t>(hg->InDegree(etype, vid));
    });
368
369

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroInDegrees")
370
371
372
373
374
375
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray vids = args[2];
      *rv = hg->InDegrees(etype, vids);
    });
376
377

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegree")
378
379
380
381
382
383
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      dgl_id_t vid = args[2];
      *rv = static_cast<int64_t>(hg->OutDegree(etype, vid));
    });
384
385

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroOutDegrees")
386
387
388
389
390
391
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      IdArray vids = args[2];
      *rv = hg->OutDegrees(etype, vids);
    });
392
393

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetAdj")
394
395
396
397
398
399
400
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_type_t etype = args[1];
      bool transpose = args[2];
      std::string fmt = args[3];
      *rv = ConvertNDArrayVectorToPackedFunc(hg->GetAdj(etype, transpose, fmt));
    });
401
402

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroVertexSubgraph")
403
404
405
406
407
408
409
410
411
412
413
414
415
416
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> vids = args[1];
      bool relabel_nodes = args[2];
      CHECK(relabel_nodes) << "Node subgraph only supports relabel_nodes=True.";
      std::vector<IdArray> vid_vec;
      vid_vec.reserve(vids.size());
      for (Value val : vids) {
        vid_vec.push_back(val->data);
      }
      std::shared_ptr<HeteroSubgraph> subg(
          new HeteroSubgraph(hg->VertexSubgraph(vid_vec)));
      *rv = HeteroSubgraphRef(subg);
    });
417
418

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroEdgeSubgraph")
419
420
421
422
423
424
425
426
427
428
429
430
431
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> eids = args[1];
      bool preserve_nodes = args[2];
      std::vector<IdArray> eid_vec;
      eid_vec.reserve(eids.size());
      for (Value val : eids) {
        eid_vec.push_back(val->data);
      }
      std::shared_ptr<HeteroSubgraph> subg(
          new HeteroSubgraph(hg->EdgeSubgraph(eid_vec, preserve_nodes)));
      *rv = HeteroSubgraphRef(subg);
    });
432
433
434
435

///////////////////////// HeteroSubgraph members /////////////////////////

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetGraph")
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroSubgraphRef subg = args[0];
      *rv = HeteroGraphRef(subg->graph);
    });

DGL_REGISTER_GLOBAL(
    "heterograph_index._CAPI_DGLHeteroSubgraphGetInducedVertices")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroSubgraphRef subg = args[0];
      List<Value> induced_verts;
      for (IdArray arr : subg->induced_vertices) {
        induced_verts.push_back(Value(MakeValue(arr)));
      }
      *rv = induced_verts;
    });
451
452

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSubgraphGetInducedEdges")
453
454
455
456
457
458
459
460
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroSubgraphRef subg = args[0];
      List<Value> induced_edges;
      for (IdArray arr : subg->induced_edges) {
        induced_edges.push_back(Value(MakeValue(arr)));
      }
      *rv = induced_edges;
    });
461

462
463
///////////////////////// Global functions and algorithms
////////////////////////////
464
465

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroAsNumBits")
466
467
468
469
470
471
472
473
474
475
476
477
478
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      int bits = args[1];
      HeteroGraphPtr bhg_ptr = hg.sptr();
      auto hg_ptr = std::dynamic_pointer_cast<HeteroGraph>(bhg_ptr);
      HeteroGraphPtr hg_new;
      if (hg_ptr) {
        hg_new = HeteroGraph::AsNumBits(hg_ptr, bits);
      } else {
        hg_new = UnitGraph::AsNumBits(bhg_ptr, bits);
      }
      *rv = HeteroGraphRef(hg_new);
    });
479
480

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyTo")
481
482
483
484
485
486
487
488
489
490
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      int device_type = args[1];
      int device_id = args[2];
      DGLContext ctx;
      ctx.device_type = static_cast<DGLDeviceType>(device_type);
      ctx.device_id = device_id;
      HeteroGraphPtr hg_new = HeteroGraph::CopyTo(hg.sptr(), ctx);
      *rv = HeteroGraphRef(hg_new);
    });
491

492
493
494
495
496
497
498
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPinMemory")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      HeteroGraphPtr hg_new = HeteroGraph::PinMemory(hg.sptr());
      *rv = HeteroGraphRef(hg_new);
    });

499
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroPinMemory_")
500
501
502
503
504
505
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
      hgindex->PinMemory_();
      *rv = hg;
    });
506
507

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroUnpinMemory_")
508
509
510
511
512
513
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
      hgindex->UnpinMemory_();
      *rv = hg;
    });
514

515
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroRecordStream")
516
517
518
519
520
521
522
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      DGLStreamHandle stream = args[1];
      auto hgindex = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
      hgindex->RecordStream(stream);
      *rv = hg;
    });
523

524
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCopyToSharedMem")
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      std::string name = args[1];
      List<Value> ntypes = args[2];
      List<Value> etypes = args[3];
      List<Value> fmts = args[4];
      auto ntypes_vec = ListValueToVector<std::string>(ntypes);
      auto etypes_vec = ListValueToVector<std::string>(etypes);
      std::set<std::string> fmts_set;
      for (const auto& fmt : fmts) {
        std::string fmt_data = fmt->data;
        fmts_set.insert(fmt_data);
      }
      auto hg_share = HeteroGraph::CopyToSharedMem(
          hg.sptr(), name, ntypes_vec, etypes_vec, fmts_set);
      *rv = HeteroGraphRef(hg_share);
    });
542
543

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFromSharedMem")
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      std::string name = args[0];
      HeteroGraphPtr hg;
      std::vector<std::string> ntypes;
      std::vector<std::string> etypes;
      std::tie(hg, ntypes, etypes) = HeteroGraph::CreateFromSharedMem(name);
      List<Value> ntypes_list;
      List<Value> etypes_list;
      for (const auto& ntype : ntypes)
        ntypes_list.push_back(Value(MakeValue(ntype)));
      for (const auto& etype : etypes)
        etypes_list.push_back(Value(MakeValue(etype)));
      List<ObjectRef> ret;
      ret.push_back(HeteroGraphRef(hg));
      ret.push_back(ntypes_list);
      ret.push_back(etypes_list);
      *rv = ret;
    });
562

563
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroJointUnion")
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      GraphRef meta_graph = args[0];
      List<HeteroGraphRef> component_graphs = args[1];
      CHECK(component_graphs.size() > 1)
          << "Expect graph list to have at least two graphs";
      std::vector<HeteroGraphPtr> component_ptrs;
      component_ptrs.reserve(component_graphs.size());
      const int64_t bits = component_graphs[0]->NumBits();
      const DGLContext ctx = component_graphs[0]->Context();
      for (const auto& component : component_graphs) {
        component_ptrs.push_back(component.sptr());
        CHECK_EQ(component->NumBits(), bits)
            << "Expect graphs to joint union have the same index dtype(int"
            << bits << "), but got int" << component->NumBits();
        CHECK_EQ(component->Context(), ctx)
            << "Expect graphs to joint union have the same context" << ctx
            << "), but got " << component->Context();
      }

      auto hgptr = JointUnionHeteroGraph(meta_graph.sptr(), component_ptrs);
      *rv = HeteroGraphRef(hgptr);
    });
586

587
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointUnion_v2")
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      GraphRef meta_graph = args[0];
      List<HeteroGraphRef> component_graphs = args[1];
      CHECK(component_graphs.size() > 0)
          << "Expect graph list has at least one graph";
      std::vector<HeteroGraphPtr> component_ptrs;
      component_ptrs.reserve(component_graphs.size());
      const int64_t bits = component_graphs[0]->NumBits();
      const DGLContext ctx = component_graphs[0]->Context();
      for (const auto& component : component_graphs) {
        component_ptrs.push_back(component.sptr());
        CHECK_EQ(component->NumBits(), bits)
            << "Expect graphs to batch have the same index dtype(int" << bits
            << "), but got int" << component->NumBits();
        CHECK_EQ(component->Context(), ctx)
            << "Expect graphs to batch have the same context" << ctx
            << "), but got " << component->Context();
      }

      auto hgptr = DisjointUnionHeteroGraph2(meta_graph.sptr(), component_ptrs);
      *rv = HeteroGraphRef(hgptr);
    });

DGL_REGISTER_GLOBAL(
    "heterograph_index._CAPI_DGLHeteroDisjointPartitionBySizes_v2")
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      const IdArray vertex_sizes = args[1];
      const IdArray edge_sizes = args[2];
      std::vector<HeteroGraphPtr> ret;
      ret = DisjointPartitionHeteroBySizes2(
          hg->meta_graph(), hg.sptr(), vertex_sizes, edge_sizes);
      List<HeteroGraphRef> ret_list;
      for (HeteroGraphPtr hgptr : ret) {
        ret_list.push_back(HeteroGraphRef(hgptr));
      }
      *rv = ret_list;
    });
626
627

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroDisjointPartitionBySizes")
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      const IdArray vertex_sizes = args[1];
      const IdArray edge_sizes = args[2];
      const int64_t bits = hg->NumBits();
      std::vector<HeteroGraphPtr> ret;
      ATEN_ID_BITS_SWITCH(bits, IdType, {
        ret = DisjointPartitionHeteroBySizes<IdType>(
            hg->meta_graph(), hg.sptr(), vertex_sizes, edge_sizes);
      });
      List<HeteroGraphRef> ret_list;
      for (HeteroGraphPtr hgptr : ret) {
        ret_list.push_back(HeteroGraphRef(hgptr));
      }
      *rv = ret_list;
    });
644

645
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroSlice")
646
647
648
649
650
651
652
653
654
655
656
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      const IdArray num_nodes_per_type = args[1];
      const IdArray start_nid_per_type = args[2];
      const IdArray num_edges_per_type = args[3];
      const IdArray start_eid_per_type = args[4];
      auto hgptr = SliceHeteroGraph(
          hg->meta_graph(), hg.sptr(), num_nodes_per_type, start_nid_per_type,
          num_edges_per_type, start_eid_per_type);
      *rv = HeteroGraphRef(hgptr);
    });
657

658
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetCreatedFormats")
659
660
661
662
663
664
665
666
667
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> format_list;
      dgl_format_code_t code = hg->GetRelationGraph(0)->GetCreatedFormats();
      for (auto format : CodeToSparseFormats(code)) {
        format_list.push_back(Value(MakeValue(ToStringSparseFormat(format))));
      }
      *rv = format_list;
    });
668

669
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetAllowedFormats")
670
671
672
673
674
675
676
677
678
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> format_list;
      dgl_format_code_t code = hg->GetRelationGraph(0)->GetAllowedFormats();
      for (auto format : CodeToSparseFormats(code)) {
        format_list.push_back(Value(MakeValue(ToStringSparseFormat(format))));
      }
      *rv = format_list;
    });
679

680
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFormat")
681
682
683
684
685
686
687
688
689
690
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      dgl_format_code_t code = hg->GetRelationGraph(0)->GetAllowedFormats();
      auto get_format_f = [&](size_t etype_b, size_t etype_e) {
        for (auto etype = etype_b; etype < etype_e; ++etype) {
          auto bg =
              std::dynamic_pointer_cast<UnitGraph>(hg->GetRelationGraph(etype));
          for (auto format : CodeToSparseFormats(code)) bg->GetFormat(format);
        }
      };
691
692

#if !(defined(DGL_USE_CUDA))
693
      runtime::parallel_for(0, hg->NumEdgeTypes(), get_format_f);
694
#else
695
      get_format_f(0, hg->NumEdgeTypes());
696
#endif
697
    });
698

699
DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFormatGraph")
700
701
702
703
704
705
706
707
708
709
710
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      List<Value> formats = args[1];
      std::vector<SparseFormat> formats_vec;
      for (Value val : formats) {
        std::string fmt = val->data;
        formats_vec.push_back(ParseSparseFormat(fmt));
      }
      auto hgptr = hg->GetGraphInFormat(SparseFormatsToCode(formats_vec));
      *rv = HeteroGraphRef(hgptr);
    });
711

712
DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLInSubgraph")
713
714
715
716
717
718
719
720
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      const auto& nodes = ListValueToVector<IdArray>(args[1]);
      bool relabel_nodes = args[2];
      std::shared_ptr<HeteroSubgraph> ret(new HeteroSubgraph);
      *ret = InEdgeGraph(hg.sptr(), nodes, relabel_nodes);
      *rv = HeteroGraphRef(ret);
    });
721

722
DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLOutSubgraph")
723
724
725
726
727
728
729
730
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      const auto& nodes = ListValueToVector<IdArray>(args[1]);
      bool relabel_nodes = args[2];
      std::shared_ptr<HeteroSubgraph> ret(new HeteroSubgraph);
      *ret = OutEdgeGraph(hg.sptr(), nodes, relabel_nodes);
      *rv = HeteroGraphRef(ret);
    });
731

732
DGL_REGISTER_GLOBAL("transform._CAPI_DGLAsImmutableGraph")
733
734
735
736
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      *rv = GraphRef(hg->AsImmutableGraph());
    });
737

738
DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortOutEdges")
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      NDArray tag = args[1];
      int64_t num_tag = args[2];

      CHECK_EQ(hg->Context().device_type, kDGLCPU)
          << "Only support sorting by tag on cpu";
      CHECK(aten::IsValidIdArray(tag));
      CHECK_EQ(tag->ctx.device_type, kDGLCPU)
          << "Only support sorting by tag on cpu";

      const auto csr = hg->GetCSRMatrix(0);

      NDArray tag_pos = aten::NullArray();
      aten::CSRMatrix output;
      std::tie(output, tag_pos) = aten::CSRSortByTag(csr, tag, num_tag);
      HeteroGraphPtr output_hg =
          CreateFromCSR(hg->NumVertexTypes(), output, ALL_CODE);
      List<ObjectRef> ret;
      ret.push_back(HeteroGraphRef(output_hg));
      ret.push_back(Value(MakeValue(tag_pos)));
      *rv = ret;
    });
762
763

DGL_REGISTER_GLOBAL("transform._CAPI_DGLHeteroSortInEdges")
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      NDArray tag = args[1];
      int64_t num_tag = args[2];

      CHECK_EQ(hg->Context().device_type, kDGLCPU)
          << "Only support sorting by tag on cpu";
      CHECK(aten::IsValidIdArray(tag));
      CHECK_EQ(tag->ctx.device_type, kDGLCPU)
          << "Only support sorting by tag on cpu";

      const auto csc = hg->GetCSCMatrix(0);

      NDArray tag_pos = aten::NullArray();
      aten::CSRMatrix output;
      std::tie(output, tag_pos) = aten::CSRSortByTag(csc, tag, num_tag);

      HeteroGraphPtr output_hg =
          CreateFromCSC(hg->NumVertexTypes(), output, ALL_CODE);
      List<ObjectRef> ret;
      ret.push_back(HeteroGraphRef(output_hg));
      ret.push_back(Value(MakeValue(tag_pos)));
      *rv = ret;
    });
788

789
DGL_REGISTER_GLOBAL("heterograph._CAPI_DGLFindSrcDstNtypes")
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      GraphRef metagraph = args[0];
      std::unordered_set<uint64_t> dst_set;
      std::unordered_set<uint64_t> src_set;

      for (uint64_t eid = 0; eid < metagraph->NumEdges(); ++eid) {
        auto edge = metagraph->FindEdge(eid);
        auto src = edge.first;
        auto dst = edge.second;
        dst_set.insert(dst);
        src_set.insert(src);
      }

      List<Value> srclist, dstlist;
      List<List<Value>> ret_list;
      for (uint64_t nid = 0; nid < metagraph->NumVertices(); ++nid) {
        auto is_dst = dst_set.count(nid);
        auto is_src = src_set.count(nid);
        if (is_dst && is_src)
          return;
        else if (is_dst)
          dstlist.push_back(Value(MakeValue(static_cast<int64_t>(nid))));
        else
          // If a node type is isolated, put it in srctype as defined in the
          // Python docstring.
          srclist.push_back(Value(MakeValue(static_cast<int64_t>(nid))));
      }
      ret_list.push_back(srclist);
      ret_list.push_back(dstlist);
      *rv = ret_list;
    });
821
822

DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroReverse")
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      HeteroGraphRef hg = args[0];
      CHECK_GT(hg->NumEdgeTypes(), 0);
      auto g = std::dynamic_pointer_cast<HeteroGraph>(hg.sptr());
      std::vector<HeteroGraphPtr> rev_ugs;
      const auto& ugs = g->relation_graphs();
      rev_ugs.resize(ugs.size());

      for (size_t i = 0; i < ugs.size(); ++i) {
        const auto& rev_ug = ugs[i]->Reverse();
        rev_ugs[i] = rev_ug;
      }
      // node types are not changed
      const auto& num_nodes = g->NumVerticesPerType();
      const auto& meta_edges = hg->meta_graph()->Edges("eid");
      // reverse the metagraph
      const auto& rev_meta = ImmutableGraph::CreateFromCOO(
          hg->meta_graph()->NumVertices(), meta_edges.dst, meta_edges.src);
      *rv = CreateHeteroGraph(rev_meta, rev_ugs, num_nodes);
    });
843
}  // namespace dgl