array.cc 23.9 KB
Newer Older
1
2
3
4
5
6
/*!
 *  Copyright (c) 2019 by Contributors
 * \file array/array.cc
 * \brief DGL array utilities implementation
 */
#include <dgl/array.h>
7
#include <dgl/graph_traversal.h>
8
9
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
10
#include <dgl/runtime/shared_mem.h>
11
12
13
14
#include "../c_api_common.h"
#include "./array_op.h"
#include "./arith.h"

15
using namespace dgl::runtime;
16

17
namespace dgl {
18
19
20
21
22
23
24
25
26
27
28
29
30
31
namespace aten {

IdArray NewIdArray(int64_t length, DLContext ctx, uint8_t nbits) {
  return IdArray::Empty({length}, DLDataType{kDLInt, nbits, 1}, ctx);
}

IdArray Clone(IdArray arr) {
  IdArray ret = NewIdArray(arr->shape[0], arr->ctx, arr->dtype.bits);
  ret.CopyFrom(arr);
  return ret;
}

IdArray Range(int64_t low, int64_t high, uint8_t nbits, DLContext ctx) {
  IdArray ret;
32
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Range", {
33
34
35
36
37
38
39
40
41
42
43
44
45
    if (nbits == 32) {
      ret = impl::Range<XPU, int32_t>(low, high, ctx);
    } else if (nbits == 64) {
      ret = impl::Range<XPU, int64_t>(low, high, ctx);
    } else {
      LOG(FATAL) << "Only int32 or int64 is supported.";
    }
  });
  return ret;
}

IdArray Full(int64_t val, int64_t length, uint8_t nbits, DLContext ctx) {
  IdArray ret;
46
  ATEN_XPU_SWITCH(ctx.device_type, XPU, "Full", {
47
48
49
50
51
52
53
54
55
56
57
58
    if (nbits == 32) {
      ret = impl::Full<XPU, int32_t>(val, length, ctx);
    } else if (nbits == 64) {
      ret = impl::Full<XPU, int64_t>(val, length, ctx);
    } else {
      LOG(FATAL) << "Only int32 or int64 is supported.";
    }
  });
  return ret;
}

IdArray AsNumBits(IdArray arr, uint8_t bits) {
59
60
61
62
63
  CHECK(bits == 32 || bits == 64)
    << "Invalid ID type. Must be int32 or int64, but got int"
    << static_cast<int>(bits) << ".";
  if (arr->dtype.bits == bits)
    return arr;
64
  IdArray ret;
65
  ATEN_XPU_SWITCH_CUDA(arr->ctx.device_type, XPU, "AsNumBits", {
66
67
68
69
70
71
72
73
74
75
76
    ATEN_ID_TYPE_SWITCH(arr->dtype, IdType, {
      ret = impl::AsNumBits<XPU, IdType>(arr, bits);
    });
  });
  return ret;
}

