array.cc 41.3 KB
Newer Older
1
/**
2
 *  Copyright (c) 2019-2022 by Contributors
3
4
 * @file array/array.cc
 * @brief DGL array utilities implementation
5
6
 */
#include <dgl/array.h>
7
#include <dgl/bcast.h>
8
#include <dgl/graph_traversal.h>
9
10
#include <dgl/packed_func_ext.h>
#include <dgl/runtime/container.h>
11
#include <dgl/runtime/device_api.h>
12
13
#include <dgl/runtime/shared_mem.h>

14
#include <sstream>
15

16
17
#include "../c_api_common.h"
#include "./arith.h"
18
#include "./array_op.h"
19
#include "./kernel_decl.h"
20

21
using namespace dgl::runtime;
22

23
namespace dgl {
24
25
namespace aten {

26
27
IdArray NewIdArray(int64_t length, DGLContext ctx, uint8_t nbits) {
  return IdArray::Empty({length}, DGLDataType{kDGLInt, nbits, 1}, ctx);
28
29
}

30
31
32
33
FloatArray NewFloatArray(int64_t length, DGLContext ctx, uint8_t nbits) {
  return FloatArray::Empty({length}, DGLDataType{kDGLFloat, nbits, 1}, ctx);
}

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

40
IdArray Range(int64_t low, int64_t high, uint8_t nbits, DGLContext ctx) {
41
  IdArray ret;
42
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Range", {
43
44
45
46
47
48
49
50
51
52
53
    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;
}

54
IdArray Full(int64_t val, int64_t length, uint8_t nbits, DGLContext ctx) {
55
  IdArray ret;
56
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Full", {
57
58
59
60
61
62
63
64
65
66
67
    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;
}

68
template <typename DType>
69
NDArray Full(DType val, int64_t length, DGLContext ctx) {
70
71
72
73
74
75
76
  NDArray ret;
  ATEN_XPU_SWITCH_CUDA(ctx.device_type, XPU, "Full", {
    ret = impl::Full<XPU, DType>(val, length, ctx);
  });
  return ret;
}

77
78
79
80
template NDArray Full<int32_t>(int32_t val, int64_t length, DGLContext ctx);
template NDArray Full<int64_t>(int64_t val, int64_t length, DGLContext ctx);
template NDArray Full<float>(float val, int64_t length, DGLContext ctx);
template NDArray Full<double>(double val, int64_t length, DGLContext ctx);
81

82
IdArray AsNumBits(IdArray arr, uint8_t bits) {
83
  CHECK(bits == 32 || bits == 64)
84
85
86
87
      << "Invalid ID type. Must be int32 or int64, but got int"
      << static_cast<int>(bits) << ".";
  if (arr->dtype.bits == bits) return arr;
  if (arr.NumElements() == 0) return NewIdArray(arr->shape[0], arr->ctx, bits);
88
  IdArray ret;
89
  ATEN_XPU_SWITCH_CUDA(arr->ctx.device_type, XPU, "AsNumBits", {
90
91
    ATEN_ID_TYPE_SWITCH(
        arr->dtype, IdType, { ret = impl::AsNumBits<XPU, IdType>(arr, bits); });
92
93
94
95
96
97
  });
  return ret;
}

IdArray HStack(IdArray lhs, IdArray rhs) {
  IdArray ret;
98
99
  CHECK_SAME_CONTEXT(lhs, rhs);
  CHECK_SAME_DTYPE(lhs, rhs);
100
101
102
103
104
105
  CHECK_EQ(lhs->shape[0], rhs->shape[0]);
  auto device = runtime::DeviceAPI::Get(lhs->ctx);
  const auto& ctx = lhs->ctx;
  ATEN_ID_TYPE_SWITCH(lhs->dtype, IdType, {
    const int64_t len = lhs->shape[0];
    ret = NewIdArray(2 * len, lhs->ctx, lhs->dtype.bits);
106
107
108
109
110
111
    device->CopyDataFromTo(
        lhs.Ptr<IdType>(), 0, ret.Ptr<IdType>(), 0, len * sizeof(IdType), ctx,
        ctx, lhs->dtype);
    device->CopyDataFromTo(
        rhs.Ptr<IdType>(), 0, ret.Ptr<IdType>(), len * sizeof(IdType),
        len * sizeof(IdType), ctx, ctx, lhs->dtype);
Jinjing Zhou's avatar
Jinjing Zhou committed
112
113
114
115
  });
  return ret;
}

116
117
NDArray IndexSelect(NDArray array, IdArray index) {
  NDArray ret;
118
  CHECK_GE(array->ndim, 1) << "Only support array with at least 1 dimension";
119
  CHECK_EQ(index->ndim, 1) << "Index array must be an 1D array.";
120
121
122
123
  // if array is not pinned, index has the same context as array
  // if array is pinned, op dispatching depends on the context of index
  CHECK_VALID_CONTEXT(array, index);
  ATEN_XPU_SWITCH_CUDA(index->ctx.device_type, XPU, "IndexSelect", {
124
125
126
127
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
        ret = impl::IndexSelect<XPU, DType, IdType>(array, index);
      });
128
129
130
131
132
    });
  });
  return ret;
}

133
template <typename ValueType>
134
ValueType IndexSelect(NDArray array, int64_t index) {
135
  CHECK_EQ(array->ndim, 1) << "Only support select values from 1D array.";
136
  CHECK(index >= 0 && index < array.NumElements())
137
      << "Index " << index << " is out of bound.";
138
  ValueType ret = 0;
139
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "IndexSelect", {
140
141
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ret = impl::IndexSelect<XPU, DType>(array, index);
142
143
144
145
    });
  });
  return ret;
}
146
147
148
149
150
151
152
153
154
155
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())
156
      << "Index " << start << " is out of bound.";
