spmat_op_impl_csr.cc 24.6 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
7
8
 */
#include <dgl/array.h>
#include <vector>
#include <unordered_set>
9
#include <numeric>
10
#include "array_utils.h"
11
12
13
14
15
16
17
18
19
20
21
22
23
24

namespace dgl {

using runtime::NDArray;

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
25
26
27
28
29
30
31
32
33
  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;
      }
34
35
36
37
38
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
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
    }
  }
  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;
  for (int64_t i = 0, j = 0; i < rowlen && j < collen; i += row_stride, j += col_stride) {
    *(rst_data++) = CSRIsNonZero<XPU, IdType>(csr, row_data[i], col_data[j])? 1 : 0;
  }
  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) {
98
  CHECK_SAME_DTYPE(csr.indices, rows);
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
  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 /////////////////////////////

129
template <DLDeviceType XPU, typename IdType>
130
131
132
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);
133
134
135
136
137
  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);
138
139
}

140
141
template NDArray CSRGetRowData<kDLCPU, int32_t>(CSRMatrix, int64_t);
template NDArray CSRGetRowData<kDLCPU, int64_t>(CSRMatrix, int64_t);
142
143
144

///////////////////////////// CSRGetData /////////////////////////////

145
146
template <DLDeviceType XPU, typename IdType>
void CollectDataFromSorted(const IdType *indices_data, const IdType *data,
Da Zheng's avatar
Da Zheng committed
147
                           const IdType start, const IdType end, const IdType col,
148
                           std::vector<IdType> *ret_vec) {
Da Zheng's avatar
Da Zheng committed
149
150
151
152
153
154
155
156
157
  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;
158
      ret_vec->push_back(data? data[idx] : idx);
Da Zheng's avatar
Da Zheng committed
159
160
161
162
163
164
165
    } else {
      // If we find a column that is different, we can stop searching now.
      break;
    }
  }
}

166
template <DLDeviceType XPU, typename IdType>
167
IdArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
168
169
170
171
172
173
174
175
176
177
178
179
180
  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);
181
  const IdType* data = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
182

183
184
185
  const int64_t retlen = std::max(rowlen, collen);
  IdArray ret = Full(-1, retlen, rows->dtype.bits, rows->ctx);
  IdType* ret_data = ret.Ptr<IdType>();
186

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
  // NOTE: In most cases, the input csr is already sorted. If not, we might need to
  //   consider sorting it especially when the number of (row, col) pairs is large.
  //   Need more benchmarks to justify the choice.

  if (csr.sorted) {
    // use binary search on each row
#pragma omp parallel for
    for (int64_t p = 0; p < retlen; ++p) {
      const IdType row_id = row_data[p * row_stride], col_id = col_data[p * col_stride];
      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;
      const IdType *start_ptr = indices_data + indptr_data[row_id];
      const IdType *end_ptr = indices_data + indptr_data[row_id + 1];
      auto it = std::lower_bound(start_ptr, end_ptr, col_id);
      if (it != end_ptr && *it == col_id) {
        const IdType idx = it - indices_data;
        ret_data[p] = data? data[idx] : idx;
      }
    }
  } else {
    // linear search on each row
#pragma omp parallel for
    for (int64_t p = 0; p < retlen; ++p) {
      const IdType row_id = row_data[p * row_stride], col_id = col_data[p * col_stride];
      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;
      for (IdType idx = indptr_data[row_id]; idx < indptr_data[row_id + 1]; ++idx) {
        if (indices_data[idx] == col_id) {
          ret_data[p] = data? data[idx] : idx;
          break;
Da Zheng's avatar
Da Zheng committed
217
        }
218
219
220
      }
    }
  }
221
  return ret;
222
223
}

224
225
template NDArray CSRGetData<kDLCPU, int32_t>(CSRMatrix csr, NDArray rows, NDArray cols);
template NDArray CSRGetData<kDLCPU, int64_t>(CSRMatrix csr, NDArray rows, NDArray cols);
226
227
228

///////////////////////////// CSRGetDataAndIndices /////////////////////////////

