array.cc 25.4 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
#include <dgl/runtime/device_api.h>
#include <sstream>
13
14
15
16
#include "../c_api_common.h"
#include "./array_op.h"
#include "./arith.h"

17
using namespace dgl::runtime;
18

19
namespace dgl {
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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;
34
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Range", {
35
36
37
38
39
40
41
42
43
44
45
46
47
    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;
48
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Full", {
49
50
51
52
53
54
55
56
57
58
59
60
    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) {
61
62
63
64
65
  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;
66
  IdArray ret;
67
  ATEN_XPU_SWITCH_CUDA(arr->ctx.device_type, XPU, "AsNumBits", {
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 HStack(IdArray lhs, IdArray rhs) {
  IdArray ret;
77
78
  CHECK_SAME_CONTEXT(lhs, rhs);
  CHECK_SAME_DTYPE(lhs, rhs);
79
  ATEN_XPU_SWITCH(lhs->ctx.device_type, XPU, "HStack", {
80
81
82
83
84
85
86
    ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
      ret = impl::HStack<XPU, IdType>(lhs, rhs);
    });
  });
  return ret;
}

87
88
NDArray IndexSelect(NDArray array, IdArray index) {
  NDArray ret;
89
  CHECK_SAME_CONTEXT(array, index);
90
91
92
  CHECK_GE(array->ndim, 1) << "Only support array with at least 1 dimension";
  CHECK_EQ(array->shape[0], array.NumElements()) << "Only support tensor"
    << " whose first dimension equals number of elements, e.g. (5,), (5, 1)";
93
94
  CHECK_EQ(index->ndim, 1) << "Index array must be an 1D array.";
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "IndexSelect", {
95
96
97
98
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
        ret = impl::IndexSelect<XPU, DType, IdType>(array, index);
      });
99
100
101
102
103
    });
  });
  return ret;
}

104
template<typename ValueType>
105
ValueType IndexSelect(NDArray array, int64_t index) {
106
  CHECK_EQ(array->ndim, 1) << "Only support select values from 1D array.";
107
108
  CHECK(index >= 0 && index < array.NumElements())
    << "Index " << index << " is out of bound.";
109
  ValueType ret = 0;
110
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "IndexSelect", {
111
112
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ret = impl::IndexSelect<XPU, DType>(array, index);
113
114
115
116
    });
  });
  return ret;
}
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
template int32_t IndexSelect<int32_t>(NDArray array, int64_t index);
template int64_t IndexSelect<int64_t>(NDArray array, int64_t index);
template uint32_t IndexSelect<uint32_t>(NDArray array, int64_t index);
template uint64_t IndexSelect<uint64_t>(NDArray array, int64_t index);
template float IndexSelect<float>(NDArray array, int64_t index);
template double IndexSelect<double>(NDArray array, int64_t index);

NDArray IndexSelect(NDArray array, int64_t start, int64_t end) {
  CHECK_EQ(array->ndim, 1) << "Only support select values from 1D array.";
  CHECK(start >= 0 && start < array.NumElements())
    << "Index " << start << " is out of bound.";
  CHECK(end >= 0 && end <= array.NumElements())
    << "Index " << end << " is out of bound.";
  CHECK_LE(start, end);
  auto device = runtime::DeviceAPI::Get(array->ctx);
  const int64_t len = end - start;
  NDArray ret = NDArray::Empty({len}, array->dtype, array->ctx);
  ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
    device->CopyDataFromTo(array->data, start * sizeof(DType),
                           ret->data, 0, len * sizeof(DType),
                           array->ctx, ret->ctx, array->dtype, nullptr);
  });
  return ret;
}
141

142
143
NDArray Scatter(NDArray array, IdArray indices) {
  NDArray ret;
144
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Scatter", {
145
146
147
148
149
150
151
152
153
154
155
    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;
156
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Repeat", {
157
158
159
160
161
162
163
164
165
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(repeats->dtype, IdType, {
        ret = impl::Repeat<XPU, DType, IdType>(array, repeats);
      });
    });
  });
  return ret;
}

166
167
IdArray Relabel_(const std::vector<IdArray>& arrays) {
  IdArray ret;
168
  ATEN_XPU_SWITCH(arrays[0]->ctx.device_type, XPU, "Relabel_", {
169
170
171
172
173
174
175
    ATEN_ID_TYPE_SWITCH(arrays[0]->dtype, IdType, {
      ret = impl::Relabel_<XPU, IdType>(arrays);
    });
  });
  return ret;
}

