graph_apis.cc 12.6 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>
Minjie Wang's avatar
Minjie Wang committed
7
#include <dgl/graph_op.h>
Lingfan Yu's avatar
Lingfan Yu committed
8
#include "../c_api_common.h"
Minjie Wang's avatar
Minjie Wang committed
9
10
11
12
13

using tvm::runtime::TVMArgs;
using tvm::runtime::TVMArgValue;
using tvm::runtime::TVMRetValue;
using tvm::runtime::PackedFunc;
Minjie Wang's avatar
Minjie Wang committed
14
using tvm::runtime::NDArray;
Minjie Wang's avatar
Minjie Wang committed
15
16

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

Minjie Wang's avatar
Minjie Wang committed
18
namespace {
Minjie Wang's avatar
Minjie Wang committed
19
// Convert EdgeArray structure to PackedFunc.
20
21
PackedFunc ConvertEdgeArrayToPackedFunc(const Graph::EdgeArray& ea) {
  auto body = [ea] (TVMArgs args, TVMRetValue* rv) {
Minjie Wang's avatar
Minjie Wang committed
22
      int which = args[0];
23
      if (which == 0) {
Minjie Wang's avatar
Minjie Wang committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
        *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);
}

// Convert Subgraph structure to PackedFunc.
PackedFunc ConvertSubgraphToPackedFunc(const Subgraph& sg) {
  auto body = [sg] (TVMArgs args, TVMRetValue* rv) {
      int which = args[0];
      if (which == 0) {
        Graph* gptr = new Graph();
        *gptr = std::move(sg.graph);
        GraphHandle ghandle = gptr;
        *rv = ghandle;
45
      } else if (which == 1) {
Minjie Wang's avatar
Minjie Wang committed
46
        *rv = std::move(sg.induced_vertices);
47
      } else if (which == 2) {
Minjie Wang's avatar
Minjie Wang committed
48
        *rv = std::move(sg.induced_edges);
49
50
51
      } else {
        LOG(FATAL) << "invalid choice";
      }
Minjie Wang's avatar
Minjie Wang committed
52
53
54
    };
  return PackedFunc(body);
}
Minjie Wang's avatar
Minjie Wang committed
55

Minjie Wang's avatar
Minjie Wang committed
56
}  // namespace
Minjie Wang's avatar
Minjie Wang committed
57

Minjie Wang's avatar
Minjie Wang committed
58
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphCreate")
Minjie Wang's avatar
Minjie Wang committed
59
.set_body([] (TVMArgs args, TVMRetValue* rv) {
60
61
    bool multigraph = static_cast<bool>(args[0]);
    GraphHandle ghandle = new Graph(multigraph);
Minjie Wang's avatar
Minjie Wang committed
62
63
    *rv = ghandle;
  });
Minjie Wang's avatar
Minjie Wang committed
64

Minjie Wang's avatar
Minjie Wang committed
65
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFree")
Minjie Wang's avatar
Minjie Wang committed
66
67
68
69
70
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    Graph* gptr = static_cast<Graph*>(ghandle);
    delete gptr;
  });
Minjie Wang's avatar
Minjie Wang committed
71

Minjie Wang's avatar
Minjie Wang committed
72
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddVertices")
Minjie Wang's avatar
Minjie Wang committed
73
74
75
76
77
78
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    Graph* gptr = static_cast<Graph*>(ghandle);
    uint64_t num_vertices = args[1];
    gptr->AddVertices(num_vertices);
  });
Minjie Wang's avatar
Minjie Wang committed
79

Minjie Wang's avatar
Minjie Wang committed
80
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdge")
Minjie Wang's avatar
Minjie Wang committed
81
82
83
84
85
86
87
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    Graph* gptr = static_cast<Graph*>(ghandle);
    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
88

Minjie Wang's avatar
Minjie Wang committed
89
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphAddEdges")
Minjie Wang's avatar
Minjie Wang committed
90
91
92
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
93
94
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
Minjie Wang's avatar
Minjie Wang committed
95
96
97
    gptr->AddEdges(src, dst);
  });

