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

[driver] added TRITON_LIBCUDA environment variable to specify libcuda

Browse Source 
path if not in LD_LIBRARY_PATH
This commit is contained in:
Philippe Tillet
2019-09-06 13:26:51 -04:00
parent b79bcbaee8
commit 0ff81badac
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/triton/tools/sys/getenv.hpp
View File

@@ -38,11 +38,11 @@ namespace tools
std::size_t sz = 0; std::size_t sz = 0;
_dupenv_s(&cache_path, &sz, name); _dupenv_s(&cache_path, &sz, name);
#else #else
const char * cache_path = std::getenv(name); const char * cstr = std::getenv(name);
#endif #endif
if(!cache_path) if(!cstr)
return ""; return "";
std::string result(cache_path); std::string result(cstr);
#ifdef _MSC_VER #ifdef _MSC_VER
free(cache_path); free(cache_path);
#endif #endif

8
lib/driver/dispatch.cc
View File

@@ -22,6 +22,7 @@
#include "triton/driver/dispatch.h" #include "triton/driver/dispatch.h"
#include "triton/driver/context.h" #include "triton/driver/context.h"
#include "triton/tools/sys/getenv.hpp"
namespace triton namespace triton
{ {
@@ -108,8 +109,13 @@ bool dispatch::clinit()
} }
bool dispatch::cuinit(){ bool dispatch::cuinit(){
if(cuda_==nullptr) if(cuda_==nullptr){
std::string libcuda = tools::getenv("TRITON_LIBCUDA");
if(libcuda.empty())
cuda_ = dlopen("libcuda.so", RTLD_LAZY); cuda_ = dlopen("libcuda.so", RTLD_LAZY);
else
cuda_ = dlopen(libcuda.c_str(), RTLD_LAZY);
}
if(cuda_ == nullptr) if(cuda_ == nullptr)
return false; return false;
CUresult (*fptr)(unsigned int); CUresult (*fptr)(unsigned int);