diff options
Diffstat (limited to 'bitsandbytes/cuda_setup.py')
| -rw-r--r-- | bitsandbytes/cuda_setup.py | 78 |
1 files changed, 46 insertions, 32 deletions
diff --git a/bitsandbytes/cuda_setup.py b/bitsandbytes/cuda_setup.py index 5ed0c89..0dd53c5 100644 --- a/bitsandbytes/cuda_setup.py +++ b/bitsandbytes/cuda_setup.py @@ -18,31 +18,36 @@ evaluation: - based on that set the default path """ +import ctypes +import shlex +import subprocess from os import environ as env from pathlib import Path from typing import Set, Union -from .utils import warn_of_missing_prerequisite, print_err -import ctypes -import shlex -import subprocess +from .utils import print_err, warn_of_missing_prerequisite + def execute_and_return(strCMD): - proc = subprocess.Popen(shlex.split(strCMD), stdout=subprocess.PIPE, stderr=subprocess.PIPE) + proc = subprocess.Popen( + shlex.split(strCMD), stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) out, err = proc.communicate() out, err = out.decode("UTF-8").strip(), err.decode("UTF-8").strip() return out, err + def check_cuda_result(cuda, result_val): if result_val != 0: cuda.cuGetErrorString(result_val, ctypes.byref(error_str)) print(f"Count not initialize CUDA - failure!") - raise Exception('CUDA exception!') + raise Exception("CUDA exception!") return result_val + # taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549 def get_compute_capability(): - libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll') + libnames = ("libcuda.so", "libcuda.dylib", "cuda.dll") for libname in libnames: try: cuda = ctypes.CDLL(libname) @@ -51,8 +56,7 @@ def get_compute_capability(): else: break else: - raise OSError("could not load any of: " + ' '.join(libnames)) - + raise OSError("could not load any of: " + " ".join(libnames)) nGpus = ctypes.c_int() cc_major = ctypes.c_int() @@ -69,39 +73,43 @@ def get_compute_capability(): ccs = [] for i in range(nGpus.value): result = check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i)) - result = check_cuda_result(cuda, cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device)) - ccs.append(f'{cc_major.value}.{cc_minor.value}') + result = check_cuda_result( + cuda, + cuda.cuDeviceComputeCapability( + ctypes.byref(cc_major), ctypes.byref(cc_minor), device + ), + ) + ccs.append(f"{cc_major.value}.{cc_minor.value}") - #TODO: handle different compute capabilities; for now, take the max + # TODO: handle different compute capabilities; for now, take the max ccs.sort() - return ccs[-1] + # return ccs[-1] + return ccs + CUDA_RUNTIME_LIB: str = "libcudart.so" + def tokenize_paths(paths: str) -> Set[Path]: - return { - Path(ld_path) for ld_path in paths.split(':') - if ld_path - } + return {Path(ld_path) for ld_path in paths.split(":") if ld_path} + def get_cuda_runtime_lib_path( # TODO: replace this with logic for all paths in env vars LD_LIBRARY_PATH: Union[str, None] = env.get("LD_LIBRARY_PATH") ) -> Union[Path, None]: - """ # TODO: add doc-string - """ + """# TODO: add doc-string""" if not LD_LIBRARY_PATH: warn_of_missing_prerequisite( - 'LD_LIBRARY_PATH is completely missing from environment!' + "LD_LIBRARY_PATH is completely missing from environment!" ) return None ld_library_paths: Set[Path] = tokenize_paths(LD_LIBRARY_PATH) - non_existent_directories: Set[Path] = { - path for path in ld_library_paths - if not path.exists() + non_existent_directories: Set[Path] = { + path for path in ld_library_paths if not path.exists() } if non_existent_directories: @@ -111,7 +119,8 @@ def get_cuda_runtime_lib_path( ) cuda_runtime_libs: Set[Path] = { - path / CUDA_RUNTIME_LIB for path in ld_library_paths + path / CUDA_RUNTIME_LIB + for path in ld_library_paths if (path / CUDA_RUNTIME_LIB).is_file() } - non_existent_directories @@ -126,26 +135,31 @@ def get_cuda_runtime_lib_path( single_cuda_runtime_lib_dir = next(iter(cuda_runtime_libs)) return single_cuda_runtime_lib_dir + def evaluate_cuda_setup(): cuda_path = get_cuda_runtime_lib_path() cc = get_compute_capability() - binary_name = 'libbitsandbytes_cpu.so' + binary_name = "libbitsandbytes_cpu.so" if not (has_gpu := bool(cc)): - print('WARNING: No GPU detected! Check our CUDA paths. Processing to load CPU-only library...') + print( + "WARNING: No GPU detected! Check our CUDA paths. Processing to load CPU-only library..." + ) return binary_name - has_cublaslt = cc in ['7.5', '8.0', '8.6'] + has_cublaslt = cc in ["7.5", "8.0", "8.6"] - # TODO: + # TODO: # (1) Model missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible) # (2) Multiple CUDA versions installed cuda_home = str(Path(cuda_path).parent.parent) - ls_output, err = execute_and_return(f'{cuda_home}/bin/nvcc --version') - cuda_version = ls_output.split('\n')[3].split(',')[-1].strip().lower().replace('v', '') - major, minor, revision = cuda_version.split('.') - cuda_version_string = f'{major}{minor}' + ls_output, err = execute_and_return(f"{cuda_home}/bin/nvcc --version") + cuda_version = ( + ls_output.split("\n")[3].split(",")[-1].strip().lower().replace("v", "") + ) + major, minor, revision = cuda_version.split(".") + cuda_version_string = f"{major}{minor}" binary_name = f'libbitsandbytes_cuda{cuda_version_string}_{("cublaslt" if has_cublaslt else "")}.so' |