157
  CHECK(end >= 0 && end <= array.NumElements())
158
      << "Index " << end << " is out of bound.";
159
160
161
162
163
  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", {
164
165
166
    device->CopyDataFromTo(
        array->data, start * sizeof(DType), ret->data, 0, len * sizeof(DType),
        array->ctx, ret->ctx, array->dtype);
167
168
169
  });
  return ret;
}
170

171
172
NDArray Scatter(NDArray array, IdArray indices) {
  NDArray ret;
173
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Scatter", {
174
175
176
177
178
179
180
181
182
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(indices->dtype, IdType, {
        ret = impl::Scatter<XPU, DType, IdType>(array, indices);
      });
    });
  });
  return ret;
}

183
184
185
186
187
void Scatter_(IdArray index, NDArray value, NDArray out) {
  CHECK_SAME_DTYPE(value, out);
  CHECK_SAME_CONTEXT(index, value);
  CHECK_SAME_CONTEXT(index, out);
  CHECK_EQ(value->shape[0], index->shape[0]);
188
  if (index->shape[0] == 0) return;
189
190
191
192
193
194
195
196
197
  ATEN_XPU_SWITCH_CUDA(value->ctx.device_type, XPU, "Scatter_", {
    ATEN_DTYPE_SWITCH(value->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(index->dtype, IdType, {
        impl::Scatter_<XPU, DType, IdType>(index, value, out);
      });
    });
  });
}

198
199
NDArray Repeat(NDArray array, IdArray repeats) {
  NDArray ret;
200
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Repeat", {
201
202
203
204
205
206
207
208
209
    ATEN_DTYPE_SWITCH(array->dtype, DType, "values", {
      ATEN_ID_TYPE_SWITCH(repeats->dtype, IdType, {
        ret = impl::Repeat<XPU, DType, IdType>(array, repeats);
      });
    });
  });
  return ret;
}

210
211
IdArray Relabel_(const std::vector<IdArray>& arrays) {
  IdArray ret;
212
  ATEN_XPU_SWITCH_CUDA(arrays[0]->ctx.device_type, XPU, "Relabel_", {
213
214
215
216
217
218
219
    ATEN_ID_TYPE_SWITCH(arrays[0]->dtype, IdType, {
      ret = impl::Relabel_<XPU, IdType>(arrays);
    });
  });
  return ret;
}

220
221
222
223
224
225
226
227
228
229
NDArray Concat(const std::vector<IdArray>& arrays) {
  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]);
  }

230
  NDArray ret_arr = NDArray::Empty({len}, arrays[0]->dtype, arrays[0]->ctx);
231
232
233
234
235

  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(
236
237
238
239
240
241
          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);

      offset += arrays[i]->shape[0] * sizeof(DType);
242
243
244
245
246
247
    });
  }

  return ret_arr;
}

248
template <typename ValueType>
249
250
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, ValueType pad_value) {
  std::tuple<NDArray, IdArray, IdArray> ret;
251
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "Pack", {
252
253
254
255
256
257
258
259
260
    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);
261
262
263
264
template std::tuple<NDArray, IdArray, IdArray> Pack<uint32_t>(
    NDArray, uint32_t);
template std::tuple<NDArray, IdArray, IdArray> Pack<uint64_t>(
    NDArray, uint64_t);
265
266
267
268
269
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;
270
  ATEN_XPU_SWITCH(array->ctx.device_type, XPU, "ConcatSlices", {
271
272
273
274
275
276
277
278
279
    ATEN_DTYPE_SWITCH(array->dtype, DType, "array", {
      ATEN_ID_TYPE_SWITCH(lengths->dtype, IdType, {
        ret = impl::ConcatSlices<XPU, DType, IdType>(array, lengths);
      });
    });
  });
  return ret;
}

280
281
282
283
284
285
286
287
288
289
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;
}

290
291
292
IdArray NonZero(NDArray array) {
  IdArray ret;
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "NonZero", {
293
294
    ATEN_ID_TYPE_SWITCH(
        array->dtype, DType, { ret = impl::NonZero<XPU, DType>(array); });
295
296
297
298
  });
  return ret;
}

299
std::pair<IdArray, IdArray> Sort(IdArray array, const int num_bits) {
300
301
302
303
304
305
306
  if (array.NumElements() == 0) {
    IdArray idx = NewIdArray(0, array->ctx, 64);
    return std::make_pair(array, idx);
  }
  std::pair<IdArray, IdArray> ret;
  ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, "Sort", {
    ATEN_ID_TYPE_SWITCH(array->dtype, IdType, {
307
      ret = impl::Sort<XPU, IdType>(array, num_bits);
308
309
310
311
312
    });
  });
  return ret;
}

313
314
std::string ToDebugString(NDArray array) {
  std::ostringstream oss;
315
  NDArray a = array.CopyTo(DGLContext{kDGLCPU, 0});
316
317
318
319
320
321
  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] << ", ";
    }
  });
322
  if (a.NumElements() > 10) oss << "...";
323
324
325
326
  oss << "], dtype=" << array->dtype << ", ctx=" << array->ctx << ")";
  return oss.str();
}

327
328
329
///////////////////////// CSR routines //////////////////////////

bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
330
331
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
  CHECK(col >= 0 && col < csr.num_cols) << "Invalid col index: " << col;
332
  bool ret = false;
333
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRIsNonZero", {
334
335
336
337
338
339
340
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
  NDArray ret;
