test_spmat_csr.cc 20.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
#include <gtest/gtest.h>
#include <dgl/array.h>
#include "./common.h"

using namespace dgl;
using namespace dgl::runtime;

namespace {

template <typename IDX>
11
aten::CSRMatrix CSR1(DLContext ctx = CTX) {
12
13
14
15
16
  // [[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]
17
18
  return aten::CSRMatrix(
      4, 5,
19
      aten::VecToIdArray(std::vector<IDX>({0, 2, 3, 5, 5}), sizeof(IDX)*8, ctx),
20
21
      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),
22
      false);
23
24
25
}

template <typename IDX>
26
aten::CSRMatrix CSR2(DLContext ctx = CTX) {
27
28
29
30
31
32
  // 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]
33
34
  return aten::CSRMatrix(
      4, 5,
35
36
37
      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),
38
      false);
39
40
}

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
template <typename IDX>
aten::CSRMatrix CSR3(DLContext ctx = CTX) {
  // 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,
      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),
      false);
}

65
template <typename IDX>
66
aten::COOMatrix COO1(DLContext ctx = CTX) {
67
68
69
70
71
72
73
  // [[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]
74
75
  return aten::COOMatrix(
      4, 5,
76
77
      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));
78
79
80
}

template <typename IDX>
81
aten::COOMatrix COO2(DLContext ctx = CTX) {
82
83
84
85
86
87
88
89
  // 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]
90
91
  return aten::COOMatrix(
      4, 5,
92
93
      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));
94
95
}

96
template <typename IDX>
97
aten::CSRMatrix SR_CSR3(DLContext ctx) {
98
99
100
101
102
103
  // [[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,
104
105
106
      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),
107
108
109
110
      false);
}

template <typename IDX>
111
aten::CSRMatrix SRC_CSR3(DLContext ctx) {
112
113
114
115
116
117
  // [[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,
118
119
120
      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),
121
122
123
124
      false);
}

template <typename IDX>
125
aten::COOMatrix COO3(DLContext ctx) {
126
127
128
129
130
131
132
133
134
  // 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,
135
136
      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));
137
138
}

139
}  // namespace
140
141

template <typename IDX>
142
void _TestCSRIsNonZero1(DLContext ctx) {
143
  auto csr = CSR1<IDX>(ctx);
144
145
  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
146
147
  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);
148
  IdArray x = aten::CSRIsNonZero(csr, r, c);
149
  IdArray tx = aten::VecToIdArray(std::vector<IDX>({0, 0, 1, 0}), sizeof(IDX)*8, ctx);
150
151
152
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

153
154
155
156
157
158
159
160
161
162
163
164
template <typename IDX>
void _TestCSRIsNonZero2(DLContext ctx) {
  auto csr = CSR3<IDX>(ctx);
  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
  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);
  IdArray x = aten::CSRIsNonZero(csr, r, c);
  IdArray tx = aten::VecToIdArray(std::vector<IDX>({0, 1, 1, 1, 0}), sizeof(IDX)*8, ctx);
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx)) << " x = " << x << ", tx = " << tx;
}

165
TEST(SpmatTest, TestCSRIsNonZero) {
166
167
168
169
  _TestCSRIsNonZero1<int32_t>(CPU);
  _TestCSRIsNonZero1<int64_t>(CPU);
  _TestCSRIsNonZero2<int32_t>(CPU);
  _TestCSRIsNonZero2<int64_t>(CPU);
170
#ifdef DGL_USE_CUDA
171
172
173
174
  _TestCSRIsNonZero1<int32_t>(GPU);
  _TestCSRIsNonZero1<int64_t>(GPU);
  _TestCSRIsNonZero2<int32_t>(GPU);
  _TestCSRIsNonZero2<int64_t>(GPU);
175
#endif
176
177
178
}

template <typename IDX>
179
180
void _TestCSRGetRowNNZ(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
181
182
  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 0), 3);
  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 3), 0);
183
  IdArray r = aten::VecToIdArray(std::vector<IDX>({0, 3}), sizeof(IDX)*8, ctx);
184
  IdArray x = aten::CSRGetRowNNZ(csr, r);
185
  IdArray tx = aten::VecToIdArray(std::vector<IDX>({3, 0}), sizeof(IDX)*8, ctx);
186
187
188
189
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

