spmat_op_impl_coo.cc 23.8 KB
Newer Older
1
2
3
4
5
/*!
 *  Copyright (c) 2019 by Contributors
 * \file array/cpu/spmat_op_impl.cc
 * \brief CPU implementation of COO sparse matrix operators
 */
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
6
#include <dmlc/omp.h>
7
#include <dgl/runtime/parallel_for.h>
8
9
10
#include <vector>
#include <unordered_set>
#include <unordered_map>
11
#include <tuple>
12
#include <numeric>
13
14
15
16
17
#include "array_utils.h"

namespace dgl {

using runtime::NDArray;
18
using runtime::parallel_for;
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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

namespace aten {
namespace impl {

/*
 * TODO(BarclayII):
 * For row-major sorted COOs, we have faster implementation with binary search,
 * sorted search, etc.  Later we should benchmark how much we can gain with
 * sorted COOs on hypersparse graphs.
 */

///////////////////////////// COOIsNonZero /////////////////////////////

template <DLDeviceType XPU, typename IdType>
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col) {
  CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;
  CHECK(col >= 0 && col < coo.num_cols) << "Invalid col index: " << col;
  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    if (coo_row_data[i] == row && coo_col_data[i] == col)
      return true;
  }
  return false;
}

template bool COOIsNonZero<kDLCPU, int32_t>(COOMatrix, int64_t, int64_t);
template bool COOIsNonZero<kDLCPU, int64_t>(COOMatrix, int64_t, int64_t);

template <DLDeviceType XPU, typename IdType>
NDArray COOIsNonZero(COOMatrix coo, 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;
  const int64_t kmax = std::max(rowlen, collen);
60
61
62
63
64
65
66
  parallel_for(0, kmax, [=](size_t b, size_t e) {
    for (auto k = b; k < e; ++k) {
      int64_t i = row_stride * k;
      int64_t j = col_stride * k;
      rst_data[k] = COOIsNonZero<XPU, IdType>(coo, row_data[i], col_data[j])? 1 : 0;
    }
  });
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
104
105
106
107
108
109
110
111
112
113
  return rst;
}

template NDArray COOIsNonZero<kDLCPU, int32_t>(COOMatrix, NDArray, NDArray);
template NDArray COOIsNonZero<kDLCPU, int64_t>(COOMatrix, NDArray, NDArray);

///////////////////////////// COOHasDuplicate /////////////////////////////

template <DLDeviceType XPU, typename IdType>
bool COOHasDuplicate(COOMatrix coo) {
  std::unordered_set<std::pair<IdType, IdType>, PairHash> hashmap;
  const IdType* src_data = static_cast<IdType*>(coo.row->data);
  const IdType* dst_data = static_cast<IdType*>(coo.col->data);
  const auto nnz = coo.row->shape[0];
  for (IdType eid = 0; eid < nnz; ++eid) {
    const auto& p = std::make_pair(src_data[eid], dst_data[eid]);
    if (hashmap.count(p)) {
      return true;
    } else {
      hashmap.insert(p);
    }
  }
  return false;
}

template bool COOHasDuplicate<kDLCPU, int32_t>(COOMatrix coo);
template bool COOHasDuplicate<kDLCPU, int64_t>(COOMatrix coo);

///////////////////////////// COOGetRowNNZ /////////////////////////////

template <DLDeviceType XPU, typename IdType>
int64_t COOGetRowNNZ(COOMatrix coo, int64_t row) {
  CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;
  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  int64_t result = 0;
  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    if (coo_row_data[i] == row)
      ++result;
  }
  return result;
}

template int64_t COOGetRowNNZ<kDLCPU, int32_t>(COOMatrix, int64_t);
template int64_t COOGetRowNNZ<kDLCPU, int64_t>(COOMatrix, int64_t);

template <DLDeviceType XPU, typename IdType>
NDArray COOGetRowNNZ(COOMatrix coo, NDArray rows) {
114
  CHECK_SAME_DTYPE(coo.col, rows);
115
116
117
118
119
  const auto len = rows->shape[0];
  const IdType* vid_data = static_cast<IdType*>(rows->data);
  NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
  IdType* rst_data = static_cast<IdType*>(rst->data);
#pragma omp parallel for
120
  for (int64_t i = 0; i < len; ++i) {
121
    rst_data[i] = COOGetRowNNZ<XPU, IdType>(coo, vid_data[i]);
122
  }
123
124
125
126
127
128
129
130
  return rst;
}