341
342
  CHECK_SAME_DTYPE(csr.indices, row);
  CHECK_SAME_DTYPE(csr.indices, col);
343
344
  CHECK_SAME_CONTEXT(row, col);
  ATEN_CSR_SWITCH_CUDA_UVA(csr, row, XPU, IdType, "CSRIsNonZero", {
345
346
347
348
349
350
351
    ret = impl::CSRIsNonZero<XPU, IdType>(csr, row, col);
  });
  return ret;
}

bool CSRHasDuplicate(CSRMatrix csr) {
  bool ret = false;
352
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRHasDuplicate", {
353
354
355
356
357
358
    ret = impl::CSRHasDuplicate<XPU, IdType>(csr);
  });
  return ret;
}

int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
359
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
360
  int64_t ret = 0;
361
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowNNZ", {
362
363
364
365
366
367
368
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray row) {
  NDArray ret;
369
  CHECK_SAME_DTYPE(csr.indices, row);
370
  ATEN_CSR_SWITCH_CUDA_UVA(csr, row, XPU, IdType, "CSRGetRowNNZ", {
371
372
373
374
375
376
    ret = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
377
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
378
  NDArray ret;
379
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowColumnIndices", {
380
381
382
383
384
385
    ret = impl::CSRGetRowColumnIndices<XPU, IdType>(csr, row);
  });
  return ret;
}

NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
386
  CHECK(row >= 0 && row < csr.num_rows) << "Invalid row index: " << row;
387
  NDArray ret;
388
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGetRowData", {
389
    ret = impl::CSRGetRowData<XPU, IdType>(csr, row);
390
391
392
393
  });
  return ret;
}

394
bool CSRIsSorted(CSRMatrix csr) {
395
  if (csr.indices->shape[0] <= 1) return true;
396
397
398
399
400
401
402
  bool ret = false;
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRIsSorted", {
    ret = impl::CSRIsSorted<XPU, IdType>(csr);
  });
  return ret;
}

403
404
NDArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
  NDArray ret;
405
406
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
407
408
  CHECK_SAME_CONTEXT(rows, cols);
  ATEN_CSR_SWITCH_CUDA_UVA(csr, rows, XPU, IdType, "CSRGetData", {
409
    ret = impl::CSRGetData<XPU, IdType>(csr, rows, cols);
410
411
412
413
  });
  return ret;
}

414
template <typename DType>
415
416
NDArray CSRGetData(
    CSRMatrix csr, NDArray rows, NDArray cols, NDArray weights, DType filler) {
417
418
419
  NDArray ret;
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
420
421
422
  CHECK_SAME_CONTEXT(rows, cols);
  CHECK_SAME_CONTEXT(rows, weights);
  ATEN_CSR_SWITCH_CUDA_UVA(csr, rows, XPU, IdType, "CSRGetData", {
423
424
    ret =
        impl::CSRGetData<XPU, IdType, DType>(csr, rows, cols, weights, filler);
425
426
427
428
429
430
431
432
433
  });
  return ret;
}

template NDArray CSRGetData<float>(
    CSRMatrix csr, NDArray rows, NDArray cols, NDArray weights, float filler);
template NDArray CSRGetData<double>(
    CSRMatrix csr, NDArray rows, NDArray cols, NDArray weights, double filler);

434
435
std::vector<NDArray> CSRGetDataAndIndices(
    CSRMatrix csr, NDArray rows, NDArray cols) {
436
437
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
438
  CHECK_SAME_CONTEXT(rows, cols);
439
  std::vector<NDArray> ret;
440
  ATEN_CSR_SWITCH_CUDA_UVA(csr, rows, XPU, IdType, "CSRGetDataAndIndices", {
441
    ret = impl::CSRGetDataAndIndices<XPU, IdType>(csr, rows, cols);
442
443
444
445
446
447
  });
  return ret;
}

CSRMatrix CSRTranspose(CSRMatrix csr) {
  CSRMatrix ret;
448
449
450
451
  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);
    });
452
453
454
455
456
457
458
  });
  return ret;
}

COOMatrix CSRToCOO(CSRMatrix csr, bool data_as_order) {
  COOMatrix ret;
  if (data_as_order) {
459
460
461
462
463
464
    ATEN_XPU_SWITCH_CUDA(
        csr.indptr->ctx.device_type, XPU, "CSRToCOODataAsOrder", {
          ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
            ret = impl::CSRToCOODataAsOrder<XPU, IdType>(csr);
          });
        });
465
  } else {
466
    ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "CSRToCOO", {
467
468
469
470
471
472
473
474
475
      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) {
476
477
478
  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);
479
  CSRMatrix ret;
480
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRSliceRows", {
481
    ret = impl::CSRSliceRows<XPU, IdType>(csr, start, end);
482
483
484
485
486
  });
  return ret;
}

CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
487
  CHECK_SAME_DTYPE(csr.indices, rows);
488
  CSRMatrix ret;
489
  ATEN_CSR_SWITCH_CUDA_UVA(csr, rows, XPU, IdType, "CSRSliceRows", {
490
    ret = impl::CSRSliceRows<XPU, IdType>(csr, rows);
491
492
493
494
495
  });
  return ret;
}

CSRMatrix CSRSliceMatrix(CSRMatrix csr, NDArray rows, NDArray cols) {
496
497
  CHECK_SAME_DTYPE(csr.indices, rows);
  CHECK_SAME_DTYPE(csr.indices, cols);
498
  CHECK_SAME_CONTEXT(rows, cols);
499
  CSRMatrix ret;
500
  ATEN_CSR_SWITCH_CUDA_UVA(csr, rows, XPU, IdType, "CSRSliceMatrix", {
501
    ret = impl::CSRSliceMatrix<XPU, IdType>(csr, rows, cols);
502
503
504
505
  });
  return ret;
}

