functor.cuh 9.05 KB
Newer Older
1
/**
2
 *  Copyright (c) 2020 by Contributors
3
4
 * @file array/cuda/functor.cuh
 * @brief Functors for template on CUDA
5
6
7
8
 */
#ifndef DGL_ARRAY_CUDA_FUNCTOR_CUH_
#define DGL_ARRAY_CUDA_FUNCTOR_CUH_

9
10
#include <cmath>
#include <limits>
11

12
#include "./atomic.cuh"
13
#include "./fp16.cuh"
14
#include "bf16.cuh"
15

16
17
18
19
namespace dgl {
namespace aten {
namespace cuda {

20
/////////////////////////// CUDA binary operators //////////////////////////////
21
22
23
24
25
26
namespace binary {
template <typename DType>
struct Add {
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = false;
27
28
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
29
30
31
    return lhs[0] + rhs[0];
  }
};
32
33
34
35
36
37
template <typename DType>
constexpr bool Add<DType>::use_lhs;
template <typename DType>
constexpr bool Add<DType>::use_rhs;
template <typename DType>
constexpr bool Add<DType>::reduce_last_dim;
38
39
40
41
42
43

template <typename DType>
struct Sub {
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = false;
44
45
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
46
47
48
    return lhs[0] - rhs[0];
  }
};
49
50
51
52
53
54
template <typename DType>
constexpr bool Sub<DType>::use_lhs;
template <typename DType>
constexpr bool Sub<DType>::use_rhs;
template <typename DType>
constexpr bool Sub<DType>::reduce_last_dim;
55
56
57
58
59
60

template <typename DType>
struct Mul {
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = false;
61
62
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
63
64
65
    return lhs[0] * rhs[0];
  }
};
66
67
68
69
70
71
template <typename DType>
constexpr bool Mul<DType>::use_lhs;
template <typename DType>
constexpr bool Mul<DType>::use_rhs;
template <typename DType>
constexpr bool Mul<DType>::reduce_last_dim;
72
73
74
75
76
77

template <typename DType>
struct Div {
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = false;
78
79
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
80
81
82
    return lhs[0] / rhs[0];
  }
};
83
84
85
86
87
88
template <typename DType>
constexpr bool Div<DType>::use_lhs;
template <typename DType>
constexpr bool Div<DType>::use_rhs;
template <typename DType>
constexpr bool Div<DType>::reduce_last_dim;
89
90

template <typename DType>
91
struct CopyLhs {
92
93
94
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = false;
  static constexpr bool reduce_last_dim = false;
95
96
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
97
98
99
    return lhs[0];
  }
};
100
101
102
103
104
105
template <typename DType>
constexpr bool CopyLhs<DType>::use_lhs;
template <typename DType>
constexpr bool CopyLhs<DType>::use_rhs;
template <typename DType>
constexpr bool CopyLhs<DType>::reduce_last_dim;
106
107

template <typename DType>
108
struct CopyRhs {
109
110
111
  static constexpr bool use_lhs = false;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = false;
112
113
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
114
115
116
    return rhs[0];
  }
};
117
118
119
120
121
122
template <typename DType>
constexpr bool CopyRhs<DType>::use_lhs;
template <typename DType>
constexpr bool CopyRhs<DType>::use_rhs;
template <typename DType>
constexpr bool CopyRhs<DType>::reduce_last_dim;
123
124
125
126
127
128

template <typename DType>
struct Dot {
  static constexpr bool use_lhs = true;
  static constexpr bool use_rhs = true;
  static constexpr bool reduce_last_dim = true;
129
130
  static __device__ __forceinline__ DType
  Call(const DType *lhs, const DType *rhs, int64_t len = 1) {
131
    DType rst = static_cast<DType>(0.0f);
132
133
134
135
136
137
    for (int64_t i = 0; i < len; ++i) {
      rst += lhs[i] * rhs[i];
    }
    return rst;
  }
};
138
139
140
141
142
143
template <typename DType>
constexpr bool Dot<DType>::use_lhs;
template <typename DType>
constexpr bool Dot<DType>::use_rhs;
template <typename DType>
constexpr bool Dot<DType>::reduce_last_dim;
144

145
}  // end of namespace binary
146

147
/////////////////////////// CUDA reduce operators //////////////////////////////
148
namespace reduce {
149
template <typename Idx, typename DType, bool atomic>
150
151
152
struct _Sum {
  static constexpr __host__ __device__ __forceinline__ DType zero() {
    return 0.;
153
  }
154
155
  static constexpr bool require_arg = false;
  static __device__ __forceinline__ void Call(
156
157
      DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
      Idx eid) {
158
159
160
161
162
163
    if (!atomic) {
      *out_buf += val;
    } else {
      cuda::AtomicAdd(out_buf, val);
    }
  }
164
  static __device__ __forceinline__ void Call(
165
      DType *out_buf, Idx *arg_buf, DType val, Idx id) {
166
167
168
169
170
171
    if (!atomic) {
      *out_buf += val;
    } else {
      cuda::AtomicAdd(out_buf, val);
    }
  }
172
173
174
  static __device__ __forceinline__ void CallArg(
      Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
      Idx uid, Idx eid) {}
175
176
};

177
178
template <typename Idx, typename DType, bool atomic = false>
struct Sum : _Sum<Idx, DType, atomic> {};
179
180

template <typename Idx, bool atomic>
181
struct Sum<Idx, half, atomic> : _Sum<Idx, half, atomic> {
182
183
  static constexpr __host__ __device__ __forceinline__ half zero() {
    return __float2half_rn(0.);
184
  }
185
};
186
187
188

#if BF16_ENABLED
template <typename Idx, bool atomic>
189
struct Sum<Idx, __nv_bfloat16, atomic> : _Sum<Idx, __nv_bfloat16, atomic> {
190
191
192
193
194
  static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
    return __float2bfloat16_rn(0.);
  }
};
#endif  // BF16_ENABLED
195

