From a79f74082e9c50a7588c8ac61df48bdd5ee2ded2 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Wed, 14 Apr 2021 10:14:54 -0700
Subject: [PATCH] Limit max channels for ring graph on single node Rome (#347)

* Limit max channels for ring graph on single node Rome
* Partially revert "Use non-temporal access for streaming data (#341)"
---
 src/collectives/device/common_kernel.h  | 9 ++++-----
 src/graph/search.cc                     | 4 ++++
 tools/rccl-prim-test/copy_kernel.h      | 9 ++++-----
 tools/rccl-prim-test/rccl_prim_test.cpp | 1 +
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/collectives/device/common_kernel.h b/src/collectives/device/common_kernel.h
index 5759c2ed7b..30c9619887 100644
--- a/src/collectives/device/common_kernel.h
+++ b/src/collectives/device/common_kernel.h
@@ -267,17 +267,16 @@ struct MULTI128 {
 
 inline __device__ void Fetch128(Pack128& v, const Pack128* p) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  v.x = __builtin_nontemporal_load(&p->x);
-  v.y = __builtin_nontemporal_load(&p->y);
+  v.x = p->x;
+  v.y = 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__)
-  __builtin_nontemporal_store(v.x, &p->x);
-  __builtin_nontemporal_store(v.y, &p->y);
+  p->x = v.x;
+  p->y = v.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/graph/search.cc b/src/graph/search.cc
index 3ea5b7c06b..55456e0073 100644
--- a/src/graph/search.cc
+++ b/src/graph/search.cc
@@ -766,6 +766,10 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
   }
   if (graph->nChannels) return ncclSuccess;
 
+  if ((graph->pattern == NCCL_TOPO_PATTERN_RING) && (system->type & RCCL_TOPO_4P2H_ROME) && (ngpus == system->nRanks)) {
+    // limit single node max channels when searching ring graph on Rome
+    graph->maxChannels = 2;
+  }
   if (ngpus == 1) if (graph->pattern != NCCL_TOPO_PATTERN_RING) graph->pattern = NCCL_TOPO_PATTERN_TREE;
 
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
diff --git a/tools/rccl-prim-test/copy_kernel.h b/tools/rccl-prim-test/copy_kernel.h
index c05251a83b..05d96d8406 100644
--- a/tools/rccl-prim-test/copy_kernel.h
+++ b/tools/rccl-prim-test/copy_kernel.h
@@ -86,17 +86,16 @@ struct MULTI128 {
 
 inline __device__ void Fetch128(Pack128& v, const Pack128* p) {
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
-  v.x = __builtin_nontemporal_load(&p->x);
-  v.y = __builtin_nontemporal_load(&p->y);
+  v.x = p->x;
+  v.y = 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__)
-  __builtin_nontemporal_store(v.x, &p->x);
-  __builtin_nontemporal_store(v.y, &p->y);
+  p->x = v.x;
+  p->y = v.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 c36b2ea78c..076387f9b9 100644
--- a/tools/rccl-prim-test/rccl_prim_test.cpp
+++ b/tools/rccl-prim-test/rccl_prim_test.cpp
@@ -29,6 +29,7 @@ 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"