copysign.h 866 B

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