From aa5f75aa9c24837b8490774588cc6dcaca2a3f60 Mon Sep 17 00:00:00 2001
From: Pedram Alizadeh <pmohamma@amd.com>
Date: Thu, 27 Apr 2023 14:06:17 -0400
Subject: [PATCH] fixing the error message for mpirun when number of requested
 GPUs exceeds the limits (#37)

[ROCm/rccl-tests commit: d16d1fb16b2abe1c1c88464097e6f1d8070d1116]
---
 projects/rccl-tests/src/common.cu | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/projects/rccl-tests/src/common.cu b/projects/rccl-tests/src/common.cu
index 332cc3f272..7107d8b4e6 100644
--- a/projects/rccl-tests/src/common.cu
+++ b/projects/rccl-tests/src/common.cu
@@ -1127,11 +1127,13 @@ int main(int argc, char* argv[]) {
   }
 
   HIPCHECK(hipGetDeviceCount(&numDevices));
+#ifndef MPI_SUPPORT
   if (nGpus > numDevices)
   {
       fprintf(stderr, "[ERROR] The number of requested GPUs (%d) is greater than the number of GPUs available (%d)\n", nGpus, numDevices);
       return testNcclError;
   }
+#endif
   if (minBytes > maxBytes) {
     fprintf(stderr, "invalid sizes for 'minbytes' and 'maxbytes': %llu > %llu\n",
            (unsigned long long)minBytes,
@@ -1163,6 +1165,7 @@ int main(int argc, char* argv[]) {
 testResult_t run() {
   int nProcs = 1, proc = 0;
   int localRank = 0;
+  int localSize = 0;
   char hostname[1024];
   getHostName(hostname, 1024);
 
@@ -1176,6 +1179,14 @@ testResult_t run() {
     if (p == proc) break;
     if (hostHashs[p] == hostHashs[proc]) localRank++;
   }
+  for (int p=0; p<nProcs; p++) {
+    if (hostHashs[p] == hostHashs[proc]) localSize++;
+  }
+  if (nGpus * localSize > numDevices)
+  {
+      fprintf(stderr, "[ERROR] The number of requested GPUs (%d) is greater than the number of GPUs available (%d) on node (%s)\n", nGpus*localSize, numDevices, hostname);
+      return testNcclError;
+  }
 #endif
   is_main_thread = (proc == 0) ? 1 : 0;