IdArray Add(IdArray lhs, IdArray rhs) {
  IdArray ret;
  CHECK_EQ(lhs->ctx, rhs->ctx) << "Both operands should have the same device context";
  CHECK_EQ(lhs->dtype, rhs->dtype) << "Both operands should have the same dtype";
77
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Add", {
78
79
80
81
82
83
84
85
86
87
88
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Add>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Sub(IdArray lhs, IdArray rhs) {
  IdArray ret;
  CHECK_EQ(lhs->ctx, rhs->ctx) << "Both operands should have the same device context";
  CHECK_EQ(lhs->dtype, rhs->dtype) << "Both operands should have the same dtype";
89
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Sub", {
90
91
92
93
94
95
96
97
98
99
100
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Sub>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Mul(IdArray lhs, IdArray rhs) {
  IdArray ret;
  CHECK_EQ(lhs->ctx, rhs->ctx) << "Both operands should have the same device context";
  CHECK_EQ(lhs->dtype, rhs->dtype) << "Both operands should have the same dtype";
101
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Mul", {
102
103
104
105
106
107
108
109
110
111
112
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Mul>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Div(IdArray lhs, IdArray rhs) {
  IdArray ret;
  CHECK_EQ(lhs->ctx, rhs->ctx) << "Both operands should have the same device context";
  CHECK_EQ(lhs->dtype, rhs->dtype) << "Both operands should have the same dtype";
113
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Div", {
114
115
116
117
118
119
120
121
122
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Div>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Add(IdArray lhs, dgl_id_t rhs) {
  IdArray ret;
123
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Add", {
124
125
126
127
128
129
130
131
132
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Add>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Sub(IdArray lhs, dgl_id_t rhs) {
  IdArray ret;
133
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Sub", {
134
135
136
137
138
139
140
141
142
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Sub>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Mul(IdArray lhs, dgl_id_t rhs) {
  IdArray ret;
143
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Mul", {
144
145
146
147
148
149
150
151
152
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Mul>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Div(IdArray lhs, dgl_id_t rhs) {
  IdArray ret;
153
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "Div", {
154
155
156
157
158
159
160
161
162
163
164
165
166
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Div>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Add(dgl_id_t lhs, IdArray rhs) {
  return Add(rhs, lhs);
}

IdArray Sub(dgl_id_t lhs, IdArray rhs) {
  IdArray ret;
167
  ATEN_XPU_SWITCH(rhs->ctx.device_type, XPU, "Sub", {
168
169
170
171
172
173
174
175
176
177
178
179
180
    ATEN_ID_TYPE_SWITCH(rhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Sub>(lhs, rhs);
    });
  });
  return ret;
}

IdArray Mul(dgl_id_t lhs, IdArray rhs) {
  return Mul(rhs, lhs);
}

IdArray Div(dgl_id_t lhs, IdArray rhs) {
  IdArray ret;
181
  ATEN_XPU_SWITCH(rhs->ctx.device_type, XPU, "Div", {
182
183
184
185
186
187
188
189
190
    ATEN_ID_TYPE_SWITCH(rhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::Div>(lhs, rhs);
    });
  });
  return ret;
}

BoolArray LT(IdArray lhs, dgl_id_t rhs) {
  BoolArray ret;
191
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "LT", {
192
193
194
195
196
197
198
199
200
201
202
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::BinaryElewise<XPU, IdType, arith::LT>(lhs, rhs);
    });
  });
  return ret;
}

IdArray HStack(IdArray lhs, IdArray rhs) {
  IdArray ret;
  CHECK_EQ(lhs->ctx, rhs->ctx) << "Both operands should have the same device context";
  CHECK_EQ(lhs->dtype, rhs->dtype) << "Both operands should have the same dtype";
203
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "HStack", {
204
205
206
207
208
209
210
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::HStack<XPU, IdType>(lhs, rhs);
    });
  });
  return ret;
}

211
212
213
NDArray IndexSelect(NDArray array, IdArray index) {
  NDArray ret;
  // TODO(BarclayII): check if array and index match in context
214
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "IndexSelect", {
215
216
217
218
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
        ret = impl::IndexSelect<XPU, DType, IdType>(array, index);
      });
219
220
221
222
223
    });
  });
  return ret;
}

224
225
226
template<typename ValueType>
ValueType IndexSelect(NDArray array, uint64_t index) {
  ValueType ret = 0;
227
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "IndexSelect", {
228
229
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ret = impl::IndexSelect<XPU, DType>(array, index);
230
231
232
233
    });
  });
  return ret;
}
234
235
236
237
238
239
template int32_t IndexSelect<int32_t>(NDArray array, uint64_t index);
template int64_t IndexSelect<int64_t>(NDArray array, uint64_t index);
template uint32_t IndexSelect<uint32_t>(NDArray array, uint64_t index);
template uint64_t IndexSelect<uint64_t>(NDArray array, uint64_t index);
template float IndexSelect<float>(NDArray array, uint64_t index);
template double IndexSelect<double>(NDArray array, uint64_t index);
240

241
242
NDArray Scatter(NDArray array, IdArray indices) {
  NDArray ret;
243
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Scatter", {
244
245
246
247
248
249
250
251
252
253
254
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(indices->dtype, IdType, {
        ret = impl::Scatter<XPU, DType, IdType>(array, indices);
      });
    });
  });
  return ret;
}

NDArray Repeat(NDArray array, IdArray repeats) {
  NDArray ret;
255
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Repeat", {
256
257
258
259
260
261
262
263
264
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(repeats->dtype, IdType, {
        ret = impl::Repeat<XPU, DType, IdType>(array, repeats);
      });
    });
  });
  return ret;
}

265
266
IdArray Relabel_(const std::vector<IdArray>& arrays) {
  IdArray ret;
267
  ATEN_XPU_SWITCH(arrays[0]->ctx.device_type, XPU, "Relabel_", {
268
269
270
271
272
273
274
    ATEN_ID_TYPE_SWITCH(arrays[0]->dtype, IdType, {
      ret = impl::Relabel_<XPU, IdType>(arrays);
    });
  });
  return ret;
}