template NDArray COOGetRowNNZ<kDLCPU, int32_t>(COOMatrix, NDArray);
template NDArray COOGetRowNNZ<kDLCPU, int64_t>(COOMatrix, NDArray);

///////////////////////////// COOGetRowDataAndIndices /////////////////////////////

131
template <DLDeviceType XPU, typename IdType>
132
133
134
135
136
137
std::pair<NDArray, NDArray> COOGetRowDataAndIndices(
    COOMatrix coo, int64_t row) {
  CHECK(row >= 0 && row < coo.num_rows) << "Invalid row index: " << row;

  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
138
  const IdType* coo_data = COOHasData(coo) ? static_cast<IdType*>(coo.data->data) : nullptr;
139
140

  std::vector<IdType> indices;
141
  std::vector<IdType> data;
142
143
144
145
146
147
148
149
150
151
152
153

  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    if (coo_row_data[i] == row) {
      indices.push_back(coo_col_data[i]);
      data.push_back(coo_data ? coo_data[i] : i);
    }
  }

  return std::make_pair(NDArray::FromVector(data), NDArray::FromVector(indices));
}

template std::pair<NDArray, NDArray>
154
COOGetRowDataAndIndices<kDLCPU, int32_t>(COOMatrix, int64_t);
155
template std::pair<NDArray, NDArray>
156
COOGetRowDataAndIndices<kDLCPU, int64_t>(COOMatrix, int64_t);
157
158
159

///////////////////////////// COOGetData /////////////////////////////

160
template <DLDeviceType XPU, typename IdType>
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
IdArray COOGetData(COOMatrix coo, IdArray rows, IdArray cols) {
  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:" << rows << " " << cols;
  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 = rows.Ptr<IdType>();
  const IdType* col_data = cols.Ptr<IdType>();

  const IdType* coo_row = coo.row.Ptr<IdType>();
  const IdType* coo_col = coo.col.Ptr<IdType>();
  const IdType* data = COOHasData(coo) ? coo.data.Ptr<IdType>() : nullptr;
  const int64_t nnz = coo.row->shape[0];

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

  // TODO(minjie): We might need to consider sorting the COO beforehand especially
  //   when the number of (row, col) pairs is large. Need more benchmarks to justify
  //   the choice.

  if (coo.row_sorted) {
185
186
187
188
189
190
191
192
193
194
    parallel_for(0, retlen, [&](size_t b, size_t e) {
      for (auto p = b; p < e; ++p) {
        const IdType row_id = row_data[p * row_stride], col_id = col_data[p * col_stride];
        auto it = std::lower_bound(coo_row, coo_row + nnz, row_id);
        for (; it < coo_row + nnz && *it == row_id; ++it) {
          const auto idx = it - coo_row;
          if (coo_col[idx] == col_id) {
            ret_data[p] = data? data[idx] : idx;
            break;
          }
195
196
        }
      }
197
    });
198
199
200
201
202
203
204
205
206
207
208
  } else {
#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];
      for (int64_t idx = 0; idx < nnz; ++idx) {
        if (coo_row[idx] == row_id && coo_col[idx] == col_id) {
          ret_data[p] = data? data[idx] : idx;
          break;
        }
      }
    }
209
  }
210
211

  return ret;
212
213
}

214
215
template IdArray COOGetData<kDLCPU, int32_t>(COOMatrix, IdArray, IdArray);
template IdArray COOGetData<kDLCPU, int64_t>(COOMatrix, IdArray, IdArray);
216
217
218

///////////////////////////// COOGetDataAndIndices /////////////////////////////

