coo_sort.cu 5.65 KB
Newer Older
1
2
3
4
5
6
/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/cuda/coo_sort.cc
 * \brief Sort COO index
 */
#include <dgl/array.h>
7

8
#include "../../c_api_common.h"
9
#include "../../runtime/cuda/cuda_common.h"
10
11
12
13
14
15
16
17
18
19
20
#include "./utils.h"

namespace dgl {

using runtime::NDArray;

namespace aten {
namespace impl {

///////////////////////////// COOSort_ /////////////////////////////

21
/**
22
23
24
25
26
27
28
29
30
31
 * @brief Encode row and column IDs into a single scalar per edge.
 *
 * @tparam IdType The type to encode as.
 * @param row The row (src) IDs per edge.
 * @param col The column (dst) IDs per edge.
 * @param nnz The number of edges.
 * @param col_bits The number of bits used to encode the destination. The row
 * information is packed into the remaining bits.
 * @param key The encoded edges (output).
 */
32
33
template <typename IdType>
__global__ void _COOEncodeEdgesKernel(
34
35
    const IdType* const row, const IdType* const col, const int64_t nnz,
    const int col_bits, IdType* const key) {
36
37
38
39
40
  int64_t tx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;

  if (tx < nnz) {
    key[tx] = row[tx] << col_bits | col[tx];
  }
41
42
}

43
/**
44
45
46
47
48
49
50
51
52
 * @brief Decode row and column IDs from the encoded edges.
 *
 * @tparam IdType The type the edges are encoded as.
 * @param key The encoded edges.
 * @param nnz The number of edges.
 * @param col_bits The number of bits used to store the column/dst ID.
 * @param row The row (src) IDs per edge (output).
 * @param col The col (dst) IDs per edge (output).
 */
53
54
55
template <typename IdType>
__global__ void _COODecodeEdgesKernel(
    const IdType* const key, const int64_t nnz, const int col_bits,
56
    IdType* const row, IdType* const col) {
57
58
59
60
61
62
  int64_t tx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;

  if (tx < nnz) {
    const IdType k = key[tx];
    row[tx] = k >> col_bits;
    col[tx] = k & ((1 << col_bits) - 1);
63
  }
64
}
65

66
template <typename T>
67
68
69
70
int _NumberOfBits(const T& range) {
  if (range <= 1) {
    // ranges of 0 or 1 require no bits to store
    return 0;
71
72
  }

73
  int bits = 1;
74
  while (bits < static_cast<int>(sizeof(T) * 8) && (1 << bits) < range) {
75
76
77
    ++bits;
  }

78
79
  CHECK_EQ((range - 1) >> bits, 0);
  CHECK_NE((range - 1) >> (bits - 1), 0);
80
81

  return bits;
82
83
}

84
template <DGLDeviceType XPU, typename IdType>
85
void COOSort_(COOMatrix* coo, bool sort_column) {
86
  cudaStream_t stream = runtime::getCurrentCUDAStream();
87
88
89
90
91
92
93
94
  const int row_bits = _NumberOfBits(coo->num_rows);

  const int64_t nnz = coo->row->shape[0];
  if (sort_column) {
    const int col_bits = _NumberOfBits(coo->num_cols);
    const int num_bits = row_bits + col_bits;

    const int nt = 256;
95
96
    const int nb = (nnz + nt - 1) / nt;
    CHECK(static_cast<int64_t>(nb) * nt >= nnz);
97
98
99

    IdArray pos = aten::NewIdArray(nnz, coo->row->ctx, coo->row->dtype.bits);

100
101
102
    CUDA_KERNEL_CALL(
        _COOEncodeEdgesKernel, nb, nt, 0, stream, coo->row.Ptr<IdType>(),
        coo->col.Ptr<IdType>(), nnz, col_bits, pos.Ptr<IdType>());
103
104
105

    auto sorted = Sort(pos, num_bits);

106
107
108
    CUDA_KERNEL_CALL(
        _COODecodeEdgesKernel, nb, nt, 0, stream, sorted.first.Ptr<IdType>(),
        nnz, col_bits, coo->row.Ptr<IdType>(), coo->col.Ptr<IdType>());
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

    if (aten::COOHasData(*coo))
      coo->data = IndexSelect(coo->data, sorted.second);
    else
      coo->data = AsNumBits(sorted.second, coo->row->dtype.bits);
    coo->row_sorted = coo->col_sorted = true;
  } else {
    const int num_bits = row_bits;

    auto sorted = Sort(coo->row, num_bits);

    coo->row = sorted.first;
    coo->col = IndexSelect(coo->col, sorted.second);

    if (aten::COOHasData(*coo))
      coo->data = IndexSelect(coo->data, sorted.second);
    else
      coo->data = AsNumBits(sorted.second, coo->row->dtype.bits);
    coo->row_sorted = true;
  }
129
130
}

131
132
template void COOSort_<kDGLCUDA, int32_t>(COOMatrix* coo, bool sort_column);
template void COOSort_<kDGLCUDA, int64_t>(COOMatrix* coo, bool sort_column);
133
134
135
136
137

///////////////////////////// COOIsSorted /////////////////////////////

template <typename IdType>
__global__ void _COOIsSortedKernel(
138
139
    const IdType* row, const IdType* col, int64_t nnz, int8_t* row_sorted,
    int8_t* col_sorted) {
140
141
142
143
144
145
146
147
  int tx = blockIdx.x * blockDim.x + threadIdx.x;
  const int stride_x = gridDim.x * blockDim.x;
  while (tx < nnz) {
    if (tx == 0) {
      row_sorted[0] = 1;
      col_sorted[0] = 1;
    } else {
      row_sorted[tx] = static_cast<int8_t>(row[tx - 1] <= row[tx]);
148
149
      col_sorted[tx] =
          static_cast<int8_t>(row[tx - 1] < row[tx] || col[tx - 1] <= col[tx]);
150
151
152
153
154
    }
    tx += stride_x;
  }
}

155
template <DGLDeviceType XPU, typename IdType>
156
157
158
std::pair<bool, bool> COOIsSorted(COOMatrix coo) {
  const int64_t nnz = coo.row->shape[0];
  const auto& ctx = coo.row->ctx;
159
  cudaStream_t stream = runtime::getCurrentCUDAStream();
160
  auto device = runtime::DeviceAPI::Get(ctx);
161
162
  // We allocate a workspace of 2*nnz bytes. It wastes a little bit memory but
  // should be fine.
163
164
165
166
  int8_t* row_flags = static_cast<int8_t*>(device->AllocWorkspace(ctx, nnz));
  int8_t* col_flags = static_cast<int8_t*>(device->AllocWorkspace(ctx, nnz));
  const int nt = cuda::FindNumThreads(nnz);
  const int nb = (nnz + nt - 1) / nt;
167
168
169
  CUDA_KERNEL_CALL(
      _COOIsSortedKernel, nb, nt, 0, stream, coo.row.Ptr<IdType>(),
      coo.col.Ptr<IdType>(), nnz, row_flags, col_flags);
170
171

  const bool row_sorted = cuda::AllTrue(row_flags, nnz, ctx);
172
173
  const bool col_sorted =
      row_sorted ? cuda::AllTrue(col_flags, nnz, ctx) : false;
174
175
176
177
178
179
180

  device->FreeWorkspace(ctx, row_flags);
  device->FreeWorkspace(ctx, col_flags);

  return {row_sorted, col_sorted};
}

181
182
template std::pair<bool, bool> COOIsSorted<kDGLCUDA, int32_t>(COOMatrix coo);
template std::pair<bool, bool> COOIsSorted<kDGLCUDA, int64_t>(COOMatrix coo);
183
184
185
186

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