spmat_op_impl_csr.cc 23 KB
Newer Older
1
2
/*!
 *  Copyright (c) 2019 by Contributors
3
4
 * \file array/cpu/spmat_op_impl_csr.cc
 * \brief CSR matrix operator CPU implementation
5
6
 */
#include <dgl/array.h>
7
#include <dgl/runtime/parallel_for.h>
8
9
#include <vector>
#include <unordered_set>
10
#include <numeric>
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
11
#include <atomic>
12
#include "array_utils.h"
13
14
15
16

namespace dgl {

using runtime::NDArray;
17
using runtime::parallel_for;
18
19
20
21
22
23
24
25
26
27

namespace aten {
namespace impl {

///////////////////////////// CSRIsNonZero /////////////////////////////

template <DLDeviceType XPU, typename IdType>
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
Da Zheng's avatar
Da Zheng committed
28
29
30
31
32
33
34
35
36
  if (csr.sorted) {
    const IdType *start = indices_data + indptr_data[row];
    const IdType *end = indices_data + indptr_data[row + 1];
    return std::binary_search(start, end, col);
  } else {
    for (IdType i = indptr_data[row]; i < indptr_data[row + 1]; ++i) {
      if (indices_data[i] == col) {
        return true;
      }
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    }
  }
  return false;
}

template bool CSRIsNonZero<kDLCPU, int32_t>(CSRMatrix, int64_t, int64_t);
template bool CSRIsNonZero<kDLCPU, int64_t>(CSRMatrix, int64_t, int64_t);

template <DLDeviceType XPU, typename IdType>
NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
  const auto rowlen = row->shape[0];
  const auto collen = col->shape[0];
  const auto rstlen = std::max(rowlen, collen);
  NDArray rst = NDArray::Empty({rstlen}, row->dtype, row->ctx);
  IdType* rst_data = static_cast<IdType*>(rst->data);
  const IdType* row_data = static_cast<IdType*>(row->data);
  const IdType* col_data = static_cast<IdType*>(col->data);
  const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
  const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
56
57
58
59
60
61
  runtime::parallel_for(0, std::max(rowlen, collen), 1, [=](int64_t b, int64_t e) {
    int64_t i = (row_stride == 0) ? 0 : b;
    int64_t j = (col_stride == 0) ? 0 : b;
    for (int64_t k = b; i < e && j < e; i += row_stride, j += col_stride, ++k)
      rst_data[k] = CSRIsNonZero<XPU, IdType>(csr, row_data[i], col_data[j]) ? 1 : 0;
  });
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
  return rst;
}

template NDArray CSRIsNonZero<kDLCPU, int32_t>(CSRMatrix, NDArray, NDArray);
template NDArray CSRIsNonZero<kDLCPU, int64_t>(CSRMatrix, NDArray, NDArray);

///////////////////////////// CSRHasDuplicate /////////////////////////////

template <DLDeviceType XPU, typename IdType>
bool CSRHasDuplicate(CSRMatrix csr) {
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
  for (IdType src = 0; src < csr.num_rows; ++src) {
    std::unordered_set<IdType> hashmap;
    for (IdType eid = indptr_data[src]; eid < indptr_data[src+1]; ++eid) {
      const IdType dst = indices_data[eid];
      if (hashmap.count(dst)) {
        return true;
      } else {
        hashmap.insert(dst);
      }
    }
  }
  return false;
}

template bool CSRHasDuplicate<kDLCPU, int32_t>(CSRMatrix csr);
template bool CSRHasDuplicate<kDLCPU, int64_t>(CSRMatrix csr);

///////////////////////////// CSRGetRowNNZ /////////////////////////////

template <DLDeviceType XPU, typename IdType>
int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  return indptr_data[row + 1] - indptr_data[row];
}

template int64_t CSRGetRowNNZ<kDLCPU, int32_t>(CSRMatrix, int64_t);
template int64_t CSRGetRowNNZ<kDLCPU, int64_t>(CSRMatrix, int64_t);

