test_spmat_csr.cc 23.8 KB
Newer Older
1
#include <dgl/array.h>
2
3
#include <gtest/gtest.h>

4
5
6
7
8
9
10
11
#include "./common.h"

using namespace dgl;
using namespace dgl::runtime;

namespace {

template <typename IDX>
12
aten::CSRMatrix CSR1(DGLContext ctx = CTX) {
13
14
15
16
17
  // [[0, 1, 1, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [0, 2, 3, 1, 4]
18
19
  return aten::CSRMatrix(
      4, 5,
20
21
22
23
24
25
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 3, 5, 5}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({1, 2, 0, 3, 2}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 3, 4, 1}), sizeof(IDX) * 8, ctx),
26
      false);
27
28
29
}

template <typename IDX>
30
aten::CSRMatrix CSR2(DGLContext ctx = CTX) {
31
32
33
34
35
36
  // has duplicate entries
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [0, 2, 5, 3, 1, 4]
37
38
  return aten::CSRMatrix(
      4, 5,
39
40
41
42
43
44
      aten::VecToIdArray(
          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
45
      false);
46
47
}

48
template <typename IDX>
49
aten::CSRMatrix CSR3(DGLContext ctx = CTX) {
50
51
52
53
54
55
56
57
58
59
60
61
62
  // has duplicate entries and the columns are not sorted
  // [[0, 1, 1, 1, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0],
  //  [1, 1, 1, 0, 0],
  //  [0, 0, 0, 1, 0]],
  //  [0, 0, 0, 0, 0]],
  //  [1, 2, 1, 1, 0]],
  //  [0, 1, 0, 0, 1]],
  // data: [5, 2, 0, 3, 1, 4, 8, 7, 6, 9, 12, 13, 11, 10, 14, 15, 16]
  return aten::CSRMatrix(
      9, 6,
63
64
65
66
67
68
69
70
71
72
      aten::VecToIdArray(
          std::vector<IDX>({0, 3, 4, 6, 6, 9, 10, 10, 15, 17}), sizeof(IDX) * 8,
          ctx),
      aten::VecToIdArray(
          std::vector<IDX>({3, 2, 1, 0, 2, 3, 1, 2, 0, 3, 1, 2, 1, 3, 0, 5, 1}),
          sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>(
              {0, 2, 5, 3, 1, 4, 6, 8, 7, 9, 13, 10, 11, 14, 12, 16, 15}),
          sizeof(IDX) * 8, ctx),
73
74
75
      false);
}

76
template <typename IDX>
77
aten::COOMatrix COO1(DGLContext ctx = CTX) {
78
79
80
81
82
83
84
  // [[0, 1, 1, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [0, 2, 3, 1, 4]
  // row : [0, 2, 0, 1, 2]
  // col : [1, 2, 2, 0, 3]
85
86
  return aten::COOMatrix(
      4, 5,
87
88
89
90
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 0, 1, 2}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({1, 2, 2, 0, 3}), sizeof(IDX) * 8, ctx));
91
92
93
}

template <typename IDX>
94
aten::COOMatrix COO2(DGLContext ctx = CTX) {
95
96
97
98
99
100
101
102
  // has duplicate entries
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [0, 2, 5, 3, 1, 4]
  // row : [0, 2, 0, 1, 2, 0]
  // col : [1, 2, 2, 0, 3, 2]
103
104
  return aten::COOMatrix(
      4, 5,
105
106
107
108
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({1, 2, 2, 0, 3, 2}), sizeof(IDX) * 8, ctx));
109
110
}