TEST(SpmatTest, TestCSRGetRowNNZ) {
190
191
192
193
194
195
  _TestCSRGetRowNNZ<int32_t>(CPU);
  _TestCSRGetRowNNZ<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetRowNNZ<int32_t>(GPU);
  _TestCSRGetRowNNZ<int64_t>(GPU);
#endif
196
197
198
}

template <typename IDX>
199
200
void _TestCSRGetRowColumnIndices(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
201
  auto x = aten::CSRGetRowColumnIndices(csr, 0);
202
  auto tx = aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX)*8, ctx);
203
204
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowColumnIndices(csr, 1);
205
  tx = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX)*8, ctx);
206
207
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowColumnIndices(csr, 3);
208
  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX)*8, ctx);
209
210
211
212
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

TEST(SpmatTest, TestCSRGetRowColumnIndices) {
213
214
215
216
217
218
  _TestCSRGetRowColumnIndices<int32_t>(CPU);
  _TestCSRGetRowColumnIndices<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetRowColumnIndices<int32_t>(GPU);
  _TestCSRGetRowColumnIndices<int64_t>(GPU);
#endif
219
220
221
}

template <typename IDX>
222
223
void _TestCSRGetRowData(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
224
  auto x = aten::CSRGetRowData(csr, 0);
225
  auto tx = aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX)*8, ctx);
226
227
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowData(csr, 1);
228
  tx = aten::VecToIdArray(std::vector<IDX>({3}), sizeof(IDX)*8, ctx);
229
230
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetRowData(csr, 3);
231
  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX)*8, ctx);
232
233
234
235
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
}

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

template <typename IDX>
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
void _TestCSRGetData(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
  // test get all data
  auto x = aten::CSRGetAllData(csr, 0, 0);
  auto tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX)*8, ctx);
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
  x = aten::CSRGetAllData(csr, 0, 2);
  tx = aten::VecToIdArray(std::vector<IDX>({2, 5}), sizeof(IDX)*8, ctx);
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));

  // test get data
  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);
  x = aten::CSRGetData(csr, r, c);
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX)*8, ctx);
260
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
261
262
263
264
265
266
267

  // test get data on sorted
  csr = aten::CSRSort(csr);
  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);
  x = aten::CSRGetData(csr, r, c);
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX)*8, ctx);
268
269
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));

270
271
272
  // test get data w/ broadcasting
  r = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX)*8, ctx);
  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX)*8, ctx);
273
  x = aten::CSRGetData(csr, r, c);
274
  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX)*8, ctx);
275
  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
276

277
278
}

279
280
281
282
283
TEST(SpmatTest, CSRGetData) {
  _TestCSRGetData<int32_t>(CPU);
  _TestCSRGetData<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetData<int32_t>(GPU);
284
  _TestCSRGetData<int64_t>(GPU);
285
#endif
286
287
288
}

template <typename IDX>
289
290
291
292
void _TestCSRGetDataAndIndices(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
  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);
293
  auto x = aten::CSRGetDataAndIndices(csr, r, c);
294
295
296
  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);
297
298
299
300
301
  ASSERT_TRUE(ArrayEQ<IDX>(x[0], tr));
  ASSERT_TRUE(ArrayEQ<IDX>(x[1], tc));
  ASSERT_TRUE(ArrayEQ<IDX>(x[2], td));
}

302
303
304
305
306
307
308
TEST(SpmatTest, CSRGetDataAndIndices) {
  _TestCSRGetDataAndIndices<int32_t>(CPU);
  _TestCSRGetDataAndIndices<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRGetDataAndIndices<int32_t>(GPU);
  _TestCSRGetDataAndIndices<int64_t>(GPU);
#endif
309
310
311
}

template <typename IDX>
312
313
void _TestCSRTranspose(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
314
315
316
317
318
319
320
321
322
  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);
323
324
325
  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);
326
327
328
329
330
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.data, td));
}

331
TEST(SpmatTest, CSRTranspose) {
332
333
334
335
  _TestCSRTranspose<int32_t>(CPU);
  _TestCSRTranspose<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRTranspose<int32_t>(GPU);
336
  _TestCSRTranspose<int64_t>(GPU);
337
#endif
338
339
340
}

template <typename IDX>
341
342
void _TestCSRToCOO(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
343
344
345
346
  {
  auto coo = CSRToCOO(csr, false);
  ASSERT_EQ(coo.num_rows, 4);
  ASSERT_EQ(coo.num_cols, 5);
347
  ASSERT_TRUE(coo.row_sorted);
348
  auto tr = aten::VecToIdArray(std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX)*8, ctx);
349
  ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tr));