template <DLDeviceType XPU, typename IdType>
NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray rows) {
104
  CHECK_SAME_DTYPE(csr.indices, rows);
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
  const auto len = rows->shape[0];
  const IdType* vid_data = static_cast<IdType*>(rows->data);
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
  IdType* rst_data = static_cast<IdType*>(rst->data);
  for (int64_t i = 0; i < len; ++i) {
    const auto vid = vid_data[i];
    rst_data[i] = indptr_data[vid + 1] - indptr_data[vid];
  }
  return rst;
}

template NDArray CSRGetRowNNZ<kDLCPU, int32_t>(CSRMatrix, NDArray);
template NDArray CSRGetRowNNZ<kDLCPU, int64_t>(CSRMatrix, NDArray);

///////////////////////////// CSRGetRowColumnIndices /////////////////////////////

template <DLDeviceType XPU, typename IdType>
NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
  const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const int64_t offset = indptr_data[row] * sizeof(IdType);
  return csr.indices.CreateView({len}, csr.indices->dtype, offset);
}

template NDArray CSRGetRowColumnIndices<kDLCPU, int32_t>(CSRMatrix, int64_t);
template NDArray CSRGetRowColumnIndices<kDLCPU, int64_t>(CSRMatrix, int64_t);

///////////////////////////// CSRGetRowData /////////////////////////////

135
template <DLDeviceType XPU, typename IdType>
136
137
138
NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
  const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
139
140
141
142
143
  const int64_t offset = indptr_data[row] * sizeof(IdType);
  if (CSRHasData(csr))
    return csr.data.CreateView({len}, csr.data->dtype, offset);
  else
    return aten::Range(offset, offset + len, csr.indptr->dtype.bits, csr.indptr->ctx);
144
145
}

146
147
template NDArray CSRGetRowData<kDLCPU, int32_t>(CSRMatrix, int64_t);
template NDArray CSRGetRowData<kDLCPU, int64_t>(CSRMatrix, int64_t);
148
149
150
151

///////////////////////////// CSRGetData /////////////////////////////
///////////////////////////// CSRGetDataAndIndices /////////////////////////////

152
153
template <DLDeviceType XPU, typename IdType>
void CollectDataIndicesFromSorted(const IdType *indices_data, const IdType *data,
Da Zheng's avatar
Da Zheng committed
154
155
                                  const IdType start, const IdType end, const IdType col,
                                  std::vector<IdType> *col_vec,
156
                                  std::vector<IdType> *ret_vec) {
Da Zheng's avatar
Da Zheng committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
  const IdType *start_ptr = indices_data + start;
  const IdType *end_ptr = indices_data + end;
  auto it = std::lower_bound(start_ptr, end_ptr, col);
  // This might be a multi-graph. We need to collect all of the matched
  // columns.
  for (; it != end_ptr; it++) {
    // If the col exist
    if (*it == col) {
      IdType idx = it - indices_data;
      col_vec->push_back(indices_data[idx]);
      ret_vec->push_back(data[idx]);
    } else {
      // If we find a column that is different, we can stop searching now.
      break;
    }
  }
}