506
void CSRSort_(CSRMatrix* csr) {
507
508
509
  if (csr->sorted) return;
  ATEN_CSR_SWITCH_CUDA(
      *csr, XPU, IdType, "CSRSort_", { impl::CSRSort_<XPU, IdType>(csr); });
Da Zheng's avatar
Da Zheng committed
510
511
}

512
std::pair<CSRMatrix, NDArray> CSRSortByTag(
513
    const CSRMatrix& csr, IdArray tag, int64_t num_tags) {
514
  CHECK_EQ(csr.indices->shape[0], tag->shape[0])
515
516
      << "The length of the tag array should be equal to the number of "
         "non-zero data.";
517
518
519
520
521
522
523
524
525
526
527
  CHECK_SAME_CONTEXT(csr.indices, tag);
  CHECK_INT(tag, "tag");
  std::pair<CSRMatrix, NDArray> ret;
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRSortByTag", {
    ATEN_ID_TYPE_SWITCH(tag->dtype, TagType, {
      ret = impl::CSRSortByTag<XPU, IdType, TagType>(csr, tag, num_tags);
    });
  });
  return ret;
}

528
529
CSRMatrix CSRReorder(
    CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids) {
Da Zheng's avatar
Da Zheng committed
530
531
532
533
534
535
536
  CSRMatrix ret;
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRReorder", {
    ret = impl::CSRReorder<XPU, IdType>(csr, new_row_ids, new_col_ids);
  });
  return ret;
}

537
538
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries) {
  CSRMatrix ret;
539
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRRemove", {
540
541
542
543
544
    ret = impl::CSRRemove<XPU, IdType>(csr, entries);
  });
  return ret;
}

545
546
547
548
549
std::pair<COOMatrix, FloatArray> CSRLaborSampling(
    CSRMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
    int importance_sampling, IdArray random_seed, IdArray NIDs) {
  std::pair<COOMatrix, FloatArray> ret;
  ATEN_CSR_SWITCH_CUDA_UVA(mat, rows, XPU, IdType, "CSRLaborSampling", {
550
551
    const auto dtype =
        IsNullArray(prob) ? DGLDataTypeTraits<float>::dtype : prob->dtype;
552
553
554
555
556
557
558
559
    ATEN_FLOAT_TYPE_SWITCH(dtype, FloatType, "probability", {
      ret = impl::CSRLaborSampling<XPU, IdType, FloatType>(
          mat, rows, num_samples, prob, importance_sampling, random_seed, NIDs);
    });
  });
  return ret;
}

560
COOMatrix CSRRowWiseSampling(
561
562
    CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
    bool replace) {
563
  COOMatrix ret;
564
  if (IsNullArray(prob_or_mask)) {
565
566
567
568
569
    ATEN_CSR_SWITCH_CUDA_UVA(
        mat, rows, XPU, IdType, "CSRRowWiseSamplingUniform", {
          ret = impl::CSRRowWiseSamplingUniform<XPU, IdType>(
              mat, rows, num_samples, replace);
        });
570
  } else {
571
572
    // prob_or_mask is pinned and rows on GPU is valid
    CHECK_VALID_CONTEXT(prob_or_mask, rows);
573
    ATEN_CSR_SWITCH_CUDA_UVA(mat, rows, XPU, IdType, "CSRRowWiseSampling", {
574
575
      CHECK(!(prob_or_mask->dtype.bits == 8 && XPU == kDGLCUDA))
          << "GPU sampling with masks is currently not supported yet.";
576
      ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(
577
          prob_or_mask->dtype, FloatType, "probability or mask", {
578
579
580
            ret = impl::CSRRowWiseSampling<XPU, IdType, FloatType>(
                mat, rows, num_samples, prob_or_mask, replace);
          });
581
582
    });
  }
583
584
585
  return ret;
}

586
COOMatrix CSRRowWisePerEtypeSampling(
587
    CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
588
589
590
    const std::vector<int64_t>& num_samples,
    const std::vector<NDArray>& prob_or_mask, bool replace,
    bool rowwise_etype_sorted) {
591
  COOMatrix ret;
592
  CHECK(prob_or_mask.size() > 0) << "probability or mask array is empty";
593
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWisePerEtypeSampling", {
594
    if (std::all_of(prob_or_mask.begin(), prob_or_mask.end(), IsNullArray)) {
595
      ret = impl::CSRRowWisePerEtypeSamplingUniform<XPU, IdType>(
596
597
          mat, rows, eid2etype_offset, num_samples, replace,
          rowwise_etype_sorted);
598
    } else {
599
600
      ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(
          prob_or_mask[0]->dtype, DType, "probability or mask", {
601
602
603
604
            ret = impl::CSRRowWisePerEtypeSampling<XPU, IdType, DType>(
                mat, rows, eid2etype_offset, num_samples, prob_or_mask, replace,
                rowwise_etype_sorted);
          });
605
606
607
608
609
    }
  });
  return ret;
}

610
COOMatrix CSRRowWiseTopk(
611
    CSRMatrix mat, IdArray rows, int64_t k, NDArray weight, bool ascending) {
612
  COOMatrix ret;
613
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseTopk", {
614
615
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::CSRRowWiseTopk<XPU, IdType, DType>(
616
617
618
619
620
621
          mat, rows, k, weight, ascending);
    });
  });
  return ret;
}

