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