spmat_op_impl_csr.hip 24.8 KB
Newer Older
sangwzh's avatar
sangwzh committed
1
2
// !!! This is a file automatically generated by hipify!!!
#include "hip/hip_runtime.h"
3
/**
4
 *  Copyright (c) 2020 by Contributors
5
6
 * @file array/cuda/spmat_op_impl_csr.cu
 * @brief CSR operator CPU implementation
7
8
 */
#include <dgl/array.h>
9
10
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
11

sangwzh's avatar
sangwzh committed
12
#include <hipcub/hipcub.hpp>
13
#include <numeric>
14
15
16
#include <unordered_set>
#include <vector>

17
#include "../../runtime/cuda/cuda_common.h"
sangwzh's avatar
sangwzh committed
18
19
#include "atomic.cuh"
#include "utils.h"
20
21
22
23

namespace dgl {

using runtime::NDArray;
24
using namespace cuda;
25
26
27
28
29
30

namespace aten {
namespace impl {

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

31
template <DGLDeviceType XPU, typename IdType>
32
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col) {
sangwzh's avatar
sangwzh committed
33
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
34
35
36
37
38
39
40
41
  const auto& ctx = csr.indptr->ctx;
  IdArray rows = aten::VecToIdArray<int64_t>({row}, sizeof(IdType) * 8, ctx);
  IdArray cols = aten::VecToIdArray<int64_t>({col}, sizeof(IdType) * 8, ctx);
  rows = rows.CopyTo(ctx);
  cols = cols.CopyTo(ctx);
  IdArray out = aten::NewIdArray(1, ctx, sizeof(IdType) * 8);
  const IdType* data = nullptr;
  // TODO(minjie): use binary search for sorted csr
42
43
44
45
46
  CUDA_KERNEL_CALL(
      dgl::cuda::_LinearSearchKernel, 1, 1, 0, stream, csr.indptr.Ptr<IdType>(),
      csr.indices.Ptr<IdType>(), data, rows.Ptr<IdType>(), cols.Ptr<IdType>(),
      1, 1, 1, static_cast<IdType*>(nullptr), static_cast<IdType>(-1),
      out.Ptr<IdType>());
47
  out = out.CopyTo(DGLContext{kDGLCPU, 0});
48
49
50
  return *out.Ptr<IdType>() != -1;
}

51
52
template bool CSRIsNonZero<kDGLCUDA, int32_t>(CSRMatrix, int64_t, int64_t);
template bool CSRIsNonZero<kDGLCUDA, int64_t>(CSRMatrix, int64_t, int64_t);
53

54
template <DGLDeviceType XPU, typename IdType>
55
56
57
NDArray CSRIsNonZero(CSRMatrix csr, NDArray row, NDArray col) {
  const auto rowlen = row->shape[0];
  const auto collen = col->shape[0];
sangwzh's avatar
sangwzh committed
58
  const auto rstlen = ::max(rowlen, collen);
59
  NDArray rst = NDArray::Empty({rstlen}, row->dtype, row->ctx);
60
  if (rstlen == 0) return rst;
61
62
  const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
  const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
sangwzh's avatar
sangwzh committed
63
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
64
  const int nt = dgl::cuda::FindNumThreads(rstlen);
65
66
  const int nb = (rstlen + nt - 1) / nt;
  const IdType* data = nullptr;
67
68
69
70
  const IdType* indptr_data =
      static_cast<IdType*>(GetDevicePointer(csr.indptr));
  const IdType* indices_data =
      static_cast<IdType*>(GetDevicePointer(csr.indices));
71
  // TODO(minjie): use binary search for sorted csr
72
73
74
75
76
  CUDA_KERNEL_CALL(
      dgl::cuda::_LinearSearchKernel, nb, nt, 0, stream, indptr_data,
      indices_data, data, row.Ptr<IdType>(), col.Ptr<IdType>(), row_stride,
      col_stride, rstlen, static_cast<IdType*>(nullptr),
      static_cast<IdType>(-1), rst.Ptr<IdType>());
77
78
79
  return rst != -1;
}

80
81
template NDArray CSRIsNonZero<kDGLCUDA, int32_t>(CSRMatrix, NDArray, NDArray);
template NDArray CSRIsNonZero<kDGLCUDA, int64_t>(CSRMatrix, NDArray, NDArray);
82
83
84

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

85
/**
86
 * @brief Check whether each row does not have any duplicate entries.
87
88
89
90
 * Assume the CSR is sorted.
 */
template <typename IdType>
__global__ void _SegmentHasNoDuplicate(
91
92
    const IdType* indptr, const IdType* indices, int64_t num_rows,
    int8_t* flags) {
93
94
95
96
97
98
99
100
101
102
103
104
  int tx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride_x = gridDim.x * blockDim.x;
  while (tx < num_rows) {
    bool f = true;
    for (IdType i = indptr[tx] + 1; f && i < indptr[tx + 1]; ++i) {
      f = (indices[i - 1] != indices[i]);
    }
    flags[tx] = static_cast<int8_t>(f);
    tx += stride_x;
  }
}

105
template <DGLDeviceType XPU, typename IdType>
106
bool CSRHasDuplicate(CSRMatrix csr) {
107
  if (!csr.sorted) csr = CSRSort(csr);
108
  const auto& ctx = csr.indptr->ctx;
sangwzh's avatar
sangwzh committed
109
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
110
  auto device = runtime::DeviceAPI::Get(ctx);
111
112
113
114
  // We allocate a workspace of num_rows bytes. It wastes a little bit memory
  // but should be fine.
  int8_t* flags =
      static_cast<int8_t*>(device->AllocWorkspace(ctx, csr.num_rows));
115
  const int nt = dgl::cuda::FindNumThreads(csr.num_rows);
116
  const int nb = (csr.num_rows + nt - 1) / nt;
117
118
119
  CUDA_KERNEL_CALL(
      _SegmentHasNoDuplicate, nb, nt, 0, stream, csr.indptr.Ptr<IdType>(),
      csr.indices.Ptr<IdType>(), csr.num_rows, flags);
120
  bool ret = dgl::cuda::AllTrue(flags, csr.num_rows, ctx);
121
122
123
124
  device->FreeWorkspace(ctx, flags);
  return !ret;
}

125
126
template bool CSRHasDuplicate<kDGLCUDA, int32_t>(CSRMatrix csr);
template bool CSRHasDuplicate<kDGLCUDA, int64_t>(CSRMatrix csr);
127
128
129

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

130
template <DGLDeviceType XPU, typename IdType>
131
132
133
134
135
136
int64_t CSRGetRowNNZ(CSRMatrix csr, int64_t row) {
  const IdType cur = aten::IndexSelect<IdType>(csr.indptr, row);
  const IdType next = aten::IndexSelect<IdType>(csr.indptr, row + 1);
  return next - cur;
}

137
138
template int64_t CSRGetRowNNZ<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
template int64_t CSRGetRowNNZ<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
139
140
141

template <typename IdType>
__global__ void _CSRGetRowNNZKernel(
142
    const IdType* vid, const IdType* indptr, IdType* out, int64_t length) {
143
144
145
146
147
148
149
150
151
  int tx = blockIdx.x * blockDim.x + threadIdx.x;
  int stride_x = gridDim.x * blockDim.x;
  while (tx < length) {
    const IdType vv = vid[tx];
    out[tx] = indptr[vv + 1] - indptr[vv];
    tx += stride_x;
  }
}

152
template <DGLDeviceType XPU, typename IdType>
153
NDArray CSRGetRowNNZ(CSRMatrix csr, NDArray rows) {
sangwzh's avatar
sangwzh committed
154
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
155
  const auto len = rows->shape[0];
156
  const IdType* vid_data = rows.Ptr<IdType>();
157
158
  const IdType* indptr_data =
      static_cast<IdType*>(GetDevicePointer(csr.indptr));
159
160
  NDArray rst = NDArray::Empty({len}, rows->dtype, rows->ctx);
  IdType* rst_data = static_cast<IdType*>(rst->data);
161
  const int nt = dgl::cuda::FindNumThreads(len);
162
  const int nb = (len + nt - 1) / nt;
163
164
165
  CUDA_KERNEL_CALL(
      _CSRGetRowNNZKernel, nb, nt, 0, stream, vid_data, indptr_data, rst_data,
      len);
166
167
168
  return rst;
}

169
170
template NDArray CSRGetRowNNZ<kDGLCUDA, int32_t>(CSRMatrix, NDArray);
template NDArray CSRGetRowNNZ<kDGLCUDA, int64_t>(CSRMatrix, NDArray);
171

172
////////////////////////// CSRGetRowColumnIndices //////////////////////////////
173

174
template <DGLDeviceType XPU, typename IdType>
175
176
NDArray CSRGetRowColumnIndices(CSRMatrix csr, int64_t row) {
  const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
177
178
  const int64_t offset =
      aten::IndexSelect<IdType>(csr.indptr, row) * sizeof(IdType);
179
180
181
  return csr.indices.CreateView({len}, csr.indices->dtype, offset);
}

182
183
template NDArray CSRGetRowColumnIndices<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
template NDArray CSRGetRowColumnIndices<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
184
185
186

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

187
template <DGLDeviceType XPU, typename IdType>
188
189
NDArray CSRGetRowData(CSRMatrix csr, int64_t row) {
  const int64_t len = impl::CSRGetRowNNZ<XPU, IdType>(csr, row);
190
191
  const int64_t offset =
      aten::IndexSelect<IdType>(csr.indptr, row) * sizeof(IdType);
192
193
194
  if (aten::CSRHasData(csr))
    return csr.data.CreateView({len}, csr.data->dtype, offset);
  else
195
196
    return aten::Range(
        offset, offset + len, csr.indptr->dtype.bits, csr.indptr->ctx);
197
198
}

199
200
template NDArray CSRGetRowData<kDGLCUDA, int32_t>(CSRMatrix, int64_t);
template NDArray CSRGetRowData<kDGLCUDA, int64_t>(CSRMatrix, int64_t);
201
202
203

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

204
template <DGLDeviceType XPU, typename IdType>
205
206
207
208
209
210
211
212
213
214
215
CSRMatrix CSRSliceRows(CSRMatrix csr, int64_t start, int64_t end) {
  const int64_t num_rows = end - start;
  const IdType st_pos = aten::IndexSelect<IdType>(csr.indptr, start);
  const IdType ed_pos = aten::IndexSelect<IdType>(csr.indptr, end);
  const IdType nnz = ed_pos - st_pos;
  IdArray ret_indptr = aten::IndexSelect(csr.indptr, start, end + 1) - st_pos;
  // indices and data can be view arrays
  IdArray ret_indices = csr.indices.CreateView(
      {nnz}, csr.indices->dtype, st_pos * sizeof(IdType));
  IdArray ret_data;
  if (CSRHasData(csr))
216
217
    ret_data =
        csr.data.CreateView({nnz}, csr.data->dtype, st_pos * sizeof(IdType));
218
  else
219
220
221
222
    ret_data =
        aten::Range(st_pos, ed_pos, csr.indptr->dtype.bits, csr.indptr->ctx);
  return CSRMatrix(
      num_rows, csr.num_cols, ret_indptr, ret_indices, ret_data, csr.sorted);
223
224
}

225
226
template CSRMatrix CSRSliceRows<kDGLCUDA, int32_t>(CSRMatrix, int64_t, int64_t);
template CSRMatrix CSRSliceRows<kDGLCUDA, int64_t>(CSRMatrix, int64_t, int64_t);
227

228
/**
229
 * @brief Copy data segment to output buffers
230
 *
231
232
233
 * For the i^th row r = row[i], copy the data from indptr[r] ~ indptr[r+1]
 * to the out_data from out_indptr[i] ~ out_indptr[i+1]
 *
234
235
 * If the provided `data` array is nullptr, write the read index to the
 * out_data.
236
237
238
239
 *
 */
template <typename IdType, typename DType>
__global__ void _SegmentCopyKernel(
240
241
    const IdType* indptr, const DType* data, const IdType* row, int64_t length,
    int64_t n_row, const IdType* out_indptr, DType* out_data) {
242
  IdType tx = static_cast<IdType>(blockIdx.x) * blockDim.x + threadIdx.x;
243
244
  const int stride_x = gridDim.x * blockDim.x;
  while (tx < length) {
245
    IdType rpos = dgl::cuda::_UpperBound(out_indptr, n_row, tx) - 1;
246
247
    IdType rofs = tx - out_indptr[rpos];
    const IdType u = row[rpos];
248
    out_data[tx] = data ? data[indptr[u] + rofs] : indptr[u] + rofs;
249
250
251
252
    tx += stride_x;
  }
}

253
template <DGLDeviceType XPU, typename IdType>
254
CSRMatrix CSRSliceRows(CSRMatrix csr, NDArray rows) {
sangwzh's avatar
sangwzh committed
255
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
256
257
258
259
  const int64_t len = rows->shape[0];
  IdArray ret_indptr = aten::CumSum(aten::CSRGetRowNNZ(csr, rows), true);
  const int64_t nnz = aten::IndexSelect<IdType>(ret_indptr, len);

260
261
262
  const int nt = 256;  // for better GPU usage of small invocations
  const int nb = (nnz + nt - 1) / nt;

263
  // Copy indices.
264
  IdArray ret_indices = NDArray::Empty({nnz}, csr.indptr->dtype, rows->ctx);
265

266
267
268
269
270
271
272
  const IdType* indptr_data =
      static_cast<IdType*>(GetDevicePointer(csr.indptr));
  const IdType* indices_data =
      static_cast<IdType*>(GetDevicePointer(csr.indices));
  const IdType* data_data =
      CSRHasData(csr) ? static_cast<IdType*>(GetDevicePointer(csr.data))
                      : nullptr;
273

274
275
276
277
  CUDA_KERNEL_CALL(
      _SegmentCopyKernel, nb, nt, 0, stream, indptr_data, indices_data,
      rows.Ptr<IdType>(), nnz, len, ret_indptr.Ptr<IdType>(),
      ret_indices.Ptr<IdType>());
278
  // Copy data.
279
  IdArray ret_data = NDArray::Empty({nnz}, csr.indptr->dtype, rows->ctx);
280
281
282
283
284
285
  CUDA_KERNEL_CALL(
      _SegmentCopyKernel, nb, nt, 0, stream, indptr_data, data_data,
      rows.Ptr<IdType>(), nnz, len, ret_indptr.Ptr<IdType>(),
      ret_data.Ptr<IdType>());
  return CSRMatrix(
      len, csr.num_cols, ret_indptr, ret_indices, ret_data, csr.sorted);
286
287
}

288
289
template CSRMatrix CSRSliceRows<kDGLCUDA, int32_t>(CSRMatrix, NDArray);
template CSRMatrix CSRSliceRows<kDGLCUDA, int64_t>(CSRMatrix, NDArray);
290
291
292

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

293
/**
294
 * @brief Generate a 0-1 mask for each index that hits the provided (row, col)
295
 *        index.
296
 *
297
298
299
300
301
302
303
304
305
306
307
 * Examples:
 * Given a CSR matrix (with duplicate entries) as follows:
 * [[0, 1, 2, 0, 0],
 *  [1, 0, 0, 0, 0],
 *  [0, 0, 1, 1, 0],
 *  [0, 0, 0, 0, 0]]
 * Given rows: [0, 1], cols: [0, 2, 3]
 * The result mask is: [0, 1, 1, 1, 0, 0]
 */
template <typename IdType>
__global__ void _SegmentMaskKernel(
308
309
310
    const IdType* indptr, const IdType* indices, const IdType* row,
    const IdType* col, int64_t row_stride, int64_t col_stride, int64_t length,
    IdType* mask) {
311
312
313
314
315
316
317
318
319
320
321
322
323
324
  int tx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride_x = gridDim.x * blockDim.x;
  while (tx < length) {
    int rpos = tx * row_stride, cpos = tx * col_stride;
    const IdType r = row[rpos], c = col[cpos];
    for (IdType i = indptr[r]; i < indptr[r + 1]; ++i) {
      if (indices[i] == c) {
        mask[i] = 1;
      }
    }
    tx += stride_x;
  }
}

325
/**
326
 * @brief Search for the insertion positions for needle in the hay.
327
328
329
330
331
 *
 * The hay is a list of sorted elements and the result is the insertion position
 * of each needle so that the insertion still gives sorted order.
 *
 * It essentially perform binary search to find lower bound for each needle
332
333
334
 * elements. Require the largest elements in the hay is larger than the given
 * needle elements. Commonly used in searching for row IDs of a given set of
 * coordinates.
335
336
337
 */
template <typename IdType>
__global__ void _SortedSearchKernel(
338
339
    const IdType* hay, int64_t hay_size, const IdType* needles,
    int64_t num_needles, IdType* pos) {
340
341
342
343
344
345
346
347
348
349
350
351
352
353
  int tx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride_x = gridDim.x * blockDim.x;
  while (tx < num_needles) {
    const IdType ele = needles[tx];
    // binary search
    IdType lo = 0, hi = hay_size - 1;
    while (lo < hi) {
      IdType mid = (lo + hi) >> 1;
      if (hay[mid] <= ele) {
        lo = mid + 1;
      } else {
        hi = mid;
      }
    }
354
    pos[tx] = (hay[hi] == ele) ? hi : hi - 1;
355
356
357
358
    tx += stride_x;
  }
}

359
template <DGLDeviceType XPU, typename IdType>
360
361
std::vector<NDArray> CSRGetDataAndIndices(
    CSRMatrix csr, NDArray row, NDArray col) {
362
363
  const auto rowlen = row->shape[0];
  const auto collen = col->shape[0];
sangwzh's avatar
sangwzh committed
364
  const auto len = ::max(rowlen, collen);
365
  if (len == 0) return {NullArray(), NullArray(), NullArray()};
366
367
368
369
370
371

  const auto& ctx = row->ctx;
  const auto nbits = row->dtype.bits;
  const int64_t nnz = csr.indices->shape[0];
  const int64_t row_stride = (rowlen == 1 && collen != 1) ? 0 : 1;
  const int64_t col_stride = (collen == 1 && rowlen != 1) ? 0 : 1;
sangwzh's avatar
sangwzh committed
372
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
373

374
375
376
377
  const IdType* indptr_data =
      static_cast<IdType*>(GetDevicePointer(csr.indptr));
  const IdType* indices_data =
      static_cast<IdType*>(GetDevicePointer(csr.indices));
378

379
380
  // Generate a 0-1 mask for matched (row, col) positions.
  IdArray mask = Full(0, nnz, nbits, ctx);
381
  const int nt = dgl::cuda::FindNumThreads(len);
382
  const int nb = (len + nt - 1) / nt;
383
384
385
  CUDA_KERNEL_CALL(
      _SegmentMaskKernel, nb, nt, 0, stream, indptr_data, indices_data,
      row.Ptr<IdType>(), col.Ptr<IdType>(), row_stride, col_stride, len,
386
387
388
389
390
391
392
393
394
      mask.Ptr<IdType>());

  IdArray idx = AsNumBits(NonZero(mask), nbits);
  if (idx->shape[0] == 0)
    // No data. Return three empty arrays.
    return {idx, idx, idx};

  // Search for row index
  IdArray ret_row = NewIdArray(idx->shape[0], ctx, nbits);
395
  const int nt2 = dgl::cuda::FindNumThreads(idx->shape[0]);
396
  const int nb2 = (idx->shape[0] + nt - 1) / nt;
397
398
399
  CUDA_KERNEL_CALL(
      _SortedSearchKernel, nb2, nt2, 0, stream, indptr_data, csr.num_rows,
      idx.Ptr<IdType>(), idx->shape[0], ret_row.Ptr<IdType>());
400
401
402

  // Column & data can be obtained by index select.
  IdArray ret_col = IndexSelect(csr.indices, idx);
403
  IdArray ret_data = CSRHasData(csr) ? IndexSelect(csr.data, idx) : idx;
404
405
406
  return {ret_row, ret_col, ret_data};
}

407
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCUDA, int32_t>(
408
    CSRMatrix csr, NDArray rows, NDArray cols);
409
template std::vector<NDArray> CSRGetDataAndIndices<kDGLCUDA, int64_t>(
410
411
412
413
    CSRMatrix csr, NDArray rows, NDArray cols);

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

414
415
416
417
418
int64_t _UpPower(int64_t numel) {
  uint64_t ret = 1 << static_cast<uint64_t>(std::log2(numel) + 1);
  return ret;
}

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
 * @brief Thomas Wang's 32 bit Mix Function.
 * Source link: https://gist.github.com/badboy/6267743
 */
__device__ inline uint32_t _Hash32Shift(uint32_t key) {
  key = ~key + (key << 15);
  key = key ^ (key >> 12);
  key = key + (key << 2);
  key = key ^ (key >> 4);
  key = key * 2057;
  key = key ^ (key >> 16);
  return key;
}

/**
 * @brief Thomas Wang's 64 bit Mix Function.
 * Source link: https://gist.github.com/badboy/6267743
 */
__device__ inline uint64_t _Hash64Shift(uint64_t key) {
  key = (~key) + (key << 21);
  key = key ^ (key >> 24);
  key = (key + (key << 3)) + (key << 8);
  key = key ^ (key >> 14);
  key = (key + (key << 2)) + (key << 4);
  key = key ^ (key >> 28);
  key = key + (key << 31);
  return key;
}

/**
 * @brief A hashmap designed for CSRSliceMatrix, similar in function to set. For
 * performance, it can only be created and called in the cuda kernel.
451
452
 */
template <typename IdType>
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
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
struct NodeQueryHashmap {
  __device__ inline NodeQueryHashmap(IdType* Kptr, size_t numel)
      : kptr_(Kptr), capacity_(numel) {}