219
template <DLDeviceType XPU, typename IdType>
220
221
222
223
std::vector<NDArray> COOGetDataAndIndices(COOMatrix coo, NDArray rows,
                                          NDArray cols) {
  CHECK_SAME_DTYPE(coo.col, rows);
  CHECK_SAME_DTYPE(coo.col, cols);
224
225
  const int64_t rowlen = rows->shape[0];
  const int64_t collen = cols->shape[0];
226
  const int64_t len = std::max(rowlen, collen);
227
228
229
230
231
232
233
234
235
236
237

  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* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
238
  const IdType* data = COOHasData(coo) ? static_cast<IdType*>(coo.data->data) : nullptr;
239
240

  std::vector<IdType> ret_rows, ret_cols;
241
  std::vector<IdType> ret_data;
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
  ret_rows.reserve(len);
  ret_cols.reserve(len);
  ret_data.reserve(len);

  // NOTE(BarclayII): With a small number of lookups, linear scan is faster.
  // The threshold 200 comes from benchmarking both algorithms on a P3.8x instance.
  // I also tried sorting plus binary search.  The speed gain is only significant for
  // medium-sized graphs and lookups, so I didn't include it.
  if (len >= 200) {
    // TODO(BarclayII) Ideally we would want to cache this object.  However I'm not sure
    // what is the best way to do so since this object is valid for CPU only.
    std::unordered_multimap<std::pair<IdType, IdType>, IdType, PairHash> pair_map;
    pair_map.reserve(coo.row->shape[0]);
    for (int64_t k = 0; k < coo.row->shape[0]; ++k)
      pair_map.emplace(std::make_pair(coo_row_data[k], coo_col_data[k]), data ? data[k]: k);

    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 < coo.num_rows) << "Invalid row index: " << row_id;
      CHECK(col_id >= 0 && col_id < coo.num_cols) << "Invalid col index: " << col_id;
      auto range = pair_map.equal_range({row_id, col_id});
      for (auto it = range.first; it != range.second; ++it) {
264
265
        ret_rows.push_back(row_id);
        ret_cols.push_back(col_id);
266
267
268
269
270
271
272
273
274
275
276
277
278
279
        ret_data.push_back(it->second);
      }
    }
  } else {
    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 < coo.num_rows) << "Invalid row index: " << row_id;
      CHECK(col_id >= 0 && col_id < coo.num_cols) << "Invalid col index: " << col_id;
      for (int64_t k = 0; k < coo.row->shape[0]; ++k) {
        if (coo_row_data[k] == row_id && coo_col_data[k] == col_id) {
          ret_rows.push_back(row_id);
          ret_cols.push_back(col_id);
          ret_data.push_back(data ? data[k] : k);
        }
280
281
282
283
284
285
286
287
288
      }
    }
  }

  return {NDArray::FromVector(ret_rows),
          NDArray::FromVector(ret_cols),
          NDArray::FromVector(ret_data)};
}

289
template std::vector<NDArray> COOGetDataAndIndices<kDLCPU, int32_t>(
290
    COOMatrix coo, NDArray rows, NDArray cols);
291
template std::vector<NDArray> COOGetDataAndIndices<kDLCPU, int64_t>(
292
293
294
295
    COOMatrix coo, NDArray rows, NDArray cols);

///////////////////////////// COOTranspose /////////////////////////////

296
template <DLDeviceType XPU, typename IdType>
297
298
299
300
COOMatrix COOTranspose(COOMatrix coo) {
  return COOMatrix{coo.num_cols, coo.num_rows, coo.col, coo.row, coo.data};
}

301
302
template COOMatrix COOTranspose<kDLCPU, int32_t>(COOMatrix coo);
template COOMatrix COOTranspose<kDLCPU, int64_t>(COOMatrix coo);
303
304
305

///////////////////////////// COOToCSR /////////////////////////////

