From 1d55fe756c87dc60abdce817118bbe0ff126c1b6 Mon Sep 17 00:00:00 2001
From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com>
Date: Fri, 19 Aug 2022 13:51:49 -0600
Subject: [PATCH] Simple tree changes (#599)

changed treebase to create basic balanced tree

[ROCm/rccl commit: 3c1b1ec8c84d0f26271a31b5c3d9d5f5c093c442]
---
 projects/rccl/src/enqueue.cc       |  7 +--
 projects/rccl/src/graph/connect.cc | 92 +++++++++---------------------
 2 files changed, 30 insertions(+), 69 deletions(-)

diff --git a/projects/rccl/src/enqueue.cc b/projects/rccl/src/enqueue.cc
index 12281cb402..4348b9dc71 100644
--- a/projects/rccl/src/enqueue.cc
+++ b/projects/rccl/src/enqueue.cc
@@ -464,13 +464,10 @@ static ncclResult_t getAlgoInfo(struct ncclInfo* info, int collNetTypeSupport, i
     info->nChannels = nc;
     if (!userTuneInput) {
       // always respect user settings
-      if (info->nBytes <= 900000) {
+      if (info->nBytes <= 2200008) {
         info->protocol = NCCL_PROTO_LL;
         info->algorithm = NCCL_ALGO_TREE;
-        info->nChannels = std::min(comm->nChannels, info->nBytes <= 16384? 4 : 24);
-      } else if (info->nBytes <= 2200008) {
-        info->protocol = NCCL_PROTO_LL;
-        info->algorithm = NCCL_ALGO_RING;
+        info->nChannels = 24;
       } else {
         info->protocol = NCCL_PROTO_SIMPLE;
         info->algorithm = NCCL_ALGO_RING;
diff --git a/projects/rccl/src/graph/connect.cc b/projects/rccl/src/graph/connect.cc
index bdf727ba3b..94e32f560b 100644
--- a/projects/rccl/src/graph/connect.cc
+++ b/projects/rccl/src/graph/connect.cc
@@ -72,80 +72,44 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm,
 ncclResult_t ncclTreeBasePostset(struct ncclComm* comm,
     struct ncclTopoGraph* treeGraph) {
   int nChannels = comm->nChannels;
-  int ring[NCCL_TOPO_MAX_NODES][NCCL_TOPO_MAX_NODES];
-  int xLimit = 0, yLimit = 0;
-  for (int j=0; j < NCCL_TOPO_MAX_NODES; j++) {
-    if (treeGraph->treeBase[0][j]  == -1) {
-      xLimit = j;
-      break;
-    }
+  int localRanks = 0;
+  for (int i=0; i<comm->topo->nodes[GPU].count; i++) {
+    localRanks += comm->topo->nodes[GPU].nodes[i].gpu.nRanksPerGpu;
   }
-  for (int k=0; k < NCCL_TOPO_MAX_NODES; k++) {
-    if (treeGraph->treeBase[k][0]  == -1) {
-      yLimit = k;
-      break;
-    }
-  }
-  for (int j=0; j < xLimit; j++) {
-    for (int k=0; k < yLimit; k++)
-      ring[k][j] = treeGraph->treeBase[k][j];
-  }
-
   //new tree
   for (int c=0; c<nChannels; c++) {
+    int* treeIntra = treeGraph->intra+c%3*localRanks;
+    int buff = ((c%6)*localRanks)/6 + ((localRanks/4)*(c/6));
+
+    int tempArray[256];
+    for (int ko=0; ko < localRanks; ko++){
+      tempArray[ko] = treeIntra[(ko+buff)%localRanks];
+    }
     struct ncclChannel* channel = comm->channels+c;
-    int ringPrev, ringNext;
-    int treeRoot = c%comm->nRanks;
+
     int curRank = comm->rank;
-    int curRankNeighborUp, curRankNeighborDown;
-    int rootRing, nextRing;
-    int rootIndex = 0;
     int arrayIndex;
-    for (int j=0; j < xLimit; j++) {
-      for (int k=0; k < yLimit; k++) {
-        if (treeRoot == ring[k][j]) {
-          rootRing = k;
-          rootIndex = j;
-        }
-        if (curRank == ring[k][j]) {
-          arrayIndex = j;
-          if (k > 0) curRankNeighborUp=ring[k-1][j];
-          else curRankNeighborUp=ring[yLimit-1][j];
-          if (k < yLimit-1) {
-            curRankNeighborDown=ring[k+1][j];
-            nextRing = k+1;
-          }
-          else {
-            curRankNeighborDown = ring[0][j];
-            nextRing = 0;
-          }
-        }
-      }
-    }
 
-    if ((curRank != ring[rootRing][arrayIndex])) {
-      channel->tree.up = curRankNeighborUp;
-      channel->tree.down[0] = nextRing==rootRing ? -1 : curRankNeighborDown;
-      channel->tree.down[1] = -1;
-    }
-    else {
-      if (arrayIndex > 0) ringPrev=ring[rootRing][arrayIndex-1];
-      else ringPrev=ring[rootRing][xLimit-1];
-      if (arrayIndex < xLimit-1) ringNext=ring[rootRing][arrayIndex+1];
-      else ringNext=ring[rootRing][0];
-      if ((c/2)%2 == 1) {
-        int temp = ringPrev;
-        ringPrev = ringNext;
-        ringNext = temp;
+    for (int i=0; i<localRanks; i++) {
+      if (tempArray[i] == curRank) {
+        if (i == 0) {
+          channel->tree.up = -1;
+          channel->tree.down[0] = tempArray[i+1];
+          channel->tree.down[1] = tempArray[localRanks-1];
+          channel->tree.down[2] = -1;
+        }
+        else {
+          channel->tree.up = i > localRanks/2 ?  tempArray[(i+1)%localRanks] : tempArray[i-1];
+          channel->tree.down[0] = i > localRanks/2 ?  tempArray[i-1] : tempArray[i+1];
+          if ((i == localRanks/2) || (i == (localRanks/2 + 1))) {
+            channel->tree.down[0] = -1;
+          }
+          channel->tree.down[1] = -1;
+          channel->tree.down[2] = -1;
+        }
       }
-      channel->tree.up = treeRoot == curRank ? -1 : ringPrev;
-      channel->tree.down[0] = curRankNeighborDown;
-      channel->tree.down[1] = ringNext == treeRoot ? -1 : ringNext;
     }
-    channel->tree.down[2] = -1; // cleanup
   }
-
-
   return ncclSuccess;
 }