From 5f96f13e75b260dd42f13178d5bb894f3c6d22d6 Mon Sep 17 00:00:00 2001
From: Wenkai Du <wenkai.du@amd.com>
Date: Tue, 4 Aug 2020 21:07:00 +0000
Subject: [PATCH] Allow setup ring through NCCL_RINGS to facilitate testing

[ROCm/rccl commit: 5b03132ace510f2536dca88cabdfb18210c9e212]
---
 projects/rccl/src/graph/connect.cc | 1 +
 projects/rccl/src/graph/search.cc  | 8 ++++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/projects/rccl/src/graph/connect.cc b/projects/rccl/src/graph/connect.cc
index 0dfde31a02..b62e7741a6 100644
--- a/projects/rccl/src/graph/connect.cc
+++ b/projects/rccl/src/graph/connect.cc
@@ -291,6 +291,7 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, struct nccl
 
   int nc = nChannels*2;
   if (comm->topo->nodes[NET].count == 0 && comm->topo->type == RCCL_TOPO_CR8G) nc = nChannels*3;
+  if (comm->topo->nodes[NET].count && comm->topo->type == RCCL_TOPO_4P2H_ROME) nc = 4*comm->topo->nodes[NET].count;
   int end = std::min((int)ncclMaxNchannels(), std::max(nc, ncclMinNchannels()));
 
   // Duplication should be complete now
diff --git a/projects/rccl/src/graph/search.cc b/projects/rccl/src/graph/search.cc
index 4d4c23a0c1..fcb0a4f864 100644
--- a/projects/rccl/src/graph/search.cc
+++ b/projects/rccl/src/graph/search.cc
@@ -931,6 +931,8 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
     if (graph->nChannels > 0) return ncclSuccess;
   }
 
+  str = getenv("NCCL_RINGS");
+  if (str) system->type = RCCL_TOPO_4P2H_ROME;
   if (!str) NCCLCHECK(parseChordalRing(system, &str));
   if (!str) NCCLCHECK(parseRome4P2H(system, &str));
   if (str) {
@@ -950,8 +952,10 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
     if (system->nodes[NET].count) {
       // do not change ring order for multi node 4P2H on Rome
       if (system->type == RCCL_TOPO_4P2H_ROME) {
-        for (int n = 0; n < graph->nChannels; n++)
-          graph->inter[n*2] = graph->inter[n*2+1] = n%system->nodes[NET].count;
+        for (int n = 0; n < graph->nChannels; n++) {
+          graph->inter[n*2] = n%system->nodes[NET].count;
+          graph->inter[n*2+1] = (n+1)%system->nodes[NET].count;
+        }
       } else {
         int *intra, *used;
         graph->nChannels = system->nodes[NET].count;