306
307
308
// complexity: time O(NNZ), space O(1) if the coo is row sorted,
// time O(NNZ/p + N), space O(NNZ + N*p) otherwise, where p is the number of
// threads.
309
template <DLDeviceType XPU, typename IdType>
310
311
312
CSRMatrix COOToCSR(COOMatrix coo) {
  const int64_t N = coo.num_rows;
  const int64_t NNZ = coo.row->shape[0];
313
314
315
  const IdType* const row_data = static_cast<IdType*>(coo.row->data);
  const IdType* const col_data = static_cast<IdType*>(coo.col->data);
  const IdType* const data = COOHasData(coo)? static_cast<IdType*>(coo.data->data) : nullptr;
316

317
  NDArray ret_indptr = NDArray::Empty({N + 1}, coo.row->dtype, coo.row->ctx);
318
319
  NDArray ret_indices;
  NDArray ret_data;
320

321
322
  const bool row_sorted = coo.row_sorted;
  const bool col_sorted = coo.col_sorted;
323

324
325
  if (row_sorted) {
    // compute indptr
326
    IdType* const Bp = static_cast<IdType*>(ret_indptr->data);
327
328
    Bp[0] = 0;

329
330
331
332
333
334
335
    if (!data) {
      // Leave empty, and populate from inside of parallel block
      coo.data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
    }
    IdType * const fill_data = data ? nullptr : static_cast<IdType*>(coo.data->data);

    if (NNZ > 0) {
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
      auto num_threads = omp_get_max_threads();
      parallel_for(0, num_threads, [&](int b, int e) {
        for (auto thread_id = b; thread_id < e; ++thread_id) {
          // We partition the set the of non-zeros among the threads
          const int64_t nz_chunk = (NNZ+num_threads-1)/num_threads;
          const int64_t nz_start = thread_id*nz_chunk;
          const int64_t nz_end = std::min(NNZ, nz_start+nz_chunk);

          // Each thread searchs the row array for a change, and marks it's
          // location in Bp. Threads, other than the first, start at the last
          // index covered by the previous, in order to detect changes in the row
          // array between thread partitions. This means that each thread after
          // the first, searches the range [nz_start-1, nz_end). That is,
          // if we had 10 non-zeros, and 4 threads, the indexes searched by each
          // thread would be:
          // 0: [0, 1, 2]
          // 1: [2, 3, 4, 5]
          // 2: [5, 6, 7, 8]
          // 3: [8, 9]
          //
          // That way, if the row array were [0, 0, 1, 2, 2, 2, 4, 5, 5, 6], each
          // change in row would be captured by one thread:
          //
          // 0: [0, 0, 1] - row 0
          // 1: [1, 2, 2, 2] - row 1
          // 2: [2, 4, 5, 5] - rows 2, 3, and 4
          // 3: [5, 6] - rows 5 and 6
          //
          int64_t row = 0;
          if (nz_start < nz_end) {
            row = nz_start == 0 ? 0 : row_data[nz_start-1];
            for (int64_t i = nz_start; i < nz_end; ++i) {
              while (row != row_data[i]) {
                ++row;
                Bp[row] = i;
              }
372
373
            }

374
375
376
377
378
379
380
381
            // We will not detect the row change for the last row, nor any empty
            // rows at the end of the matrix, so the last active thread needs
            // mark all remaining rows in Bp with NNZ.
            if (nz_end == NNZ) {
              while (row < N) {
                ++row;
                Bp[row] = NNZ;
              }
382
383
            }

384
385
386
387
388
389
390
391
392
            if (fill_data) {
              // TODO(minjie): Many of our current implementation assumes that CSR must have
              //   a data array. This is a temporary workaround. Remove this after:
              //   - The old immutable graph implementation is deprecated.
              //   - The old binary reduce kernel is deprecated.
              std::iota(fill_data+nz_start,
                        fill_data+nz_end,
                        nz_start);
            }
393
394
          }
        }
395
      });
396
397
398
    } else {
      std::fill(Bp, Bp+N+1, 0);
    }
399

400
    // compute indices and data
401
402
403
    ret_indices = coo.col;
    ret_data = coo.data;
  } else {
404
    // compute indptr
405
406
    IdType* const Bp = static_cast<IdType*>(ret_indptr->data);
    Bp[0] = 0;
407
408

    // compute indices and data
409
410
    ret_indices = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
    ret_data = NDArray::Empty({NNZ}, coo.row->dtype, coo.row->ctx);
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
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
480
    IdType* const Bi = static_cast<IdType*>(ret_indices->data);
    IdType* const Bx = static_cast<IdType*>(ret_data->data);

    // the offset within each row, that each thread will write to
    std::vector<std::vector<IdType>> local_ptrs;
    std::vector<int64_t> thread_prefixsum;

#pragma omp parallel
    {
      const int num_threads = omp_get_num_threads();
      const int thread_id = omp_get_thread_num();
      CHECK_LT(thread_id, num_threads);

      const int64_t nz_chunk = (NNZ+num_threads-1)/num_threads;
      const int64_t nz_start = thread_id*nz_chunk;
      const int64_t nz_end = std::min(NNZ, nz_start+nz_chunk);

      const int64_t n_chunk = (N+num_threads-1)/num_threads;
      const int64_t n_start = thread_id*n_chunk;
      const int64_t n_end = std::min(N, n_start+n_chunk);

#pragma omp master
      {
        local_ptrs.resize(num_threads);
        thread_prefixsum.resize(num_threads+1);
      }

#pragma omp barrier
      local_ptrs[thread_id].resize(N, 0);

      for (int64_t i = nz_start; i < nz_end; ++i) {
        ++local_ptrs[thread_id][row_data[i]];
      }

#pragma omp barrier
      // compute prefixsum in parallel
      int64_t sum = 0;
      for (int64_t i = n_start; i < n_end; ++i) {
        IdType tmp = 0;
        for (int j = 0; j < num_threads; ++j) {
          std::swap(tmp, local_ptrs[j][i]);
          tmp += local_ptrs[j][i];
        }
        sum += tmp;
        Bp[i+1] = sum;
      }
      thread_prefixsum[thread_id+1] = sum;

#pragma omp barrier
#pragma omp master
      {
        for (int64_t i = 0; i < num_threads; ++i) {
          thread_prefixsum[i+1] += thread_prefixsum[i];
        }
        CHECK_EQ(thread_prefixsum[num_threads], NNZ);
      }
#pragma omp barrier

      sum = thread_prefixsum[thread_id];
      for (int64_t i = n_start; i < n_end; ++i) {
        Bp[i+1] += sum;
      }

#pragma omp barrier
      for (int64_t i = nz_start; i < nz_end; ++i) {
        const IdType r = row_data[i];
        const int64_t index = Bp[r] + local_ptrs[thread_id][r]++;
        Bi[index] = col_data[i];
        Bx[index] = data ? data[i] : i;
      }
481
    }
482
    CHECK_EQ(Bp[N], NNZ);
483
484
  }

485
486
  return CSRMatrix(coo.num_rows, coo.num_cols,
                   ret_indptr, ret_indices, ret_data,
487
                   col_sorted);
488
489
}

