GPU (WSL): support nvml too

Author: CarterLi

src/detection/gpu/gpu_linux.c

@@ -251,7 +251,7 @@ static void pciHandleDevice(FF_MAYBE_UNUSED const FFGPUOptions* options, FFlist*
             .temp = options->temp ? &gpu->temperature : NULL,
             .memory = options->useNvml ? &gpu->dedicated : NULL,
             .coreCount = options->useNvml ? (uint32_t*) &gpu->coreCount : NULL,
-        });
+        }, "libnvidia-ml.so");
 
         if (gpu->dedicated.total != FF_GPU_VMEM_SIZE_UNSET)
             gpu->type = gpu->dedicated.total > 1024 * 1024 * 1024 ? FF_GPU_TYPE_DISCRETE : FF_GPU_TYPE_INTEGRATED;

src/detection/gpu/gpu_nvidia.c

@@ -15,19 +15,12 @@ struct FFNvmlData {
     bool inited;
 } nvmlData;
 
-const char* ffDetectNvidiaGpuInfo(FFGpuNvidiaCondition cond, FFGpuNvidiaResult result)
+const char* ffDetectNvidiaGpuInfo(FFGpuNvidiaCondition cond, FFGpuNvidiaResult result, const char* soName)
 {
     if (!nvmlData.inited)
     {
         nvmlData.inited = true;
-        FF_LIBRARY_LOAD(libnvml, NULL, "dlopen nvml failed",
-#ifdef _WIN32
-            "nvml"
-#else
-            "libnvidia-ml"
-#endif
-            FF_LIBRARY_EXTENSION, -1
-        );
+        FF_LIBRARY_LOAD(libnvml, NULL, "dlopen nvml failed", soName , 1);
         FF_LIBRARY_LOAD_SYMBOL_MESSAGE(libnvml, nvmlInit_v2)
         FF_LIBRARY_LOAD_SYMBOL_MESSAGE(libnvml, nvmlShutdown)
         FF_LIBRARY_LOAD_SYMBOL_VAR_MESSAGE(libnvml, nvmlData, nvmlDeviceGetCount_v2)

src/detection/gpu/gpu_nvidia.h

@@ -3,11 +3,11 @@
 #include "gpu.h"
 
 // Use pciBusId if not NULL; use pciDeviceId and pciSubSystemId otherwise
-// pciBusId = "domain:bus:device.function"
-// pciDeviceId = (deviceId << 16) | vendorId
 typedef struct FFGpuNvidiaCondition
 {
+    // "domain:bus:device.function"
     const char* pciBusId;
+    // (deviceId << 16) | vendorId
     uint32_t pciDeviceId;
     uint32_t pciSubSystemId;
 } FFGpuNvidiaCondition;
@@ -20,4 +20,4 @@ typedef struct FFGpuNvidiaResult
     uint32_t* coreCount;
 } FFGpuNvidiaResult;
 
-const char* ffDetectNvidiaGpuInfo(FFGpuNvidiaCondition cond, FFGpuNvidiaResult result);
+const char* ffDetectNvidiaGpuInfo(FFGpuNvidiaCondition cond, FFGpuNvidiaResult result, const char* soName);

src/detection/gpu/gpu_windows.c

@@ -99,7 +99,7 @@ const char* ffDetectGPUImpl(FF_MAYBE_UNUSED const FFGPUOptions* options, FFlist*
                     .temp = options->temp ? &gpu->temperature : NULL,
                     .memory = options->useNvml ? &gpu->dedicated : NULL,
                     .coreCount = options->useNvml ? (uint32_t*) &gpu->coreCount : NULL,
-                });
+                }, "nvml.dll");
             }
         }
     }

src/detection/gpu/gpu_wsl.cpp

@@ -4,6 +4,7 @@
 extern "C" {
 #include "common/library.h"
 #include "detection/gpu/gpu.h"
+#include "detection/gpu/gpu_nvidia.h"
 }
 
 #include <wsl/winadapter.h>
@@ -92,7 +93,20 @@ const char* ffGPUDetectByDirectX(FF_MAYBE_UNUSED const FFGPUOptions* options, FF
         if (SUCCEEDED(adapter->GetProperty(DXCoreAdapterProperty::HardwareID, sizeof(hardwareId), &hardwareId)))
         {
             const char* vendorStr = ffGetGPUVendorString((unsigned) hardwareId.vendorID);
-            ffStrbufAppendS(&gpu->vendor, vendorStr);
+            ffStrbufSetStatic(&gpu->vendor, vendorStr);
+
+            if (vendorStr == FF_GPU_VENDOR_NAME_NVIDIA && options->useNvml)
+            {
+                ffDetectNvidiaGpuInfo((FFGpuNvidiaCondition) {
+                    .pciBusId = nullptr,
+                    .pciDeviceId = (hardwareId.deviceID << 16) | hardwareId.vendorID,
+                    .pciSubSystemId = hardwareId.subSysID
+                }, (FFGpuNvidiaResult) {
+                    .temp = options->temp ? &gpu->temperature : NULL,
+                    .memory = options->useNvml ? &gpu->dedicated : NULL,
+                    .coreCount = options->useNvml ? (uint32_t*) &gpu->coreCount : NULL,
+                }, "/usr/lib/wsl/lib/libnvidia-ml.so");
+            }
         }
     }