229
230
template <DLDeviceType XPU, typename IdType>
void CollectDataIndicesFromSorted(const IdType *indices_data, const IdType *data,
Da Zheng's avatar
Da Zheng committed
231
232
                                  const IdType start, const IdType end, const IdType col,
                                  std::vector<IdType> *col_vec,
233
                                  std::vector<IdType> *ret_vec) {
Da Zheng's avatar
Da Zheng committed
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
  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;
    }
  }
}

252
template <DLDeviceType XPU, typename IdType>
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
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);
268
  const IdType* data = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
269
270

  std::vector<IdType> ret_rows, ret_cols;
271
  std::vector<IdType> ret_data;
272
273
274
275
276

  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
277
278
    if (csr.sorted) {
      // Here we collect col indices and data.
279
280
281
282
283
      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
284
285
286
287
288
289
290
      // 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) {
291
292
          ret_rows.push_back(row_id);
          ret_cols.push_back(col_id);
293
          ret_data.push_back(data? data[i] : i);
Da Zheng's avatar
Da Zheng committed
294
        }
295
296
297
298
      }
    }
  }

299
300
301
  return {NDArray::FromVector(ret_rows, csr.indptr->ctx),
          NDArray::FromVector(ret_cols, csr.indptr->ctx),
          NDArray::FromVector(ret_data, csr.data->ctx)};
302
303
}

304
template std::vector<NDArray> CSRGetDataAndIndices<kDLCPU, int32_t>(
305
    CSRMatrix csr, NDArray rows, NDArray cols);
306
template std::vector<NDArray> CSRGetDataAndIndices<kDLCPU, int64_t>(
307
308
309
310
311
312
    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)
313
template <DLDeviceType XPU, typename IdType>
314
315
316
317
318
319
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);
320
  const IdType* Ax = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
321
322
  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);
323
  NDArray ret_data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
324
325
  IdType* Bp = static_cast<IdType*>(ret_indptr->data);
  IdType* Bi = static_cast<IdType*>(ret_indices->data);
326
  IdType* Bx = static_cast<IdType*>(ret_data->data);
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345

  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;
346
      Bx[Bp[dst]] = Ax? Ax[j] : j;
347
348
349
350
351
352
353
354
355
356
357
358
359
360
      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};
}

361
362
template CSRMatrix CSRTranspose<kDLCPU, int32_t>(CSRMatrix csr);
template CSRMatrix CSRTranspose<kDLCPU, int64_t>(CSRMatrix csr);
363
364
365
366
367
368
369
370
371
372
373
374
375

///////////////////////////// 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);
  }
376
377
378
  return COOMatrix(csr.num_rows, csr.num_cols,
                   ret_row, csr.indices, csr.data,
                   true, csr.sorted);
379
380
381
382
383
384
385
386
387
388
389
390
391
392
}

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
393
  const IdType* data = CSRHasData(csr) ? static_cast<IdType*>(csr.data->data) : nullptr;
394
395
396
397
398
399
400
401
  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
402
403
      ret_row_data[data ? data[j] : j] = row;
      ret_col_data[data ? data[j] : j] = col;
404
405
    }
  }
406
  return COOMatrix(N, M, ret_row, ret_col);
407
408
409
410
411
412
413
}

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

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

414
template <DLDeviceType XPU, typename IdType>
415
416
417
418
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];
419
420
  IdArray ret_indptr = IdArray::Empty({num_rows + 1}, csr.indptr->dtype, csr.indices->ctx);
  IdType* r_indptr = static_cast<IdType*>(ret_indptr->data);
421
422
423
424
  for (int64_t i = start; i < end + 1; ++i) {
    r_indptr[i - start] = indptr[i] - indptr[start];
  }
  // indices and data can be view arrays
425
426
427
428
429
430
431
432
433
434
435
  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);
436
437
}

438
439
template CSRMatrix CSRSliceRows<kDLCPU, int32_t>(CSRMatrix, int64_t, int64_t);
template CSRMatrix CSRSliceRows<kDLCPU, int64_t>(CSRMatrix, int64_t, int64_t);
440

441
template <DLDeviceType XPU, typename IdType>
442
CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
443
  CHECK_SAME_DTYPE(csr.indices, rows);
