diff --git a/hipamd/src/hip_hcc.cpp b/hipamd/src/hip_hcc.cpp index ed2dd79f2e..aa9759fc2b 100644 --- a/hipamd/src/hip_hcc.cpp +++ b/hipamd/src/hip_hcc.cpp @@ -28,10 +28,13 @@ THE SOFTWARE. #include #include #include +#include #include #include #include #include +#include +#include #include #include @@ -73,6 +76,8 @@ int HIP_STAGING_SIZE = 64; /* size of staging buffers, in KB */ int HIP_STAGING_BUFFERS = 2; // TODO - remove, two buffers should be enough. int HIP_PININPLACE = 0; int HIP_STREAM_SIGNALS = 2; /* number of signals to allocate at stream creation */ +int HIP_VISIBLE_DEVICES = 0; /* Contains a comma-separated sequence of GPU identifiers */ +std::vector g_hip_visible_devices; /* vector of integers that contains the visible device IDs */ //--- @@ -851,17 +856,42 @@ void ihipReadEnv_I(int *var_ptr, const char *var_name1, const char *var_name2, c env = getenv(var_name2); } - // Default is set when variable is initialized (at top of this file), so only override if we find - // an environment variable. - if (env) { - long int v = strtol(env, NULL, 0); - *var_ptr = (int) (v); + // Check if the environment variable is either HIP_VISIBLE_DEVICES or CUDA_LAUNCH_BLOCKING, which + // contains a sequence of comma-separated device IDs + if (!(strcmp(var_name1,"HIP_VISIBLE_DEVICES") && strcmp(var_name2, "CUDA_VISIBLE_DEVICES")) && env){ + // Parse the string stream of env and store the device ids to g_hip_visible_devices global variable + std::string str = env; + std::istringstream ss(str); + std::string device_id; + // Clean up the defult value + g_hip_visible_devices.clear(); + // Read the visible device numbers + while (std::getline(ss, device_id, ',')) { + if (atoi(device_id.c_str()) >= 0) { + g_hip_visible_devices.push_back(atoi(device_id.c_str())); + }else// Any device number after invalid number will not present + break; + } + // Print out the number of ids + if (HIP_PRINT_ENV) { + printf ("%-30s = ", var_name1); + for(int i=0;i= deviceCnt){ + // Make sure any DeviceID after invalid DeviceID will be erased. + g_hip_visible_devices.resize(i); + break; } } - assert(deviceCnt == g_deviceCnt); + g_devices = new ihipDevice_t[deviceCnt]; + g_deviceCnt = 0; + for (int i=0; i"); @@ -941,6 +996,12 @@ INLINE bool ihipIsValidDevice(unsigned deviceIndex) return (deviceIndex < g_deviceCnt); } +// check if the device ID is set as visible +INLINE bool ihipIsVisibleDevice(unsigned deviceIndex) +{ + return std::find(g_hip_visible_devices.begin(), g_hip_visible_devices.end(), + (int)deviceIndex) != g_hip_visible_devices.end(); +} //--- INLINE ihipDevice_t *ihipGetTlsDefaultDevice()