622
COOMatrix CSRRowWiseSamplingBiased(
623
624
    CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray tag_offset,
    FloatArray bias, bool replace) {
625
626
627
  COOMatrix ret;
  ATEN_CSR_SWITCH(mat, XPU, IdType, "CSRRowWiseSamplingBiased", {
    ATEN_FLOAT_TYPE_SWITCH(bias->dtype, FloatType, "bias", {
628
      ret = impl::CSRRowWiseSamplingBiased<XPU, IdType, FloatType>(
629
630
631
632
633
634
          mat, rows, num_samples, tag_offset, bias, replace);
    });
  });
  return ret;
}

635
std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling(
636
637
    const CSRMatrix& csr, int64_t num_samples, int num_trials,
    bool exclude_self_loops, bool replace, double redundancy) {
638
639
640
641
642
643
644
645
646
647
  CHECK_GT(num_samples, 0) << "Number of samples must be positive";
  CHECK_GT(num_trials, 0) << "Number of sampling trials must be positive";
  std::pair<IdArray, IdArray> result;
  ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, "CSRGlobalUniformNegativeSampling", {
    result = impl::CSRGlobalUniformNegativeSampling<XPU, IdType>(
        csr, num_samples, num_trials, exclude_self_loops, replace, redundancy);
  });
  return result;
}

648
649
CSRMatrix UnionCsr(const std::vector<CSRMatrix>& csrs) {
  CSRMatrix ret;
650
651
  CHECK_GT(csrs.size(), 1)
      << "UnionCsr creates a union of multiple CSRMatrixes";
652
653
  // sanity check
  for (size_t i = 1; i < csrs.size(); ++i) {
654
655
656
657
    CHECK_EQ(csrs[0].num_rows, csrs[i].num_rows)
        << "UnionCsr requires both CSRMatrix have same number of rows";
    CHECK_EQ(csrs[0].num_cols, csrs[i].num_cols)
        << "UnionCsr requires both CSRMatrix have same number of cols";
658
659
660
661
662
663
664
665
666
667
    CHECK_SAME_CONTEXT(csrs[0].indptr, csrs[i].indptr);
    CHECK_SAME_DTYPE(csrs[0].indptr, csrs[i].indptr);
  }

  ATEN_CSR_SWITCH(csrs[0], XPU, IdType, "UnionCsr", {
    ret = impl::UnionCsr<XPU, IdType>(csrs);
  });
  return ret;
}

668
std::tuple<CSRMatrix, IdArray, IdArray> CSRToSimple(const CSRMatrix& csr) {
669
670
671
672
673
674
675
676
677
  std::tuple<CSRMatrix, IdArray, IdArray> ret;

  CSRMatrix sorted_csr = (CSRIsSorted(csr)) ? csr : CSRSort(csr);
  ATEN_CSR_SWITCH(csr, XPU, IdType, "CSRToSimple", {
    ret = impl::CSRToSimple<XPU, IdType>(sorted_csr);
  });
  return ret;
}

678
679
///////////////////////// COO routines //////////////////////////

680
681
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
  bool ret = false;
682
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
683
684
685
686
687
688
689
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

NDArray COOIsNonZero(COOMatrix coo, NDArray row, NDArray col) {
  NDArray ret;
690
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOIsNonZero", {
691
692
693
694
695
    ret = impl::COOIsNonZero<XPU, IdType>(coo, row, col);
  });
  return ret;
}

696
697
bool COOHasDuplicate(COOMatrix coo) {
  bool ret = false;
698
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOHasDuplicate", {
699
700
701
702
703
    ret = impl::COOHasDuplicate<XPU, IdType>(coo);
  });
  return ret;
}

704
705
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
  int64_t ret = 0;
706
  ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, "COOGetRowNNZ", {
707
708
709
710
711
712
713
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

NDArray COOGetRowNNZ(COOMatrix coo, NDArray row) {
  NDArray ret;
714
  ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, "COOGetRowNNZ", {
715
716
717
718
719
    ret = impl::COOGetRowNNZ<XPU, IdType>(coo, row);
  });
  return ret;
}

720
721
std::pair<NDArray, NDArray> COOGetRowDataAndIndices(
    COOMatrix coo, int64_t row) {
722
  std::pair<NDArray, NDArray> ret;
723
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetRowDataAndIndices", {
724
    ret = impl::COOGetRowDataAndIndices<XPU, IdType>(coo, row);
725
726
727
728
729
730
731
  });
  return ret;
}

std::vector<NDArray> COOGetDataAndIndices(
    COOMatrix coo, NDArray rows, NDArray cols) {
  std::vector<NDArray> ret;
732
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetDataAndIndices", {
733
    ret = impl::COOGetDataAndIndices<XPU, IdType>(coo, rows, cols);
734
735
736
737
  });
  return ret;
}

738
739
740
741
742
743
744
745
NDArray COOGetData(COOMatrix coo, NDArray rows, NDArray cols) {
  NDArray ret;
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOGetData", {
    ret = impl::COOGetData<XPU, IdType>(coo, rows, cols);
  });
  return ret;
}

746
COOMatrix COOTranspose(COOMatrix coo) {
747
  return COOMatrix(coo.num_cols, coo.num_rows, coo.col, coo.row, coo.data);
748
749
}

750
751
CSRMatrix COOToCSR(COOMatrix coo) {
  CSRMatrix ret;
752
  ATEN_XPU_SWITCH_CUDA(coo.row->ctx.device_type, XPU, "COOToCSR", {
753
754
    ATEN_ID_TYPE_SWITCH(
        coo.row->dtype, IdType, { ret = impl::COOToCSR<XPU, IdType>(coo); });
755
756
757
758
  });
  return ret;
}

