fp16.cuh 4.51 KB
Newer Older
1
/**
2
3
4
5
6
7
8
9
10
11
12
13
14
15
 *  Copyright (c) 2020-2022 by Contributors
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
16
17
18
 * @file array/cuda/fp16.cuh
 * @brief float16 related functions.
 * @note this file is modified from TVM project:
19
20
 *       https://github.com/apache/tvm/blob/e561007f0c330e3d14c2bc8a3ef40fb741db9004/src/target/source/literal/cuda_half_t.h.
 */
21
22
#ifndef DGL_ARRAY_CUDA_FP16_CUH_
#define DGL_ARRAY_CUDA_FP16_CUH_
23
24

#include <cuda_fp16.h>
25

26
#include <algorithm>
27

28
static __device__ __forceinline__ half max(half a, half b) {
29
30
31
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
  return __hgt(__half(a), __half(b)) ? a : b;
#else
32
  return __half(max(float(a), float(b)));  // NOLINT
33
34
35
#endif
}

36
static __device__ __forceinline__ half min(half a, half b) {
37
38
39
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
  return __hlt(__half(a), __half(b)) ? a : b;
#else
40
  return __half(min(float(a), float(b)));  // NOLINT
41
42
#endif
}
43
44

#ifdef __CUDACC__
45
46
// Arithmetic FP16 operations for architecture >= 5.3 are already defined in
// cuda_fp16.h
47
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)
48
49
// CUDA 12.2 adds "emulated" support for older architectures.
#if defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
50
51
__device__ __forceinline__ __half
operator+(const __half& lh, const __half& rh) {
52
53
  return __half(float(lh) + float(rh));  // NOLINT
}
54
55
__device__ __forceinline__ __half
operator-(const __half& lh, const __half& rh) {
56
57
  return __half(float(lh) - float(rh));  // NOLINT
}
58
59
__device__ __forceinline__ __half
operator*(const __half& lh, const __half& rh) {
60
61
  return __half(float(lh) * float(rh));  // NOLINT
}
62
63
__device__ __forceinline__ __half
operator/(const __half& lh, const __half& rh) {
64
65
  return __half(float(lh) / float(rh));  // NOLINT
}
66

67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
__device__ __forceinline__ __half& operator+=(
    __half& lh, const __half& rh) {    // NOLINT
  lh = __half(float(lh) + float(rh));  // NOLINT
  return lh;
}
__device__ __forceinline__ __half& operator-=(
    __half& lh, const __half& rh) {    // NOLINT
  lh = __half(float(lh) - float(rh));  // NOLINT
  return lh;
}
__device__ __forceinline__ __half& operator*=(
    __half& lh, const __half& rh) {    // NOLINT
  lh = __half(float(lh) * float(rh));  // NOLINT
  return lh;
}
__device__ __forceinline__ __half& operator/=(
    __half& lh, const __half& rh) {    // NOLINT
  lh = __half(float(lh) / float(rh));  // NOLINT
  return lh;
86
}
87

88
__device__ __forceinline__ __half& operator++(__half& h) {  // NOLINT
89
90
  h = __half(float(h) + 1.0f);                              // NOLINT
  return h;
91
92
}
__device__ __forceinline__ __half& operator--(__half& h) {  // NOLINT
93
94
  h = __half(float(h) - 1.0f);                              // NOLINT
  return h;
95
}
96
97
98
99
__device__ __forceinline__ __half operator++(__half& h, int) {  // NOLINT
  __half ret = h;
  h = __half(float(h) + 1.0f);  // NOLINT
  return ret;
100
}
101
102
103
104
__device__ __forceinline__ __half operator--(__half& h, int) {  // NOLINT
  __half ret = h;
  h = __half(float(h) - 1.0f);  // NOLINT
  return ret;
105
}
106
107

__device__ __forceinline__ __half operator+(const __half& h) { return h; }
108
109
110
__device__ __forceinline__ __half operator-(const __half& h) {
  return __half(-float(h));  // NOLINT
}
111

112
113
114
115
116
117
__device__ __forceinline__ bool operator==(const __half& lh, const __half& rh) {
  return float(lh) == float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator!=(const __half& lh, const __half& rh) {
  return float(lh) != float(rh);  // NOLINT
}
118
119
__device__ __forceinline__ bool operator>(const __half& lh, const __half& rh) {
  return float(lh) > float(rh);  // NOLINT
120
}
121
122
__device__ __forceinline__ bool operator<(const __half& lh, const __half& rh) {
  return float(lh) < float(rh);  // NOLINT
123
124
125
126
127
128
129
}
__device__ __forceinline__ bool operator>=(const __half& lh, const __half& rh) {
  return float(lh) >= float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator<=(const __half& lh, const __half& rh) {
  return float(lh) <= float(rh);  // NOLINT
}
130
131
#endif  // defined(CUDART_VERSION) && (CUDART_VERSION < 12020)
#endif  // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)
132
133
#endif  // __CUDACC__

134
#endif  // DGL_ARRAY_CUDA_FP16_CUH_