176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
NDArray Concat(const std::vector<IdArray>& arrays) {
  CHECK(arrays.size() > 1) << "Number of arrays should larger than 1";
  IdArray ret;

  int64_t len = 0, offset = 0;
  for (size_t i = 0; i < arrays.size(); ++i) {
    len += arrays[i]->shape[0];
    CHECK_SAME_DTYPE(arrays[0], arrays[i]);
    CHECK_SAME_CONTEXT(arrays[0], arrays[i]);
  }

  NDArray ret_arr = NDArray::Empty({len},
                                   arrays[0]->dtype,
                                   arrays[0]->ctx);

  auto device = runtime::DeviceAPI::Get(arrays[0]->ctx);
  for (size_t i = 0; i < arrays.size(); ++i) {
    ATEN_DTYPE_SWITCH(arrays[i]->dtype, DType, "array", {
      device->CopyDataFromTo(
        static_cast<DType*>(arrays[i]->data),
        0,
        static_cast<DType*>(ret_arr->data),
        offset,
        arrays[i]->shape[0] * sizeof(DType),
        arrays[i]->ctx,
        ret_arr->ctx,
        arrays[i]->dtype,
        nullptr);

        offset += arrays[i]->shape[0] * sizeof(DType);
    });
  }

  return ret_arr;
}

212
213
214
template<typename ValueType>
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, ValueType pad_value) {
  std::tuple<NDArray, IdArray, IdArray> ret;
215
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Pack", {
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
    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;
232
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "ConcatSlices", {
233
234
235
236
237
238
239
240
241
    ATEN_DTYPE_SWITCH(array->dtype, DType, "array", {
      ATEN_ID_TYPE_SWITCH(lengths->dtype, IdType, {
        ret = impl::ConcatSlices<XPU, DType, IdType>(array, lengths);
      });
    });
  });
  return ret;
}

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
IdArray CumSum(IdArray array, bool prepend_zero) {
  IdArray ret;
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "CumSum", {
    ATEN_ID_TYPE_SWITCH(array->dtype, IdType, {
      ret = impl::CumSum<XPU, IdType>(array, prepend_zero);
    });
  });
  return ret;
}

std::string ToDebugString(NDArray array) {
  std::ostringstream oss;
  NDArray a = array.CopyTo(DLContext{kDLCPU, 0});
  oss << "array([";
  ATEN_DTYPE_SWITCH(a->dtype, DType, "array", {
    for (int64_t i = 0; i < std::min<int64_t>(a.NumElements(), 10L); ++i) {
      oss << a.Ptr<DType>()[i] << ", ";
    }
  });
  if (a.NumElements() > 10)
    oss << "...";
  oss << "], dtype=" << array->dtype << ", ctx=" << array->ctx << ")";
  return oss.str();
}

267
268
269
///////////////////////// CSR routines //////////////////////////

bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
270
271
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
  CHECK(col >= 0 && col < csr.num_cols) << "Invalid col index: " << col;
272
  bool ret = false;
273
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRIsNonZero", {
274
275
276
277
278
279
280
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
  NDArray ret;
281
282
283
284
285
  CHECK_SAME_DTYPE(csr.indices, row);
  CHECK_SAME_DTYPE(csr.indices, col);
  CHECK_SAME_CONTEXT(csr.indices, row);
  CHECK_SAME_CONTEXT(csr.indices, col);
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRIsNonZero", {
286
287
288
289
290
291
292
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

bool CSRHasDuplicate(CSRMatrix csr) {
  bool ret = false;
293
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRHasDuplicate", {
294
295
296
297
298
299
    ret = impl::CSRHasDuplicate<XPU, IdType>(csr);
  });
  return ret;
}

int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
300
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
301
  int64_t ret = 0;
302
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowNNZ", {
303
304
305
306
307
308
309
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray row) {
  NDArray ret;
310
311
312
  CHECK_SAME_DTYPE(csr.indices, row);
  CHECK_SAME_CONTEXT(csr.indices, row);
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowNNZ", {
313
314
315
316
317
318
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
319
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
320
  NDArray ret;
321
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowColumnIndices", {
322
323
324
325
326
327
    ret = impl::CSRGetRowColumnIndices<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
328
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
329
  NDArray ret;
330
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowData", {
331
    ret = impl::CSRGetRowData<XPU, IdType>(csr, row);
332
333
334
335
  });
  return ret;
}

336
337
338
339
340
341
342
343
344
345
bool CSRIsSorted(CSRMatrix csr) {
  if (csr.indices->shape[0] <= 1)
    return true;
  bool ret = false;
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRIsSorted", {
    ret = impl::CSRIsSorted<XPU, IdType>(csr);
  });
  return ret;
}