759
760
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end) {
  COOMatrix ret;
761
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
762
    ret = impl::COOSliceRows<XPU, IdType>(coo, start, end);
763
764
765
766
767
768
  });
  return ret;
}

COOMatrix COOSliceRows(COOMatrix coo, NDArray rows) {
  COOMatrix ret;
769
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceRows", {
770
    ret = impl::COOSliceRows<XPU, IdType>(coo, rows);
771
772
773
774
775
776
  });
  return ret;
}

COOMatrix COOSliceMatrix(COOMatrix coo, NDArray rows, NDArray cols) {
  COOMatrix ret;
777
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOSliceMatrix", {
778
779
780
781
782
    ret = impl::COOSliceMatrix<XPU, IdType>(coo, rows, cols);
  });
  return ret;
}

783
void COOSort_(COOMatrix* mat, bool sort_column) {
784
  if ((mat->row_sorted && !sort_column) || mat->col_sorted) return;
785
786
787
  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);
788
    });
789
  });
790
791
792
}

std::pair<bool, bool> COOIsSorted(COOMatrix coo) {
793
  if (coo.row->shape[0] <= 1) return {true, true};
794
795
796
797
  std::pair<bool, bool> ret;
  ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, "COOIsSorted", {
    ret = impl::COOIsSorted<XPU, IdType>(coo);
  });
798
799
800
  return ret;
}

801
802
COOMatrix COOReorder(
    COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids) {
803
804
805
806
807
808
809
  COOMatrix ret;
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOReorder", {
    ret = impl::COOReorder<XPU, IdType>(coo, new_row_ids, new_col_ids);
  });
  return ret;
}

810
811
COOMatrix COORemove(COOMatrix coo, IdArray entries) {
  COOMatrix ret;
812
  ATEN_COO_SWITCH(coo, XPU, IdType, "COORemove", {
813
814
815
816
817
    ret = impl::COORemove<XPU, IdType>(coo, entries);
  });
  return ret;
}

818
819
820
821
822
std::pair<COOMatrix, FloatArray> COOLaborSampling(
    COOMatrix mat, IdArray rows, int64_t num_samples, FloatArray prob,
    int importance_sampling, IdArray random_seed, IdArray NIDs) {
  std::pair<COOMatrix, FloatArray> ret;
  ATEN_COO_SWITCH(mat, XPU, IdType, "COOLaborSampling", {
823
824
    const auto dtype =
        IsNullArray(prob) ? DGLDataTypeTraits<float>::dtype : prob->dtype;
825
826
827
828
829
830
831
832
    ATEN_FLOAT_TYPE_SWITCH(dtype, FloatType, "probability", {
      ret = impl::COOLaborSampling<XPU, IdType, FloatType>(
          mat, rows, num_samples, prob, importance_sampling, random_seed, NIDs);
    });
  });
  return ret;
}

833
COOMatrix COORowWiseSampling(
834
835
    COOMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
    bool replace) {
836
  COOMatrix ret;
837
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseSampling", {
838
    if (IsNullArray(prob_or_mask)) {
839
840
      ret = impl::COORowWiseSamplingUniform<XPU, IdType>(
          mat, rows, num_samples, replace);
841
    } else {
842
      ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(
843
          prob_or_mask->dtype, DType, "probability or mask", {
844
845
846
            ret = impl::COORowWiseSampling<XPU, IdType, DType>(
                mat, rows, num_samples, prob_or_mask, replace);
          });
847
848
849
850
851
    }
  });
  return ret;
}

852
COOMatrix COORowWisePerEtypeSampling(
853
    COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
854
855
    const std::vector<int64_t>& num_samples,
    const std::vector<NDArray>& prob_or_mask, bool replace) {
856
  COOMatrix ret;
857
  CHECK(prob_or_mask.size() > 0) << "probability or mask array is empty";
858
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWisePerEtypeSampling", {
859
    if (std::all_of(prob_or_mask.begin(), prob_or_mask.end(), IsNullArray)) {
860
      ret = impl::COORowWisePerEtypeSamplingUniform<XPU, IdType>(
861
          mat, rows, eid2etype_offset, num_samples, replace);
862
    } else {
863
864
      ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(
          prob_or_mask[0]->dtype, DType, "probability or mask", {
865
866
867
868
            ret = impl::COORowWisePerEtypeSampling<XPU, IdType, DType>(
                mat, rows, eid2etype_offset, num_samples, prob_or_mask,
                replace);
          });
869
870
871
872
873
    }
  });
  return ret;
}

874
875
876
COOMatrix COORowWiseTopk(
    COOMatrix mat, IdArray rows, int64_t k, FloatArray weight, bool ascending) {
  COOMatrix ret;
877
  ATEN_COO_SWITCH(mat, XPU, IdType, "COORowWiseTopk", {
878
879
    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
      ret = impl::COORowWiseTopk<XPU, IdType, DType>(
880
881
          mat, rows, k, weight, ascending);
    });
882
883
884
885
  });
  return ret;
}

886
887
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo) {
  std::pair<COOMatrix, IdArray> ret;
888
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOCoalesce", {
889
890
891
892
893
    ret = impl::COOCoalesce<XPU, IdType>(coo);
  });
  return ret;
}

894
895
896
897
898
899
900
901
902
903
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) {
  COOMatrix ret;
  ATEN_XPU_SWITCH_CUDA(coos[0].row->ctx.device_type, XPU, "DisjointUnionCoo", {
    ATEN_ID_TYPE_SWITCH(coos[0].row->dtype, IdType, {
      ret = impl::DisjointUnionCoo<XPU, IdType>(coos);
    });
  });
  return ret;
}