  /**
   * @brief Insert a key. It must be called by cuda threads.
   *
   * @param key The key to be inserted.
   */
  __device__ inline void Insert(IdType key) {
    uint32_t delta = 1;
    uint32_t pos = Hash(key);
    IdType prev = dgl::aten::cuda::AtomicCAS(&kptr_[pos], kEmptyKey_, key);
    while (prev != key && prev != kEmptyKey_) {
      pos = Hash(pos + delta);
      delta += 1;
      prev = dgl::aten::cuda::AtomicCAS(&kptr_[pos], kEmptyKey_, key);
    }
  }

  /**
   * @brief Check whether a key exists within the hashtable. It must be called
   * by cuda threads.
   *
   * @param key The key to check for.
   * @return True if the key exists in the hashtable.
   */
  __device__ inline bool Query(IdType key) {
    uint32_t delta = 1;
    uint32_t pos = Hash(key);
    while (true) {
      if (kptr_[pos] == key) return true;
      if (kptr_[pos] == kEmptyKey_) return false;
      pos = Hash(pos + delta);
      delta += 1;
    }
    return false;
  }

  __device__ inline uint32_t Hash(int32_t key) {
    return _Hash32Shift(key) & (capacity_ - 1);
  }

  __device__ inline uint32_t Hash(uint32_t key) {
    return _Hash32Shift(key) & (capacity_ - 1);
  }

