graph_apis.cc 19.3 KB
Newer Older
1
2
3
4
5
/*!
 *  Copyright (c) 2018 by Contributors
 * \file graph/graph.cc
 * \brief DGL graph index APIs
 */
Minjie Wang's avatar
Minjie Wang committed
6
#include <dgl/graph.h>
7
#include <dgl/immutable_graph.h>
Minjie Wang's avatar
Minjie Wang committed
8
#include <dgl/graph_op.h>
Da Zheng's avatar
Da Zheng committed
9
#include <dgl/sampler.h>
10
#include <dgl/nodeflow.h>
Lingfan Yu's avatar
Lingfan Yu committed
11
#include "../c_api_common.h"
Minjie Wang's avatar
Minjie Wang committed
12

13
14
15
16
17
using dgl::runtime::DGLArgs;
using dgl::runtime::DGLArgValue;
using dgl::runtime::DGLRetValue;
using dgl::runtime::PackedFunc;
using dgl::runtime::NDArray;
Minjie Wang's avatar
Minjie Wang committed
18
19

namespace dgl {
Minjie Wang's avatar
Minjie Wang committed
20

Minjie Wang's avatar
Minjie Wang committed
21
namespace {
Minjie Wang's avatar
Minjie Wang committed
22
// Convert EdgeArray structure to PackedFunc.
23
24
template<class EdgeArray>
PackedFunc ConvertEdgeArrayToPackedFunc(const EdgeArray& ea) {
25
  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
26
      const int which = args[0];
27
      if (which == 0) {
Minjie Wang's avatar
Minjie Wang committed
28
29
30
31
32
33
34
35
36
37
38
39
        *rv = std::move(ea.src);
      } else if (which == 1) {
        *rv = std::move(ea.dst);
      } else if (which == 2) {
        *rv = std::move(ea.id);
      } else {
        LOG(FATAL) << "invalid choice";
      }
    };
  return PackedFunc(body);
}

40
41
42
43
44
45
46
47
48
49
50
51
52
// Convert CSRArray structure to PackedFunc.
PackedFunc ConvertAdjToPackedFunc(const std::vector<IdArray>& ea) {
  auto body = [ea] (DGLArgs args, DGLRetValue* rv) {
      const int which = args[0];
      if ((size_t) which < ea.size()) {
        *rv = std::move(ea[which]);
      } else {
        LOG(FATAL) << "invalid choice";
      }
    };
  return PackedFunc(body);
}

Minjie Wang's avatar
Minjie Wang committed
53
54
// Convert Subgraph structure to PackedFunc.
PackedFunc ConvertSubgraphToPackedFunc(const Subgraph& sg) {
55
  auto body = [sg] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
56
      const int which = args[0];
Minjie Wang's avatar
Minjie Wang committed
57
      if (which == 0) {
58
        GraphInterface* gptr = sg.graph->Reset();
Minjie Wang's avatar
Minjie Wang committed
59
60
        GraphHandle ghandle = gptr;
        *rv = ghandle;
61
      } else if (which == 1) {
Minjie Wang's avatar
Minjie Wang committed
62
        *rv = std::move(sg.induced_vertices);
63
      } else if (which == 2) {
Minjie Wang's avatar
Minjie Wang committed
64
        *rv = std::move(sg.induced_edges);
65
66
67
      } else {
        LOG(FATAL) << "invalid choice";
      }
Minjie Wang's avatar
Minjie Wang committed
68
69
70
    };
  return PackedFunc(body);
}
Minjie Wang's avatar
Minjie Wang committed
71

Minjie Wang's avatar
Minjie Wang committed
72
}  // namespace
Minjie Wang's avatar
Minjie Wang committed
73