904
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking) {
905
906
907
908
909
910
  COOMatrix ret;
  ATEN_COO_SWITCH(coo, XPU, IdType, "COOLineGraph", {
    ret = impl::COOLineGraph<XPU, IdType>(coo, backtracking);
  });
  return ret;
}
911
912
913

COOMatrix UnionCoo(const std::vector<COOMatrix>& coos) {
  COOMatrix ret;
914
915
  CHECK_GT(coos.size(), 1)
      << "UnionCoo creates a union of multiple COOMatrixes";
916
917
  // sanity check
  for (size_t i = 1; i < coos.size(); ++i) {
918
919
920
921
    CHECK_EQ(coos[0].num_rows, coos[i].num_rows)
        << "UnionCoo requires both COOMatrix have same number of rows";
    CHECK_EQ(coos[0].num_cols, coos[i].num_cols)
        << "UnionCoo requires both COOMatrix have same number of cols";
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
    CHECK_SAME_CONTEXT(coos[0].row, coos[i].row);
    CHECK_SAME_DTYPE(coos[0].row, coos[i].row);
  }

  // we assume the number of coos is not large in common cases
  std::vector<IdArray> coo_row;
  std::vector<IdArray> coo_col;
  bool has_data = false;

  for (size_t i = 0; i < coos.size(); ++i) {
    coo_row.push_back(coos[i].row);
    coo_col.push_back(coos[i].col);
    has_data |= COOHasData(coos[i]);
  }

  IdArray row = Concat(coo_row);
  IdArray col = Concat(coo_col);
  IdArray data = NullArray();

  if (has_data) {
    std::vector<IdArray> eid_data;
943
944
945
946
947
    eid_data.push_back(
        COOHasData(coos[0]) ? coos[0].data
                            : Range(
                                  0, coos[0].row->shape[0],
                                  coos[0].row->dtype.bits, coos[0].row->ctx));
948
949
    int64_t num_edges = coos[0].row->shape[0];
    for (size_t i = 1; i < coos.size(); ++i) {
950
951
952
953
954
955
      eid_data.push_back(
          COOHasData(coos[i])
              ? coos[i].data + num_edges
              : Range(
                    num_edges, num_edges + coos[i].row->shape[0],
                    coos[i].row->dtype.bits, coos[i].row->ctx));
956
957
958
959
960
961
962
      num_edges += coos[i].row->shape[0];
    }

    data = Concat(eid_data);
  }

  return COOMatrix(
963
      coos[0].num_rows, coos[0].num_cols, row, col, data, false, false);
964
965
}

966
std::tuple<COOMatrix, IdArray, IdArray> COOToSimple(const COOMatrix& coo) {
967
968
  // coo column sorted
  const COOMatrix sorted_coo = COOSort(coo, true);
969
970
971
972
973
974
975
976
977
  const IdArray eids_shuffled =
      COOHasData(sorted_coo)
          ? sorted_coo.data
          : Range(
                0, sorted_coo.row->shape[0], sorted_coo.row->dtype.bits,
                sorted_coo.row->ctx);
  const auto& coalesced_result = COOCoalesce(sorted_coo);
  const COOMatrix& coalesced_adj = coalesced_result.first;
  const IdArray& count = coalesced_result.second;
978

979
  /**
980
981
   * eids_shuffled actually already contains the mapping from old edge space to
   * the new one:
982
   *
983
984
985
986
987
988
   * * eids_shuffled[0:count[0]] indicates the original edge IDs that coalesced
   * into new edge #0.
   * * eids_shuffled[count[0]:count[0] + count[1]] indicates those that
   * coalesced into new edge #1.
   * * eids_shuffled[count[0] + count[1]:count[0] + count[1] + count[2]]
   * indicates those that coalesced into new edge #2.
989
990
   * * etc.
   *
991
992
993
   * Here, we need to translate eids_shuffled to an array "eids_remapped" such
   * that eids_remapped[i] indicates the new edge ID the old edge #i is mapped
   * to.  The translation can simply be achieved by (in numpy code):
994
995
996
997
998
999
   *
   *     new_eid_for_eids_shuffled = np.range(len(count)).repeat(count)
   *     eids_remapped = np.zeros_like(new_eid_for_eids_shuffled)
   *     eids_remapped[eids_shuffled] = new_eid_for_eids_shuffled
   */
  const IdArray new_eids = Range(
1000
1001
      0, coalesced_adj.row->shape[0], coalesced_adj.row->dtype.bits,
      coalesced_adj.row->ctx);
1002
1003
1004
  const IdArray eids_remapped = Scatter(Repeat(new_eids, count), eids_shuffled);

  COOMatrix ret = COOMatrix(
1005
1006
      coalesced_adj.num_rows, coalesced_adj.num_cols, coalesced_adj.row,
      coalesced_adj.col, NullArray(), true, true);
1007
1008
1009
  return std::make_tuple(ret, count, eids_remapped);
}

1010
///////////////////////// Graph Traverse routines //////////////////////////
1011
1012
Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source) {
  Frontiers ret;
1013
1014
1015
1016
1017
1018
  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.";
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
  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;
1029
1030
1031
1032
1033
1034
  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.";
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
  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;
1045
1046
1047
1048
1049
1050
1051
1052
  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);
        });
      });
1053
1054
1055
1056
1057
  return ret;
}

Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source) {
  Frontiers ret;
1058
1059
1060
1061
1062
1063
  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.";
1064
1065
1066
1067
1068
1069
1070
  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;
}
1071

