diff --git a/src/nvapi.cpp b/src/nvapi.cpp index d30a9ad4f..e0ec359a3 100644 --- a/src/nvapi.cpp +++ b/src/nvapi.cpp @@ -26,8 +26,41 @@ static NvAPI_Unload_t NvAPI_Unload; ///// +static std::set parse_visible_devices() { + // set to store device IDs + std::set devices; + + // retrieve the value of the environment variable "CUDA_VISIBLE_DEVICES" + const char * env_p = std::getenv("CUDA_VISIBLE_DEVICES"); + if (!env_p) { + return devices; + } + + // create a string stream from the environment variable value + std::stringstream ss(env_p); + + // string to hold each device ID from the stream + std::string item; + + // iterate over the comma-separated device IDs in the environment variable + while (std::getline(ss, item, ",")) { + try { + // convert the current item to an integer and insert it into the set + devices.insert(std::stoi(item)); + } catch (...) { + // ignore any exceptions that occur during the conversion + continue; + } + } + + // return the set of device IDs + return devices; +} + +///// + void nvapi_init() { - // load library + // attempt to load the NVAPI library #ifdef _WIN32 if (!lib) { lib = LoadLibrary("nvapi64.dll"); @@ -46,7 +79,7 @@ void nvapi_init() { } #endif - // lookup QueryInterface + // obtain the address of the nvapi_QueryInterface function if (lib) { #ifdef _WIN32 nvapi_QueryInterface = (nvapi_QueryInterface_t) GetProcAddress(lib, "nvapi_QueryInterface"); @@ -55,7 +88,7 @@ void nvapi_init() { #endif } - // resolve functions + // retrieve function pointers for various NVAPI functions if (nvapi_QueryInterface) { NvAPI_EnumPhysicalGPUs = (NvAPI_EnumPhysicalGPUs_t) nvapi_QueryInterface(0xe5ac921f); NvAPI_GPU_SetForcePstate = (NvAPI_GPU_SetForcePstate_t) nvapi_QueryInterface(0x025bfb10); @@ -63,19 +96,19 @@ void nvapi_init() { NvAPI_Unload = (NvAPI_Unload_t) nvapi_QueryInterface(0xd22bdd7e); } - // initialize library + // initialize the NVAPI library if (NvAPI_Initialize()) { load_success = true; } } void nvapi_free() { - // deinitialize library + // if the library was successfully initialized, unload it if (load_success) { NvAPI_Unload(); } - // free library + // release the library resources if (lib) { #ifdef _WIN32 FreeLibrary(lib); @@ -84,27 +117,42 @@ void nvapi_free() { #endif } - // invalidate pointers + // reset the pointers and flags lib = nullptr; load_success = false; } -void nvapi_set_pstate(int ids[], int ids_size, int pstate) { +void nvapi_set_pstate(int pstate) { + // check if the library initialization was successful before proceeding if (!load_success) { return; } - // TODO - (void) ids; - (void) ids_size; - (void) pstate; - printf("nvapi_set_pstate: %d", pstate); -} + // array to hold GPU handles + void *gpu_array[64] = {0}; -void nvapi_set_pstate_high() { - nvapi_set_pstate({}, 0, 16); -} + // integer to hold GPU count + int gpu_count = 0; -void nvapi_set_pstate_low() { - nvapi_set_pstate({}, 0, 8); + // attempt to enumerate GPUs + if (NvAPI_EnumPhysicalGPUs(gpu_array, &gpu_count) != 0) { + fprintf(stderr, "Failed to enumerate GPUs\n"); + return; + } + + // try to retrieve the set of visible CUDA devices + std::set devices = parse_visible_devices(); + + // iterate over each GPU + for (int i = 0; i < gpu_count; i++) { + // if the set of visible devices is not empty and the current GPU ID is not in this set, skip this iteration + if (!devices.empty() && !devices.find(i)) { + continue; + } + + // attempt to set the performance state for the current GPU + if (NvAPI_GPU_SetForcePstate(gpu_array[gpu_id], pstate, 2) != 0) { + fprintf(stderr, "Failed to set performance state for gpu #%d\n", gpu_id); + } + } } diff --git a/src/nvapi.h b/src/nvapi.h index 4a6e4d530..1d7d8431e 100644 --- a/src/nvapi.h +++ b/src/nvapi.h @@ -7,9 +7,7 @@ extern "C" { void nvapi_init(); void nvapi_free(); -void nvapi_set_pstate(int ids[], int ids_size, int pstate); -void nvapi_set_pstate_high(); -void nvapi_set_pstate_low(); +void nvapi_set_pstate(int pstate); #ifdef __cplusplus }