74
75
76
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
namespace {
// This namespace contains template functions for batching
// and unbatching over graph and immutable graph
template<typename T>
void DGLDisjointPartitionByNum(const T *gptr, DGLArgs args, DGLRetValue *rv) {
  int64_t num = args[1];
  const std::vector<T> &&rst = GraphOp::DisjointPartitionByNum(gptr, num);
  // return the pointer array as an integer array
  const int64_t len = rst.size();
  NDArray ptr_array = NDArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t *ptr_array_data = static_cast<int64_t *>(ptr_array->data);
  for (size_t i = 0; i < rst.size(); ++i) {
    T *ptr = new T();
    *ptr = std::move(rst[i]);
    ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
  }
  *rv = ptr_array;
}

template<typename T>
void DGLDisjointUnion(GraphHandle *inhandles, int list_size, DGLRetValue *rv) {
  std::vector<const T *> graphs;
  for (int i = 0; i < list_size; ++i) {
    const GraphInterface *ptr = static_cast<const GraphInterface *>(inhandles[i]);
    const T *gr = dynamic_cast<const T *>(ptr);
    CHECK(gr) << "Error: Attempted to batch MutableGraph with ImmutableGraph";
    graphs.push_back(gr);
  }

  T *gptr = new T();
  *gptr = GraphOp::DisjointUnion(std::move(graphs));
  GraphHandle ghandle = gptr;
  *rv = ghandle;
}

template<typename T>
void DGLDisjointPartitionBySizes(const T *gptr, const IdArray sizes, DGLRetValue *rv) {
  std::vector<T> &&rst = GraphOp::DisjointPartitionBySizes(gptr, sizes);
  // return the pointer array as an integer array
  const int64_t len = rst.size();
  NDArray ptr_array = NDArray::Empty({len}, DLDataType{kDLInt, 64, 1}, DLContext{kDLCPU, 0});
  int64_t *ptr_array_data = static_cast<int64_t *>(ptr_array->data);
  for (size_t i = 0; i < rst.size(); ++i) {
    T *ptr = new T();
    *ptr = std::move(rst[i]);
    ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
  }
  *rv = ptr_array;
}
}  // namespace

125
126
127
///////////////////////////// Graph API ///////////////////////////////////

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreateMutable")
128
.set_body([] (DGLArgs args, DGLRetValue* rv) {
129
130
    bool multigraph = static_cast<bool>(args[0]);
    GraphHandle ghandle = new Graph(multigraph);
Minjie Wang's avatar
Minjie Wang committed
131
132
    *rv = ghandle;
  });
Minjie Wang's avatar
Minjie Wang committed
133

134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const IdArray src_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray dst_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray edge_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
    const bool multigraph = static_cast<bool>(args[3]);
    const int64_t num_nodes = static_cast<int64_t>(args[4]);
    const bool readonly = static_cast<bool>(args[5]);
    GraphHandle ghandle;
    if (readonly)
      ghandle = new ImmutableGraph(src_ids, dst_ids, edge_ids, num_nodes, multigraph);
    else
      ghandle = new Graph(src_ids, dst_ids, edge_ids, num_nodes, multigraph);
    *rv = ghandle;
  });

150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreate")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const IdArray indptr = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray indices = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray edge_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
    const std::string shared_mem_name = args[3];
    const bool multigraph = static_cast<bool>(args[4]);
    const std::string edge_dir = args[5];
    ImmutableGraph::CSR::Ptr csr;
    if (shared_mem_name.empty())
      csr.reset(new ImmutableGraph::CSR(indptr, indices, edge_ids));
    else
      csr.reset(new ImmutableGraph::CSR(indptr, indices, edge_ids, shared_mem_name));

    GraphHandle ghandle;
    if (edge_dir == "in")
      ghandle = new ImmutableGraph(csr, nullptr, multigraph);
    else
      ghandle = new ImmutableGraph(nullptr, csr, multigraph);
    *rv = ghandle;
  });

DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCSRCreateMMap")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const std::string shared_mem_name = args[0];
    const int64_t num_vertices = args[1];
    const int64_t num_edges = args[2];
    const bool multigraph = static_cast<bool>(args[3]);
    const std::string edge_dir = args[4];
    ImmutableGraph::CSR::Ptr csr(new ImmutableGraph::CSR(shared_mem_name,
                                                         num_vertices, num_edges));
    GraphHandle ghandle;
    if (edge_dir == "in")
      ghandle = new ImmutableGraph(csr, nullptr, multigraph);
    else
      ghandle = new ImmutableGraph(nullptr, csr, multigraph);
    *rv = ghandle;
  });