346
NDArray CSRGetData(CSRMatrix csr, int64_t row, int64_t col) {
347
348
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
  CHECK(col >= 0 && col < csr.num_cols) << "Invalid col index: " << col;
349
  NDArray ret;
350
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetData", {
351
    ret = impl::CSRGetData<XPU, IdType>(csr, row, col);
352
353
354
355
356
357
  });
  return ret;
}

NDArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
  NDArray ret;
358
359
360
361
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
  CHECK_SAME_CONTEXT(csr.indices, rows);
  CHECK_SAME_CONTEXT(csr.indices, cols);
362
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetData", {
363
    ret = impl::CSRGetData<XPU, IdType>(csr, rows, cols);
364
365
366
367
368
369
  });
  return ret;
}

std::vector<NDArray> CSRGetDataAndIndices(
    CSRMatrix csr, NDArray rows, NDArray cols) {
370
371
372
373
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
  CHECK_SAME_CONTEXT(csr.indices, rows);
  CHECK_SAME_CONTEXT(csr.indices, cols);
374
  std::vector<NDArray> ret;
375
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRGetDataAndIndices", {
376
    ret = impl::CSRGetDataAndIndices<XPU, IdType>(csr, rows, cols);
377
378
379
380
381
382
  });
  return ret;
}

CSRMatrix CSRTranspose(CSRMatrix csr) {
  CSRMatrix ret;
383
384
385
386
  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);
    });
387
388
389
390
391
392
393
  });
  return ret;
}

COOMatrix CSRToCOO(CSRMatrix csr, bool data_as_order) {
  COOMatrix ret;
  if (data_as_order) {
394
    ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRToCOODataAsOrder", {
395
396
397
398
399
      ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
        ret = impl::CSRToCOODataAsOrder<XPU, IdType>(csr);
      });
    });
  } else {
400
    ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRToCOO", {
401
402
403
404
405
406
407
408
409
      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) {
410
411
412
  CHECK(start >= 0 && start < csr.num_rows) << "Invalid start index: " << start;
  CHECK(end >= 0 && end <= csr.num_rows) << "Invalid end index: " << end;
  CHECK_GE(end, start);
413
  CSRMatrix ret;
414
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRSliceRows", {
415
    ret = impl::CSRSliceRows<XPU, IdType>(csr, start, end);
416
417
418
419
420
  });
  return ret;
}

CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
421
422
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_CONTEXT(csr.indices, rows);
423
  CSRMatrix ret;
424
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRSliceRows", {
425
    ret = impl::CSRSliceRows<XPU, IdType>(csr, rows);
426
427
428
429
430
  });
  return ret;
}

CSRMatrix CSRSliceMatrix(CSRMatrix csr, NDArray rows, NDArray cols) {
431
432
433
434
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
  CHECK_SAME_CONTEXT(csr.indices, rows);
  CHECK_SAME_CONTEXT(csr.indices, cols);
435
  CSRMatrix ret;
436
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRSliceMatrix", {
437
    ret = impl::CSRSliceMatrix<XPU, IdType>(csr, rows, cols);
438
439
440
441
  });
  return ret;
}

442
void CSRSort_(CSRMatrix* csr) {
443
444
445
  if (csr->sorted)
    return;
  ATEN_CSR_SWITCH_CUDA(*csr, XPU, IdType, "CSRSort_", {
446
    impl::CSRSort_<XPU, IdType>(csr);
Da Zheng's avatar
Da Zheng committed
447
448
449
  });
}

Da Zheng's avatar
Da Zheng committed
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
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;
}

466
467
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries) {
  CSRMatrix ret;
468
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRRemove", {
469
470
471
472
473
    ret = impl::CSRRemove<XPU, IdType>(csr, entries);
  });
  return ret;
}

474
475
476
COOMatrix CSRRowWiseSampling(
    CSRMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob, bool replace) {
  COOMatrix ret;
477
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseSampling", {
478
    if (IsNullArray(prob)) {
479
480
481
482
483
484
485
486
487
488
489
490
      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(
491
    CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending) {
492
  COOMatrix ret;
493
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseTopk", {
494
495
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::CSRRowWiseTopk<XPU, IdType, DType>(
496
497
498
499
500
501
          mat, rows, k, weight, ascending);
    });
  });
  return ret;
}

502
503
///////////////////////// COO routines //////////////////////////

504
505
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
  bool ret = false;
506
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
507
508
509
510
511
512
513
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

NDArray COOIsNonZero(COOMatrix coo, NDArray row, NDArray col) {
  NDArray ret;
514
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
515
516
517
518
519
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

520
521
bool COOHasDuplicate(COOMatrix coo) {
  bool ret = false;
522
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOHasDuplicate", {
523
524
525
526
527
    ret = impl::COOHasDuplicate<XPU, IdType>(coo);
  });
  return ret;
}

