From 3809236428e704f9a7e22232701a651aafa5ca1b Mon Sep 17 00:00:00 2001 From: Titus von Koeller Date: Tue, 2 Aug 2022 07:42:27 -0700 Subject: move cuda_setup code into subpackage --- bitsandbytes/__init__.py | 1 + bitsandbytes/cextension.py | 2 +- bitsandbytes/cuda_setup.py | 173 -------------------------- bitsandbytes/cuda_setup/__init__.py | 0 bitsandbytes/cuda_setup/compute_capability.py | 65 ++++++++++ bitsandbytes/cuda_setup/main.py | 173 ++++++++++++++++++++++++++ 6 files changed, 240 insertions(+), 174 deletions(-) delete mode 100644 bitsandbytes/cuda_setup.py create mode 100644 bitsandbytes/cuda_setup/__init__.py create mode 100644 bitsandbytes/cuda_setup/compute_capability.py create mode 100644 bitsandbytes/cuda_setup/main.py (limited to 'bitsandbytes') diff --git a/bitsandbytes/__init__.py b/bitsandbytes/__init__.py index 6e5b6ac..76a5b48 100644 --- a/bitsandbytes/__init__.py +++ b/bitsandbytes/__init__.py @@ -12,6 +12,7 @@ from .autograd._functions import ( ) from .cextension import COMPILED_WITH_CUDA from .nn import modules +from . import cuda_setup if COMPILED_WITH_CUDA: from .optim import adam diff --git a/bitsandbytes/cextension.py b/bitsandbytes/cextension.py index bc11474..f5b97fd 100644 --- a/bitsandbytes/cextension.py +++ b/bitsandbytes/cextension.py @@ -2,7 +2,7 @@ import ctypes as ct import os from warnings import warn -from bitsandbytes.cuda_setup import evaluate_cuda_setup +from bitsandbytes.cuda_setup.main import evaluate_cuda_setup class CUDALibrary_Singleton(object): diff --git a/bitsandbytes/cuda_setup.py b/bitsandbytes/cuda_setup.py deleted file mode 100644 index e68cd5e..0000000 --- a/bitsandbytes/cuda_setup.py +++ /dev/null @@ -1,173 +0,0 @@ -""" -extract factors the build is dependent on: -[X] compute capability - [ ] TODO: Q - What if we have multiple GPUs of different makes? -- CUDA version -- Software: - - CPU-only: only CPU quantization functions (no optimizer, no matrix multipl) - - 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: - - determine CUDA version - - determine capabilities - - based on that set the default path -""" - -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 - - -def check_cuda_result(cuda, result_val): - # 3. Check for CUDA errors - if result_val != 0: - error_str = ctypes.c_char_p() - cuda.cuGetErrorString(result_val, ctypes.byref(error_str)) - 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 - - # 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)) - - nGpus = ctypes.c_int() - cc_major = ctypes.c_int() - cc_minor = ctypes.c_int() - - result = ctypes.c_int() - device = ctypes.c_int() - - check_cuda_result(cuda, cuda.cuInit(0)) - - check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus))) - ccs = [] - for i in range(nGpus.value): - 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}" - ) - - cuda_runtime_libs: Set[Path] = { - path / CUDA_RUNTIME_LIB - for path in paths - if (path / CUDA_RUNTIME_LIB).is_file() - } - non_existent_directories - - 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 evaluate_cuda_setup(): - cuda_path = get_cuda_runtime_lib_path() - print(f'CUDA SETUP: CUDA path found: {cuda_path}') - cc = get_compute_capability() - binary_name = "libbitsandbytes_cpu.so" - - if not (has_gpu := bool(cc)): - print( - "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..." - ) - return binary_name - - 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 - - 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('.') - cuda_version_string = f"{major}{minor}" - - binary_name = f'libbitsandbytes_cuda{cuda_version_string}{("" if has_cublaslt else "_nocublaslt")}.so' - - return binary_name diff --git a/bitsandbytes/cuda_setup/__init__.py b/bitsandbytes/cuda_setup/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bitsandbytes/cuda_setup/compute_capability.py b/bitsandbytes/cuda_setup/compute_capability.py new file mode 100644 index 0000000..19ceb3b --- /dev/null +++ b/bitsandbytes/cuda_setup/compute_capability.py @@ -0,0 +1,65 @@ +import ctypes +from dataclasses import dataclass, field + + +CUDA_SUCCESS = 0 + +@dataclass +class CudaLibVals: + # code bits taken from + # https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549 + + nGpus = ctypes.c_int() + cc_major = ctypes.c_int() + cc_minor = ctypes.c_int() + device = ctypes.c_int() + error_str = ctypes.c_char_p() + cuda: ctypes.CDLL = field(init=False, repr=False) + ccs: List[str, ...] = field(init=False) + + def load_cuda_lib(self): + """ + 1. find libcuda.so library (GPU driver) (/usr/lib) + init_device -> init variables -> call function by reference + """ + libnames = ("libcuda.so") + for libname in libnames: + try: + self.cuda = ctypes.CDLL(libname) + except OSError: + continue + else: + break + else: + raise OSError("could not load any of: " + " ".join(libnames)) + + def check_cuda_result(self, result_val): + """ + 2. call extern C function to determine CC + (see https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html) + """ + cls_fields: Tuple[Field, ...] = fields(self.__class__) + + if result_val != 0: + self.cuda.cuGetErrorString(result_val, ctypes.byref(self.error_str)) + print("Count not initialize CUDA - failure!") + raise Exception("CUDA exception!") + return result_val + + def __post_init__(self): + self.load_cuda_lib() + self.check_cuda_result(self.cuda.cuInit(0)) + self.check_cuda_result(self.cuda, self.cuda.cuDeviceGetCount(ctypes.byref(self.nGpus))) + tmp_ccs = [] + for gpu_index in range(self.nGpus.value): + check_cuda_result( + self.cuda, self.cuda.cuDeviceGet(ctypes.byref(self.device), gpu_index) + ) + check_cuda_result( + self.cuda, + self.cuda.cuDeviceComputeCapability( + ctypes.byref(self.cc_major), ctypes.byref(self.cc_minor), self.device + ), + ) + tmp_ccs.append(f"{self.cc_major.value}.{self.cc_minor.value}") + self.ccs = sorted(tmp_ccs, reverse=True) diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py new file mode 100644 index 0000000..6d70c92 --- /dev/null +++ b/bitsandbytes/cuda_setup/main.py @@ -0,0 +1,173 @@ +""" +extract factors the build is dependent on: +[X] compute capability + [ ] TODO: Q - What if we have multiple GPUs of different makes? +- CUDA version +- Software: + - CPU-only: only CPU quantization functions (no optimizer, no matrix multipl) + - 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: + - determine CUDA version + - determine capabilities + - based on that set the default path +""" + +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 + + +def check_cuda_result(cuda, result_val): + # 3. Check for CUDA errors + if result_val != 0: + error_str = ctypes.c_char_p() + cuda.cuGetErrorString(result_val, ctypes.byref(error_str)) + 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 + + # 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)) + + nGpus = ctypes.c_int() + cc_major = ctypes.c_int() + cc_minor = ctypes.c_int() + + result = ctypes.c_int() + device = ctypes.c_int() + + check_cuda_result(cuda, cuda.cuInit(0)) + + check_cuda_result(cuda, cuda.cuDeviceGetCount(ctypes.byref(nGpus))) + ccs = [] + for i in range(nGpus.value): + 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}" + ) + + cuda_runtime_libs: Set[Path] = { + path / CUDA_RUNTIME_LIB + for path in paths + if (path / CUDA_RUNTIME_LIB).is_file() + } - non_existent_directories + + 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 evaluate_cuda_setup(): + cuda_path = get_cuda_runtime_lib_path() + print(f'CUDA SETUP: CUDA path found: {cuda_path}') + cc = get_compute_capability() + binary_name = "libbitsandbytes_cpu.so" + + if not (has_gpu := bool(cc)): + print( + "WARNING: No GPU detected! Check your CUDA paths. Processing to load CPU-only library..." + ) + return binary_name + + 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 + + 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('.') + cuda_version_string = f"{major}{minor}" + + binary_name = f'libbitsandbytes_cuda{cuda_version_string}{("" if has_cublaslt else "_nocublaslt")}.so' + + return binary_name -- cgit v1.2.3