111
template <typename IDX>
112
aten::CSRMatrix SR_CSR3(DGLContext ctx) {
113
114
115
116
117
118
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  return aten::CSRMatrix(
      4, 5,
119
120
121
122
123
124
      aten::VecToIdArray(
          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({2, 1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
125
126
127
128
      false);
}

template <typename IDX>
129
aten::CSRMatrix SRC_CSR3(DGLContext ctx) {
130
131
132
133
134
135
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  return aten::CSRMatrix(
      4, 5,
136
137
138
139
140
141
      aten::VecToIdArray(
          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({2, 0, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
142
143
144
145
      false);
}

template <typename IDX>
146
aten::COOMatrix COO3(DGLContext ctx) {
147
148
149
150
151
152
153
154
155
  // has duplicate entries
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // row : [0, 2, 0, 1, 2, 0]
  // col : [2, 2, 1, 0, 3, 2]
  return aten::COOMatrix(
      4, 5,
156
157
158
159
      aten::VecToIdArray(
          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
      aten::VecToIdArray(
          std::vector<IDX>({2, 2, 1, 0, 3, 2}), sizeof(IDX) * 8, ctx));
160
161
}

162
}  // namespace
163
164

template <typename IDX>
165
void _TestCSRIsNonZero1(DGLContext ctx) {
166
  auto csr = CSR1<IDX>(ctx);
167
168
  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
169
170
171
172
  IdArray r =
      aten::VecToIdArray(std::vector<IDX>({2, 2, 0, 0}), sizeof(IDX) * 8, ctx);
  IdArray c =
      aten::VecToIdArray(std::vector<IDX>({1, 1, 1, 3}), sizeof(IDX) * 8, ctx);
173
  IdArray x = aten::CSRIsNonZero(csr, r, c);
174
175
  IdArray tx =
      aten::VecToIdArray(std::vector<IDX>({0, 0, 1, 0}), sizeof(IDX) * 8, ctx);
176
177
178
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

179
template <typename IDX>
180
void _TestCSRIsNonZero2(DGLContext ctx) {
181
182
183
  auto csr = CSR3<IDX>(ctx);
  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
  IdArray r = aten::VecToIdArray(
      std::vector<IDX>({
          0,
          0,
          0,
          0,
          0,
      }),
      sizeof(IDX) * 8, ctx);
  IdArray c = aten::VecToIdArray(
      std::vector<IDX>({
          0,
          1,
          2,
          3,
          4,
      }),
      sizeof(IDX) * 8, ctx);
202
  IdArray x = aten::CSRIsNonZero(csr, r, c);
203
204
  IdArray tx = aten::VecToIdArray(
      std::vector<IDX>({0, 1, 1, 1, 0}), sizeof(IDX) * 8, ctx);
205
206
207
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx)) << " x = " << x << ", tx = " << tx;
}

208
TEST(SpmatTest, TestCSRIsNonZero) {
209
210
211
212
  _TestCSRIsNonZero1<int32_t>(CPU);
  _TestCSRIsNonZero1<int64_t>(CPU);
  _TestCSRIsNonZero2<int32_t>(CPU);
  _TestCSRIsNonZero2<int64_t>(CPU);
213
#ifdef DGL_USE_CUDA
214
215
216
217
  _TestCSRIsNonZero1<int32_t>(GPU);
  _TestCSRIsNonZero1<int64_t>(GPU);
  _TestCSRIsNonZero2<int32_t>(GPU);
  _TestCSRIsNonZero2<int64_t>(GPU);
218
#endif
219
220
221
}

template <typename IDX>
222
void _TestCSRGetRowNNZ(DGLContext ctx) {
223
  auto csr = CSR2<IDX>(ctx);
224
225
  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 0), 3);
  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 3), 0);
226
227
  IdArray r =
      aten::VecToIdArray(std::vector<IDX>({0, 3}), sizeof(IDX) * 8, ctx);
228
  IdArray x = aten::CSRGetRowNNZ(csr, r);
229
230
  IdArray tx =
      aten::VecToIdArray(std::vector<IDX>({3, 0}), sizeof(IDX) * 8, ctx);
231
232
233
234
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

TEST(SpmatTest, TestCSRGetRowNNZ) {
235
236
237
238
239
240
  _TestCSRGetRowNNZ<int32_t>(CPU);
  _TestCSRGetRowNNZ<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetRowNNZ<int32_t>(GPU);
  _TestCSRGetRowNNZ<int64_t>(GPU);
#endif
241
242
243
}

template <typename IDX>
244
void _TestCSRGetRowColumnIndices(DGLContext ctx) {
245
  auto csr = CSR2<IDX>(ctx);
246
  auto x = aten::CSRGetRowColumnIndices(csr, 0);
247
248
  auto tx =
      aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX) * 8, ctx);
249
250
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowColumnIndices(csr, 1);
251
  tx = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX) * 8, ctx);
