Adding support for hipMallocManaged() in unit tests (#375)
* Adding HMM support for unit tests * Fixing HMM opt-in check
This commit is contained in:
+80
-12
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user