AtomicAddFloat.h 857 B

12345678910111213141516171819202122232425262728293031323334353637
  1. #ifndef ATOMIC_ADD_FLOAT
  2. #define ATOMIC_ADD_FLOAT
  3. #if (defined(__x86_64__) || defined(__i386__) || defined(__aarch64__))
  4. #include <ATen/native/cpu/Intrinsics.h>
  5. #else
  6. #define _mm_pause()
  7. #endif
  8. #include <atomic>
  9. static inline void cpu_atomic_add_float(float* dst, float fvalue)
  10. {
  11. typedef union {
  12. unsigned intV;
  13. float floatV;
  14. } uf32_t;
  15. uf32_t new_value, old_value;
  16. std::atomic<unsigned>* dst_intV = (std::atomic<unsigned>*)(dst);
  17. old_value.floatV = *dst;
  18. new_value.floatV = old_value.floatV + fvalue;
  19. unsigned* old_intV = (unsigned*)(&old_value.intV);
  20. while (!std::atomic_compare_exchange_strong(dst_intV, old_intV, new_value.intV)) {
  21. #ifdef __aarch64__
  22. __asm__ __volatile__("yield;" : : : "memory");
  23. #else
  24. _mm_pause();
  25. #endif
  26. old_value.floatV = *dst;
  27. new_value.floatV = old_value.floatV + fvalue;
  28. }
  29. }
  30. #endif