175
template <DLDeviceType XPU, typename IdType>
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
std::vector<NDArray> CSRGetDataAndIndices(CSRMatrix csr, NDArray rows, NDArray cols) {
  // TODO(minjie): more efficient implementation for matrix without duplicate entries
  const int64_t rowlen = rows->shape[0];
  const int64_t collen = cols->shape[0];

  CHECK((rowlen == collen) || (rowlen == 1) || (collen == 1))
    << "Invalid row and col id array.";

  const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
  const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
  const IdType* row_data = static_cast<IdType*>(rows->data);
  const IdType* col_data = static_cast<IdType*>(cols->data);

  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
191
  const IdType* data = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
192
193

  std::vector<IdType> ret_rows, ret_cols;
194
  std::vector<IdType> ret_data;
195
196
197
198
199

  for (int64_t i = 0, j = 0; i < rowlen && j < collen; i += row_stride, j += col_stride) {
    const IdType row_id = row_data[i], col_id = col_data[j];
    CHECK(row_id >= 0 && row_id < csr.num_rows) << "Invalid row index: " << row_id;
    CHECK(col_id >= 0 && col_id < csr.num_cols) << "Invalid col index: " << col_id;
Da Zheng's avatar
Da Zheng committed
200
201
    if (csr.sorted) {
      // Here we collect col indices and data.
202
203
204
205
206
      CollectDataIndicesFromSorted<XPU, IdType>(indices_data, data,
                                                indptr_data[row_id],
                                                indptr_data[row_id + 1],
                                                col_id, &ret_cols,
                                                &ret_data);
Da Zheng's avatar
Da Zheng committed
207
208
209
210
211
212
213
      // We need to add row Ids.
      while (ret_rows.size() < ret_data.size()) {
        ret_rows.push_back(row_id);
      }
    } else {
      for (IdType i = indptr_data[row_id]; i < indptr_data[row_id+1]; ++i) {
        if (indices_data[i] == col_id) {
214
215
          ret_rows.push_back(row_id);
          ret_cols.push_back(col_id);
216
          ret_data.push_back(data? data[i] : i);
Da Zheng's avatar
Da Zheng committed
217
        }
218
219
220
221
      }
    }
  }

222
223
224
  return {NDArray::FromVector(ret_rows, csr.indptr->ctx),
          NDArray::FromVector(ret_cols, csr.indptr->ctx),
          NDArray::FromVector(ret_data, csr.data->ctx)};
225
226
}

227
template std::vector<NDArray> CSRGetDataAndIndices<kDLCPU, int32_t>(
228
    CSRMatrix csr, NDArray rows, NDArray cols);
229
template std::vector<NDArray> CSRGetDataAndIndices<kDLCPU, int64_t>(
230
231
232
233
234
235
    CSRMatrix csr, NDArray rows, NDArray cols);

///////////////////////////// CSRTranspose /////////////////////////////

// for a matrix of shape (N, M) and NNZ
// complexity: time O(NNZ + max(N, M)), space O(1)
236
template <DLDeviceType XPU, typename IdType>
237
238
239
240
241
242
CSRMatrix CSRTranspose(CSRMatrix csr) {
  const int64_t N = csr.num_rows;
  const int64_t M = csr.num_cols;
  const int64_t nnz = csr.indices->shape[0];
  const IdType* Ap = static_cast<IdType*>(csr.indptr->data);
  const IdType* Aj = static_cast<IdType*>(csr.indices->data);
243
  const IdType* Ax = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
244
245
  NDArray ret_indptr = NDArray::Empty({M + 1}, csr.indptr->dtype, csr.indptr->ctx);
  NDArray ret_indices = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
246
  NDArray ret_data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
247
248
  IdType* Bp = static_cast<IdType*>(ret_indptr->data);
  IdType* Bi = static_cast<IdType*>(ret_indices->data);
249
  IdType* Bx = static_cast<IdType*>(ret_data->data);
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268

  std::fill(Bp, Bp + M, 0);

  for (int64_t j = 0; j < nnz; ++j) {
    Bp[Aj[j]]++;
  }

  // cumsum
  for (int64_t i = 0, cumsum = 0; i < M; ++i) {
    const IdType temp = Bp[i];
    Bp[i] = cumsum;
    cumsum += temp;
  }
  Bp[M] = nnz;

  for (int64_t i = 0; i < N; ++i) {
    for (IdType j = Ap[i]; j < Ap[i+1]; ++j) {
      const IdType dst = Aj[j];
      Bi[Bp[dst]] = i;
269
      Bx[Bp[dst]] = Ax? Ax[j] : j;
270
271
272
273
274
275
276
277
278
279
280
281
282
283
      Bp[dst]++;
    }
  }

  // correct the indptr
  for (int64_t i = 0, last = 0; i <= M; ++i) {
    IdType temp = Bp[i];
    Bp[i] = last;
    last = temp;
  }

  return CSRMatrix{csr.num_cols, csr.num_rows, ret_indptr, ret_indices, ret_data};
}