275
276
277
template<typename ValueType>
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, ValueType pad_value) {
  std::tuple<NDArray, IdArray, IdArray> ret;
278
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Pack", {
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
    ATEN_DTYPE_SWITCH(array->dtype, DType, "array", {
      ret = impl::Pack<XPU, DType>(array, static_cast<DType>(pad_value));
    });
  });
  return ret;
}

template std::tuple<NDArray, IdArray, IdArray> Pack<int32_t>(NDArray, int32_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<int64_t>(NDArray, int64_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<uint32_t>(NDArray, uint32_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<uint64_t>(NDArray, uint64_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<float>(NDArray, float);
template std::tuple<NDArray, IdArray, IdArray> Pack<double>(NDArray, double);

std::pair<NDArray, IdArray> ConcatSlices(NDArray array, IdArray lengths) {
  std::pair<NDArray, IdArray> ret;
295
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "ConcatSlices", {
296
297
298
299
300
301
302
303
304
    ATEN_DTYPE_SWITCH(array->dtype, DType, "array", {
      ATEN_ID_TYPE_SWITCH(lengths->dtype, IdType, {
        ret = impl::ConcatSlices<XPU, DType, IdType>(array, lengths);
      });
    });
  });
  return ret;
}

305
306
307
308
///////////////////////// CSR routines //////////////////////////

bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
  bool ret = false;
309
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRIsNonZero", {
310
311
312
313
314
315
316
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
  NDArray ret;
317
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRIsNonZero", {
318
319
320
321
322
323
324
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

bool CSRHasDuplicate(CSRMatrix csr) {
  bool ret = false;
325
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRHasDuplicate", {
326
327
328
329
330
331
332
    ret = impl::CSRHasDuplicate<XPU, IdType>(csr);
  });
  return ret;
}

int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
  int64_t ret = 0;
333
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetRowNNZ", {
334
335
336
337
338
339
340
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray row) {
  NDArray ret;
341
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetRowNNZ", {
342
343
344
345
346
347
348
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
  NDArray ret;
349
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetRowColumnIndices", {
350
351
352
353
354
355
356
    ret = impl::CSRGetRowColumnIndices<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
  NDArray ret;
357
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetRowData", {
358
    ret = impl::CSRGetRowData<XPU, IdType>(csr, row);
359
360
361
362
363
364
  });
  return ret;
}

NDArray CSRGetData(CSRMatrix csr, int64_t row, int64_t col) {
  NDArray ret;
365
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetData", {
366
    ret = impl::CSRGetData<XPU, IdType>(csr, row, col);
367
368
369
370
371
372
  });
  return ret;
}

NDArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
  NDArray ret;
373
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetData", {
374
    ret = impl::CSRGetData<XPU, IdType>(csr, rows, cols);
375
376
377
378
379
380
381
  });
  return ret;
}

std::vector<NDArray> CSRGetDataAndIndices(
    CSRMatrix csr, NDArray rows, NDArray cols) {
  std::vector<NDArray> ret;
382
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetDataAndIndices", {
383
    ret = impl::CSRGetDataAndIndices<XPU, IdType>(csr, rows, cols);
384
385
386
387
388
389
  });
  return ret;
}

CSRMatrix CSRTranspose(CSRMatrix csr) {
  CSRMatrix ret;
390
391
392
393
  ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRTranspose", {
    ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
      ret = impl::CSRTranspose<XPU, IdType>(csr);
    });
394
395
396
397
398
399
400
  });
  return ret;
}

COOMatrix CSRToCOO(CSRMatrix csr, bool data_as_order) {
  COOMatrix ret;
  if (data_as_order) {
401
    ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRToCOODataAsOrder", {
402
403
404
405
406
      ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
        ret = impl::CSRToCOODataAsOrder<XPU, IdType>(csr);
      });
    });
  } else {
407
    ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRToCOO", {
408
409
410
411
412
413
414
415
416
417
      ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
        ret = impl::CSRToCOO<XPU, IdType>(csr);
      });
    });
  }
  return ret;
}

CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end) {
  CSRMatrix ret;
418
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRSliceRows", {
419
    ret = impl::CSRSliceRows<XPU, IdType>(csr, start, end);
420
421
422
423
424
425
  });
  return ret;
}

CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
  CSRMatrix ret;