490
491
template CSRMatrix COOToCSR<kDLCPU, int32_t>(COOMatrix coo);
template CSRMatrix COOToCSR<kDLCPU, int64_t>(COOMatrix coo);
492
493
494

///////////////////////////// COOSliceRows /////////////////////////////

495
template <DLDeviceType XPU, typename IdType>
496
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end) {
497
  // TODO(minjie): use binary search when coo.row_sorted is true
498
499
500
501
502
  CHECK(start >= 0 && start < coo.num_rows) << "Invalid start row " << start;
  CHECK(end > 0 && end <= coo.num_rows) << "Invalid end row " << end;

  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
503
  const IdType* coo_data = COOHasData(coo) ? static_cast<IdType*>(coo.data->data) : nullptr;
504
505

  std::vector<IdType> ret_row, ret_col;
506
  std::vector<IdType> ret_data;
507
508
509
510
511
512
513
514
515
516

  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    const IdType row_id = coo_row_data[i];
    const IdType col_id = coo_col_data[i];
    if (row_id < end && row_id >= start) {
      ret_row.push_back(row_id - start);
      ret_col.push_back(col_id);
      ret_data.push_back(coo_data ? coo_data[i] : i);
    }
  }
517
  return COOMatrix(
518
519
520
521
    end - start,
    coo.num_cols,
    NDArray::FromVector(ret_row),
    NDArray::FromVector(ret_col),
522
523
524
    NDArray::FromVector(ret_data),
    coo.row_sorted,
    coo.col_sorted);
525
526
}

527
528
template COOMatrix COOSliceRows<kDLCPU, int32_t>(COOMatrix, int64_t, int64_t);
template COOMatrix COOSliceRows<kDLCPU, int64_t>(COOMatrix, int64_t, int64_t);
529

530
template <DLDeviceType XPU, typename IdType>
531
532
533
COOMatrix COOSliceRows(COOMatrix coo, NDArray rows) {
  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
534
  const IdType* coo_data = COOHasData(coo) ? static_cast<IdType*>(coo.data->data) : nullptr;
535
536

  std::vector<IdType> ret_row, ret_col;
537
  std::vector<IdType> ret_data;
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556

  IdHashMap<IdType> hashmap(rows);

  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    const IdType row_id = coo_row_data[i];
    const IdType col_id = coo_col_data[i];
    const IdType mapped_row_id = hashmap.Map(row_id, -1);
    if (mapped_row_id != -1) {
      ret_row.push_back(mapped_row_id);
      ret_col.push_back(col_id);
      ret_data.push_back(coo_data ? coo_data[i] : i);
    }
  }

  return COOMatrix{
    rows->shape[0],
    coo.num_cols,
    NDArray::FromVector(ret_row),
    NDArray::FromVector(ret_col),
557
558
    NDArray::FromVector(ret_data),
    coo.row_sorted, coo.col_sorted};