284
285
template CSRMatrix CSRTranspose<kDLCPU, int32_t>(CSRMatrix csr);
template CSRMatrix CSRTranspose<kDLCPU, int64_t>(CSRMatrix csr);
286
287
288
289
290
291
292
293
294
295
296
297
298

///////////////////////////// CSRToCOO /////////////////////////////
template <DLDeviceType XPU, typename IdType>
COOMatrix CSRToCOO(CSRMatrix csr) {
  const int64_t nnz = csr.indices->shape[0];
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  NDArray ret_row = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
  IdType* ret_row_data = static_cast<IdType*>(ret_row->data);
  for (IdType i = 0; i < csr.indptr->shape[0] - 1; ++i) {
    std::fill(ret_row_data + indptr_data[i],
              ret_row_data + indptr_data[i + 1],
              i);
  }
299
300
301
  return COOMatrix(csr.num_rows, csr.num_cols,
                   ret_row, csr.indices, csr.data,
                   true, csr.sorted);
302
303
304
305
306
307
308
309
310
311
312
313
314
315
}

template COOMatrix CSRToCOO<kDLCPU, int32_t>(CSRMatrix csr);
template COOMatrix CSRToCOO<kDLCPU, int64_t>(CSRMatrix csr);

// complexity: time O(NNZ), space O(1)
template <DLDeviceType XPU, typename IdType>
COOMatrix CSRToCOODataAsOrder(CSRMatrix csr) {
  const int64_t N = csr.num_rows;
  const int64_t M = csr.num_cols;
  const int64_t nnz = csr.indices->shape[0];
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
  // data array should have the same type as the indices arrays
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
316
  const IdType* data = CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
317
318
319
320
321
322
323
324
  NDArray ret_row = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
  NDArray ret_col = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
  IdType* ret_row_data = static_cast<IdType*>(ret_row->data);
  IdType* ret_col_data = static_cast<IdType*>(ret_col->data);
  // scatter using the indices in the data array
  for (IdType row = 0; row < N; ++row) {
    for (IdType j = indptr_data[row]; j < indptr_data[row + 1]; ++j) {
      const IdType col = indices_data[j];
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
325
326
      ret_row_data[data ? data[j] : j] = row;
      ret_col_data[data ? data[j] : j] = col;
327
328
    }
  }
329
  return COOMatrix(N, M, ret_row, ret_col);
330
331
332
333
334
335
336
}

template COOMatrix CSRToCOODataAsOrder<kDLCPU, int32_t>(CSRMatrix csr);
template COOMatrix CSRToCOODataAsOrder<kDLCPU, int64_t>(CSRMatrix csr);

///////////////////////////// CSRSliceRows /////////////////////////////

337
template <DLDeviceType XPU, typename IdType>
338
339
340
341
CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end) {
  const IdType* indptr = static_cast<IdType*>(csr.indptr->data);
  const int64_t num_rows = end - start;
  const int64_t nnz = indptr[end] - indptr[start];
342
343
  IdArray ret_indptr = IdArray::Empty({num_rows + 1}, csr.indptr->dtype, csr.indices->ctx);
  IdType* r_indptr = static_cast<IdType*>(ret_indptr->data);
344
345
346
347
  for (int64_t i = start; i < end + 1; ++i) {
    r_indptr[i - start] = indptr[i] - indptr[start];
  }
  // indices and data can be view arrays
348
349
350
351
352
353
354
355
356
357
358
  IdArray ret_indices = csr.indices.CreateView(
      {nnz}, csr.indices->dtype, indptr[start] * sizeof(IdType));
  IdArray ret_data;
  if (CSRHasData(csr))
    ret_data = csr.data.CreateView({nnz}, csr.data->dtype, indptr[start] * sizeof(IdType));
  else
    ret_data = aten::Range(indptr[start], indptr[end],
                           csr.indptr->dtype.bits, csr.indptr->ctx);
  return CSRMatrix(num_rows, csr.num_cols,
                   ret_indptr, ret_indices, ret_data,
                   csr.sorted);
359
360
}

