implement pstate switching

src/nvapi.cpp

@@ -26,8 +26,41 @@ static NvAPI_Unload_t             NvAPI_Unload;
 
 /////
 
+static std::set<int> parse_visible_devices() {
+    // set to store device IDs
+    std::set<int> 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<int> 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);
+        }
+    }
 }

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
 }