auto cuda install boost

python/jittor/__init__.py

@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.2.3.12'
+__version__ = '1.2.3.13'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int

python/jittor/compile_extern.py

@@ -182,7 +182,7 @@ def setup_cuda_lib(lib_name, link=True, extra_flags=""):
             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)
 
-        if lib_name == "cublas":
+        if lib_name == "cublas" and nvcc_version[0] >= 10:
             # 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)
             ctypes.CDLL(cublas_lt_lib_path, dlopen_flags)

python/jittor/compiler.py

@@ -761,7 +761,7 @@ def compile_extern():
     LOG.vv(f"Compile extern llvm passes: {str(files)}")
 
 def check_cuda():
-    if nvcc_path == "":
+    if not nvcc_path:
         return
     global cc_flags, has_cuda, core_link_flags, cuda_dir, cuda_lib, cuda_include, cuda_home
     cuda_dir = os.path.dirname(get_full_path_of_executable(nvcc_path))
@@ -867,12 +867,24 @@ ex_python_path = python_path + '.' + str(sys.version_info.minor)
 if os.path.isfile(ex_python_path):
     python_path = ex_python_path
 py3_config_path = jit_utils.py3_config_path
-nvcc_path = env_or_try_find('nvcc_path', 'nvcc') or try_find_exe('/usr/local/cuda/bin/nvcc') or try_find_exe('/usr/bin/nvcc')
-if not nvcc_path:
-    cuda_driver = install_cuda.get_cuda_driver()
+
+# if jtcuda is already installed
+nvcc_path = None
+if install_cuda.has_installation():
     nvcc_path = install_cuda.install_cuda()
     if nvcc_path:
         nvcc_path = try_find_exe(nvcc_path)
+# check system installed cuda
+if not nvcc_path:
+    nvcc_path = env_or_try_find('nvcc_path', 'nvcc') or try_find_exe('/usr/local/cuda/bin/nvcc') or try_find_exe('/usr/bin/nvcc')
+# if system has no cuda, install jtcuda
+if not nvcc_path:
+    nvcc_path = install_cuda.install_cuda()
+    if nvcc_path:
+        nvcc_path = try_find_exe(nvcc_path)
+if not nvcc_path:
+    nvcc_path = ""
+
 gdb_path = try_find_exe('gdb')
 addr2line_path = try_find_exe('addr2line')
 has_pybt = check_pybt(gdb_path, python_path)

python/jittor_utils/install_cuda.py

@@ -26,13 +26,30 @@ def get_cuda_driver():
     except:
         return None
 
+def has_installation():
+    jtcuda_path = os.path.join(pathlib.Path.home(), ".cache", "jittor", "jtcuda")
+    return os.path.isdir(jtcuda_path)
 
 def install_cuda():
     cuda_driver_version = get_cuda_driver()
+    if not cuda_driver_version:
+        return None
     LOG.i("cuda_driver_version: ", cuda_driver_version)
 
-    cuda_tgz = "cuda10.2_cudnn7_linux.tgz"
-    md5 = "a78f296746d97e9d76615289c2fe98ac"
+    if cuda_driver_version >= [11,2]:
+        cuda_tgz = "cuda11.2_cudnn8_linux.tgz"
+        md5 = "b93a1a5d19098e93450ee080509e9836"
+    elif cuda_driver_version >= [11,]:
+        cuda_tgz = "cuda11.0_cudnn8_linux.tgz"
+        md5 = "5dbdb43e35b4db8249027997720bf1ca"
+    elif cuda_driver_version >= [10,2]:
+        cuda_tgz = "cuda10.2_cudnn7_linux.tgz"
+        md5 = "a78f296746d97e9d76615289c2fe98ac"
+    elif cuda_driver_version >= [10,]:
+        cuda_tgz = "cuda10.0_cudnn7_linux.tgz"
+        md5 = "f16d3ff63f081031d21faec3ec8b7dac"
+    else:
+        raise RuntimeError(f"Unsupport cuda driver version: {cuda_driver_version}")
     jtcuda_path = os.path.join(pathlib.Path.home(), ".cache", "jittor", "jtcuda")
     nvcc_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "bin", "nvcc")
     nvcc_lib_path = os.path.join(jtcuda_path, cuda_tgz[:-4], "lib64")