From 59a615b3869eb8488a748e2aa51224a5e3d366bb Mon Sep 17 00:00:00 2001 From: Titus von Koeller Date: Tue, 2 Aug 2022 21:26:50 -0700 Subject: factored cuda_setup.main out into smaller modules and functions --- bitsandbytes/cuda_setup/main.py | 142 ++++++++++++++-------------------------- 1 file changed, 48 insertions(+), 94 deletions(-) (limited to 'bitsandbytes/cuda_setup/main.py') diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py index 6d70c92..e96ac70 100644 --- a/bitsandbytes/cuda_setup/main.py +++ b/bitsandbytes/cuda_setup/main.py @@ -8,8 +8,6 @@ extract factors the build is dependent on: - CuBLAS-LT: full-build 8-bit optimizer - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`) -alle Binaries packagen - evaluation: - if paths faulty, return meaningful error - else: @@ -19,11 +17,10 @@ evaluation: """ import ctypes -import os from pathlib import Path -from typing import Set, Union -from ..utils import print_err, warn_of_missing_prerequisite, execute_and_return +from ..utils import execute_and_return +from .paths import determine_cuda_runtime_lib_path def check_cuda_result(cuda, result_val): @@ -34,26 +31,23 @@ def check_cuda_result(cuda, result_val): raise Exception(f"CUDA exception! ERROR: {error_str}") -# taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549 -def get_compute_capability(): - # 1. find libcuda.so library (GPU driver) (/usr/lib) - # init_device -> init variables -> call function by reference - # 2. call extern C function to determine CC - # (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html) - # 3. Check for CUDA errors - # https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api +def get_compute_capabilities(): + """ + 1. find libcuda.so library (GPU driver) (/usr/lib) + init_device -> init variables -> call function by reference + 2. call extern C function to determine CC + (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html) + 3. Check for CUDA errors + https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api + # bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549 + """ # 1. find libcuda.so library (GPU driver) (/usr/lib) - libnames = ("libcuda.so",) - for libname in libnames: - try: - cuda = ctypes.CDLL(libname) - except OSError: - continue - else: - break - else: - raise OSError("could not load any of: " + " ".join(libnames)) + try: + cuda = ctypes.CDLL("libcuda.so") + except OSError: + # TODO: shouldn't we error or at least warn here? + return None nGpus = ctypes.c_int() cc_major = ctypes.c_int() @@ -70,104 +64,64 @@ def get_compute_capability(): check_cuda_result(cuda, cuda.cuDeviceGet(ctypes.byref(device), i)) ref_major = ctypes.byref(cc_major) ref_minor = ctypes.byref(cc_minor) - # 2. call extern C function to determine CC - check_cuda_result(cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device)) - ccs.append(f"{cc_major.value}.{cc_minor.value}") - - # TODO: handle different compute capabilities; for now, take the max - ccs.sort() - max_cc = ccs[-1] - return max_cc - - -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} - - -def resolve_env_variable(env_var): - '''Searches a given envirionmental library or path for the CUDA runtime library (libcudart.so)''' - paths: Set[Path] = tokenize_paths(env_var) - - non_existent_directories: Set[Path] = { - path for path in paths if not path.exists() - } - - if non_existent_directories: - print_err( - "WARNING: The following directories listed your path were found to " - f"be non-existent: {non_existent_directories}" + # 2. call extern C function to determine CC + check_cuda_result( + cuda, cuda.cuDeviceComputeCapability(ref_major, ref_minor, device) ) + ccs.append(f"{cc_major.value}.{cc_minor.value}") - cuda_runtime_libs: Set[Path] = { - path / CUDA_RUNTIME_LIB - for path in paths - if (path / CUDA_RUNTIME_LIB).is_file() - } - non_existent_directories + return ccs.sort() - if len(cuda_runtime_libs) > 1: - err_msg = ( - f"Found duplicate {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.." - ) - raise FileNotFoundError(err_msg) - elif len(cuda_runtime_libs) == 0: return None # this is not en error, since other envs can contain CUDA - else: return next(iter(cuda_runtime_libs)) # for now just return the first - -def get_cuda_runtime_lib_path() -> Union[Path, None]: - '''Searches conda installation and environmental paths for a cuda installations.''' - - cuda_runtime_libs = [] - # CONDA_PREFIX/lib is the default location for a default conda - # install of pytorch. This location takes priortiy over all - # other defined variables - if 'CONDA_PREFIX' in os.environ: - lib_conda_path = f'{os.environ["CONDA_PREFIX"]}/lib/' - print(lib_conda_path) - cuda_runtime_libs.append(resolve_env_variable(lib_conda_path)) - - if len(cuda_runtime_libs) == 1: return cuda_runtime_libs[0] - - # if CONDA_PREFIX does not have the library, search the environment - # (in particualr LD_LIBRARY PATH) - for var in os.environ: - cuda_runtime_libs.append(resolve_env_variable(var)) - - if len(cuda_runtime_libs) < 1: - err_msg = ( - f"Did not find {CUDA_RUNTIME_LIB} files: {cuda_runtime_libs}.." - ) - raise FileNotFoundError(err_msg) - return cuda_runtime_libs.pop() +# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error +def get_compute_capability(): + """ + Extracts the highest compute capbility from all available GPUs, as compute + capabilities are downwards compatible. If no GPUs are detected, it returns + None. + """ + if ccs := get_compute_capabilities() is not None: + # TODO: handle different compute capabilities; for now, take the max + return ccs[-1] + return None def evaluate_cuda_setup(): - cuda_path = get_cuda_runtime_lib_path() - print(f'CUDA SETUP: CUDA path found: {cuda_path}') + cuda_path = determine_cuda_runtime_lib_path() + print(f"CUDA SETUP: CUDA path found: {cuda_path}") cc = get_compute_capability() binary_name = "libbitsandbytes_cpu.so" + # FIXME: has_gpu is still unused if not (has_gpu := bool(cc)): print( "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..." ) return binary_name + # 7.5 is the minimum CC vor cublaslt has_cublaslt = cc in ["7.5", "8.0", "8.6"] # TODO: # (1) CUDA missing cases (no CUDA installed by CUDA driver (nvidia-smi accessible) # (2) Multiple CUDA versions installed + # FIXME: cuda_home is still unused cuda_home = str(Path(cuda_path).parent.parent) # we use ls -l instead of nvcc to determine the cuda version # since most installations will have the libcudart.so installed, but not the compiler ls_output, err = execute_and_return(f"ls -l {cuda_path}") - major, minor, revision = ls_output.split(' ')[-1].replace('libcudart.so.', '').split('.') + major, minor, revision = ( + ls_output.split(" ")[-1].replace("libcudart.so.", "").split(".") + ) cuda_version_string = f"{major}{minor}" - binary_name = f'libbitsandbytes_cuda{cuda_version_string}{("" if has_cublaslt else "_nocublaslt")}.so' + def get_binary_name(): + "if not has_cublaslt (CC < 7.5), then we have to choose _nocublaslt.so" + bin_base_name = "libbitsandbytes_cuda" + if has_cublaslt: + return f"{bin_base_name}{cuda_version_string}.so" + else: + return f"{bin_base_name}_nocublaslt.so" return binary_name -- cgit v1.2.3