Use non-temporal access for slow path (#602)

src/collectives/device/common_kernel.h

@@ -425,12 +425,12 @@ struct MULTI<FUNC, int64_t> {
 
 template<typename T> inline __device__
 T vFetch(const volatile T* ptr) {
-  return *ptr;
+  return __builtin_nontemporal_load(ptr);
 }
 
 template<typename T> 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<half>(volatile half* ptr, const half val) {
 #else
 template<> inline __device__
 half vFetch<half>(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<half>(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<rccl_bfloat16>(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<rccl_bfloat16>(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