559
560
}

561
562
template COOMatrix COOSliceRows<kDLCPU, int32_t>(COOMatrix , NDArray);
template COOMatrix COOSliceRows<kDLCPU, int64_t>(COOMatrix , NDArray);
563
564
565

///////////////////////////// COOSliceMatrix /////////////////////////////

566
template <DLDeviceType XPU, typename IdType>
567
568
569
COOMatrix COOSliceMatrix(COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols) {
  const IdType* coo_row_data = static_cast<IdType*>(coo.row->data);
  const IdType* coo_col_data = static_cast<IdType*>(coo.col->data);
570
  const IdType* coo_data = COOHasData(coo) ? static_cast<IdType*>(coo.data->data) : nullptr;
571
572
573
574

  IdHashMap<IdType> row_map(rows), col_map(cols);

  std::vector<IdType> ret_row, ret_col;
575
  std::vector<IdType> ret_data;
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590

  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
    const IdType row_id = coo_row_data[i];
    const IdType col_id = coo_col_data[i];
    const IdType mapped_row_id = row_map.Map(row_id, -1);
    if (mapped_row_id != -1) {
      const IdType mapped_col_id = col_map.Map(col_id, -1);
      if (mapped_col_id != -1) {
        ret_row.push_back(mapped_row_id);
        ret_col.push_back(mapped_col_id);
        ret_data.push_back(coo_data ? coo_data[i] : i);
      }
    }
  }

591
592
593
594
595
  return COOMatrix(rows->shape[0], cols->shape[0],
                   NDArray::FromVector(ret_row),
                   NDArray::FromVector(ret_col),
                   NDArray::FromVector(ret_data),
                   coo.row_sorted, coo.col_sorted);
596
597
}

598
template COOMatrix COOSliceMatrix<kDLCPU, int32_t>(
599
    COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
600
template COOMatrix COOSliceMatrix<kDLCPU, int64_t>(
601
602
    COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);

Da Zheng's avatar
Da Zheng committed
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633

///////////////////////////// COOReorder /////////////////////////////

template <DLDeviceType XPU, typename IdType>
COOMatrix COOReorder(COOMatrix coo, runtime::NDArray new_row_id_arr,
                     runtime::NDArray new_col_id_arr) {
  CHECK_SAME_DTYPE(coo.row, new_row_id_arr);
  CHECK_SAME_DTYPE(coo.col, new_col_id_arr);

  // Input COO
  const IdType* in_rows = static_cast<IdType*>(coo.row->data);
  const IdType* in_cols = static_cast<IdType*>(coo.col->data);
  int64_t num_rows = coo.num_rows;
  int64_t num_cols = coo.num_cols;
  int64_t nnz = coo.row->shape[0];
  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 COO";
  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 COO";

  // 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 COO
  NDArray out_row_arr = NDArray::Empty({nnz}, coo.row->dtype, coo.row->ctx);
  NDArray out_col_arr = NDArray::Empty({nnz}, coo.col->dtype, coo.col->ctx);
  NDArray out_data_arr = COOHasData(coo) ? coo.data : NullArray();
  IdType *out_row = static_cast<IdType*>(out_row_arr->data);
  IdType *out_col = static_cast<IdType*>(out_col_arr->data);

634
635
636
637
638
639
  parallel_for(0, nnz, [=](size_t b, size_t e) {
    for (auto i = b; i < e; ++i) {
      out_row[i] = new_row_ids[in_rows[i]];
      out_col[i] = new_col_ids[in_cols[i]];
    }
  });
Da Zheng's avatar
Da Zheng committed
640
641
642
643
644
645
646
647
  return COOMatrix(num_rows, num_cols, out_row_arr, out_col_arr, out_data_arr);
}

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

648
649
650
}  // namespace impl
}  // namespace aten
}  // namespace dgl