252
253
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowColumnIndices(csr, 3);
254
  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
255
256
257
258
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

TEST(SpmatTest, TestCSRGetRowColumnIndices) {
259
260
261
262
263
264
  _TestCSRGetRowColumnIndices<int32_t>(CPU);
  _TestCSRGetRowColumnIndices<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetRowColumnIndices<int32_t>(GPU);
  _TestCSRGetRowColumnIndices<int64_t>(GPU);
#endif
265
266
267
}

template <typename IDX>
268
void _TestCSRGetRowData(DGLContext ctx) {
269
  auto csr = CSR2<IDX>(ctx);
270
  auto x = aten::CSRGetRowData(csr, 0);
271
272
  auto tx =
      aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
273
274
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowData(csr, 1);
275
  tx = aten::VecToIdArray(std::vector<IDX>({3}), sizeof(IDX) * 8, ctx);
276
277
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowData(csr, 3);
278
  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
279
280
281
282
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

TEST(SpmatTest, TestCSRGetRowData) {
283
284
285
286
287
288
  _TestCSRGetRowData<int32_t>(CPU);
  _TestCSRGetRowData<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetRowData<int32_t>(GPU);
  _TestCSRGetRowData<int64_t>(GPU);
#endif
289
290
291
}

template <typename IDX>
292
void _TestCSRGetData(DGLContext ctx) {
293
294
295
  auto csr = CSR2<IDX>(ctx);
  // test get all data
  auto x = aten::CSRGetAllData(csr, 0, 0);
296
  auto tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
297
298
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetAllData(csr, 0, 2);
299
  tx = aten::VecToIdArray(std::vector<IDX>({2, 5}), sizeof(IDX) * 8, ctx);
300
301
302
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));

  // test get data
303
304
305
306
  auto r =
      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
  auto c =
      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
307
  x = aten::CSRGetData(csr, r, c);
308
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
309
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
310
311
312

  // test get data on sorted
  csr = aten::CSRSort(csr);
313
314
  r = aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
315
  x = aten::CSRGetData(csr, r, c);
316
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
317
318
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));

319
  // test get data w/ broadcasting
320
321
  r = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX) * 8, ctx);
  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
322
  x = aten::CSRGetData(csr, r, c);
323
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
324
325
326
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

327
328
329
330
331
TEST(SpmatTest, CSRGetData) {
  _TestCSRGetData<int32_t>(CPU);
  _TestCSRGetData<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetData<int32_t>(GPU);
332
  _TestCSRGetData<int64_t>(GPU);
333
#endif
334
335
336
}

template <typename IDX>
337
void _TestCSRGetDataAndIndices(DGLContext ctx) {
338
  auto csr = CSR2<IDX>(ctx);
339
340
341
342
  auto r =
      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
  auto c =
      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
343
  auto x = aten::CSRGetDataAndIndices(csr, r, c);
344
345
346
347
348
349
  auto tr =
      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
  auto tc =
      aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX) * 8, ctx);
  auto td =
      aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
350
351
352
353
354
  ASSERT_TRUE(ArrayEQ<IDX>(x[0], tr));
  ASSERT_TRUE(ArrayEQ<IDX>(x[1], tc));
  ASSERT_TRUE(ArrayEQ<IDX>(x[2], td));
}

355
356
357
358
359
360
361
TEST(SpmatTest, CSRGetDataAndIndices) {
  _TestCSRGetDataAndIndices<int32_t>(CPU);
  _TestCSRGetDataAndIndices<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetDataAndIndices<int32_t>(GPU);
  _TestCSRGetDataAndIndices<int64_t>(GPU);
#endif
362
363
364
}