Minjie Wang's avatar
Minjie Wang committed
98
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphClear")
Minjie Wang's avatar
Minjie Wang committed
99
100
101
102
103
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    Graph* gptr = static_cast<Graph*>(ghandle);
    gptr->Clear();
  });
Minjie Wang's avatar
Minjie Wang committed
104

105
106
107
108
109
110
111
112
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphIsMultigraph")
.set_body([] (TVMArgs args, TVMRetValue *rv) {
    GraphHandle ghandle = args[0];
    // NOTE: not const since we have caches
    const Graph* gptr = static_cast<Graph*>(ghandle);
    *rv = gptr->IsMultigraph();
  });

Minjie Wang's avatar
Minjie Wang committed
113
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumVertices")
Minjie Wang's avatar
Minjie Wang committed
114
115
116
117
118
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    *rv = static_cast<int64_t>(gptr->NumVertices());
  });
Minjie Wang's avatar
Minjie Wang committed
119

Minjie Wang's avatar
Minjie Wang committed
120
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphNumEdges")
Minjie Wang's avatar
Minjie Wang committed
121
122
123
124
125
126
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    *rv = static_cast<int64_t>(gptr->NumEdges());
  });

Minjie Wang's avatar
Minjie Wang committed
127
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertex")
Minjie Wang's avatar
Minjie Wang committed
128
129
130
131
132
133
134
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
    *rv = gptr->HasVertex(vid);
  });

Minjie Wang's avatar
Minjie Wang committed
135
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasVertices")
Minjie Wang's avatar
Minjie Wang committed
136
137
138
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
139
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
140
141
142
    *rv = gptr->HasVertices(vids);
  });

143
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgeBetween")
Minjie Wang's avatar
Minjie Wang committed
144
145
146
147
148
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
149
    *rv = gptr->HasEdgeBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
150
151
  });

152
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphHasEdgesBetween")
Minjie Wang's avatar
Minjie Wang committed
153
154
155
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
156
157
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
158
    *rv = gptr->HasEdgesBetween(src, dst);
Minjie Wang's avatar
Minjie Wang committed
159
160
  });

Minjie Wang's avatar
Minjie Wang committed
161
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphPredecessors")
Minjie Wang's avatar
Minjie Wang committed
162
163
164
165
166
167
168
169
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Predecessors(vid, radius);
  });

Minjie Wang's avatar
Minjie Wang committed
170
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphSuccessors")
Minjie Wang's avatar
Minjie Wang committed
171
172
173
174
175
176
177
178
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
    const uint64_t radius = args[2];
    *rv = gptr->Successors(vid, radius);
  });

Minjie Wang's avatar
Minjie Wang committed
179
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeId")
Minjie Wang's avatar
Minjie Wang committed
180
181
182
183
184
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t src = args[1];
    const dgl_id_t dst = args[2];
185
    *rv = gptr->EdgeId(src, dst);
Minjie Wang's avatar
Minjie Wang committed
186
187
  });

Minjie Wang's avatar
Minjie Wang committed
188
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeIds")
Minjie Wang's avatar
Minjie Wang committed
189
190
191
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
192
193
    const IdArray src = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    const IdArray dst = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
194
195
196
197
198
199
200
201
202
    *rv = ConvertEdgeArrayToPackedFunc(gptr->EdgeIds(src, dst));
  });

TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphFindEdges")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertEdgeArrayToPackedFunc(gptr->FindEdges(eids));
Minjie Wang's avatar
Minjie Wang committed
203
204
  });

Minjie Wang's avatar
Minjie Wang committed
205
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_1")
Minjie Wang's avatar
Minjie Wang committed
206
207
208
209
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
210
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
211
212
  });

Minjie Wang's avatar
Minjie Wang committed
213
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInEdges_2")
Minjie Wang's avatar
Minjie Wang committed
214
215
216
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
217
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
218
    *rv = ConvertEdgeArrayToPackedFunc(gptr->InEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
219
220
  });

Minjie Wang's avatar
Minjie Wang committed
221
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_1")
Minjie Wang's avatar
Minjie Wang committed
222
223
224
225
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
226
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vid));
Minjie Wang's avatar
Minjie Wang committed
227
228
  });

