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