From 90cbef704275b8504156328723f6afac4aa965bd Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Tue, 6 Jun 2023 08:41:38 -0700
Subject: [PATCH] Add NCCL_NCHANNELS_PER_PEER override (#767)

Also fix topol_expl build issue

[ROCm/rccl commit: 3af90902c82796082fbb3d6a43211e7124e036ab]
---
 projects/rccl/src/graph/paths.cc        | 11 ++++++++---
 projects/rccl/tools/topo_expl/Makefile  |  1 +
 projects/rccl/tools/topo_expl/utils.cpp | 25 ++++++++-----------------
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/projects/rccl/src/graph/paths.cc b/projects/rccl/src/graph/paths.cc
index cbd529bda0..8fc65f9748 100644
--- a/projects/rccl/src/graph/paths.cc
+++ b/projects/rccl/src/graph/paths.cc
@@ -767,6 +767,7 @@ void ncclTopoFree(struct ncclTopoSystem* system) {
 }
 
 NCCL_PARAM(NChannelsPerNetPeer, "NCHANNELS_PER_NET_PEER", 1);
+NCCL_PARAM(NChannelsPerPeer, "NCHANNELS_PER_PEER", -2);
 
 static ncclResult_t ncclTopoGetNchannels(struct ncclTopoSystem* system, int g /*local gpu index*/, int peerRank, int* nChannels) {
   int peer;
@@ -780,8 +781,12 @@ static ncclResult_t ncclTopoGetNchannels(struct ncclTopoSystem* system, int g /*
     // Local rank
     path = system->nodes[GPU].nodes[peer].paths[GPU]+g;
     if (path->type == PATH_NVL) {
-      float nvlBw = ncclTopoXGMISpeed(system->nodes[GPU].nodes[g].gpu.gcn);
-      *nChannels = 2*std::max(1, (int)(path->bw / nvlBw));
+      if (ncclParamNChannelsPerPeer() == -2) {
+        float nvlBw = ncclTopoXGMISpeed(system->nodes[GPU].nodes[g].gpu.gcn);
+        *nChannels = 2*std::max(1, (int)(path->bw / nvlBw));
+      } else {
+        *nChannels = ncclParamNChannelsPerPeer();
+      }
     } else {
       *nChannels = 2;
     }
@@ -829,7 +834,7 @@ ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm) {
   else {
     // Round to next pow2 nChannelsPerPeer and nChannels
     comm->p2pnChannelsPerPeer = nextPow2(minChannels);
-    comm->p2pnChannels = std::min(comm->topo->nodes[GPU].count == comm->topo->nRanks ? 2*comm->nRanks : comm->nRanks, nextPow2(comm->p2pnChannels));
+    comm->p2pnChannels = nextPow2(comm->p2pnChannels);
   }
 
   // Init channels that weren't used so far
diff --git a/projects/rccl/tools/topo_expl/Makefile b/projects/rccl/tools/topo_expl/Makefile
index 6a0d5d62a5..4a51a427ce 100644
--- a/projects/rccl/tools/topo_expl/Makefile
+++ b/projects/rccl/tools/topo_expl/Makefile
@@ -17,6 +17,7 @@ $(EXE): $(files)
 	$(HIPCC) $(CXXFLAGS) $^ -o $@
 
 hipify:
+	rm -rf hipify_rccl
 	mkdir -p hipify_rccl
 	cp -a ../../src/include/ hipify_rccl/
 	cp -a ../../src/graph/ hipify_rccl/
diff --git a/projects/rccl/tools/topo_expl/utils.cpp b/projects/rccl/tools/topo_expl/utils.cpp
index 70b70fa640..46e4f3a568 100644
--- a/projects/rccl/tools/topo_expl/utils.cpp
+++ b/projects/rccl/tools/topo_expl/utils.cpp
@@ -943,35 +943,23 @@ ncclResult_t initTransportsRank_3(struct ncclComm* comm, struct allGather3Data_t
   NCCLCHECKGOTO(ncclCalloc(&rings, nranks*MAXCHANNELS), ret, fail);
   NCCLCHECKGOTO(ncclTopoPostset(comm, nodesFirstRank, nodesTreePatterns, allTopoRanks, rings, &collNetGraph, nc), ret, fail);
 
-  if (comm->topo->pivotA2ANumBiRings == 3) {
-    NCCLCHECK(ncclTreeBasePostset(comm, &treeGraph));
-    NCCLCHECK(ncclBinaryTreePostset(comm, &treeGraph));
-  }
-
+  if (comm->topo->pivotA2ANumBiRings == 3) NCCLCHECK(ncclTreeBasePostset(comm, &treeGraph));
 
   // AllGather3 - end
 
   TRACE(NCCL_INIT, "rank %d nranks %d - BUILT %d TREES/RINGS", rank, nranks, comm->nChannels);
 
-  char line[1024], binline[1024];
+  char line[1024];
   line[0]='\0';
-  binline[0]='\0';
   for (int c=0; c<comm->nChannels; c++) {
     struct ncclTree* tree = &comm->channels[c].tree;
-    struct ncclTree* binTree = &comm->channels[c].binTree;
     snprintf(line+strlen(line), 1023-strlen(line), " [%d] %d/%d/%d->%d->%d",
         c, tree->down[0], tree->down[1], tree->down[2], rank, tree->up);
-    if (comm->topo->pivotA2ANumBiRings == 3)
-      snprintf(binline+strlen(binline), 1023-strlen(binline), " [%d] %d/%d/%d->%d->%d",
-	       c, binTree->down[0], binTree->down[1], binTree->down[2], rank, binTree->up);
-    INFO(NCCL_GRAPH, "Ring %d : %d -> %d -> %d", c, comm->channels[c].ring.prev, comm->rank, comm->channels[c].ring.next);
+    INFO(NCCL_GRAPH, "Ring %d : %d -> %d -> %d comm %p nRanks %02d busId %lx", c, comm->channels[c].ring.prev,
+         comm->rank, comm->channels[c].ring.next, comm, comm->nRanks, comm->busId);
   }
   line[1023] = '\0';
-  INFO(NCCL_INIT, "Trees%s", line);
-  if (comm->topo->pivotA2ANumBiRings == 3) {
-    binline[1023] = '\0';
-    INFO(NCCL_INIT, "BinTrees%s", binline);
-  }
+  INFO(NCCL_INIT, "Trees%s comm %p nRanks %02d busId %lx", line, comm, comm->nRanks, comm->busId);
 
   //NCCLCHECKGOTO(computeBuffSizes(comm), ret, fail);
 
@@ -1023,6 +1011,9 @@ ncclResult_t initTransportsRank_3(struct ncclComm* comm, struct allGather3Data_t
 
   // Compute nChannels per peer for p2p
   NCCLCHECKGOTO(ncclTopoComputeP2pChannels(comm), ret, fail);
+
+  INFO(NCCL_INIT, "%d coll channels, %d nvls channels, %d p2p channels, %d p2p channels per peer", comm->nChannels, comm->nvlsChannels, comm->p2pnChannels, comm->p2pnChannelsPerPeer);
+
 #if 0
   do { // Setup p2p structures in comm->tasks
     struct ncclTasks* tasks = &comm->tasks;