Minjie Wang's avatar
Minjie Wang committed
229
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutEdges_2")
Minjie Wang's avatar
Minjie Wang committed
230
231
232
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
233
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
234
    *rv = ConvertEdgeArrayToPackedFunc(gptr->OutEdges(vids));
Minjie Wang's avatar
Minjie Wang committed
235
236
  });

Minjie Wang's avatar
Minjie Wang committed
237
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdges")
Minjie Wang's avatar
Minjie Wang committed
238
239
240
241
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const bool sorted = args[1];
242
    *rv = ConvertEdgeArrayToPackedFunc(gptr->Edges(sorted));
Minjie Wang's avatar
Minjie Wang committed
243
244
  });

Minjie Wang's avatar
Minjie Wang committed
245
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegree")
Minjie Wang's avatar
Minjie Wang committed
246
247
248
249
250
251
252
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->InDegree(vid));
  });

Minjie Wang's avatar
Minjie Wang committed
253
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphInDegrees")
Minjie Wang's avatar
Minjie Wang committed
254
255
256
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
257
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
258
259
260
    *rv = gptr->InDegrees(vids);
  });

Minjie Wang's avatar
Minjie Wang committed
261
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegree")
Minjie Wang's avatar
Minjie Wang committed
262
263
264
265
266
267
268
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const dgl_id_t vid = args[1];
    *rv = static_cast<int64_t>(gptr->OutDegree(vid));
  });

Minjie Wang's avatar
Minjie Wang committed
269
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphOutDegrees")
Minjie Wang's avatar
Minjie Wang committed
270
271
272
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
Minjie Wang's avatar
Minjie Wang committed
273
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
Minjie Wang's avatar
Minjie Wang committed
274
275
    *rv = gptr->OutDegrees(vids);
  });
Minjie Wang's avatar
Minjie Wang committed
276

Minjie Wang's avatar
Minjie Wang committed
277
278
279
280
281
282
283
284
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphVertexSubgraph")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->VertexSubgraph(vids));
  });

285
286
287
288
289
290
291
292
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphEdgeSubgraph")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph *gptr = static_cast<Graph*>(ghandle);
    const IdArray eids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    *rv = ConvertSubgraphToPackedFunc(gptr->EdgeSubgraph(eids));
  });

Minjie Wang's avatar
Minjie Wang committed
293
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointUnion")
294
295
296
297
298
299
300
301
302
303
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    void* list = args[0];
    GraphHandle* inhandles = static_cast<GraphHandle*>(list);
    int list_size = args[1];
    std::vector<const Graph*> graphs;
    for (int i = 0; i < list_size; ++i) {
      const Graph* gr = static_cast<const Graph*>(inhandles[i]);
      graphs.push_back(gr);
    }
    Graph* gptr = new Graph();
Minjie Wang's avatar
Minjie Wang committed
304
    *gptr = GraphOp::DisjointUnion(std::move(graphs));
305
306
307
308
    GraphHandle ghandle = gptr;
    *rv = ghandle;
  });

Minjie Wang's avatar
Minjie Wang committed
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionByNum")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    int64_t num = args[1];
    std::vector<Graph>&& 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) {
      Graph* ptr = new Graph();
      *ptr = std::move(rst[i]);
      ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
    }
    *rv = ptr_array;
  });

TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLDisjointPartitionBySizes")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    const IdArray sizes = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
    std::vector<Graph>&& 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) {
      Graph* ptr = new Graph();
      *ptr = std::move(rst[i]);
      ptr_array_data[i] = reinterpret_cast<std::intptr_t>(ptr);
    }
    *rv = ptr_array;
  });
GaiYu0's avatar
cpp lg  
GaiYu0 committed
344
345
346
347
348
349
350
351
352
353
354

TVM_REGISTER_GLOBAL("graph_index._CAPI_DGLGraphLineGraph")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
    GraphHandle ghandle = args[0];
    bool backtracking = args[1];
    const Graph* gptr = static_cast<Graph*>(ghandle);
    Graph* lgptr = new Graph();
    *lgptr = GraphOp::LineGraph(gptr, backtracking);
    GraphHandle lghandle = lgptr;
    *rv = lghandle;
  });
Minjie Wang's avatar
Minjie Wang committed
355
}  // namespace dgl