From 9dfc2c183e0e0fc2ed19a5b2bdd29aa7c177dbc3 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 7 Apr 2021 17:34:35 -0700
Subject: [PATCH] Use non-temporal access for streaming data (#341)

* Use non-temporal access for streaming data

* Revert to ulong2 after fixing compiling issue
---
 src/collectives/device/common_kernel.h  |  9 +++++----
 src/collectives/device/prims_ll.h       | 26 ++++++++++++-------------
 src/nccl.h.in                           |  2 +-
 tools/rccl-prim-test/copy_kernel.h      |  9 +++++----
 tools/rccl-prim-test/rccl_prim_test.cpp |  1 -
 5 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/src/collectives/device/common_kernel.h b/src/collectives/device/common_kernel.h
index 30c9619887..5759c2ed7b 100644
--- a/src/collectives/device/common_kernel.h
+++ b/src/collectives/device/common_kernel.h
@@ -267,16 +267,17 @@ struct MULTI128 {
 
 inline __device__ void Fetch128(Pack128& v, const Pack128* p) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  v.x = p->x;
-  v.y = p->y;
+  v.x = __builtin_nontemporal_load(&p->x);
+  v.y = __builtin_nontemporal_load(&p->y);
 #else
   asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];" : "=l"(v.x), "=l"(v.y) : "l"(p) : "memory");
 #endif
 }
+
 inline __device__ void Store128(Pack128* p, Pack128& v) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  p->x = v.x;
-  p->y = v.y;
+  __builtin_nontemporal_store(v.x, &p->x);
+  __builtin_nontemporal_store(v.y, &p->y);
 #else
   asm volatile("st.volatile.global.v2.u64 [%0], {%1,%2};" :: "l"(p), "l"(v.x), "l"(v.y) : "memory");
 #endif
diff --git a/src/collectives/device/prims_ll.h b/src/collectives/device/prims_ll.h
index 44263e928e..ce6ea697ee 100644
--- a/src/collectives/device/prims_ll.h
+++ b/src/collectives/device/prims_ll.h
@@ -96,14 +96,13 @@ class ncclLLPrimitives {
     uint32_t data1, flag1, data2, flag2;
     spins = 0;
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-    using Vec = uint32_t __attribute__((ext_vector_type(4)));
-    Vec i4;
+    union ncclLLFifoLine i4;
     do {
-      asm volatile ("flat_load_dwordx4 %0, %1, glc, slc\n"
-        "s_waitcnt vmcnt(0)\n" : "=v"(i4) : "v"(src));
+      i4.v[0] = __builtin_nontemporal_load(src->v);
+      i4.v[1] = __builtin_nontemporal_load(src->v+1);
       if (checkAbort(i, 0)) break;
-    } while ((i4[1] != flag) || (i4[3] != flag));
-    uint64_t val64 = (uint64_t)(i4[0]) + (((uint64_t)i4[2]) << 32);
+    } while ((i4.flag1 != flag) || (i4.flag2 != flag));
+    uint64_t val64 = (uint64_t)(i4.data1) + (((uint64_t)i4.data2) << 32);
 #else
     do {
       asm volatile("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(data1), "=r"(flag1), "=r"(data2), "=r"(flag2) : "l"(&src->i4));
@@ -116,14 +115,13 @@ class ncclLLPrimitives {
 
   __device__ void storeLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-    using Vec = uint32_t __attribute__((ext_vector_type(4)));
-    Vec i4;
-    i4[0] = val & 0xffffffff;
-    i4[1] = flag;
-    i4[2] = (val >> 32);
-    i4[3] = flag;
-    asm volatile ("flat_store_dwordx4 %0, %1, glc, slc\n"
-      "s_waitcnt vmcnt(0)\n" : : "v"(dst), "v"(i4));
+    union ncclLLFifoLine i4;
+    i4.data1 = val & 0xffffffff;
+    i4.flag1 = flag;
+    i4.data2 = (val >> 32);
+    i4.flag2 = flag;
+    __builtin_nontemporal_store(i4.v[0], dst->v);
+    __builtin_nontemporal_store(i4.v[1], dst->v+1);
 #else
     asm volatile("st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};" :: "l"(&dst->i4), "r"((uint32_t)val), "r"(flag), "r"((uint32_t)(val >> 32)), "r"(flag));
 #endif
diff --git a/src/nccl.h.in b/src/nccl.h.in
index 0b6a74cb57..06b51d1871 100644
--- a/src/nccl.h.in
+++ b/src/nccl.h.in
@@ -8,7 +8,7 @@
 #ifndef NCCL_H_
 #define NCCL_H_
 
-#include <hip/hip_runtime_api.h>
+#include <hip/hip_runtime.h>
 #include <hip/hip_fp16.h>
 
 #define NCCL_MAJOR ${NCCL_MAJOR}
diff --git a/tools/rccl-prim-test/copy_kernel.h b/tools/rccl-prim-test/copy_kernel.h
index 05d96d8406..c05251a83b 100644
--- a/tools/rccl-prim-test/copy_kernel.h
+++ b/tools/rccl-prim-test/copy_kernel.h
@@ -86,16 +86,17 @@ struct MULTI128 {
 
 inline __device__ void Fetch128(Pack128& v, const Pack128* p) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  v.x = p->x;
-  v.y = p->y;
+  v.x = __builtin_nontemporal_load(&p->x);
+  v.y = __builtin_nontemporal_load(&p->y);
 #else
   asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];" : "=l"(v.x), "=l"(v.y) : "l"(p) : "memory");
 #endif
 }
+
 inline __device__ void Store128(Pack128* p, Pack128& v) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  p->x = v.x;
-  p->y = v.y;
+  __builtin_nontemporal_store(v.x, &p->x);
+  __builtin_nontemporal_store(v.y, &p->y);
 #else
   asm volatile("st.volatile.global.v2.u64 [%0], {%1,%2};" :: "l"(p), "l"(v.x), "l"(v.y) : "memory");
 #endif
diff --git a/tools/rccl-prim-test/rccl_prim_test.cpp b/tools/rccl-prim-test/rccl_prim_test.cpp
index 076387f9b9..c36b2ea78c 100644
--- a/tools/rccl-prim-test/rccl_prim_test.cpp
+++ b/tools/rccl-prim-test/rccl_prim_test.cpp
@@ -29,7 +29,6 @@ THE SOFTWARE.
 #include <iostream> //cerr
 #include <unistd.h> //usleep
 #include <cstring>
-#include <hip/hip_runtime_api.h>
 #include <hip/hip_runtime.h>
 #include "copy_kernel.h"