From eb077fc993787d7473c33eb8c345fa4cc6745ede Mon Sep 17 00:00:00 2001
From: Victor <lostfreeman@gmail.com>
Date: Thu, 16 Dec 2021 22:09:52 -0800
Subject: [PATCH] [RUNTIME] fixed NVidia DLL names on Windows (#392)

---
 lib/driver/dispatch.cc | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/lib/driver/dispatch.cc b/lib/driver/dispatch.cc
index 4059ac235..9e2aca432 100755
--- a/lib/driver/dispatch.cc
+++ b/lib/driver/dispatch.cc
@@ -91,9 +91,13 @@ void* dispatch::fname ## _;
 
 bool dispatch::cuinit(){
   if(cuda_==nullptr){
+    #ifdef _WIN32
+    cuda_ = dlopen("cudart64_110.dll", RTLD_LAZY);
+    #else
     cuda_ = dlopen("libcuda.so", RTLD_LAZY);
     if(!cuda_)
       cuda_ = dlopen("libcuda.so.1", RTLD_LAZY);
+    #endif
     if(!cuda_)
       throw std::runtime_error("Could not find `libcuda.so`. Make sure it is in your LD_LIBRARY_PATH.");
   }
@@ -176,8 +180,13 @@ CUDA_DEFINE1(CUresult, cuEventDestroy_v2, CUevent)
  * NVML
  * ------------------- */
 bool dispatch::nvmlinit(){
+  #ifdef _WIN32
+  if(nvml_==nullptr)
+    nvml_ = dlopen("nvml.dll", RTLD_LAZY);
+  #else
   if(nvml_==nullptr)
     nvml_ = dlopen("libnvidia-ml.so", RTLD_LAZY);
+  #endif
   nvmlReturn_t (*fptr)();
   nvmlInit_v2_ = dlsym(nvml_, "nvmlInit_v2");
   *reinterpret_cast<void **>(&fptr) = nvmlInit_v2_;