361
362
template CSRMatrix CSRSliceRows<kDLCPU, int32_t>(CSRMatrix, int64_t, int64_t);
template CSRMatrix CSRSliceRows<kDLCPU, int64_t>(CSRMatrix, int64_t, int64_t);
363

364
template <DLDeviceType XPU, typename IdType>
365
CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
366
  CHECK_SAME_DTYPE(csr.indices, rows);
367
368
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
369
  const IdType* data = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
370
371
372
373
374
375
376
377
  const auto len = rows->shape[0];
  const IdType* rows_data = static_cast<IdType*>(rows->data);
  int64_t nnz = 0;

  CSRMatrix ret;
  ret.num_rows = len;
  ret.num_cols = csr.num_cols;
  ret.indptr = NDArray::Empty({len + 1}, csr.indptr->dtype, csr.indices->ctx);
378
379
380
381
382
383

  IdType* ret_indptr_data = static_cast<IdType*>(ret.indptr->data);
  ret_indptr_data[0] = 0;

  std::vector<IdType> sums;

Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
384
385
386
387
  std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
  bool err = false;
  std::stringstream err_msg_stream;

388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
  // Perform two-round parallel prefix sum using OpenMP
  #pragma omp parallel
  {
    int64_t tid = omp_get_thread_num();
    int64_t num_threads = omp_get_num_threads();

    #pragma omp single
    {
        sums.resize(num_threads + 1);
        sums[0] = 0;
    }

    int64_t sum = 0;

    // First round of parallel prefix sum. All threads perform local prefix sums.
    #pragma omp for schedule(static) nowait
    for (int64_t i = 0; i < len; ++i) {
      int64_t rid = rows_data[i];
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
406
407
408
409
410
411
412
413
414
415
      if (rid >= csr.num_rows) {
        if (!err_flag.test_and_set()) {
          err_msg_stream << "expect row ID " << rid << " to be less than number of rows "
            << csr.num_rows;
          err = true;
        }
      } else {
        sum += indptr_data[rid + 1] - indptr_data[rid];
        ret_indptr_data[i + 1] = sum;
      }
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
    }
    sums[tid + 1] = sum;
    #pragma omp barrier

    #pragma omp single
    {
      for (int64_t i = 1; i < num_threads; ++i)
        sums[i] += sums[i - 1];
    }

    int64_t offset = sums[tid];

    // Second round of parallel prefix sum. Update the local prefix sums.
    #pragma omp for schedule(static)
    for (int64_t i = 0; i < len; ++i)
      ret_indptr_data[i + 1] += offset;
  }
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
433
434
435
436
  if (err) {
    LOG(FATAL) << err_msg_stream.str();
    return ret;
  }
437
438
439
440
441

  // After the prefix sum, the last element of ret_indptr_data holds the
  // sum of all elements
  nnz = ret_indptr_data[len];

442
  ret.indices = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
443
444
  ret.data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
  ret.sorted = csr.sorted;
445
446

  IdType* ret_indices_data = static_cast<IdType*>(ret.indices->data);
447
  IdType* ret_data = static_cast<IdType*>(ret.data->data);
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462

  parallel_for(0, len, [=](int64_t b, int64_t e) {
    for (auto i = b; i < e; ++i) {
      const IdType rid = rows_data[i];
      // note: zero is allowed
      std::copy(indices_data + indptr_data[rid], indices_data + indptr_data[rid + 1],
                ret_indices_data + ret_indptr_data[i]);
      if (data)
        std::copy(data + indptr_data[rid], data + indptr_data[rid + 1],
                  ret_data + ret_indptr_data[i]);
      else
        std::iota(ret_data + ret_indptr_data[i], ret_data + ret_indptr_data[i + 1],
                  indptr_data[rid]);
    }
  });