template <typename IDX>
365
void _TestCSRTranspose(DGLContext ctx) {
366
  auto csr = CSR2<IDX>(ctx);
367
368
369
370
371
372
373
374
375
  auto csr_t = aten::CSRTranspose(csr);
  // [[0, 1, 0, 0],
  //  [1, 0, 0, 0],
  //  [2, 0, 1, 0],
  //  [0, 0, 1, 0],
  //  [0, 0, 0, 0]]
  // data: [3, 0, 2, 5, 1, 4]
  ASSERT_EQ(csr_t.num_rows, 5);
  ASSERT_EQ(csr_t.num_cols, 4);
376
377
378
379
380
381
  auto tp = aten::VecToIdArray(
      std::vector<IDX>({0, 1, 2, 5, 6, 6}), sizeof(IDX) * 8, ctx);
  auto ti = aten::VecToIdArray(
      std::vector<IDX>({1, 0, 0, 0, 2, 2}), sizeof(IDX) * 8, ctx);
  auto td = aten::VecToIdArray(
      std::vector<IDX>({3, 0, 2, 5, 1, 4}), sizeof(IDX) * 8, ctx);
382
383
384
385
386
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.data, td));
}

387
TEST(SpmatTest, CSRTranspose) {
388
389
390
391
  _TestCSRTranspose<int32_t>(CPU);
  _TestCSRTranspose<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRTranspose<int32_t>(GPU);
392
  _TestCSRTranspose<int64_t>(GPU);
393
#endif
394
395
396
}

template <typename IDX>
397
void _TestCSRToCOO(DGLContext ctx) {
398
  auto csr = CSR2<IDX>(ctx);
399
  {
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    auto coo = CSRToCOO(csr, false);
    ASSERT_EQ(coo.num_rows, 4);
    ASSERT_EQ(coo.num_cols, 5);
    ASSERT_TRUE(coo.row_sorted);
    auto tr = aten::VecToIdArray(
        std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tr));
    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, csr.indices));
    ASSERT_TRUE(ArrayEQ<IDX>(coo.data, csr.data));

    // convert from sorted csr
    auto s_csr = CSRSort(csr);
    coo = CSRToCOO(s_csr, false);
    ASSERT_EQ(coo.num_rows, 4);
    ASSERT_EQ(coo.num_cols, 5);
    ASSERT_TRUE(coo.row_sorted);
    ASSERT_TRUE(coo.col_sorted);
    tr = aten::VecToIdArray(
        std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tr));
    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, s_csr.indices));
    ASSERT_TRUE(ArrayEQ<IDX>(coo.data, s_csr.data));
422
423
  }
  {
424
425
426
427
428
429
    auto coo = CSRToCOO(csr, true);
    ASSERT_EQ(coo.num_rows, 4);
    ASSERT_EQ(coo.num_cols, 5);
    auto tcoo = COO2<IDX>(ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tcoo.row));
    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, tcoo.col));
430
431
432
  }
}

433
TEST(SpmatTest, CSRToCOO) {
434
435
436
437
  _TestCSRToCOO<int32_t>(CPU);
  _TestCSRToCOO<int64_t>(CPU);
#if DGL_USE_CUDA
  _TestCSRToCOO<int32_t>(GPU);
438
  _TestCSRToCOO<int64_t>(GPU);
439
#endif
440
441
442
}

template <typename IDX>
443
void _TestCSRSliceRows(DGLContext ctx) {
444
  auto csr = CSR2<IDX>(ctx);
445
446
447
448
449
450
451
  auto x = aten::CSRSliceRows(csr, 1, 4);
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [3, 1, 4]
  ASSERT_EQ(x.num_rows, 3);
  ASSERT_EQ(x.num_cols, 5);
452
453
454
455
456
457
  auto tp =
      aten::VecToIdArray(std::vector<IDX>({0, 1, 3, 3}), sizeof(IDX) * 8, ctx);
  auto ti =
      aten::VecToIdArray(std::vector<IDX>({0, 2, 3}), sizeof(IDX) * 8, ctx);
  auto td =
      aten::VecToIdArray(std::vector<IDX>({3, 1, 4}), sizeof(IDX) * 8, ctx);
458
459
460
461
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));

