diff --git a/src/NvidiaGPU.h b/src/NvidiaGPU.h index 8952cbe..6ab6c20 100644 --- a/src/NvidiaGPU.h +++ b/src/NvidiaGPU.h @@ -7,10 +7,30 @@ #ifdef WITH_NVIDIA namespace NvidiaGPU { + struct GPUUtilization + { + uint32_t gpu; + uint32_t vram; + }; + struct VRAM + { + uint64_t totalB; + uint64_t freeB; + uint64_t usedB; + }; + static void* nvmldl; static void* nvmlGPUHandle; + // Preloaded Query functions + typedef int (*PFN_nvmlDeviceGetUtilizationRates)(void*, struct GPUUtilization*); + static PFN_nvmlDeviceGetUtilizationRates nvmlDeviceGetUtilizationRates; + typedef int (*PFN_nvmlDeviceGetTemperature)(void*, uint32_t, uint32_t*); + static PFN_nvmlDeviceGetTemperature nvmlDeviceGetTemperature; + typedef int (*PFN_nvmlDeviceGetMemoryInfo)(void*, struct VRAM*); + static PFN_nvmlDeviceGetMemoryInfo nvmlDeviceGetMemoryInfo; + inline void Init() { if (nvmldl || !RuntimeConfig::Get().hasNvidia) @@ -40,6 +60,37 @@ namespace NvidiaGPU auto nvmlDeviceGetHandle = (PFN_nvmlDeviceGetHandle)dlsym(nvmldl, "nvmlDeviceGetHandleByIndex"); res = nvmlDeviceGetHandle(0, &nvmlGPUHandle); ASSERT(res == 0, "Failed getting device (Error: " << res << ")!"); + + // Dynamically load functions + nvmlDeviceGetUtilizationRates = (PFN_nvmlDeviceGetUtilizationRates)dlsym(nvmldl, "nvmlDeviceGetUtilizationRates"); + nvmlDeviceGetTemperature = (PFN_nvmlDeviceGetTemperature)dlsym(nvmldl, "nvmlDeviceGetTemperature"); + nvmlDeviceGetMemoryInfo = (PFN_nvmlDeviceGetMemoryInfo)dlsym(nvmldl, "nvmlDeviceGetMemoryInfo"); + + // Check if all information is available + GPUUtilization util; + int res = nvmlDeviceGetUtilizationRates(nvmlGPUHandle, &util); + if (res != 0) + { + LOG("Failed querying utilization rates (Error: " << res << "), disabling Nvidia GPU!"); + RuntimeConfig::Get().hasNvidia = false; + return; + } + uint32_t temp; + res = nvmlDeviceGetTemperature(nvmlGPUHandle, 0, &temp); + if (res != 0) + { + LOG("Failed querying temperature (Error: " << res << "), disabling Nvidia GPU!"); + RuntimeConfig::Get().hasNvidia = false; + return; + } + VRAM mem; + res = nvmlDeviceGetMemoryInfo(nvmlGPUHandle, &mem); + if (res != 0) + { + LOG("Failed querying VRAM (Error: " << res << "), disabling Nvidia GPU!"); + RuntimeConfig::Get().hasNvidia = false; + return; + } } inline void Shutdown() @@ -48,18 +99,6 @@ namespace NvidiaGPU dlclose(nvmldl); } - struct GPUUtilization - { - uint32_t gpu; - uint32_t vram; - }; - struct VRAM - { - uint64_t totalB; - uint64_t freeB; - uint64_t usedB; - }; - inline GPUUtilization GetUtilization() { if (!RuntimeConfig::Get().hasNvidia)