1072
1073
1074
Frontiers DGLDFSLabeledEdges(
    const CSRMatrix& csr, IdArray source, const bool has_reverse_edge,
    const bool has_nontree_edge, const bool return_labels) {
1075
  Frontiers ret;
1076
1077
1078
1079
1080
1081
  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.";
1082
1083
  ATEN_XPU_SWITCH(source->ctx.device_type, XPU, "DGLDFSLabeledEdges", {
    ATEN_ID_TYPE_SWITCH(source->dtype, IdType, {
1084
1085
      ret = impl::DGLDFSLabeledEdges<XPU, IdType>(
          csr, source, has_reverse_edge, has_nontree_edge, return_labels);
1086
1087
1088
1089
1090
    });
  });
  return ret;
}

1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
void CSRSpMM(
    const std::string& op, const std::string& reduce, const CSRMatrix& csr,
    NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux) {
  const auto& bcast = CalcBcastOff(op, ufeat, efeat);

  ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "SpMM", {
    ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
      ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
        SpMMCsr<XPU, IdType, Dtype>(
            op, reduce, bcast, csr, ufeat, efeat, out, out_aux);
      });
    });
  });
}

void CSRSpMM(
    const char* op, const char* reduce, const CSRMatrix& csr, NDArray ufeat,
    NDArray efeat, NDArray out, std::vector<NDArray> out_aux) {
  CSRSpMM(
      std::string(op), std::string(reduce), csr, ufeat, efeat, out, out_aux);
}

void CSRSDDMM(
    const std::string& op, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out, int lhs_target, int rhs_target) {
  const auto& bcast = CalcBcastOff(op, ufeat, efeat);

  ATEN_XPU_SWITCH_CUDA(csr.indptr->ctx.device_type, XPU, "SDDMM", {
    ATEN_ID_TYPE_SWITCH(csr.indptr->dtype, IdType, {
      ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
        SDDMMCsr<XPU, IdType, Dtype>(
            op, bcast, csr, ufeat, efeat, out, lhs_target, rhs_target);
      });
    });
  });
}

void CSRSDDMM(
    const char* op, const CSRMatrix& csr, NDArray ufeat, NDArray efeat,
    NDArray out, int lhs_target, int rhs_target) {
  return CSRSDDMM(
      std::string(op), csr, ufeat, efeat, out, lhs_target, rhs_target);
}

void COOSpMM(
    const std::string& op, const std::string& reduce, const COOMatrix& coo,
    NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux) {
  const auto& bcast = CalcBcastOff(op, ufeat, efeat);

  ATEN_XPU_SWITCH_CUDA(coo.row->ctx.device_type, XPU, "SpMM", {
    ATEN_ID_TYPE_SWITCH(coo.row->dtype, IdType, {
      ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
        SpMMCoo<XPU, IdType, Dtype>(
            op, reduce, bcast, coo, ufeat, efeat, out, out_aux);
      });
    });
  });
}

void COOSpMM(
    const char* op, const char* reduce, const COOMatrix& coo, NDArray ufeat,
    NDArray efeat, NDArray out, std::vector<NDArray> out_aux) {
  COOSpMM(
      std::string(op), std::string(reduce), coo, ufeat, efeat, out, out_aux);
}

void COOSDDMM(
    const std::string& op, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
    NDArray out, int lhs_target, int rhs_target) {
  const auto& bcast = CalcBcastOff(op, ufeat, efeat);

  ATEN_XPU_SWITCH_CUDA(coo.row->ctx.device_type, XPU, "SDDMM", {
    ATEN_ID_TYPE_SWITCH(coo.row->dtype, IdType, {
      ATEN_FLOAT_TYPE_SWITCH_16BITS(out->dtype, Dtype, XPU, "Feature data", {
        SDDMMCoo<XPU, IdType, Dtype>(
            op, bcast, coo, ufeat, efeat, out, lhs_target, rhs_target);
      });
    });
  });
}

void COOSDDMM(
    const char* op, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
    NDArray out, int lhs_target, int rhs_target) {
  COOSDDMM(std::string(op), coo, ufeat, efeat, out, lhs_target, rhs_target);
}

1178
1179
///////////////////////// C APIs /////////////////////////
DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetFormat")
1180
1181
1182
1183
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      SparseMatrixRef spmat = args[0];
      *rv = spmat->format;
    });
1184
1185

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

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

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

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLSparseMatrixGetFlags")
1205
1206
1207
1208
1209
1210
1211
1212
    .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;
    });
1213
1214

DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLCreateSparseMatrix")
1215
1216
1217
1218
1219
1220
1221
1222
    .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),
1223
          ListValueToVector<bool>(flags)));
1224
1225
      *rv = SparseMatrixRef(spmat);
    });
1226

1227
DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLExistSharedMemArray")
1228
1229
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      const std::string name = args[0];
1230
#ifndef _WIN32
1231
      *rv = SharedMemory::Exist(name);
1232
#else
1233
      *rv = false;
1234
#endif  // _WIN32
1235
    });
1236

1237
DGL_REGISTER_GLOBAL("ndarray._CAPI_DGLArrayCastToSigned")
1238
1239
1240
1241
1242
1243
1244
1245
    .set_body([](DGLArgs args, DGLRetValue* rv) {
      NDArray array = args[0];
      CHECK_EQ(array->dtype.code, kDGLUInt);
      std::vector<int64_t> shape(array->shape, array->shape + array->ndim);
      DGLDataType dtype = array->dtype;
      dtype.code = kDGLInt;
      *rv = array.CreateView(shape, dtype, 0);
    });
1246

1247
1248
}  // namespace aten
}  // namespace dgl
1249

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