mirror of https://github.com/Jittor/Jittor
add /usr/lib as cuda lib search path
This commit is contained in:
parent
2c2f5b156d
commit
fd7d68e6aa
|
@ -9,7 +9,7 @@
|
||||||
# file 'LICENSE.txt', which is part of this source code package.
|
# file 'LICENSE.txt', which is part of this source code package.
|
||||||
# ***************************************************************
|
# ***************************************************************
|
||||||
|
|
||||||
__version__ = '1.2.3.17'
|
__version__ = '1.2.3.18'
|
||||||
from jittor_utils import lock
|
from jittor_utils import lock
|
||||||
with lock.lock_scope():
|
with lock.lock_scope():
|
||||||
ori_int = int
|
ori_int = int
|
||||||
|
|
|
@ -180,11 +180,11 @@ def setup_cuda_lib(lib_name, link=True, extra_flags=""):
|
||||||
prefer_version = ()
|
prefer_version = ()
|
||||||
if nvcc_version[0] == 11:
|
if nvcc_version[0] == 11:
|
||||||
prefer_version = ("8",)
|
prefer_version = ("8",)
|
||||||
culib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu"], f"lib{lib_name}.so", prefer_version)
|
culib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu", "/usr/lib"], f"lib{lib_name}.so", prefer_version)
|
||||||
|
|
||||||
if lib_name == "cublas" and nvcc_version[0] >= 10:
|
if lib_name == "cublas" and nvcc_version[0] >= 10:
|
||||||
# manual link libcublasLt.so
|
# manual link libcublasLt.so
|
||||||
cublas_lt_lib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu"], f"libcublasLt.so", nvcc_version)
|
cublas_lt_lib_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu", "/usr/lib"], f"libcublasLt.so", nvcc_version)
|
||||||
ctypes.CDLL(cublas_lt_lib_path, dlopen_flags)
|
ctypes.CDLL(cublas_lt_lib_path, dlopen_flags)
|
||||||
|
|
||||||
|
|
||||||
|
@ -193,7 +193,7 @@ def setup_cuda_lib(lib_name, link=True, extra_flags=""):
|
||||||
if nvcc_version >= (11,0,0):
|
if nvcc_version >= (11,0,0):
|
||||||
libs = ["libcudnn_ops_infer.so", "libcudnn_ops_train.so", "libcudnn_cnn_infer.so", "libcudnn_cnn_train.so"]
|
libs = ["libcudnn_ops_infer.so", "libcudnn_ops_train.so", "libcudnn_cnn_infer.so", "libcudnn_cnn_train.so"]
|
||||||
for l in libs:
|
for l in libs:
|
||||||
ex_cudnn_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu"], l, prefer_version)
|
ex_cudnn_path = search_file([cuda_lib, extra_lib_path, "/usr/lib/x86_64-linux-gnu", "/usr/lib"], l, prefer_version)
|
||||||
ctypes.CDLL(ex_cudnn_path, dlopen_flags)
|
ctypes.CDLL(ex_cudnn_path, dlopen_flags)
|
||||||
|
|
||||||
# dynamic link cuda library
|
# dynamic link cuda library
|
||||||
|
|
Loading…
Reference in New Issue