462
463
  auto r =
      aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
464
465
466
467
468
  x = aten::CSRSliceRows(csr, r);
  // [[0, 1, 2, 0, 0],
  //  [1, 0, 0, 0, 0],
  //  [0, 0, 0, 0, 0]]
  // data: [0, 2, 5, 3]
469
470
471
  tp = aten::VecToIdArray(std::vector<IDX>({0, 3, 4, 4}), sizeof(IDX) * 8, ctx);
  ti = aten::VecToIdArray(std::vector<IDX>({1, 2, 2, 0}), sizeof(IDX) * 8, ctx);
  td = aten::VecToIdArray(std::vector<IDX>({0, 2, 5, 3}), sizeof(IDX) * 8, ctx);
472
473
474
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
475
476

  // Testing non-increasing row id based slicing
477
  r = aten::VecToIdArray(std::vector<IDX>({3, 2, 1}), sizeof(IDX) * 8, ctx);
478
479
480
481
482
  x = aten::CSRSliceRows(csr, r);
  // [[0, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0],
  //  [1, 0, 0, 0, 0]]
  // data: [1, 4, 3]
483
484
485
  tp = aten::VecToIdArray(std::vector<IDX>({0, 0, 2, 3}), sizeof(IDX) * 8, ctx);
  ti = aten::VecToIdArray(std::vector<IDX>({2, 3, 0}), sizeof(IDX) * 8, ctx);
  td = aten::VecToIdArray(std::vector<IDX>({1, 4, 3}), sizeof(IDX) * 8, ctx);
486
487
488
489
490
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));

  // Testing zero-degree row slicing with different rows
491
492
  r = aten::VecToIdArray(
      std::vector<IDX>({1, 3, 0, 3, 2}), sizeof(IDX) * 8, ctx);
493
494
495
496
497
498
499
  x = aten::CSRSliceRows(csr, r);
  // [[1, 0, 0, 0, 0],
  //  [0, 0, 0, 0, 0],
  //  [0, 1, 2, 0, 0],
  //  [0, 0, 0, 0, 0],
  //  [0, 0, 1, 1, 0]]
  // data: [3, 0, 2, 5, 1, 4]
500
501
502
503
504
505
  tp = aten::VecToIdArray(
      std::vector<IDX>({0, 1, 1, 4, 4, 6}), sizeof(IDX) * 8, ctx);
  ti = aten::VecToIdArray(
      std::vector<IDX>({0, 1, 2, 2, 2, 3}), sizeof(IDX) * 8, ctx);
  td = aten::VecToIdArray(
      std::vector<IDX>({3, 0, 2, 5, 1, 4}), sizeof(IDX) * 8, ctx);
506
507
508
509
510
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));

  // Testing empty output (i.e. sliced rows will be zero-degree)
511
  r = aten::VecToIdArray(std::vector<IDX>({3, 3, 3}), sizeof(IDX) * 8, ctx);
512
513
514
515
516
  x = aten::CSRSliceRows(csr, r);
  // [[0, 0, 0, 0, 0],
  //  [0, 0, 0, 0, 0],
  //  [0, 0, 0, 0, 0]]
  // data: []
517
518
519
  tp = aten::VecToIdArray(std::vector<IDX>({0, 0, 0, 0}), sizeof(IDX) * 8, ctx);
  ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
  td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
520
521
522
523
524
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));

  // Testing constant output: we pick last row with at least one nnz
525
  r = aten::VecToIdArray(std::vector<IDX>({2, 2, 2}), sizeof(IDX) * 8, ctx);
526
527
528
529
530
  x = aten::CSRSliceRows(csr, r);
  // [[0, 0, 1, 1, 0],
  //  [0, 0, 1, 1, 0],
  //  [0, 0, 1, 1, 0]]
  // data: [1, 4, 1, 4, 1, 4]