350
351
352
353
354
355
356
357
358
359
360
361
362
363
  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));
364
365
366
367
368
  }
  {
  auto coo = CSRToCOO(csr, true);
  ASSERT_EQ(coo.num_rows, 4);
  ASSERT_EQ(coo.num_cols, 5);
369
  auto tcoo = COO2<IDX>(ctx);
370
371
372
373
374
  ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tcoo.row));
  ASSERT_TRUE(ArrayEQ<IDX>(coo.col, tcoo.col));
  }
}

375
TEST(SpmatTest, CSRToCOO) {
376
377
378
379
  _TestCSRToCOO<int32_t>(CPU);
  _TestCSRToCOO<int64_t>(CPU);
#if DGL_USE_CUDA
  _TestCSRToCOO<int32_t>(GPU);
380
  _TestCSRToCOO<int64_t>(GPU);
381
#endif
382
383
384
}

template <typename IDX>
385
386
void _TestCSRSliceRows(DLContext ctx) {
  auto csr = CSR2<IDX>(ctx);
387
388
389
390
391
392
393
  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);
394
395
396
  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);
397
398
399
400
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));

401
  auto r = aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX)*8, ctx);
402
403
404
405
406
  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]
407
408
409
  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);
410
411
412
413
414
415
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
}

TEST(SpmatTest, TestCSRSliceRows) {
416
417
418
419
420
421
  _TestCSRSliceRows<int32_t>(CPU);
  _TestCSRSliceRows<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRSliceRows<int32_t>(GPU);
  _TestCSRSliceRows<int64_t>(GPU);
#endif
422
423
424
}

template <typename IDX>
425
void _TestCSRSliceMatrix1(DLContext ctx) {
426
427
428
429
430
  auto csr = CSR2<IDX>(ctx);
  {
  // 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);
431
432
433
434
435
436
437
  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);
438
439
440
  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);
441
442
443
  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
  }
  {
  // 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));
  }
  {
  // 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));
  }
480
481
}

482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
template <typename IDX>
void _TestCSRSliceMatrix2(DLContext ctx) {
  auto csr = CSR3<IDX>(ctx);
  {
  // 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));
  }
  {
  // 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));
  }
  {
  // 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));
  }
}


542
TEST(SpmatTest, CSRSliceMatrix) {
543
544
545
546
  _TestCSRSliceMatrix1<int32_t>(CPU);
  _TestCSRSliceMatrix1<int64_t>(CPU);
  _TestCSRSliceMatrix2<int32_t>(CPU);
  _TestCSRSliceMatrix2<int64_t>(CPU);
547
#ifdef DGL_USE_CUDA
548
549
550
551
  _TestCSRSliceMatrix1<int32_t>(GPU);
  _TestCSRSliceMatrix1<int64_t>(GPU);
  _TestCSRSliceMatrix2<int32_t>(GPU);
  _TestCSRSliceMatrix2<int64_t>(GPU);
552
#endif
553
554
555
}

template <typename IDX>
556
557
void _TestCSRHasDuplicate(DLContext ctx) {
  auto csr = CSR1<IDX>(ctx);
558
  ASSERT_FALSE(aten::CSRHasDuplicate(csr));
559
  csr = CSR2<IDX>(ctx);
560
561
562
  ASSERT_TRUE(aten::CSRHasDuplicate(csr));
}

563
564
565
566
567
TEST(SpmatTest, CSRHasDuplicate) {
  _TestCSRHasDuplicate<int32_t>(CPU);
  _TestCSRHasDuplicate<int64_t>(CPU);
#ifdef DGL_USE_CUDA
  _TestCSRHasDuplicate<int32_t>(GPU);
568
  _TestCSRHasDuplicate<int64_t>(GPU);
569
#endif
570
571
}

572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
template <typename IDX>
void _TestCSRSort(DLContext ctx) {
  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);
592
  _TestCSRSort<int64_t>(GPU);
593
594
595
#endif
}

Da Zheng's avatar
Da Zheng committed
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
template <typename IDX>
void _TestCSRReorder() {
  auto csr = CSR2<IDX>();
  auto new_row = aten::VecToIdArray(
    std::vector<IDX>({2, 0, 3, 1}), sizeof(IDX)*8, CTX);
  auto new_col = aten::VecToIdArray(
    std::vector<IDX>({2, 0, 4, 3, 1}), sizeof(IDX)*8, CTX);
  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>();
}