1 | #pragma once |
2 | |
3 | #include <c10/util/BFloat16.h> |
4 | #include <c10/util/Half.h> |
5 | #include <c10/util/math_compat.h> |
6 | |
7 | namespace c10 { |
8 | |
9 | // Note: Explicit implementation of copysign for Half and BFloat16 |
10 | // is needed to workaround g++-7/8 crash on aarch64, but also makes |
11 | // copysign faster for the half-precision types |
12 | template <typename T, typename U> |
13 | inline auto copysign(const T& a, const U& b) { |
14 | return std::copysign(a, b); |
15 | } |
16 | |
17 | // Implement copysign for half precision floats using bit ops |
18 | // Sign is the most significant bit for both half and bfloat16 types |
19 | inline c10::Half copysign(c10::Half a, c10::Half b) { |
20 | return c10::Half((a.x & 0x7fff) | (b.x & 0x8000), c10::Half::from_bits()); |
21 | } |
22 | |
23 | inline c10::BFloat16 copysign(c10::BFloat16 a, c10::BFloat16 b) { |
24 | return c10::BFloat16( |
25 | (a.x & 0x7fff) | (b.x & 0x8000), c10::BFloat16::from_bits()); |
26 | } |
27 | |
28 | } // namespace c10 |
29 | |