Introduce multi-rank support per device.

This is a single commit of the source code changes required to
introduce support for multiple ranks per device.
A new interface (ncclCommRankInitMulti) has to be used to make use of
this new feature.

src/graph/connect.cc

@@ -19,8 +19,11 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm,
     struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph,
     struct ncclTopoRanks* topoRanks) {
   int rank = comm->rank;
-  int localRanks = comm->topo->nodes[GPU].count;
   int nChannels = comm->nChannels;
+  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 c=0; c<nChannels; c++) {
     struct ncclChannel* channel = comm->channels+c;

src/graph/paths.cc

@@ -268,7 +268,7 @@ ncclResult_t ncclTopoCheckP2p(struct ncclTopoSystem* system, int64_t id1, int64_
     struct ncclTopoNode* intermediateNode = path->list[0]->remNode;
     if (intermediateNode->type == GPU) {
       intermediateIndex = intermediateNode - system->nodes[GPU].nodes;
-      if (intermediateRank) *intermediateRank = intermediateNode->gpu.rank;
+      if (intermediateRank) *intermediateRank = intermediateNode->gpu.rank[0];
     }
   }
 
@@ -404,7 +404,7 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* system, int64_t busId, int
   if (distance == PATH_PXN) {
     // In case of PXN, use the intermediate GPU distance instead
     int proxyRank, g;
-    NCCLCHECK(ncclTopoGetIntermediateRank(system, gpu->gpu.rank, netDev, &proxyRank));
+    NCCLCHECK(ncclTopoGetIntermediateRank(system, gpu->gpu.rank[0], netDev, &proxyRank));
     NCCLCHECK(ncclTopoRankToIndex(system, proxyRank, &g));
     struct ncclTopoNode* proxyGpu = system->nodes[GPU].nodes+g;
     distance = proxyGpu->paths[NET][n].type;
@@ -437,7 +437,7 @@ ncclResult_t ncclTopoGetIntermediateRank(struct ncclTopoSystem* system, int rank
       WARN("Could not find intermediate GPU between GPU rank %d and NIC %d\n", rank, netDev);
       return ncclInternalError;
     }
-    *intermediateRank = node->gpu.rank;
+    *intermediateRank = node->gpu.rank[0];
   } else {
     *intermediateRank = rank;
   }
@@ -520,10 +520,10 @@ ncclResult_t ncclTopoComputePaths(struct ncclTopoSystem* system, struct ncclPeer
 
     if (peerInfos == NULL) continue;
     // Remove GPUs we can't talk to because of containers.
-    struct ncclPeerInfo* dstInfo = peerInfos+system->nodes[GPU].nodes[g].gpu.rank;
+    struct ncclPeerInfo* dstInfo = peerInfos+system->nodes[GPU].nodes[g].gpu.rank[0];
     for (int p=0; p<system->nodes[GPU].count; p++) {
       if (p == g) continue;
-      struct ncclPeerInfo* srcInfo = peerInfos+system->nodes[GPU].nodes[p].gpu.rank;
+      struct ncclPeerInfo* srcInfo = peerInfos+system->nodes[GPU].nodes[p].gpu.rank[0];
       int shm;
       NCCLCHECK(ncclTransports[TRANSPORT_SHM].canConnect(&shm, system, NULL, srcInfo, dstInfo));
       int p2p;
@@ -556,7 +556,8 @@ ncclResult_t ncclTopoComputePaths(struct ncclTopoSystem* system, struct ncclPeer
           pxnGpu = p;
 
           int netDev;
-          NCCLCHECK(ncclTopoGetLocalNet(system, peerNode->gpu.rank, &netDev));
+
+          NCCLCHECK(ncclTopoGetLocalNet(system, peerNode->gpu.rank[0], &netDev));
           // To ensure proper balancing, use preferably a local GPU which advertised that NIC as its preferred one.
           if (netDev == netNode->id) break;
         }
@@ -599,7 +600,12 @@ ncclResult_t ncclTopoTrimSystem(struct ncclTopoSystem* system, struct ncclComm*
         domains[g] = std::min(domains[g], domains[p]);
       }
     }
-    if (gpu->gpu.rank == comm->rank) myDomain = domains[g];
+    for (int j=0; j<gpu->gpu.nRanksPerGpu; j++ ) {
+      if (gpu->gpu.rank[j] == comm->rank) {
+	myDomain = domains[g];
+	break;
+      }
+    }
   }
 
   int ngpus = system->nodes[GPU].count;
@@ -650,7 +656,7 @@ ncclResult_t ncclTopoTrimSystem(struct ncclTopoSystem* system, struct ncclComm*
     int gdr, ret = 1;
     int net;
     for (int g = 0; g < system->nodes[GPU].count; g++) {
-      NCCLCHECK(ncclTopoGetLocalNet(system, system->nodes[GPU].nodes[g].gpu.rank, &net));
+      NCCLCHECK(ncclTopoGetLocalNet(system, system->nodes[GPU].nodes[g].gpu.rank[0], &net));
       NCCLCHECK(ncclTopoCheckGdr(system, system->nodes[GPU].nodes[g].id, net, 1, &gdr));
       if (!gdr) {
         ret = 0;
@@ -677,12 +683,16 @@ ncclResult_t ncclTopoTrimSystem(struct ncclTopoSystem* system, struct ncclComm*
       }
     }
   }
+
   if (rcclParamEnableIntranet()) {
     remove = 0;
     system->type |= RCCL_TOPO_FORCE_INTRA;
   }
-  comm->localRanks = system->nodes[GPU].count;
-  if (system->nodes[GPU].count == comm->nRanks && remove) {
+  comm->localRanks = 0;
+  for (int n=0; n<system->nodes[GPU].count; n++ ) {
+    comm->localRanks += system->nodes[GPU].nodes[n].gpu.nRanksPerGpu;
+  }
+  if (comm->localRanks == comm->nRanks && remove) {
     for (int n=system->nodes[NET].count-1; n>=0; n--)
       NCCLCHECK(ncclTopoRemoveNode(system, NET, n));
   }
@@ -777,10 +787,14 @@ ncclResult_t ncclTopoGetNvbGpus(struct ncclTopoSystem* system, int rank, int* nr
   int nvbGpus = 0;
   for (int g=0; g<ngpus; g++) {
     struct ncclTopoNode* gpu = system->nodes[GPU].nodes+g;
-    if (gpu->gpu.rank != rank) continue;
+    int j=0;
+    for ( ; j<gpu->gpu.nRanksPerGpu; j++ ){
+      if (gpu->gpu.rank[j] == rank) break;
+    }
+    if ( j == gpu->gpu.nRanksPerGpu ) continue;
     for (int p=0; p<ngpus; p++) {
       if (gpu->paths[GPU][p].type == PATH_NVB) {
-        (*ranks)[nvbGpus++] = system->nodes[GPU].nodes[p].gpu.rank;
+        (*ranks)[nvbGpus++] = system->nodes[GPU].nodes[p].gpu.rank[j];
       }
     }
   }

src/graph/rome_models.cc

@@ -691,7 +691,7 @@ ncclResult_t parseGraph(const char* str, struct ncclTopoSystem* system, struct n
               if (g == system->nodes[GPU].nodes[j].gpu.dev)
                 break;
             if (j < ngpus)
-              graph->intra[nChannels*ngpus+r] = system->nodes[GPU].nodes[j].gpu.rank;
+              graph->intra[nChannels*ngpus+r] = system->nodes[GPU].nodes[j].gpu.rank[0];
             else
               return ncclInternalError;
           }
@@ -781,7 +781,7 @@ ncclResult_t parseGraphLight(const char* str, struct ncclTopoSystem* system, str
               break;
           if (j < ngpus)
           {
-            graph->treeBase[r][x] = system->nodes[GPU].nodes[j].gpu.rank;
+            graph->treeBase[r][x] = system->nodes[GPU].nodes[j].gpu.rank[0];
             y=r;
           }
           else
@@ -913,15 +913,15 @@ ncclResult_t parseChordalRing(struct ncclTopoSystem* system, struct ncclTopoGrap
       // find the first unsed GPU that is closest to NIC
       int f, m;
       for (f = 0; f < ngpus; f++) {
-        int j = 0; for (j = 0; j < n; j++) if(used[j] == system->nodes[GPU].nodes[f].gpu.rank) break;
+        int j = 0; for (j = 0; j < n; j++) if(used[j] == system->nodes[GPU].nodes[f].gpu.rank[0]) break;
         if(j >= n) break;
       }
       for (int i = 0; i < ngpus; i++) {
-        int j = 0; for (j = 0; j < n; j++) if(used[j] == system->nodes[GPU].nodes[i].gpu.rank) break;
+        int j = 0; for (j = 0; j < n; j++) if(used[j] == system->nodes[GPU].nodes[i].gpu.rank[0]) break;
         if (j < n) continue;
         if (paths[i].count < paths[f].count) f = i;
       }
-      for (m = 0; m<ngpus; m++) if (graph->intra[n*ngpus+m] == system->nodes[GPU].nodes[f].gpu.rank) break;
+      for (m = 0; m<ngpus; m++) if (graph->intra[n*ngpus+m] == system->nodes[GPU].nodes[f].gpu.rank[0]) break;
       used[n] = graph->intra[n*ngpus+m];
       for (int i = 0; i < ngpus; i++) intra[i] = graph->intra[n*ngpus+((i+m)%ngpus)];
       for (int i = 0; i < ngpus; i++) graph->intra[n*ngpus+i] = intra[i];

src/graph/search.cc

@@ -186,9 +186,11 @@ static int cmpIntraScores(struct ncclGpuScore* scores, int count) {
 
 static ncclResult_t getGpuIndex(struct ncclTopoSystem* system, int rank, int* index) {
   for (int g=0; g<system->nodes[GPU].count; g++) {
-    if (system->nodes[GPU].nodes[g].gpu.rank == rank) {
-      *index = g;
-      return ncclSuccess;
+    for (int j=0; j<system->nodes[GPU].nodes[g].gpu.nRanksPerGpu; j++) {
+      if (system->nodes[GPU].nodes[g].gpu.rank[j] == rank) {
+	*index = g;
+	return ncclSuccess;
+      }
     }
   }
   WARN("Could not find gpu rank %d", rank);
@@ -270,9 +272,13 @@ ncclResult_t ncclTopoReplayGetGpu(struct ncclTopoSystem* system, struct ncclTopo
   if (graph->nChannels == 0) return ncclInternalError;
   int ngpus = system->nodes[GPU].count;
   int nextRank = graph->intra[(graph->nChannels-1)*ngpus+step+1];
-  for (int i=0; i<ngpus; i++) if (system->nodes[GPU].nodes[i].gpu.rank == nextRank) {
-    *g = i;
-    return ncclSuccess;
+  for (int i=0; i<ngpus; i++) {
+    for (int j=0; j<system->nodes[GPU].nodes[i].gpu.nRanksPerGpu; j++ ) {
+      if (system->nodes[GPU].nodes[i].gpu.rank[j] == nextRank) {
+	*g = i;
+	return ncclSuccess;
+      }
+    }
   }
   if (*g == -1) return ncclInternalError;
   return ncclSuccess;
@@ -302,18 +308,26 @@ static int ncclTopoCountXGMI(struct ncclTopoSystem* system, struct ncclTopoGraph
       int n = graph->intra[ngpus*c+((i+1)%ngpus)];
       struct ncclTopoNode *node;
       int j;
-      for (j=0; j<ngpus; j++)
-        if (system->nodes[GPU].nodes[j].gpu.rank == g) break;
+      for (j=0; j<ngpus; j++) {
+	bool found=false;
+	for (int k=0; k<system->nodes[GPU].nodes[j].gpu.nRanksPerGpu; k++) {
+	  if (system->nodes[GPU].nodes[j].gpu.rank[k] == g)
+	    found = true;
+	}
+	if (found) break;
+      }
       if (j<ngpus) {
         node = system->nodes[GPU].nodes+j;
         for (int k = 0; k<system->nodes[GPU].count; k++) {
           if (node->paths[GPU][k].count == 1) {
             struct ncclTopoLink* link = node->paths[GPU][k].list[0];
             struct ncclTopoNode* remNode = link->remNode;
-            if (remNode->gpu.rank == n) {
-              if (link->type == LINK_NVL)
-                count ++;
-            }
+	    for (int l=0; l<remNode->gpu.nRanksPerGpu; l++) {
+	      if (remNode->gpu.rank[l] == n) {
+		if (link->type == LINK_NVL)
+		  count ++;
+	      }
+	    }
           }
         }
       }
@@ -412,7 +426,7 @@ ncclResult_t ncclTopoSearchRecGpu(struct ncclTopoSystem* system, struct ncclTopo
     graph->nChannels--;
     return ncclSuccess;
   }
-  graph->intra[graph->nChannels*ngpus+step] = gpu->gpu.rank;
+  graph->intra[graph->nChannels*ngpus+step] = gpu->gpu.rank[0];
   int g = gpu - system->nodes[GPU].nodes;
   if (step == backToNet) {
     // first get back to NIC
@@ -669,7 +683,7 @@ ncclResult_t ncclTopoGetChannelFromXml(struct ncclXmlNode *xmlChannel, int c, st
     } else if (strcmp(sub->name, "gpu") == 0) {
       int rank = -1;
       for (int g=0; g<ngpus; g++) {
-        if (system->nodes[GPU].nodes[g].gpu.dev == dev) rank = system->nodes[GPU].nodes[g].gpu.rank;
+        if (system->nodes[GPU].nodes[g].gpu.dev == dev) rank = system->nodes[GPU].nodes[g].gpu.rank[0];
       }
       if (rank == -1) {
         WARN("XML Import Channel : dev %d not found.", dev);
@@ -730,7 +744,9 @@ ncclResult_t ncclTopoGetXmlFromChannel(struct ncclTopoGraph* graph, int c, struc
     NCCLCHECK(xmlAddNode(xml, xmlChannel, "gpu", &node));
     int dev = -1;
     for (int i=0; i<ngpus; i++) {
-      if (system->nodes[GPU].nodes[i].gpu.rank == intra[g]) dev = system->nodes[GPU].nodes[i].gpu.dev;
+       for ( int j=0; j<system->nodes[GPU].nodes[i].gpu.nRanksPerGpu; j++ ) {
+	 if (system->nodes[GPU].nodes[i].gpu.rank[j] == intra[g]) dev = system->nodes[GPU].nodes[i].gpu.dev;
+       }
     }
     if (dev == -1) {
       WARN("XML Export Channel : rank %d not found.", intra[g]);
@@ -789,6 +805,27 @@ float speedArrayInter[] = { 48.0, 30.0, 24.0, 22.0, 18.0, 15.0, 12.0, 10.0, 9.0,
 RCCL_PARAM(ModelMatchingDisable, "MODEL_MATCHING_DISABLE", 0);
 NCCL_PARAM(CrossNic, "CROSS_NIC", 2);
 
+static void ncclExpandMultiRank(ncclTopoSystem* system, struct ncclTopoGraph* graph)
+{
+  // Expand the intra array to the multi-ranks per node scenario
+  int ngpus = system->nodes[GPU].count;
+  int intraCpy[MAXCHANNELS*NCCL_TOPO_MAX_NODES];
+  TRACE(NCCL_GRAPH, "TopoCompute: expanding intra array for multi-rank per GPU scenarios nChannels %d", graph->nChannels);
+  memcpy(intraCpy, graph->intra, ngpus*sizeof(int)*graph->nChannels);
+  int tk=0;
+  for (int n=0; n<graph->nChannels; n++ ) {
+    for (int i=0; i<ngpus; i++) {
+      for (int j=0; j<ngpus; j++) {
+	if (intraCpy[n*ngpus+i] == system->nodes[GPU].nodes[j].gpu.rank[0] ) {
+	  for (int k=0; k<system->nodes[GPU].nodes[j].gpu.nRanksPerGpu; k++) {
+	    graph->intra[tk++] = system->nodes[GPU].nodes[j].gpu.rank[k];
+	  }
+	}
+      }
+    }
+  }
+}
+
 ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph) {
   int ngpus = system->nodes[GPU].count;
   graph->crossNic = ncclParamCrossNic();
@@ -813,7 +850,10 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
     NCCLCHECK(ncclTopoGetGraphFromXml(xml->nodes, system, graph, &nChannels));
     INFO(NCCL_GRAPH, "Search %d : %d channels loaded from XML graph", graph->id, nChannels);
     free(xml);
-    if (graph->nChannels > 0) return ncclSuccess;
+    if (graph->nChannels > 0) {
+      ncclExpandMultiRank(system, graph);
+      return ncclSuccess;
+    }
   }
 
   str = getenv("NCCL_RINGS");
@@ -826,17 +866,29 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
   } else if (!rcclParamModelMatchingDisable() && !graph->collNet) {
     // try to match 8P6L
     NCCLCHECK(parseChordalRing(system, graph));
-    if (graph->nChannels) return ncclSuccess;
+    if (graph->nChannels) {
+      ncclExpandMultiRank(system, graph);
+      return ncclSuccess;
+    }
     // try to match Rome 4P2H
     NCCLCHECK(parseRome4P2H(system, graph));
-    if (graph->nChannels) return ncclSuccess;
+    if (graph->nChannels) {
+      ncclExpandMultiRank(system, graph);
+      return ncclSuccess;
+    }
     // try to match 1H16P
     NCCLCHECK(parse1H16P(system, graph));
-    if (graph->nChannels) return ncclSuccess;
+    if (graph->nChannels) {
+      ncclExpandMultiRank(system, graph);
+      return ncclSuccess;
+    }
     // try to match 4H4P
     NCCLCHECK(parse4H4P(system, graph));
   }
-  if (graph->nChannels) return ncclSuccess;
+  if (graph->nChannels) {
+    ncclExpandMultiRank(system, graph);
+      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
@@ -866,7 +918,6 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
   while (speedArray[speedIndex] > system->maxWidth && speedIndex < nspeeds-1) speedIndex++;
   tmpGraph.speedIntra = tmpGraph.speedInter = speedArray[speedIndex];
   int64_t globalTimeout = NCCL_SEARCH_GLOBAL_TIMEOUT;
-
 search:
   int time = tmpGraph.sameChannels ? NCCL_SEARCH_TIMEOUT_SAMECHANNELS :
     tmpGraph.pattern == NCCL_TOPO_PATTERN_TREE ? NCCL_SEARCH_TIMEOUT_TREE : NCCL_SEARCH_TIMEOUT;
@@ -902,7 +953,6 @@ search:
     if (time != -1) globalTimeout += time;
     else globalTimeout = NCCL_SEARCH_GLOBAL_TIMEOUT;
     if (globalTimeout < 0 && graph->nChannels) goto done;
-
     int maxTypeIntra = system->nodes[NET].count > 0 ? tmpGraph.typeInter : PATH_SYS;
     if (tmpGraph.typeIntra < maxTypeIntra && (graph->nChannels == 0 || tmpGraph.typeIntra < graph->typeIntra)) {
       tmpGraph.typeIntra += 1;
@@ -967,13 +1017,12 @@ done:
 
   if (graph->nChannels == 0 && graph->collNet == 0) {
     WARN("Could not find a path for pattern %d, falling back to simple order", graph->pattern);
-    for (int i=0; i<ngpus; i++) graph->intra[i] = system->nodes[GPU].nodes[i].gpu.rank;
+    for (int i=0; i<ngpus; i++) graph->intra[i] = system->nodes[GPU].nodes[i].gpu.rank[0];
     graph->inter[0] = graph->inter[1] = 0;
     graph->speedIntra = graph->speedInter = 0.1;
     graph->typeIntra = graph->typeInter = PATH_SYS;
     graph->nChannels = 1;
   }
-
   if (graph->speedIntra >= 25.0) {
     int dupChannels = std::min(graph->nChannels*2, graph->maxChannels);
     memcpy(graph->intra+graph->nChannels*ngpus, graph->intra, (dupChannels-graph->nChannels)*ngpus*sizeof(int));
@@ -983,6 +1032,7 @@ done:
     graph->speedInter /= DIVUP(dupChannels, graph->nChannels);
     graph->nChannels = dupChannels;
   }
+  ncclExpandMultiRank(system, graph);
   return ncclSuccess;
 }
 

src/graph/topo.cc

@@ -117,7 +117,10 @@ ncclResult_t ncclTopoCreateNode(struct ncclTopoSystem* system, struct ncclTopoNo
     n->links[0].remNode = n;
     n->links[0].width = LOC_WIDTH;
     n->gpu.dev = NCCL_TOPO_UNDEF;
-    n->gpu.rank = NCCL_TOPO_UNDEF;
+    for (int i=0; i<RCCL_TOPO_MAX_RANKS_PER_GPU; i++) {
+      n->gpu.rank[i] = NCCL_TOPO_UNDEF;
+    }
+    n->gpu.nRanksPerGpu = NCCL_TOPO_UNDEF;
     n->gpu.cudaCompCap = NCCL_TOPO_UNDEF;
   } else if (type == CPU) {
     n->cpu.arch = NCCL_TOPO_UNDEF;
@@ -253,7 +256,15 @@ ncclResult_t ncclTopoConnectCpus(struct ncclTopoSystem* system) {
 
 static ncclResult_t ncclTopoPrintRec(struct ncclTopoNode* node, struct ncclTopoNode* prevNode, char* line, int offset) {
   if (node->type == GPU) {
-    sprintf(line+offset, "%s/%lX (%d)", topoNodeTypeStr[node->type], node->id, node->gpu.rank);
+    sprintf(line+offset, "%s/%lX (%d", topoNodeTypeStr[node->type], node->id, node->gpu.rank[0]);
+    int nextOffset;
+    int nextRank = 1;
+    while ( nextRank < node->gpu.nRanksPerGpu ) {
+      nextOffset = strlen(line);
+      sprintf(line+nextOffset, "/%d", node->gpu.rank[nextRank++]);
+    }
+    nextOffset = strlen(line);
+    sprintf(line+nextOffset, ")");
   } else if (node->type == CPU) {
     sprintf(line+offset, "%s/%lX (%d/%d/%d)", topoNodeTypeStr[node->type], node->id, node->cpu.arch, node->cpu.vendor, node->cpu.model);
   } else if (node->type == PCI) {
@@ -373,7 +384,17 @@ ncclResult_t ncclTopoAddGpu(struct ncclXmlNode* xmlGpu, struct ncclTopoSystem* s
   rcclHipDeviceArch_t arch;
   NCCLCHECK(xmlGetAttrInt(xmlGpu, "arch", &arch.value));
   memcpy(&gpu->gpu.arch, &arch.arch, sizeof(hipDeviceArch_t));
-  NCCLCHECK(xmlGetAttrInt(xmlGpu, "rank", &gpu->gpu.rank));
+
+  //NCCLCHECK(xmlGetAttrInt(xmlGpu, "rank", &gpu->gpu.rank));
+  const char *rankStr;
+  NCCLCHECK(xmlGetAttrStr(xmlGpu, "rank", &rankStr));
+  char *tmpStr;
+  char *token = strtok_r ( (char *)rankStr, ",", &tmpStr);
+  gpu->gpu.nRanksPerGpu = 0;
+  while (token != NULL && gpu->gpu.nRanksPerGpu < RCCL_TOPO_MAX_RANKS_PER_GPU) {
+    gpu->gpu.rank[gpu->gpu.nRanksPerGpu++] = atoi(token);
+    token = strtok_r(NULL, ",", &tmpStr);
+  }
   NCCLCHECK(xmlGetAttrInt(xmlGpu, "dev", &gpu->gpu.dev));
   NCCLCHECK(xmlGetAttrInt(xmlGpu, "gdr", &gpu->gpu.gdrSupport));
   // Do not go any further, nvlinks will be added in a second pass
@@ -385,6 +406,7 @@ struct kvDict kvDictPciGen[] = {
   { "2.5 GT/s", 15 }, { "5 GT/s", 30 }, { "8 GT/s", 60 }, { "16 GT/s", 120 }, { "32 GT/s", 240 }, /* Kernel 5.6 and earlier */
   { "2.5 GT/s PCIe", 15 }, { "5.0 GT/s PCIe", 30 }, { "8.0 GT/s PCIe", 60 }, { "16.0 GT/s PCIe", 120 }, { "32.0 GT/s PCIe", 240 }, { "64.0 GT/s PCIe", 480 },
   { NULL, 60 /* Default fallback */ } }; // x100 Mbps per lane
+
 ncclResult_t ncclTopoAddPci(struct ncclXmlNode* xmlPci, struct ncclTopoSystem* system, struct ncclTopoNode* parent) {
   const char* str;
 
@@ -694,7 +716,8 @@ ncclResult_t ncclTopoGetSystem(struct ncclComm* comm, struct ncclTopoSystem** sy
       NCCLCHECK(ncclTopoFillGpu(xml, busId, &node));
       if (node == NULL) continue;
       NCCLCHECK(xmlSetAttrInt(node, "keep", 1));
-      NCCLCHECK(xmlSetAttrInt(node, "rank", r));
+      //NCCLCHECK(xmlSetAttrInt(node, "rank", r));
+      NCCLCHECK(xmlSetOrAppendAttrInt(node, "rank", r));
       NCCLCHECK(xmlInitAttrInt(node, "gdr", comm->peerInfo[r].gdrSupport));
     }
   }
@@ -795,18 +818,20 @@ NCCL_PARAM(IgnoreCpuAffinity, "IGNORE_CPU_AFFINITY", 0);
 ncclResult_t ncclTopoGetCpuAffinity(struct ncclTopoSystem* system, int rank, cpu_set_t* affinity) {
   struct ncclTopoNode* cpu = NULL, *gpu = NULL;
   for (int g=0; g<system->nodes[GPU].count; g++) {
-    if (system->nodes[GPU].nodes[g].gpu.rank == rank) {
-      gpu = system->nodes[GPU].nodes+g;
-      // Find closer CPU
-      int cpuIndex = -1, minHops = 0;
-      for (int c=0; c<system->nodes[CPU].count; c++) {
-        int nHops = system->nodes[GPU].nodes[g].paths[CPU][c].count;
-        if (cpuIndex == -1 || nHops < minHops) {
-          cpuIndex = c;
-          minHops = nHops;
-        }
+    for (int j=0; j<system->nodes[GPU].nodes[g].gpu.nRanksPerGpu; j++) {
+      if (system->nodes[GPU].nodes[g].gpu.rank[j] == rank) {
+	gpu = system->nodes[GPU].nodes+g;
+	// Find closer CPU
+	int cpuIndex = -1, minHops = 0;
+	for (int c=0; c<system->nodes[CPU].count; c++) {
+	  int nHops = system->nodes[GPU].nodes[g].paths[CPU][c].count;
+	  if (cpuIndex == -1 || nHops < minHops) {
+	    cpuIndex = c;
+	    minHops = nHops;
+	  }
+	}
+	cpu = system->nodes[CPU].nodes+cpuIndex;
       }
-      cpu = system->nodes[CPU].nodes+cpuIndex;
     }
   }
   if (cpu == NULL) {
@@ -876,9 +901,11 @@ ncclResult_t ncclTopoGetCompCap(struct ncclTopoSystem* system, int* ccMin, int*
 
 ncclResult_t ncclTopoGetLocalRank(struct ncclTopoSystem* system, int rank, int* localRank) {
   for (int g=0; g<system->nodes[GPU].count; g++) {
-    if (system->nodes[GPU].nodes[g].gpu.rank == rank) {
-      *localRank = g;
-      return ncclSuccess;
+    for ( int j=0; j<system->nodes[GPU].nodes[g].gpu.nRanksPerGpu; j++ ){
+      if (system->nodes[GPU].nodes[g].gpu.rank[j] == rank) {
+	*localRank = g;
+	return ncclSuccess;
+      }
     }
   }
   WARN("Could not find local GPU with rank %d\n", rank);

src/graph/topo.h

@@ -105,6 +105,7 @@ struct ncclTopoLinkList {
 #define RCCL_TOPO_16P1H     8
 #define RCCL_TOPO_FORCE_INTRA 16
 
+#define RCCL_TOPO_MAX_RANKS_PER_GPU 8
 struct ncclTopoNode {
   int type;
   int64_t id;
@@ -112,7 +113,8 @@ struct ncclTopoNode {
   union {
     struct {
       int dev; // NVML dev number
-      int rank;
+      int rank[RCCL_TOPO_MAX_RANKS_PER_GPU];
+      int nRanksPerGpu;
       int cudaCompCap;
       int gdrSupport;
       int gcn;
@@ -192,9 +194,11 @@ static ncclResult_t ncclTopoIdToIndex(struct ncclTopoSystem* system, int type, i
 static ncclResult_t ncclTopoRankToIndex(struct ncclTopoSystem* system, int rank, int* index) {
   *index = -1;
   for (int i=0; i<system->nodes[GPU].count; i++) {
-    if (system->nodes[GPU].nodes[i].gpu.rank == rank) {
-      *index = i;
-      return ncclSuccess;
+    for (int j=0; j<system->nodes[GPU].nodes[i].gpu.nRanksPerGpu; j++ ) {
+      if (system->nodes[GPU].nodes[i].gpu.rank[j] == rank) {
+	*index = i;
+	return ncclSuccess;
+      }
     }
   }
   return ncclInternalError;

src/graph/xml.h

@@ -176,6 +176,25 @@ static ncclResult_t xmlSetAttrInt(struct ncclXmlNode* node, const char* attrName
   return ncclSuccess;
 }
 
+static ncclResult_t xmlSetOrAppendAttrInt(struct ncclXmlNode* node, const char* attrName, const int value) {
+  int index;
+  NCCLCHECK(xmlGetAttrIndex(node, attrName, &index));
+  if (index == -1) {
+    index = node->nAttrs++;
+    strncpy(node->attrs[index].key, attrName, MAX_STR_LEN);
+    node->attrs[index].key[MAX_STR_LEN] = '\0';
+    snprintf(node->attrs[index].value, MAX_STR_LEN, "%d", value);
+    node->attrs[index].value[MAX_STR_LEN] = '\0';
+    return ncclSuccess;
+  }
+  char *tmp = strdup(node->attrs[index].value);
+  snprintf(node->attrs[index].value, MAX_STR_LEN, "%s,%d", tmp, value);
+  node->attrs[index].value[MAX_STR_LEN] = '\0';
+  free (tmp);
+  return ncclSuccess;
+}
+
+
 static ncclResult_t xmlSetAttrFloat(struct ncclXmlNode* node, const char* attrName, const float value) {
   int index;
   NCCLCHECK(xmlGetAttrIndex(node, attrName, &index));