diff --git a/projects/rccl/src/collectives/device/common_kernel.h b/projects/rccl/src/collectives/device/common_kernel.h index dd108edaf6..c6a4f7a088 100644 --- a/projects/rccl/src/collectives/device/common_kernel.h +++ b/projects/rccl/src/collectives/device/common_kernel.h @@ -425,12 +425,12 @@ struct MULTI { template inline __device__ T vFetch(const volatile T* ptr) { - return *ptr; + return __builtin_nontemporal_load(ptr); } template inline __device__ void vStore(volatile T* ptr, const T val) { - *ptr = val; + __builtin_nontemporal_store(val, ptr); } #if CUDART_VERSION < 9000 && !(defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)) @@ -448,26 +448,34 @@ void vStore(volatile half* ptr, const half val) { #else template<> inline __device__ half vFetch(const volatile half* ptr) { - half r; - r = ((half*)ptr)[0]; - return r; + half h; + uint8_t *pr = (uint8_t *)&h; + pr[0] = __builtin_nontemporal_load((uint8_t*)ptr); + pr[1] = __builtin_nontemporal_load((uint8_t*)ptr+1); + return h; } template<> inline __device__ void vStore(volatile half* ptr, const half val) { - ((half*)ptr)[0] = val; + uint8_t *pr = (uint8_t *)&val; + __builtin_nontemporal_store(pr[0], (uint8_t*)ptr); + __builtin_nontemporal_store(pr[1], ((uint8_t*)ptr)+1); } template<> inline __device__ rccl_bfloat16 vFetch(const volatile rccl_bfloat16* ptr) { rccl_bfloat16 r; - r.data = ptr->data; + uint8_t *pr = (uint8_t *)&r.data; + pr[0] = __builtin_nontemporal_load((uint8_t*)&ptr->data); + pr[1] = __builtin_nontemporal_load(((uint8_t*)&ptr->data)+1); return r; } template<> inline __device__ void vStore(volatile rccl_bfloat16* ptr, const rccl_bfloat16 val) { - ptr->data = val.data; + uint8_t *pr = (uint8_t *)&val.data; + __builtin_nontemporal_store(pr[0], (uint8_t*)&ptr->data); + __builtin_nontemporal_store(pr[1], ((uint8_t*)&ptr->data)+1); } #endif