463
464
465
  return ret;
}

466
467
template CSRMatrix CSRSliceRows<kDLCPU, int32_t>(CSRMatrix , NDArray);
template CSRMatrix CSRSliceRows<kDLCPU, int64_t>(CSRMatrix , NDArray);
468
469
470

///////////////////////////// CSRSliceMatrix /////////////////////////////

471
template <DLDeviceType XPU, typename IdType>
472
473
474
475
476
CSRMatrix CSRSliceMatrix(CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols) {
  IdHashMap<IdType> hashmap(cols);
  const int64_t new_nrows = rows->shape[0];
  const int64_t new_ncols = cols->shape[0];
  const IdType* rows_data = static_cast<IdType*>(rows->data);
477
  const bool has_data = CSRHasData(csr);
478
479
480

  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
481
  const IdType* data = has_data? static_cast<IdType*>(csr.data->data) : nullptr;
482
483

  std::vector<IdType> sub_indptr, sub_indices;
484
  std::vector<IdType> sub_data;
485
486
487
488
489
490
491
492
493
494
495
496
  sub_indptr.resize(new_nrows + 1, 0);
  const IdType kInvalidId = new_ncols + 1;
  for (int64_t i = 0; i < new_nrows; ++i) {
    // NOTE: newi == i
    const IdType oldi = rows_data[i];
    CHECK(oldi >= 0 && oldi < csr.num_rows) << "Invalid row index: " << oldi;
    for (IdType p = indptr_data[oldi]; p < indptr_data[oldi + 1]; ++p) {
      const IdType oldj = indices_data[p];
      const IdType newj = hashmap.Map(oldj, kInvalidId);
      if (newj != kInvalidId) {
        ++sub_indptr[i];
        sub_indices.push_back(newj);
497
        sub_data.push_back(has_data? data[p] : p);
498
499
500
501
502
503
504
505
506
507
508
509
510
      }
    }
  }

  // cumsum sub_indptr
  for (int64_t i = 0, cumsum = 0; i < new_nrows; ++i) {
    const IdType temp = sub_indptr[i];
    sub_indptr[i] = cumsum;
    cumsum += temp;
  }
  sub_indptr[new_nrows] = sub_indices.size();

  const int64_t nnz = sub_data.size();
511
512
  NDArray sub_data_arr = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
  IdType* ptr = static_cast<IdType*>(sub_data_arr->data);
513
514
  std::copy(sub_data.begin(), sub_data.end(), ptr);
  return CSRMatrix{new_nrows, new_ncols,
515
516
    NDArray::FromVector(sub_indptr, csr.indptr->ctx),
    NDArray::FromVector(sub_indices, csr.indptr->ctx),
517
518
519
    sub_data_arr};
}