426
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRSliceRows", {
427
    ret = impl::CSRSliceRows<XPU, IdType>(csr, rows);
428
429
430
431
432
433
  });
  return ret;
}

CSRMatrix CSRSliceMatrix(CSRMatrix csr, NDArray rows, NDArray cols) {
  CSRMatrix ret;
434
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRSliceMatrix", {
435
    ret = impl::CSRSliceMatrix<XPU, IdType>(csr, rows, cols);
436
437
438
439
  });
  return ret;
}

440
void CSRSort_(CSRMatrix* csr) {
441
  ATEN_CSR_SWITCH(*csr, XPU, IdType, "CSRSort_", {
442
    impl::CSRSort_<XPU, IdType>(csr);
Da Zheng's avatar
Da Zheng committed
443
444
445
  });
}

Da Zheng's avatar
Da Zheng committed
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
CSRMatrix CSRReorder(CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids) {
  CSRMatrix ret;
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRReorder", {
    ret = impl::CSRReorder<XPU, IdType>(csr, new_row_ids, new_col_ids);
  });
  return ret;
}

COOMatrix COOReorder(COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids) {
  COOMatrix ret;
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOReorder", {
    ret = impl::COOReorder<XPU, IdType>(coo, new_row_ids, new_col_ids);
  });
  return ret;
}

462
463
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries) {
  CSRMatrix ret;
464
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRRemove", {
465
466
467
468
469
    ret = impl::CSRRemove<XPU, IdType>(csr, entries);
  });
  return ret;
}

470
471
472
COOMatrix CSRRowWiseSampling(
    CSRMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob, bool replace) {
  COOMatrix ret;
473
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseSampling", {
474
    if (IsNullArray(prob)) {
475
476
477
478
479
480
481
482
483
484
485
486
      ret = impl::CSRRowWiseSamplingUniform<XPU, IdType>(mat, rows, num_samples, replace);
    } else {
      ATEN_FLOAT_TYPE_SWITCH(prob->dtype, FloatType, "probability", {
        ret = impl::CSRRowWiseSampling<XPU, IdType, FloatType>(
            mat, rows, num_samples, prob, replace);
      });
    }
  });
  return ret;
}

COOMatrix CSRRowWiseTopk(
487
    CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending) {
488
  COOMatrix ret;
489
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseTopk", {
490
491
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::CSRRowWiseTopk<XPU, IdType, DType>(
492
493
494
495
496
497
          mat, rows, k, weight, ascending);
    });
  });
  return ret;
}

498
499
///////////////////////// COO routines //////////////////////////

500
501
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
  bool ret = false;
502
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
503
504
505
506
507
508
509
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

NDArray COOIsNonZero(COOMatrix coo, NDArray row, NDArray col) {
  NDArray ret;
510
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
511
512
513
514
515
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

516
517
bool COOHasDuplicate(COOMatrix coo) {
  bool ret = false;
518
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOHasDuplicate", {
519
520
521
522
523
    ret = impl::COOHasDuplicate<XPU, IdType>(coo);
  });
  return ret;
}

524
525
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
  int64_t ret = 0;
526
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowNNZ", {
527
528
529
530
531
532
533
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

NDArray COOGetRowNNZ(COOMatrix coo, NDArray row) {
  NDArray ret;
534
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowNNZ", {
535
536
537
538
539
540
541
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

std::pair<NDArray, NDArray> COOGetRowDataAndIndices(COOMatrix coo, int64_t row) {
  std::pair<NDArray, NDArray> ret;
542
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowDataAndIndices", {
543
    ret = impl::COOGetRowDataAndIndices<XPU, IdType>(coo, row);
544
545
546
547
548
549
  });
  return ret;
}

NDArray COOGetData(COOMatrix coo, int64_t row, int64_t col) {
  NDArray ret;
550
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetData", {
551
    ret = impl::COOGetData<XPU, IdType>(coo, row, col);
552
553
554
555
556
557
558
  });
  return ret;
}

std::vector<NDArray> COOGetDataAndIndices(
    COOMatrix coo, NDArray rows, NDArray cols) {
  std::vector<NDArray> ret;
559
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetDataAndIndices", {
560
    ret = impl::COOGetDataAndIndices<XPU, IdType>(coo, rows, cols);
561
562
563
564
565
  });
  return ret;
}

COOMatrix COOTranspose(COOMatrix coo) {
566
  return COOMatrix(coo.num_cols, coo.num_rows, coo.col, coo.row, coo.data);
567
568
}

