copysign.h 832 B

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