189
190
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
191
    GraphHandle ghandle = args[0];
192
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
193
194
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
195

196
197
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
198
    GraphHandle ghandle = args[0];
199
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
200
201
202
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
203

204
205
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
206
    GraphHandle ghandle = args[0];
207
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
208
209
210
211
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
    gptr->AddEdge(src, dst);
  });
Minjie Wang's avatar
Minjie Wang committed
212

213
214
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
215
    GraphHandle ghandle = args[0];
216
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
217
218
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
Minjie Wang's avatar
Minjie Wang committed
219
220
221
    gptr->AddEdges(src, dst);
  });

222
223
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
224
    GraphHandle ghandle = args[0];
225
    GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
226
227
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
228

229
230
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
231
232
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
233
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
234
235
236
    *rv = gptr->IsMultigraph();
  });

237
238
239
240
241
242
243
244
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsReadonly")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    *rv = gptr->IsReadonly();
  });

245
246
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
247
    GraphHandle ghandle = args[0];
248
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
249
250
    *rv = static_cast<int64_t>(gptr->NumVertices());
  });
Minjie Wang's avatar
Minjie Wang committed
251

252
253
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
254
    GraphHandle ghandle = args[0];
255
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
256
257
258
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

259
260
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
261
    GraphHandle ghandle = args[0];
262
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
263
264
265
266
    const dgl_id_t vid = args[1];
    *rv = gptr->HasVertex(vid);
  });

267
268
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
269
    GraphHandle ghandle = args[0];
270
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
271
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
272
273
274
    *rv = gptr->HasVertices(vids);
  });

275
276
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLMapSubgraphNID")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
277
278
279
280
281
    const IdArray parent_vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
    const IdArray query = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = GraphOp::MapParentIdToSubgraphId(parent_vids, query);
  });

282
283
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
284
    GraphHandle ghandle = args[0];
285
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
286
287
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
288
    *rv = gptr->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
289
290
  });

291
292
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
293
    GraphHandle ghandle = args[0];
294
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
295
296
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
297
    *rv = gptr->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
298
299
  });

300
301
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
302
    GraphHandle ghandle = args[0];
303
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
304
305
306
307
308
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Predecessors(vid, radius);
  });

309
310
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
311
    GraphHandle ghandle = args[0];
312
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
313
314
315
316
317
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

318
319
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
320
    GraphHandle ghandle = args[0];
321
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
322
323
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
324
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
325
326
  });

327
328
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
329
    GraphHandle ghandle = args[0];
330
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
331
332
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
333
334
335
    *rv = ConvertEdgeArrayToPackedFunc(gptr->EdgeIds(src, dst));
  });

336
337
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
338
    GraphHandle ghandle = args[0];
339
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
340
341
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertEdgeArrayToPackedFunc(gptr->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
342
343
  });

344
345
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
346
    GraphHandle ghandle = args[0];
347
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
348
    const dgl_id_t vid = args[1];
349
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
350
351
  });

352
353
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
354
    GraphHandle ghandle = args[0];
355
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
356
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
357
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
358
359
  });

360
361
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
362
    GraphHandle ghandle = args[0];
363
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
364
    const dgl_id_t vid = args[1];
365
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
366
367
  });

368
369
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
370
    GraphHandle ghandle = args[0];
371
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
372
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
373
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
374
375
  });

376
377
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
378
    GraphHandle ghandle = args[0];
379
380
381
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
    std::string order = args[1];
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(order));
Minjie Wang's avatar
Minjie Wang committed
382
383
  });

384
385
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
386
    GraphHandle ghandle = args[0];
387
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
388
389
390
391
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->InDegree(vid));
  });

392
393
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
394
    GraphHandle ghandle = args[0];
