lzhengning
/
JittorMirror
mirror of https://github.com/Jittor/Jittor
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add new cuda search path

This commit is contained in:
Dun Liang 2020-09-09 20:27:43 +08:00
parent e079a8813e
commit 839f354daf
3 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/jittor/__init__.py
View File

@ -7,7 +7,7 @@
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
__version__ = '1.1.7.13'
__version__ = '1.1.7.14'
from . import lock
with lock.lock_scope():
from . import compiler

6
python/jittor/compile_extern.py
View File

@ -143,8 +143,10 @@ def setup_cuda_lib(lib_name, link=True, extra_flags=""):
link_flags = ""
if link:
cuda_include_name = search_file([cuda_include, "/usr/include"], lib_name+".h")
culib_path = search_file([cuda_lib, "/usr/lib/x86_64-linux-gnu"], f"lib{lib_name}.so")
extra_include_path = os.path.abspath(os.path.join(cuda_include, "..", "targets/x86_64-linux/include"))
extra_lib_path = os.path.abspath(os.path.join(cuda_lib, "..", "targets/x86_64-linux/lib"))
cuda_include_name = search_file([cuda_include, extra_include_path, "/usr/include"], lib_name+".h")
culib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu"], f"lib{lib_name}.so")
# dynamic link cuda library
ctypes.CDLL(culib_path, dlopen_flags)
link_flags = f"-l{lib_name} -L'{cuda_lib}'"

5
src/misc/cuda_flags.cc
View File

@ -21,8 +21,9 @@ void setter_use_cuda(int value) {
cudaGetDeviceCount(&count);
CHECK(count>0) << "No device found.";
LOGi << "CUDA enabled.";
} else
LOGi << "CUDA disabled.";
} else {
LOGv << "CUDA disabled.";
}
#else
CHECK(value==0) << "No CUDA found.";
#endif