gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[RUNTIME] fixed NVidia DLL names on Windows (#392)

Browse Source
This commit is contained in:
Victor
2021-12-16 22:09:52 -08:00
committed by GitHub
parent 73b04d71b2
commit eb077fc993
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
lib/driver/dispatch.cc
View File

@@ -91,9 +91,13 @@ void* dispatch::fname ## _;
bool dispatch::cuinit(){ bool dispatch::cuinit(){
if(cuda_==nullptr){ if(cuda_==nullptr){
#ifdef _WIN32
cuda_ = dlopen("cudart64_110.dll", RTLD_LAZY);
#else
cuda_ = dlopen("libcuda.so", RTLD_LAZY); cuda_ = dlopen("libcuda.so", RTLD_LAZY);
if(!cuda_) if(!cuda_)
cuda_ = dlopen("libcuda.so.1", RTLD_LAZY); cuda_ = dlopen("libcuda.so.1", RTLD_LAZY);
#endif
if(!cuda_) if(!cuda_)
throw std::runtime_error("Could not find `libcuda.so`. Make sure it is in your LD_LIBRARY_PATH."); 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 * NVML
* ------------------- */ * ------------------- */
bool dispatch::nvmlinit(){ bool dispatch::nvmlinit(){
#ifdef _WIN32
if(nvml_==nullptr)
nvml_ = dlopen("nvml.dll", RTLD_LAZY);
#else
if(nvml_==nullptr) if(nvml_==nullptr)
nvml_ = dlopen("libnvidia-ml.so", RTLD_LAZY); nvml_ = dlopen("libnvidia-ml.so", RTLD_LAZY);
#endif
nvmlReturn_t (*fptr)(); nvmlReturn_t (*fptr)();
nvmlInit_v2_ = dlsym(nvml_, "nvmlInit_v2"); nvmlInit_v2_ = dlsym(nvml_, "nvmlInit_v2");
*reinterpret_cast<void **>(&fptr) = nvmlInit_v2_; *reinterpret_cast<void **>(&fptr) = nvmlInit_v2_;