196
template <typename Idx, typename DType, bool atomic>
197
198
199
struct _Max {
  static constexpr __host__ __device__ __forceinline__ DType zero() {
    return -std::numeric_limits<DType>::infinity();
200
  }
201
202
  static constexpr bool require_arg = true;
  static __device__ __forceinline__ void Call(
203
204
      DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
      Idx eid) {
205
206
207
208
209
210
211
212
213
214
    if (!atomic) {
      if (*out_buf < val) {
        *out_buf = val;
        *arg_u_buf = uid;
        *arg_e_buf = eid;
      }
    } else {
      cuda::AtomicMax(out_buf, val);
    }
  }
215
  static __device__ __forceinline__ void Call(
216
      DType *out_buf, Idx *arg_buf, DType val, Idx id) {
217
218
219
220
221
222
223
224
225
    if (!atomic) {
      if (*out_buf < val) {
        *out_buf = val;
        *arg_buf = id;
      }
    } else {
      cuda::AtomicMax(out_buf, val);
    }
  }
226
227
228
  static __device__ __forceinline__ void CallArg(
      Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
      Idx uid, Idx eid) {
229
230
    if (atomic) {
      if (val == val_ref) {
231
232
        if (arg_u_buf) arg_u_buf[fid] = uid;
        if (arg_e_buf) arg_e_buf[fid] = eid;
233
234
235
236
237
      }
    }
  }
};

238
239
template <typename Idx, typename DType, bool atomic = false>
struct Max : _Max<Idx, DType, atomic> {};
240

241
template <typename Idx, bool atomic>
242
243
244
struct Max<Idx, half, atomic> : _Max<Idx, half, atomic> {
  static constexpr __host__ __device__ __forceinline__ half zero() {
    return __float2half_rn(-6.550400e+04f);
245
  }
246
};
247
248

#if BF16_ENABLED
249
template <typename Idx, bool atomic>
250
251
252
253
254
255
struct Max<Idx, __nv_bfloat16, atomic> : _Max<Idx, __nv_bfloat16, atomic> {
  static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
    return __float2bfloat16_rn(-std::numeric_limits<float>::infinity());
  }
};
#endif  // BF16_ENABLED
256

257
template <typename Idx, typename DType, bool atomic>
258
259
260
struct _Min {
  static constexpr __host__ __device__ __forceinline__ DType zero() {
    return std::numeric_limits<DType>::infinity();
261
  }
262
263
  static constexpr bool require_arg = true;
  static __device__ __forceinline__ void Call(
264
265
      DType *out_buf, Idx *arg_u_buf, Idx *arg_e_buf, DType val, Idx uid,
      Idx eid) {
266
267
268
269
270
271
272
273
274
275
    if (!atomic) {
      if (*out_buf > val) {
        *out_buf = val;
        *arg_u_buf = uid;
        *arg_e_buf = eid;
      }
    } else {
      cuda::AtomicMin(out_buf, val);
    }
  }
276
  static __device__ __forceinline__ void Call(
277
      DType *out_buf, Idx *arg_buf, DType val, Idx id) {
278
279
280
281
282
283
284
285
286
    if (!atomic) {
      if (*out_buf > val) {
        *out_buf = val;
        *arg_buf = id;
      }
    } else {
      cuda::AtomicMin(out_buf, val);
    }
  }
287
288
289
  static __device__ __forceinline__ void CallArg(
      Idx fid, Idx *arg_u_buf, Idx *arg_e_buf, DType val, DType val_ref,
      Idx uid, Idx eid) {
290
291
    if (atomic) {
      if (val == val_ref) {
292
293
        if (arg_u_buf) arg_u_buf[fid] = uid;
        if (arg_e_buf) arg_e_buf[fid] = eid;
294
295
296
297
      }
    }
  }
};
298

299
300
template <typename Idx, typename DType, bool atomic = false>
struct Min : _Min<Idx, DType, atomic> {};
301

302
template <typename Idx, bool atomic>
303
304
305
struct Min<Idx, half, atomic> : _Min<Idx, half, atomic> {
  static constexpr __host__ __device__ __forceinline__ half zero() {
    return __float2half_rn(6.550400e+04f);
306
  }
307
};
308
309

#if BF16_ENABLED
310
template <typename Idx, bool atomic>
311
312
313
314
315
316
struct Min<Idx, __nv_bfloat16, atomic> : _Min<Idx, __nv_bfloat16, atomic> {
  static constexpr __host__ __device__ __forceinline__ __nv_bfloat16 zero() {
    return __float2bfloat16_rn(std::numeric_limits<float>::infinity());
  }
};
#endif  // BF16_ENABLED
317
318
319
320
321
322
323
324

}  // namespace reduce

}  // namespace cuda
}  // namespace aten
}  // namespace dgl

#endif  // DGL_ARRAY_CUDA_FUNCTOR_CUH_