569
570
CSRMatrix COOToCSR(COOMatrix coo) {
  CSRMatrix ret;
571
572
573
574
  ATEN_XPU_SWITCH_CUDA(coo.row->ctx.device_type, XPU, "COOToCSR", {
    ATEN_ID_TYPE_SWITCH(coo.row->dtype, IdType, {
      ret = impl::COOToCSR<XPU, IdType>(coo);
    });
575
576
577
578
  });
  return ret;
}

579
580
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end) {
  COOMatrix ret;
581
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
582
    ret = impl::COOSliceRows<XPU, IdType>(coo, start, end);
583
584
585
586
587
588
  });
  return ret;
}

COOMatrix COOSliceRows(COOMatrix coo, NDArray rows) {
  COOMatrix ret;
589
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
590
    ret = impl::COOSliceRows<XPU, IdType>(coo, rows);
591
592
593
594
595
596
  });
  return ret;
}

COOMatrix COOSliceMatrix(COOMatrix coo, NDArray rows, NDArray cols) {
  COOMatrix ret;
597
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceMatrix", {
598
599
600
601
602
603
604
    ret = impl::COOSliceMatrix<XPU, IdType>(coo, rows, cols);
  });
  return ret;
}

COOMatrix COOSort(COOMatrix mat, bool sort_column) {
  COOMatrix ret;
605
606
607
608
  ATEN_XPU_SWITCH_CUDA(mat.row->ctx.device_type, XPU, "COOSort", {
    ATEN_ID_TYPE_SWITCH(mat.row->dtype, IdType, {
      ret = impl::COOSort<XPU, IdType>(mat, sort_column);
    });
609
610
611
612
  });
  return ret;
}

613
614
COOMatrix COORemove(COOMatrix coo, IdArray entries) {
  COOMatrix ret;
615
  ATEN_COO_SWITCH(coo, XPU, IdType, "COORemove", {
616
617
618
619
620
    ret = impl::COORemove<XPU, IdType>(coo, entries);
  });
  return ret;
}

621
622
623
COOMatrix COORowWiseSampling(
    COOMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob, bool replace) {
  COOMatrix ret;
624
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseSampling", {
625
    if (IsNullArray(prob)) {
626
627
628
629
630
631
632
633
634
635
636
637
638
639
      ret = impl::COORowWiseSamplingUniform<XPU, IdType>(mat, rows, num_samples, replace);
    } else {
      ATEN_FLOAT_TYPE_SWITCH(prob->dtype, FloatType, "probability", {
        ret = impl::COORowWiseSampling<XPU, IdType, FloatType>(
            mat, rows, num_samples, prob, replace);
      });
    }
  });
  return ret;
}

COOMatrix COORowWiseTopk(
    COOMatrix mat, IdArray rows, int64_t k, FloatArray weight, bool ascending) {
  COOMatrix ret;
640
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseTopk", {
641
642
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::COORowWiseTopk<XPU, IdType, DType>(
643
644
          mat, rows, k, weight, ascending);
    });
645
646
647
648
  });
  return ret;
}

649
650
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
  std::pair<COOMatrix, IdArray> ret;
651
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOCoalesce", {
652
653
654
655
656
    ret = impl::COOCoalesce<XPU, IdType>(coo);
  });
  return ret;
}

657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
///////////////////////// Graph Traverse  routines //////////////////////////
Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source) {
  Frontiers ret;
  CHECK_EQ(csr.indptr->ctx.device_type, source->ctx.device_type) <<
    "Graph and source should in the same device context";
  CHECK_EQ(csr.indices->dtype, source->dtype) <<
    "Graph and source should in the same dtype";
  CHECK_EQ(csr.num_rows, csr.num_cols) <<
    "Graph traversal can only work on square-shaped CSR.";
  ATEN_XPU_SWITCH(source->ctx.device_type, XPU, "BFSNodesFrontiers", {
    ATEN_ID_TYPE_SWITCH(source->dtype, IdType, {
      ret = impl::BFSNodesFrontiers<XPU, IdType>(csr, source);
    });
  });
  return ret;
}

