|
1 | 1 | #pragma once |
2 | 2 |
|
3 | | -#include <c10/util/BFloat16.h> |
4 | | -#include <c10/util/Half.h> |
5 | | - |
6 | | -C10_CLANG_DIAGNOSTIC_PUSH() |
7 | | -#if C10_CLANG_HAS_WARNING("-Wimplicit-float-conversion") |
8 | | -C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-float-conversion") |
9 | | -#endif |
10 | | - |
11 | | -namespace c10 { |
12 | | -template <typename T> |
13 | | -struct is_reduced_floating_point |
14 | | - : std::integral_constant< |
15 | | - bool, |
16 | | - std::is_same_v<T, c10::Half> || std::is_same_v<T, c10::BFloat16>> {}; |
17 | | - |
18 | | -template <typename T> |
19 | | -constexpr bool is_reduced_floating_point_v = |
20 | | - is_reduced_floating_point<T>::value; |
21 | | -} // namespace c10 |
22 | | - |
23 | | -namespace std { |
24 | | - |
25 | | -#if !defined(FBCODE_CAFFE2) && !defined(C10_NODEPRECATED) |
26 | | -using c10::is_reduced_floating_point; |
27 | | -using c10::is_reduced_floating_point_v; |
28 | | -#endif // !defined(FBCODE_CAFFE2) && !defined(C10_NODEPRECATED) |
29 | | - |
30 | | -template < |
31 | | - typename T, |
32 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
33 | | -inline T acos(T a) { |
34 | | - return std::acos(float(a)); |
35 | | -} |
36 | | -template < |
37 | | - typename T, |
38 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
39 | | -inline T asin(T a) { |
40 | | - return std::asin(float(a)); |
41 | | -} |
42 | | -template < |
43 | | - typename T, |
44 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
45 | | -inline T atan(T a) { |
46 | | - return std::atan(float(a)); |
47 | | -} |
48 | | -template < |
49 | | - typename T, |
50 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
51 | | -inline T atanh(T a) { |
52 | | - return std::atanh(float(a)); |
53 | | -} |
54 | | -template < |
55 | | - typename T, |
56 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
57 | | -inline T erf(T a) { |
58 | | - return std::erf(float(a)); |
59 | | -} |
60 | | -template < |
61 | | - typename T, |
62 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
63 | | -inline T erfc(T a) { |
64 | | - return std::erfc(float(a)); |
65 | | -} |
66 | | -template < |
67 | | - typename T, |
68 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
69 | | -inline T exp(T a) { |
70 | | - return std::exp(float(a)); |
71 | | -} |
72 | | -template < |
73 | | - typename T, |
74 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
75 | | -inline T expm1(T a) { |
76 | | - return std::expm1(float(a)); |
77 | | -} |
78 | | -template < |
79 | | - typename T, |
80 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
81 | | -inline bool isfinite(T a) { |
82 | | - return std::isfinite(float(a)); |
83 | | -} |
84 | | -template < |
85 | | - typename T, |
86 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
87 | | -inline T log(T a) { |
88 | | - return std::log(float(a)); |
89 | | -} |
90 | | -template < |
91 | | - typename T, |
92 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
93 | | -inline T log10(T a) { |
94 | | - return std::log10(float(a)); |
95 | | -} |
96 | | -template < |
97 | | - typename T, |
98 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
99 | | -inline T log1p(T a) { |
100 | | - return std::log1p(float(a)); |
101 | | -} |
102 | | -template < |
103 | | - typename T, |
104 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
105 | | -inline T log2(T a) { |
106 | | - return std::log2(float(a)); |
107 | | -} |
108 | | -template < |
109 | | - typename T, |
110 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
111 | | -inline T ceil(T a) { |
112 | | - return std::ceil(float(a)); |
113 | | -} |
114 | | -template < |
115 | | - typename T, |
116 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
117 | | -inline T cos(T a) { |
118 | | - return std::cos(float(a)); |
119 | | -} |
120 | | -template < |
121 | | - typename T, |
122 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
123 | | -inline T floor(T a) { |
124 | | - return std::floor(float(a)); |
125 | | -} |
126 | | -template < |
127 | | - typename T, |
128 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
129 | | -inline T nearbyint(T a) { |
130 | | - return std::nearbyint(float(a)); |
131 | | -} |
132 | | -template < |
133 | | - typename T, |
134 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
135 | | -inline T sin(T a) { |
136 | | - return std::sin(float(a)); |
137 | | -} |
138 | | -template < |
139 | | - typename T, |
140 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
141 | | -inline T tan(T a) { |
142 | | - return std::tan(float(a)); |
143 | | -} |
144 | | -template < |
145 | | - typename T, |
146 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
147 | | -inline T sinh(T a) { |
148 | | - return std::sinh(float(a)); |
149 | | -} |
150 | | -template < |
151 | | - typename T, |
152 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
153 | | -inline T cosh(T a) { |
154 | | - return std::cosh(float(a)); |
155 | | -} |
156 | | -template < |
157 | | - typename T, |
158 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
159 | | -inline T tanh(T a) { |
160 | | - return std::tanh(float(a)); |
161 | | -} |
162 | | -template < |
163 | | - typename T, |
164 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
165 | | -inline T trunc(T a) { |
166 | | - return std::trunc(float(a)); |
167 | | -} |
168 | | -template < |
169 | | - typename T, |
170 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
171 | | -inline T lgamma(T a) { |
172 | | - return std::lgamma(float(a)); |
173 | | -} |
174 | | -template < |
175 | | - typename T, |
176 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
177 | | -inline T sqrt(T a) { |
178 | | - return std::sqrt(float(a)); |
179 | | -} |
180 | | -template < |
181 | | - typename T, |
182 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
183 | | -inline T rsqrt(T a) { |
184 | | - return 1.0f / std::sqrt(float(a)); |
185 | | -} |
186 | | -template < |
187 | | - typename T, |
188 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
189 | | -inline T abs(T a) { |
190 | | - return std::abs(float(a)); |
191 | | -} |
192 | | -#if defined(_MSC_VER) && defined(__CUDACC__) |
193 | | -template < |
194 | | - typename T, |
195 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
196 | | -inline T pow(T a, double b) { |
197 | | - return std::pow(float(a), float(b)); |
198 | | -} |
199 | | -#else |
200 | | -template < |
201 | | - typename T, |
202 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
203 | | -inline T pow(T a, double b) { |
204 | | - return std::pow(float(a), b); |
205 | | -} |
206 | | -#endif |
207 | | -template < |
208 | | - typename T, |
209 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
210 | | -inline T pow(T a, T b) { |
211 | | - return std::pow(float(a), float(b)); |
212 | | -} |
213 | | -template < |
214 | | - typename T, |
215 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
216 | | -inline T fmod(T a, T b) { |
217 | | - return std::fmod(float(a), float(b)); |
218 | | -} |
219 | | - |
220 | | -/* |
221 | | - The following function is inspired from the implementation in `musl` |
222 | | - Link to License: https://git.musl-libc.org/cgit/musl/tree/COPYRIGHT |
223 | | - ---------------------------------------------------------------------- |
224 | | - Copyright © 2005-2020 Rich Felker, et al. |
225 | | -
|
226 | | - Permission is hereby granted, free of charge, to any person obtaining |
227 | | - a copy of this software and associated documentation files (the |
228 | | - "Software"), to deal in the Software without restriction, including |
229 | | - without limitation the rights to use, copy, modify, merge, publish, |
230 | | - distribute, sublicense, and/or sell copies of the Software, and to |
231 | | - permit persons to whom the Software is furnished to do so, subject to |
232 | | - the following conditions: |
233 | | -
|
234 | | - The above copyright notice and this permission notice shall be |
235 | | - included in all copies or substantial portions of the Software. |
236 | | -
|
237 | | - THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
238 | | - EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
239 | | - MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. |
240 | | - IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY |
241 | | - CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, |
242 | | - TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE |
243 | | - SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. |
244 | | - ---------------------------------------------------------------------- |
245 | | - */ |
246 | | -template < |
247 | | - typename T, |
248 | | - typename std::enable_if_t<c10::is_reduced_floating_point_v<T>, int> = 0> |
249 | | -C10_HOST_DEVICE inline T nextafter(T from, T to) { |
250 | | - // Reference: |
251 | | - // https://git.musl-libc.org/cgit/musl/tree/src/math/nextafter.c |
252 | | - using int_repr_t = uint16_t; |
253 | | - constexpr uint8_t bits = 16; |
254 | | - union { |
255 | | - T f; |
256 | | - int_repr_t i; |
257 | | - } ufrom = {from}, uto = {to}; |
258 | | - |
259 | | - // get a mask to get the sign bit i.e. MSB |
260 | | - int_repr_t sign_mask = int_repr_t{1} << (bits - 1); |
261 | | - |
262 | | - // short-circuit: if either is NaN, return NaN |
263 | | - if (from != from || to != to) { |
264 | | - return from + to; |
265 | | - } |
266 | | - |
267 | | - // short-circuit: if they are exactly the same. |
268 | | - if (ufrom.i == uto.i) { |
269 | | - return from; |
270 | | - } |
271 | | - |
272 | | - // mask the sign-bit to zero i.e. positive |
273 | | - // equivalent to abs(x) |
274 | | - int_repr_t abs_from = ufrom.i & ~sign_mask; |
275 | | - int_repr_t abs_to = uto.i & ~sign_mask; |
276 | | - if (abs_from == 0) { |
277 | | - // if both are zero but with different sign, |
278 | | - // preserve the sign of `to`. |
279 | | - if (abs_to == 0) { |
280 | | - return to; |
281 | | - } |
282 | | - // smallest subnormal with sign of `to`. |
283 | | - ufrom.i = (uto.i & sign_mask) | int_repr_t{1}; |
284 | | - return ufrom.f; |
285 | | - } |
286 | | - |
287 | | - // if abs(from) > abs(to) or sign(from) != sign(to) |
288 | | - if (abs_from > abs_to || ((ufrom.i ^ uto.i) & sign_mask)) { |
289 | | - ufrom.i--; |
290 | | - } else { |
291 | | - ufrom.i++; |
292 | | - } |
293 | | - |
294 | | - return ufrom.f; |
295 | | -} |
296 | | - |
297 | | -} // namespace std |
298 | | - |
299 | | -C10_CLANG_DIAGNOSTIC_POP() |
| 3 | +#include <torch/headeronly/util/BFloat16-math.h> |
0 commit comments