520
template CSRMatrix CSRSliceMatrix<kDLCPU, int32_t>(
521
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
522
template CSRMatrix CSRSliceMatrix<kDLCPU, int64_t>(
523
524
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);

Da Zheng's avatar
Da Zheng committed
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
///////////////////////////// CSRReorder /////////////////////////////

template <DLDeviceType XPU, typename IdType>
CSRMatrix CSRReorder(CSRMatrix csr, runtime::NDArray new_row_id_arr,
                     runtime::NDArray new_col_id_arr) {
  CHECK_SAME_DTYPE(csr.indices, new_row_id_arr);
  CHECK_SAME_DTYPE(csr.indices, new_col_id_arr);

  // Input CSR
  const IdType* in_indptr = static_cast<IdType*>(csr.indptr->data);
  const IdType* in_indices = static_cast<IdType*>(csr.indices->data);
  const IdType* in_data = static_cast<IdType*>(csr.data->data);
  int64_t num_rows = csr.num_rows;
  int64_t num_cols = csr.num_cols;
  int64_t nnz = csr.indices->shape[0];
  CHECK_EQ(nnz, in_indptr[num_rows]);
  CHECK_EQ(num_rows, new_row_id_arr->shape[0])
      << "The new row Id array needs to be the same as the number of rows of CSR";
  CHECK_EQ(num_cols, new_col_id_arr->shape[0])
      << "The new col Id array needs to be the same as the number of cols of CSR";

  // New row/col Ids.
  const IdType* new_row_ids = static_cast<IdType*>(new_row_id_arr->data);
  const IdType* new_col_ids = static_cast<IdType*>(new_col_id_arr->data);

  // Output CSR
  NDArray out_indptr_arr = NDArray::Empty({num_rows + 1}, csr.indptr->dtype, csr.indptr->ctx);
  NDArray out_indices_arr = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
  NDArray out_data_arr = NDArray::Empty({nnz}, csr.data->dtype, csr.data->ctx);
  IdType *out_indptr = static_cast<IdType*>(out_indptr_arr->data);
  IdType *out_indices = static_cast<IdType*>(out_indices_arr->data);
  IdType *out_data = static_cast<IdType*>(out_data_arr->data);

  // Compute the length of rows for the new matrix.
  std::vector<IdType> new_row_lens(num_rows, -1);
560
561
562
563
564
565
  parallel_for(0, num_rows, [=, &new_row_lens](size_t b, size_t e) {
    for (auto i = b; i < e; ++i) {
      int64_t new_row_id = new_row_ids[i];
      new_row_lens[new_row_id] = in_indptr[i + 1] - in_indptr[i];
    }
  });
Da Zheng's avatar
Da Zheng committed
566
567
568
569
570
571
572
573
574
575
  // Compute the starting location of each row in the new matrix.
  out_indptr[0] = 0;
  // This is sequential. It should be pretty fast.
  for (int64_t i = 0; i < num_rows; i++) {
    CHECK_GE(new_row_lens[i], 0);
    out_indptr[i + 1] = out_indptr[i] + new_row_lens[i];
  }
  CHECK_EQ(out_indptr[num_rows], nnz);
  // Copy indieces and data with the new order.
  // Here I iterate rows in the order of the old matrix.
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
  parallel_for(0, num_rows, [=](size_t b, size_t e) {
    for (auto i = b; i < e; ++i) {
      const IdType *in_row = in_indices + in_indptr[i];
      const IdType *in_row_data = in_data + in_indptr[i];

      int64_t new_row_id = new_row_ids[i];
      IdType *out_row = out_indices + out_indptr[new_row_id];
      IdType *out_row_data = out_data + out_indptr[new_row_id];

      int64_t row_len = new_row_lens[new_row_id];
      // Here I iterate col indices in a row in the order of the old matrix.
      for (int64_t j = 0; j < row_len; j++) {
        out_row[j] = new_col_ids[in_row[j]];
        out_row_data[j] = in_row_data[j];
      }
      // TODO(zhengda) maybe we should sort the column indices.
Da Zheng's avatar
Da Zheng committed
592
    }
593
  });
Da Zheng's avatar
Da Zheng committed
594
595
596
597
598
599
600
601
602
  return CSRMatrix(num_rows, num_cols,
    out_indptr_arr, out_indices_arr, out_data_arr);
}

template CSRMatrix CSRReorder<kDLCPU, int64_t>(CSRMatrix csr, runtime::NDArray new_row_ids,
                                               runtime::NDArray new_col_ids);
template CSRMatrix CSRReorder<kDLCPU, int32_t>(CSRMatrix csr, runtime::NDArray new_row_ids,
                                               runtime::NDArray new_col_ids);

603
604
605
}  // namespace impl
}  // namespace aten
}  // namespace dgl