  __device__ inline uint32_t Hash(int64_t key) {
    return static_cast<uint32_t>(_Hash64Shift(key)) & (capacity_ - 1);
  }

  __device__ inline uint32_t Hash(uint64_t key) {
    return static_cast<uint32_t>(_Hash64Shift(key)) & (capacity_ - 1);
  }

  IdType kEmptyKey_{-1};
  IdType* kptr_;
  uint32_t capacity_{0};
};

/**
 * @brief Generate a 0-1 mask for each index whose column is in the provided
 * hashmap. It also counts the number of masked values per row.
 *
 * @tparam IdType The ID type used for matrices.
 * @tparam WARP_SIZE The number of cuda threads in a cuda warp.
 * @tparam BLOCK_WARPS The number of warps in a cuda block.
 * @tparam TILE_SIZE The number of rows covered by each threadblock.
 */
template <typename IdType, int WARP_SIZE, int BLOCK_WARPS, int TILE_SIZE>
523
__global__ void _SegmentMaskColKernel(
524
    const IdType* indptr, const IdType* indices, int64_t num_rows,
525
526
527
528
529
530
531
532
533
534
535
536
    IdType* hashmap_buffer, int64_t buffer_size, IdType* mask, IdType* count) {
  assert(blockDim.x == WARP_SIZE);
  assert(blockDim.y == BLOCK_WARPS);

  int warp_id = threadIdx.y;
  int laneid = threadIdx.x;
  IdType out_row = blockIdx.x * TILE_SIZE + threadIdx.y;
  IdType last_row =
      min(static_cast<IdType>((blockIdx.x + 1) * TILE_SIZE),
          static_cast<IdType>(num_rows));

  NodeQueryHashmap<IdType> hashmap(hashmap_buffer, buffer_size);
sangwzh's avatar
sangwzh committed
537
  typedef hipcub::WarpReduce<IdType> WarpReduce;
538
539
540
541
542
543
544
545
546
547
548
549
  __shared__ typename WarpReduce::TempStorage temp_storage[BLOCK_WARPS];

  while (out_row < last_row) {
    IdType local_count = 0;
    IdType in_row_start = indptr[out_row];
    IdType in_row_end = indptr[out_row + 1];
    for (int idx = in_row_start + laneid; idx < in_row_end; idx += WARP_SIZE) {
      bool is_in = hashmap.Query(indices[idx]);
      if (is_in) {
        local_count += 1;
        mask[idx] = 1;
      }
550
    }
551
    IdType reduce_count = WarpReduce(temp_storage[warp_id]).Sum(local_count);
sangwzh's avatar
sangwzh committed
552
    printf("out_row = %d , reduce_count = %d \n", out_row, reduce_count);
553
554
555
556
    if (laneid == 0) {
      count[out_row] = reduce_count;
    }
    out_row += BLOCK_WARPS;
557
558
559
  }
}

560
template <DGLDeviceType XPU, typename IdType>
561
562
CSRMatrix CSRSliceMatrix(
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols) {
sangwzh's avatar
sangwzh committed
563
  hipStream_t stream = runtime::getCurrentHIPStreamMasqueradingAsCUDA();
564
565
566
567
568
569
  const auto& ctx = rows->ctx;
  const auto& dtype = rows->dtype;
  const auto nbits = dtype.bits;
  const int64_t new_nrows = rows->shape[0];
  const int64_t new_ncols = cols->shape[0];

sangwzh's avatar
sangwzh committed
570
571
572
  std::cout << "new_nrows : " << new_nrows << std::endl;
  std::cout << "new_ncols : " << new_ncols << std::endl;

573
  if (new_nrows == 0 || new_ncols == 0)
574
575
576
    return CSRMatrix(
        new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
        NullArray(dtype, ctx), NullArray(dtype, ctx));
577
578
579
580

  // First slice rows
  csr = CSRSliceRows(csr, rows);

sangwzh's avatar
sangwzh committed
581
  std::cout << "csr.indices->shape[0] : " << csr.indices->shape[0] << std::endl;
582
  if (csr.indices->shape[0] == 0)
583
584
585
    return CSRMatrix(
        new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
        NullArray(dtype, ctx), NullArray(dtype, ctx));
586
587
588
589
590

  // Generate a 0-1 mask for matched (row, col) positions.
  IdArray mask = Full(0, csr.indices->shape[0], nbits, ctx);
  // A count for how many masked values per row.
  IdArray count = NewIdArray(csr.num_rows, ctx, nbits);
sangwzh's avatar
sangwzh committed
591
  std::cout << "1 IdArray count : " << count << std::endl;
592
  CUDA_CALL(
sangwzh's avatar
sangwzh committed
593
      hipMemset(count.Ptr<IdType>(), 0, sizeof(IdType) * (csr.num_rows)));
594

sangwzh's avatar
sangwzh committed
595
  std::cout << "2 IdArray count : " << count << std::endl;
596
597
598
599
600
601
602
603
604
  // Generate a NodeQueryHashmap buffer. The key of the hashmap is col.
  // For performance, the load factor of the hashmap is in (0.25, 0.5);
  // Because num_cols is usually less than 1 Million (on GPU), the
  // memory overhead is not significant (less than 31MB) at a low load factor.
  int64_t buffer_size = _UpPower(new_ncols) * 2;
  IdArray hashmap_buffer = Full(-1, buffer_size, nbits, ctx);

  using it = thrust::counting_iterator<int64_t>;
  runtime::CUDAWorkspaceAllocator allocator(ctx);
sangwzh's avatar
sangwzh committed
605
  const auto exec_policy = thrust::hip::par_nosync(allocator).on(stream);
606
607
608
609
610
611
612
  thrust::for_each(
      exec_policy, it(0), it(new_ncols),
      [key = cols.Ptr<IdType>(), buffer = hashmap_buffer.Ptr<IdType>(),
       buffer_size] __device__(int64_t i) {
        NodeQueryHashmap<IdType> hashmap(buffer, buffer_size);
        hashmap.Insert(key[i]);
      });
613

614
615
616
617
  const IdType* indptr_data =
      static_cast<IdType*>(GetDevicePointer(csr.indptr));
  const IdType* indices_data =
      static_cast<IdType*>(GetDevicePointer(csr.indices));
618

619
  // Execute SegmentMaskColKernel
620
621
622
623
624
625
626
627
628
  const int64_t num_rows = csr.num_rows;
  constexpr int WARP_SIZE = 32;
  // With a simple fine-tuning, TILE_SIZE=16 gives a good performance.
  constexpr int TILE_SIZE = 16;
  constexpr int BLOCK_WARPS = CUDA_MAX_NUM_THREADS / WARP_SIZE;
  IdType nb =
      dgl::cuda::FindNumBlocks<'x'>((num_rows + TILE_SIZE - 1) / TILE_SIZE);
  const dim3 nthrs(WARP_SIZE, BLOCK_WARPS);
  const dim3 nblks(nb);
sangwzh's avatar
sangwzh committed
629
630
631
632
633
634
635
636
637
638
639
640
641

  std::cout << "nthrs.x : " << nthrs.x << " nthrs.y : " << nthrs.y << " nthrs.z : " << nthrs.z << std::endl;
  std::cout << "nblks.x : " << nblks.x << " nblks.y : " << nblks.y << " nblks.z : " << nblks.z << std::endl;
  std::cout << "WARP_SIZE : " << WARP_SIZE << " BLOCK_WARPS : " << BLOCK_WARPS << "TILE_SIZE : " << std::endl;
  std::cout << "indptr_data : " << indptr_data << std::endl;
  std::cout << "indices_data : " << indices_data << std::endl;
  std::cout << "num_rows : " << num_rows << std::endl;
  std::cout << "buffer_size  : " << buffer_size << std::endl;
  std::cout << "mask  : " << mask << std::endl;
  std::cout << "count  : " << count << std::endl;
  std::cout << "hashmap_buffer  : " << hashmap_buffer << std::endl;


642
  CUDA_KERNEL_CALL(
643
644
645
      (_SegmentMaskColKernel<IdType, WARP_SIZE, BLOCK_WARPS, TILE_SIZE>), nblks,
      nthrs, 0, stream, indptr_data, indices_data, num_rows,
      hashmap_buffer.Ptr<IdType>(), buffer_size, mask.Ptr<IdType>(),
646
      count.Ptr<IdType>());
647

sangwzh's avatar
sangwzh committed
648
  std::cout << "3 IdArray count : " << count << std::endl;
649
  IdArray idx = AsNumBits(NonZero(mask), nbits);
sangwzh's avatar
sangwzh committed
650
  std::cout << "idx->shape[0] : " << idx->shape[0] << std::endl;
651
  if (idx->shape[0] == 0)
652
653
654
    return CSRMatrix(
        new_nrows, new_ncols, Full(0, new_nrows + 1, nbits, ctx),
        NullArray(dtype, ctx), NullArray(dtype, ctx));
655
656

  // Indptr needs to be adjusted according to the new nnz per row.
sangwzh's avatar
sangwzh committed
657
  std::cout << " count  : " << count << std::endl;  
658
  IdArray ret_indptr = CumSum(count, true);
sangwzh's avatar
sangwzh committed
659
  std::cout << " IdArray ret_indptr  : " << ret_indptr << std::endl;
660
661
662

  // Column & data can be obtained by index select.
  IdArray ret_col = IndexSelect(csr.indices, idx);
663
  IdArray ret_data = CSRHasData(csr) ? IndexSelect(csr.data, idx) : idx;
664
665
666
667
668
669

  // Relabel column
  IdArray col_hash = NewIdArray(csr.num_cols, ctx, nbits);
  Scatter_(cols, Range(0, cols->shape[0], nbits, ctx), col_hash);
  ret_col = IndexSelect(col_hash, ret_col);

sangwzh's avatar
sangwzh committed
670
671
  // std::cout << "new_nrows : " << new_nrows << " new_ncols : " << new_ncols << " ret_indptr : " << ret_indptr << " ret_col : " << ret_col << " ret_data : " << std::endl;

672
  return CSRMatrix(new_nrows, new_ncols, ret_indptr, ret_col, ret_data);
673
674
}

675
template CSRMatrix CSRSliceMatrix<kDGLCUDA, int32_t>(
676
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
677
template CSRMatrix CSRSliceMatrix<kDGLCUDA, int64_t>(
678
679
680
681
682
    CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);

}  // namespace impl
}  // namespace aten
}  // namespace dgl