From b46260260a69c6d79c5e38894dc8239b03c81e00 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Tue, 16 Mar 2021 16:51:32 -0700
Subject: [PATCH] Sort GPUs by HIP device ID (#329)

* Sort GPUs by HIP device ID

* Remove extra space
---
 src/graph/search.cc | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/graph/search.cc b/src/graph/search.cc
index d52dbe1315..2c00068eb9 100644
--- a/src/graph/search.cc
+++ b/src/graph/search.cc
@@ -864,15 +864,15 @@ static ncclResult_t parseChordalRing(struct ncclTopoSystem* system, struct ncclT
   return ncclSuccess;
 }
 
-static ncclResult_t ncclGpuIdToIndex(struct ncclTopoSystem* system, int id, int* index) {
-  *index = -1;
-  for (int i=0; i<system->nodes[GPU].count; i++) {
-    if (system->nodes[GPU].nodes[i].gpu.dev == id) {
-      *index = i;
-      return ncclSuccess;
-    }
-  }
-  return ncclInternalError;
+struct ncclGpuIdHIP {
+  int g;
+  int dev;
+};
+
+static int cmpIds(const void * g1, const void * g2) {
+  struct ncclGpuIdHIP *s1 = (struct ncclGpuIdHIP*)g1;
+  struct ncclGpuIdHIP *s2 = (struct ncclGpuIdHIP*)g2;
+  return s1->dev - s2->dev;
 }
 
 static ncclResult_t parseRomeSystem(struct ncclTopoSystem* system, struct rcclRomeModel* romeTopo, char *pattern, int *net_map) {
@@ -881,9 +881,16 @@ static ncclResult_t parseRomeSystem(struct ncclTopoSystem* system, struct rcclRo
   romeTopo->nCpus = system->nodes[CPU].count;
   romeTopo->nNics = 0;
   romeTopo->nLinks = 0;
+  // sort GPU devices by HIP device ID
+  struct ncclGpuIdHIP scores[MAX_ROME_GPUS];
+  for (int i = 0; i < romeTopo->nGpus; i ++) {
+    scores[i].g = i;
+    scores[i].dev = system->nodes[GPU].nodes[i].gpu.dev;
+  }
+  qsort(scores, romeTopo->nGpus, sizeof(struct ncclGpuIdHIP), cmpIds);
   for (int i = 0; i < romeTopo->nGpus; i ++) {
     int gpu, n, m, distance;
-    NCCLCHECK(ncclGpuIdToIndex(system, i, &gpu));
+    gpu = scores[i].g;
     romeTopo->gpuIds[i] = system->nodes[GPU].nodes[gpu].id;
     m = 0;
     distance = system->nodes[GPU].nodes[gpu].paths[CPU][m].count;