444
445
  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
446
  const IdType* data = CSRHasData(csr)? static_cast<IdType*>(csr.data->data) : nullptr;
447
448
449
450
451
452
453
454
455
456
457
458
459
  const auto len = rows->shape[0];
  const IdType* rows_data = static_cast<IdType*>(rows->data);
  int64_t nnz = 0;
  for (int64_t i = 0; i < len; ++i) {
    IdType vid = rows_data[i];
    nnz += impl::CSRGetRowNNZ<XPU, IdType>(csr, vid);
  }

  CSRMatrix ret;
  ret.num_rows = len;
  ret.num_cols = csr.num_cols;
  ret.indptr = NDArray::Empty({len + 1}, csr.indptr->dtype, csr.indices->ctx);
  ret.indices = NDArray::Empty({nnz}, csr.indices->dtype, csr.indices->ctx);
460
461
  ret.data = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
  ret.sorted = csr.sorted;
462
463
464

  IdType* ret_indptr_data = static_cast<IdType*>(ret.indptr->data);
  IdType* ret_indices_data = static_cast<IdType*>(ret.indices->data);
465
  IdType* ret_data = static_cast<IdType*>(ret.data->data);
466
467
468
469
470
471
472
  ret_indptr_data[0] = 0;
  for (int64_t i = 0; i < len; ++i) {
    const IdType rid = rows_data[i];
    // note: zero is allowed
    ret_indptr_data[i + 1] = ret_indptr_data[i] + indptr_data[rid + 1] - indptr_data[rid];
    std::copy(indices_data + indptr_data[rid], indices_data + indptr_data[rid + 1],
              ret_indices_data + ret_indptr_data[i]);
473
474
475
476
477
478
    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]);
479
480
481
482
  }
  return ret;
}

483
484
template CSRMatrix CSRSliceRows<kDLCPU, int32_t>(CSRMatrix , NDArray);
template CSRMatrix CSRSliceRows<kDLCPU, int64_t>(CSRMatrix , NDArray);
485
486
487

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

488
template <DLDeviceType XPU, typename IdType>
489
490
491
492
493
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);
494
  const bool has_data = CSRHasData(csr);
495
496
497

  const IdType* indptr_data = static_cast<IdType*>(csr.indptr->data);
  const IdType* indices_data = static_cast<IdType*>(csr.indices->data);
498
  const IdType* data = has_data? static_cast<IdType*>(csr.data->data) : nullptr;
499
500

  std::vector<IdType> sub_indptr, sub_indices;
501
  std::vector<IdType> sub_data;
502
503
504
505
506
507
508
509
510
511
512
513
  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);
514
        sub_data.push_back(has_data? data[p] : p);
515
516
517
518
519
520
521
522
523
524
525
526
527
      }
    }
  }

  // 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();
528
529
  NDArray sub_data_arr = NDArray::Empty({nnz}, csr.indptr->dtype, csr.indptr->ctx);
  IdType* ptr = static_cast<IdType*>(sub_data_arr->data);
530
531
  std::copy(sub_data.begin(), sub_data.end(), ptr);
  return CSRMatrix{new_nrows, new_ncols,
532
533
    NDArray::FromVector(sub_indptr, csr.indptr->ctx),
    NDArray::FromVector(sub_indices, csr.indptr->ctx),
534
535
536
    sub_data_arr};
}

537
template CSRMatrix CSRSliceMatrix<kDLCPU, int32_t>(
538
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
539
template CSRMatrix CSRSliceMatrix<kDLCPU, int64_t>(
540
541
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);

Da Zheng's avatar
Da Zheng committed
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
///////////////////////////// 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);
#pragma omp parallel for
  for (int64_t i = 0; i < num_rows; i++) {
    int64_t new_row_id = new_row_ids[i];
    new_row_lens[new_row_id] = in_indptr[i + 1] - in_indptr[i];
  }
  // 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.
#pragma omp parallel for
  for (int64_t i = 0; i < num_rows; 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.
  }
  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);

618
619
620
}  // namespace impl
}  // namespace aten
}  // namespace dgl