528
529
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
  int64_t ret = 0;
530
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowNNZ", {
531
532
533
534
535
536
537
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

NDArray COOGetRowNNZ(COOMatrix coo, NDArray row) {
  NDArray ret;
538
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowNNZ", {
539
540
541
542
543
544
545
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

std::pair<NDArray, NDArray> COOGetRowDataAndIndices(COOMatrix coo, int64_t row) {
  std::pair<NDArray, NDArray> ret;
546
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowDataAndIndices", {
547
    ret = impl::COOGetRowDataAndIndices<XPU, IdType>(coo, row);
548
549
550
551
552
553
  });
  return ret;
}

NDArray COOGetData(COOMatrix coo, int64_t row, int64_t col) {
  NDArray ret;
554
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetData", {
555
    ret = impl::COOGetData<XPU, IdType>(coo, row, col);
556
557
558
559
560
561
562
  });
  return ret;
}

std::vector<NDArray> COOGetDataAndIndices(
    COOMatrix coo, NDArray rows, NDArray cols) {
  std::vector<NDArray> ret;
563
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetDataAndIndices", {
564
    ret = impl::COOGetDataAndIndices<XPU, IdType>(coo, rows, cols);
565
566
567
568
569
  });
  return ret;
}

COOMatrix COOTranspose(COOMatrix coo) {
570
  return COOMatrix(coo.num_cols, coo.num_rows, coo.col, coo.row, coo.data);
571
572
}

573
574
CSRMatrix COOToCSR(COOMatrix coo) {
  CSRMatrix ret;
575
576
577
578
  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);
    });
579
580
581
582
  });
  return ret;
}

583
584
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end) {
  COOMatrix ret;
585
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
586
    ret = impl::COOSliceRows<XPU, IdType>(coo, start, end);
587
588
589
590
591
592
  });
  return ret;
}

COOMatrix COOSliceRows(COOMatrix coo, NDArray rows) {
  COOMatrix ret;
593
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
594
    ret = impl::COOSliceRows<XPU, IdType>(coo, rows);
595
596
597
598
599
600
  });
  return ret;
}

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

607
608
609
610
611
612
void COOSort_(COOMatrix* mat, bool sort_column) {
  if ((mat->row_sorted && !sort_column) || mat->col_sorted)
    return;
  ATEN_XPU_SWITCH_CUDA(mat->row->ctx.device_type, XPU, "COOSort_", {
    ATEN_ID_TYPE_SWITCH(mat->row->dtype, IdType, {
      impl::COOSort_<XPU, IdType>(mat, sort_column);
613
    });
614
  });
615
616
617
618
619
620
621
622
623
}

std::pair<bool, bool> COOIsSorted(COOMatrix coo) {
  if (coo.row->shape[0] <= 1)
    return {true, true};
  std::pair<bool, bool> ret;
  ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, "COOIsSorted", {
    ret = impl::COOIsSorted<XPU, IdType>(coo);
  });
624
625
626
  return ret;
}

627
628
COOMatrix COORemove(COOMatrix coo, IdArray entries) {
  COOMatrix ret;
629
  ATEN_COO_SWITCH(coo, XPU, IdType, "COORemove", {
630
631
632
633
634
    ret = impl::COORemove<XPU, IdType>(coo, entries);
  });
  return ret;
}

635
636
637
COOMatrix COORowWiseSampling(
    COOMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob, bool replace) {
  COOMatrix ret;
638
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseSampling", {
639
    if (IsNullArray(prob)) {
640
641
642
643
644
645
646
647
648
649
650
651
652
653
      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;
654
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseTopk", {
655
656
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::COORowWiseTopk<XPU, IdType, DType>(
657
658
          mat, rows, k, weight, ascending);
    });
659
660
661
662
  });
  return ret;
}

663
664
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
  std::pair<COOMatrix, IdArray> ret;
665
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOCoalesce", {
666
667
668
669
670
    ret = impl::COOCoalesce<XPU, IdType>(coo);
  });
  return ret;
}

671
///////////////////////// Graph Traverse routines //////////////////////////
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
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;
}
731

732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
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;
}

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
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
///////////////////////// 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);
  });

807
808
809
810
811
812
813
814
815
816
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
  });

817
818
}  // namespace aten
}  // namespace dgl
819
820
821
822

std::ostream& operator << (std::ostream& os, dgl::runtime::NDArray array) {
  return os << dgl::aten::ToDebugString(array);
}