531
532
533
534
535
  tp = aten::VecToIdArray(std::vector<IDX>({0, 2, 4, 6}), sizeof(IDX) * 8, ctx);
  ti = aten::VecToIdArray(
      std::vector<IDX>({2, 3, 2, 3, 2, 3}), sizeof(IDX) * 8, ctx);
  td = aten::VecToIdArray(
      std::vector<IDX>({1, 4, 1, 4, 1, 4}), sizeof(IDX) * 8, ctx);
536
537
538
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
539
540
541
}

TEST(SpmatTest, TestCSRSliceRows) {
542
543
544
545
546
547
  _TestCSRSliceRows<int32_t>(CPU);
  _TestCSRSliceRows<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRSliceRows<int32_t>(GPU);
  _TestCSRSliceRows<int64_t>(GPU);
#endif
548
549
550
}

template <typename IDX>
551
void _TestCSRSliceMatrix1(DGLContext ctx) {
552
553
  auto csr = CSR2<IDX>(ctx);
  {
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
    // square
    auto r =
        aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
    auto c =
        aten::VecToIdArray(std::vector<IDX>({1, 2, 3}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[1, 2, 0],
    //  [0, 0, 0],
    //  [0, 0, 0]]
    // data: [0, 2, 5]
    ASSERT_EQ(x.num_rows, 3);
    ASSERT_EQ(x.num_cols, 3);
    auto tp = aten::VecToIdArray(
        std::vector<IDX>({0, 3, 3, 3}), sizeof(IDX) * 8, ctx);
    auto ti =
        aten::VecToIdArray(std::vector<IDX>({0, 1, 1}), sizeof(IDX) * 8, ctx);
    auto td =
        aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
575
576
  }
  {
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
    // non-square
    auto r =
        aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[0, 1],
    //  [1, 0],
    //  [0, 0]]
    // data: [0, 3]
    ASSERT_EQ(x.num_rows, 3);
    ASSERT_EQ(x.num_cols, 2);
    auto tp = aten::VecToIdArray(
        std::vector<IDX>({0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
    auto ti =
        aten::VecToIdArray(std::vector<IDX>({1, 0}), sizeof(IDX) * 8, ctx);
    auto td =
        aten::VecToIdArray(std::vector<IDX>({0, 3}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
597
598
  }
  {
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
    // empty slice
    auto r = aten::VecToIdArray(std::vector<IDX>({2, 3}), sizeof(IDX) * 8, ctx);
    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[0, 0],
    //  [0, 0]]
    // data: []
    ASSERT_EQ(x.num_rows, 2);
    ASSERT_EQ(x.num_cols, 2);
    auto tp =
        aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
    auto ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
    auto td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
615
  }
616
617
}

618
template <typename IDX>
619
void _TestCSRSliceMatrix2(DGLContext ctx) {
620
621
  auto csr = CSR3<IDX>(ctx);
  {
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
    // square
    auto r =
        aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
    auto c =
        aten::VecToIdArray(std::vector<IDX>({1, 2, 3}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[1, 1, 1],
    //  [0, 0, 0],
    //  [0, 0, 0]]
    // data: [5, 2, 0]
    ASSERT_EQ(x.num_rows, 3);
    ASSERT_EQ(x.num_cols, 3);
    auto tp = aten::VecToIdArray(
        std::vector<IDX>({0, 3, 3, 3}), sizeof(IDX) * 8, ctx);
    // indexes are in reverse order in CSR3
    auto ti =
        aten::VecToIdArray(std::vector<IDX>({2, 1, 0}), sizeof(IDX) * 8, ctx);
    auto td =
        aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
644
645
  }
  {
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
    // non-square
    auto r =
        aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[0, 1],
    //  [1, 0],
    //  [0, 0]]
    // data: [0, 3]
    ASSERT_EQ(x.num_rows, 3);
    ASSERT_EQ(x.num_cols, 2);
    auto tp = aten::VecToIdArray(
        std::vector<IDX>({0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
    auto ti =
        aten::VecToIdArray(std::vector<IDX>({1, 0}), sizeof(IDX) * 8, ctx);
    auto td =
        aten::VecToIdArray(std::vector<IDX>({5, 3}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
666
667
  }
  {
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
    // empty slice
    auto r = aten::VecToIdArray(std::vector<IDX>({2, 3}), sizeof(IDX) * 8, ctx);
    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
    auto x = aten::CSRSliceMatrix(csr, r, c);
    // [[0, 0],
    //  [0, 0]]
    // data: []
    ASSERT_EQ(x.num_rows, 2);
    ASSERT_EQ(x.num_cols, 2);
    auto tp =
        aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
    auto ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
    auto td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
684
685
686
  }
}

687
TEST(SpmatTest, CSRSliceMatrix) {
688
689
690
691
  _TestCSRSliceMatrix1<int32_t>(CPU);
  _TestCSRSliceMatrix1<int64_t>(CPU);
  _TestCSRSliceMatrix2<int32_t>(CPU);
  _TestCSRSliceMatrix2<int64_t>(CPU);
692
#ifdef DGL_USE_CUDA
693
694
695
696
  _TestCSRSliceMatrix1<int32_t>(GPU);
  _TestCSRSliceMatrix1<int64_t>(GPU);
  _TestCSRSliceMatrix2<int32_t>(GPU);
  _TestCSRSliceMatrix2<int64_t>(GPU);
697
#endif
698
699
700
}

template <typename IDX>
701
void _TestCSRHasDuplicate(DGLContext ctx) {
702
  auto csr = CSR1<IDX>(ctx);
703
  ASSERT_FALSE(aten::CSRHasDuplicate(csr));
704
  csr = CSR2<IDX>(ctx);
705
706
707
  ASSERT_TRUE(aten::CSRHasDuplicate(csr));
}

708
709
710
711
712
TEST(SpmatTest, CSRHasDuplicate) {
  _TestCSRHasDuplicate<int32_t>(CPU);
  _TestCSRHasDuplicate<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRHasDuplicate<int32_t>(GPU);
713
  _TestCSRHasDuplicate<int64_t>(GPU);
714
#endif
715
716
}

717
template <typename IDX>
718
void _TestCSRSort(DGLContext ctx) {
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
  auto csr = CSR1<IDX>(ctx);
  ASSERT_FALSE(aten::CSRIsSorted(csr));
  auto csr1 = aten::CSRSort(csr);
  ASSERT_FALSE(aten::CSRIsSorted(csr));
  ASSERT_TRUE(aten::CSRIsSorted(csr1));
  ASSERT_TRUE(csr1.sorted);
  aten::CSRSort_(&csr);
  ASSERT_TRUE(aten::CSRIsSorted(csr));
  ASSERT_TRUE(csr.sorted);
  csr = CSR2<IDX>(ctx);
  ASSERT_TRUE(aten::CSRIsSorted(csr));
}

TEST(SpmatTest, CSRSort) {
  _TestCSRSort<int32_t>(CPU);
  _TestCSRSort<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRSort<int32_t>(GPU);
737
  _TestCSRSort<int64_t>(GPU);
738
739
740
#endif
}

Da Zheng's avatar
Da Zheng committed
741
742
743
template <typename IDX>
void _TestCSRReorder() {
  auto csr = CSR2<IDX>();
744
745
  auto new_row =
      aten::VecToIdArray(std::vector<IDX>({2, 0, 3, 1}), sizeof(IDX) * 8, CTX);
Da Zheng's avatar
Da Zheng committed
746
  auto new_col = aten::VecToIdArray(
747
      std::vector<IDX>({2, 0, 4, 3, 1}), sizeof(IDX) * 8, CTX);
Da Zheng's avatar
Da Zheng committed
748
749
750
751
752
753
754
755
756
  auto new_csr = CSRReorder(csr, new_row, new_col);
  ASSERT_EQ(new_csr.num_rows, csr.num_rows);
  ASSERT_EQ(new_csr.num_cols, csr.num_cols);
}

TEST(SpmatTest, TestCSRReorder) {
  _TestCSRReorder<int32_t>();
  _TestCSRReorder<int64_t>();
}