From 78fc48f1c80686eacee31fc5b44aa69c3fcc1416 Mon Sep 17 00:00:00 2001
From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com>
Date: Wed, 14 Dec 2022 16:10:14 -0700
Subject: [PATCH] added a different way for getting device count, by running it
 in a child process (#665)

[ROCm/rccl commit: 24aa8bd8025f3abc27295dab812f500ebac6d0a4]
---
 projects/rccl/test/common/EnvVars.cpp | 43 ++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/projects/rccl/test/common/EnvVars.cpp b/projects/rccl/test/common/EnvVars.cpp
index 3e557cec77..f6815c554e 100644
--- a/projects/rccl/test/common/EnvVars.cpp
+++ b/projects/rccl/test/common/EnvVars.cpp
@@ -7,30 +7,51 @@
 #include "EnvVars.hpp"
 #include "CollectiveArgs.hpp"
 #include <cstdlib>
+#include <unistd.h>
+#include <sys/wait.h>
 
 namespace RcclUnitTesting
 {
   int const UT_SINGLE_PROCESS = (1<<0);
   int const UT_MULTI_PROCESS  = (1<<1);
 
-  hsa_status_t CountGpus(hsa_agent_t agent, void* data)
+  int getDeviceCount(int *devices)
   {
-    int* currCount = (int*)data;
-    hsa_device_type_t device;
-    hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &device);
-    if (device == HSA_DEVICE_TYPE_GPU)
-      *currCount = *currCount + 1;
-    return HSA_STATUS_SUCCESS;
+    // Prepare parent->child pipe
+    int pipefd[2];
+    if (pipe(pipefd) == -1)
+    {
+      ERROR("Unable to create parent->child pipe for getting number of devices\n");
+      return TEST_FAIL;
+    }
+    pid_t pid = fork();
+    if (0 == pid)
+    {
+      int dev;
+      hipGetDeviceCount(&dev);
+      if (write(pipefd[1], &dev, sizeof(dev)) != sizeof(dev)) return TEST_FAIL;
+      close(pipefd[0]);
+      close(pipefd[1]);
+      exit(EXIT_SUCCESS);
+    }
+    else
+    {
+      int status;
+      if (read(pipefd[0], devices, sizeof(*devices)) != sizeof(*devices)) return TEST_FAIL;
+      waitpid(pid, &status, 0);
+      assert(!status);
+      close(pipefd[0]);
+      close(pipefd[1]);
+    }
+    return TEST_SUCCESS;
   }
 
   EnvVars::EnvVars()
   {
     // Collect number of GPUs available
-    // NOTE: Cannot use HIP call prior to launching child processes via fork so use HSA
+    // NOTE: Cannot use HIP call prior to launching unless it is inside another child process
     int numDevicesAvailable = 0;
-    hsa_init();
-    hsa_iterate_agents(CountGpus, &numDevicesAvailable);
-    hsa_shut_down();
+    getDeviceCount(&numDevicesAvailable);
 
     showNames      = GetEnvVar("UT_SHOW_NAMES"  , 1);
     minGpus        = GetEnvVar("UT_MIN_GPUS"    , 2);