From 256403d4f0685c9159e6064cd2e3e00972071d70 Mon Sep 17 00:00:00 2001
From: Stanley Tsang <stanley.tsang@amd.com>
Date: Tue, 25 May 2021 17:07:12 -0600
Subject: [PATCH] Adding support for hipMallocManaged() in unit tests (#375)

* Adding HMM support for unit tests

* Fixing HMM opt-in check
---
 test/CorrectnessTest.hpp | 92 ++++++++++++++++++++++++++++++++++------
 1 file changed, 80 insertions(+), 12 deletions(-)

diff --git a/test/CorrectnessTest.hpp b/test/CorrectnessTest.hpp
index 8894ee3767..ed7cfb44e9 100644
--- a/test/CorrectnessTest.hpp
+++ b/test/CorrectnessTest.hpp
@@ -100,6 +100,49 @@ namespace CorrectnessTests
             return numElements * DataTypeToBytes(dataType);
         }
 
+        // Checks if the current HIP Runtime and GPU support managed memory
+        bool SupportsHmm()
+        {
+            hipDeviceProp_t device_prop;
+            int device_id;
+            hipGetDevice(&device_id);
+            hipGetDeviceProperties(&device_prop, device_id);
+            if (device_prop.managedMemory == 1) return true;
+
+            return false;
+        }
+
+        // Check if user has opted-in to use managed memory
+        static bool UseHmm()
+        {
+            if (getenv("RCCL_USE_HMM") == nullptr)
+            {
+                return false;
+            }
+
+            if (strcmp(getenv("RCCL_USE_HMM"), "1") == 0)
+            {
+                return true;
+            }
+            return false;
+        }
+
+        // Helper for HMM allocations: if device supports managedMemory, and HMM is requested through
+        // RCCL_USE_HMM environment variable
+        template <class T>
+        hipError_t hipMallocHelper(T** devPtr, size_t size)
+        {
+            if (SupportsHmm() && UseHmm())
+            {
+                return hipMallocManaged((void**)devPtr, size);
+            }
+            else
+            {
+                return hipMalloc((void**)devPtr, size);
+            }
+            return hipSuccess;
+        }
+
         // To be used in multi-process tests, in the parent process before forking children.
         void InitializeRootProcess(int            const numDevices_,
                                    size_t         const numElements_,
@@ -155,22 +198,22 @@ namespace CorrectnessTests
             if (multiProcessRank_ > -1)
             {
                 HIP_CALL(hipSetDevice(multiProcessRank_));
-                HIP_CALL(hipMalloc((void **)&inputs[multiProcessRank_], NumBytes(ncclInputBuffer)));
+                HIP_CALL(hipMallocHelper((void **)&inputs[multiProcessRank_], NumBytes(ncclInputBuffer)));
                 if (inPlace)
                     outputs[multiProcessRank_] = inputs[multiProcessRank_];
                 else
-                    HIP_CALL(hipMalloc((void **)&outputs[multiProcessRank_], NumBytes(ncclOutputBuffer)));
+                    HIP_CALL(hipMallocHelper((void **)&outputs[multiProcessRank_], NumBytes(ncclOutputBuffer)));
             }
             else
             {
               for (int i = 0; i < numDevices; i++)
               {
                   HIP_CALL(hipSetDevice(i));
-                  HIP_CALL(hipMalloc((void **)&inputs[i], NumBytes(ncclInputBuffer)));
+                  HIP_CALL(hipMallocHelper((void **)&inputs[i], NumBytes(ncclInputBuffer)));
                   if (inPlace)
                       outputs[i] = inputs[i];
                   else
-                      HIP_CALL(hipMalloc((void **)&outputs[i], NumBytes(ncclOutputBuffer)));
+                      HIP_CALL(hipMallocHelper((void **)&outputs[i], NumBytes(ncclOutputBuffer)));
 
                   expected[i] = malloc(NumBytes(ncclOutputBuffer));
               }
@@ -547,16 +590,17 @@ dropback:
             // Only proceed with testing if there are enough GPUs
             if (numDevices > numDevicesAvailable)
             {
-              GTEST_SKIP();
-              return;
+                fprintf(stdout, "[  SKIPPED ] Test requires %d devices (only %d available)\n",
+                        numDevices, numDevicesAvailable);
+                GTEST_SKIP();
             }
 
+            bool enableClique = false;
             envString = 0;
             numTokens = 0;
             setenv("RCCL_TEST_ENV_VARS", "ENABLE", 1);
             if (strcmp(envVals, "")) {
                 // enable RCCL env vars testing
-
                 envString = strdup(envVals);
                 tokens[numTokens] = strtok(envString, "=, ");
                 numTokens++;
@@ -570,9 +614,22 @@ dropback:
                         savedEnv[i] = 0;
                     setenv(tokens[i*2], tokens[i*2+1], 1);
                     fprintf(stdout, "[          ] setting environmental variable %s to %s\n", tokens[i*2], getenv(tokens[i*2]));
+                    if (strcmp(tokens[i*2], "RCCL_ENABLE_CLIQUE") == 0)
+                    {
+                        if (strcmp(getenv(tokens[i*2]), "1") == 0)
+                        {
+                            enableClique = true;
+                        }
+                    }
                 }
             }
 
+            if (Dataset::UseHmm() && enableClique)
+            {
+                fprintf(stdout, "[  SKIPPED ] Clique mode and unified memory together not supported\n");
+                GTEST_SKIP();
+            }
+
             // Initialize communicators
             comms.resize(numDevices);
             NCCL_CALL(ncclCommInitAll(comms.data(), numDevices, NULL));
@@ -803,6 +860,8 @@ dropback:
 
             envString = 0;
             numTokens = 0;
+            bool enableClique = false;
+
             setenv("RCCL_TEST_ENV_VARS", "ENABLE", 1);
             if (strcmp(envVals, "")) {
                 // enable RCCL env vars testing
@@ -819,9 +878,22 @@ dropback:
                         savedEnv[i] = 0;
                     setenv(tokens[i*2], tokens[i*2+1], 1);
                     fprintf(stdout, "[          ] setting environmental variable %s to %s\n", tokens[i*2], getenv(tokens[i*2]));
+                    if (strcmp(tokens[i*2], "RCCL_ENABLE_CLIQUE") == 0)
+                    {
+                        if (strcmp(getenv(tokens[i*2]), "1") == 0)
+                        {
+                            enableClique = true;
+                        }
+                    }
                 }
             }
 
+            if (Dataset::UseHmm() && enableClique)
+            {
+                fprintf(stdout, "[  SKIPPED ] Clique mode and unified memory together not supported\n");
+                GTEST_SKIP();
+            }
+
             comms.resize(numDevices);
             streams.resize(numDevices);
             dataset = (Dataset*)mmap(NULL, sizeof(Dataset), PROT_READ|PROT_WRITE, MAP_SHARED|MAP_ANONYMOUS, -1, 0);
@@ -869,11 +941,7 @@ dropback:
                     fprintf(stdout, "[  SKIPPED ] Test requires %d devices (only %d available)\n",
                             numDevices, numDevicesAvailable);
                 }
-                // Modify the number of devices so that tear-down doesn't occur
-                // This is temporary until GTEST_SKIP() becomes available
-                numDevices = 0;
-                numDevicesAvailable = -1;
-                return;
+                GTEST_SKIP();
             }
 
             HIP_CALL(hipSetDevice(rank));