Frontiers BFSEdgesFrontiers(const CSRMatrix& csr, IdArray source) {
  Frontiers ret;
  CHECK_EQ(csr.indptr->ctx.device_type, source->ctx.device_type) <<
    "Graph and source should in the same device context";
  CHECK_EQ(csr.indices->dtype, source->dtype) <<
    "Graph and source should in the same dtype";
  CHECK_EQ(csr.num_rows, csr.num_cols) <<
    "Graph traversal can only work on square-shaped CSR.";
  ATEN_XPU_SWITCH(source->ctx.device_type, XPU, "BFSEdgesFrontiers", {
    ATEN_ID_TYPE_SWITCH(source->dtype, IdType, {
      ret = impl::BFSEdgesFrontiers<XPU, IdType>(csr, source);
    });
  });
  return ret;
}

Frontiers TopologicalNodesFrontiers(const CSRMatrix& csr) {
  Frontiers ret;
  CHECK_EQ(csr.num_rows, csr.num_cols) <<
    "Graph traversal can only work on square-shaped CSR.";
  ATEN_XPU_SWITCH(csr.indptr->ctx.device_type, XPU, "TopologicalNodesFrontiers", {
    ATEN_ID_TYPE_SWITCH(csr.indices->dtype, IdType, {
      ret = impl::TopologicalNodesFrontiers<XPU, IdType>(csr);
    });
  });
  return ret;
}

Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source) {
  Frontiers ret;
  CHECK_EQ(csr.indptr->ctx.device_type, source->ctx.device_type) <<
    "Graph and source should in the same device context";
  CHECK_EQ(csr.indices->dtype, source->dtype) <<
    "Graph and source should in the same dtype";
  CHECK_EQ(csr.num_rows, csr.num_cols) <<
    "Graph traversal can only work on square-shaped CSR.";
  ATEN_XPU_SWITCH(source->ctx.device_type, XPU, "DGLDFSEdges", {
    ATEN_ID_TYPE_SWITCH(source->dtype, IdType, {
      ret = impl::DGLDFSEdges<XPU, IdType>(csr, source);
    });
  });
  return ret;
}
Frontiers DGLDFSLabeledEdges(const CSRMatrix& csr,
                             IdArray source,
                             const bool has_reverse_edge,
                             const bool has_nontree_edge,
                             const bool return_labels) {
  Frontiers ret;
  CHECK_EQ(csr.indptr->ctx.device_type, source->ctx.device_type) <<
    "Graph and source should in the same device context";
  CHECK_EQ(csr.indices->dtype, source->dtype) <<
    "Graph and source should in the same dtype";
  CHECK_EQ(csr.num_rows, csr.num_cols) <<
    "Graph traversal can only work on square-shaped CSR.";
  ATEN_XPU_SWITCH(source->ctx.device_type, XPU, "DGLDFSLabeledEdges", {
    ATEN_ID_TYPE_SWITCH(source->dtype, IdType, {
      ret = impl::DGLDFSLabeledEdges<XPU, IdType>(csr,
                                                  source,
                                                  has_reverse_edge,
                                                  has_nontree_edge,
                                                  return_labels);
    });
  });
  return ret;
}

741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
///////////////////////// C APIs /////////////////////////
DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetFormat")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SparseMatrixRef spmat = args[0];
    *rv = spmat->format;
  });

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetNumRows")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SparseMatrixRef spmat = args[0];
    *rv = spmat->num_rows;
  });

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetNumCols")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SparseMatrixRef spmat = args[0];
    *rv = spmat->num_cols;
  });

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetIndices")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SparseMatrixRef spmat = args[0];
    const int64_t i = args[1];
    *rv = spmat->indices[i];
  });

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetFlags")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    SparseMatrixRef spmat = args[0];
    List<Value> flags;
    for (bool flg : spmat->flags) {
      flags.push_back(Value(MakeValue(flg)));
    }
    *rv = flags;
  });

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLCreateSparseMatrix")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const int32_t format = args[0];
    const int64_t nrows = args[1];
    const int64_t ncols = args[2];
    const List<Value> indices = args[3];
    const List<Value> flags = args[4];
    std::shared_ptr<SparseMatrix> spmat(new SparseMatrix(
          format, nrows, ncols,
          ListValueToVector<IdArray>(indices),
          ListValueToVector<bool>(flags)));
    *rv = SparseMatrixRef(spmat);
  });

791
792
793
794
795
796
797
798
799
800
DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLExistSharedMemArray")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    const std::string name = args[0];
#ifndef _WIN32
    *rv = SharedMemory::Exist(name);
#else
    *rv = false;
#endif  // _WIN32
  });

801
802
}  // namespace aten
}  // namespace dgl