395
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
396
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
397
398
399
    *rv = gptr->InDegrees(vids);
  });

400
401
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
402
    GraphHandle ghandle = args[0];
403
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
404
405
406
407
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

408
409
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
410
    GraphHandle ghandle = args[0];
411
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
412
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
413
414
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
415

416
417
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
418
    GraphHandle ghandle = args[0];
419
    const GraphInterface* gptr = static_cast<GraphInterface*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
420
421
422
423
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

424
425
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
426
    GraphHandle ghandle = args[0];
427
    const GraphInterface *gptr = static_cast<GraphInterface*>(ghandle);
428
429
430
431
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids));
  });

432
433
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
434
435
436
    void* list = args[0];
    GraphHandle* inhandles = static_cast<GraphHandle*>(list);
    int list_size = args[1];
437
438
439
440
441
442
443
444
    const GraphInterface *ptr = static_cast<const GraphInterface *>(inhandles[0]);
    const ImmutableGraph *im_gr = dynamic_cast<const ImmutableGraph *>(ptr);
    const Graph *gr = dynamic_cast<const Graph *>(ptr);
    if (gr) {
      DGLDisjointUnion<Graph>(inhandles, list_size, rv);
    } else {
      CHECK(im_gr) << "Args[0] is not a list of valid DGLGraph";
      DGLDisjointUnion<ImmutableGraph>(inhandles, list_size, rv);
445
446
447
    }
  });

448
449
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
450
    GraphHandle ghandle = args[0];
451
452
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
453
454
455
456
457
458
    const ImmutableGraph* im_gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    if (gptr) {
      DGLDisjointPartitionByNum(gptr, args, rv);
    } else {
      CHECK(im_gptr) << "Args[0] is not a valid DGLGraph";
      DGLDisjointPartitionByNum(im_gptr, args, rv);
Minjie Wang's avatar
Minjie Wang committed
459
460
461
    }
  });

462
463
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
464
    GraphHandle ghandle = args[0];
465
    const IdArray sizes = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
466
467
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
468
469
470
471
472
473
    const ImmutableGraph* im_gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    if (gptr) {
      DGLDisjointPartitionBySizes(gptr, sizes, rv);
    } else {
      CHECK(im_gptr) << "Args[0] is not a valid DGLGraph";
      DGLDisjointPartitionBySizes(im_gptr, sizes, rv);
Minjie Wang's avatar
Minjie Wang committed
474
    }
475
});
GaiYu0's avatar
cpp lg  
GaiYu0 committed
476

477
478
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
GaiYu0's avatar
cpp lg  
GaiYu0 committed
479
480
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
481
482
483
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const Graph* gptr = dynamic_cast<const Graph*>(ptr);
    CHECK(gptr) << "_CAPI_DGLGraphLineGraph isn't implemented in immutable graph";
GaiYu0's avatar
cpp lg  
GaiYu0 committed
484
485
486
487
488
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
GaiYu0's avatar
GaiYu0 committed
489

490
491
492
493
494
495
496
497
498
499
DGL_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphGetAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    bool transpose = args[1];
    std::string format = args[2];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    auto res = ptr->GetAdj(transpose, format);
    *rv = ConvertAdjToPackedFunc(res);
  });

500
501
502
503
504
505
506
507
508
509
510
511
512
DGL_REGISTER_GLOBAL("nodeflow._CAPI_NodeFlowGetBlockAdj")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    GraphHandle ghandle = args[0];
    std::string format = args[1];
    int64_t layer0_size = args[2];
    int64_t start = args[3];
    int64_t end = args[4];
    const GraphInterface *ptr = static_cast<const GraphInterface *>(ghandle);
    const ImmutableGraph* gptr = dynamic_cast<const ImmutableGraph*>(ptr);
    auto res = GetNodeFlowSlice(*gptr, format, layer0_size, start, end, true);
    *rv = ConvertAdjToPackedFunc(res);
  });

Minjie Wang's avatar
Minjie Wang committed
513
}  // namespace dgl