summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--bitsandbytes/cuda_setup/main.py69
-rw-r--r--bitsandbytes/cuda_setup/paths.py2
2 files changed, 47 insertions, 24 deletions
diff --git a/bitsandbytes/cuda_setup/main.py b/bitsandbytes/cuda_setup/main.py
index 1e52f89..f1c845c 100644
--- a/bitsandbytes/cuda_setup/main.py
+++ b/bitsandbytes/cuda_setup/main.py
@@ -28,10 +28,40 @@ def check_cuda_result(cuda, result_val):
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}")
+ raise Exception(f"CUDA exception! Error code: {error_str.value.decode()}")
+def get_cuda_version(cuda, cudart_path):
+ # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION
+ try:
+ cudart = ctypes.CDLL(cudart_path)
+ except OSError:
+ # TODO: shouldn't we error or at least warn here?
+ print(f'ERROR: libcudart.so could not be read from path: {cudart_path}!')
+ return None
+
+ version = ctypes.c_int()
+ check_cuda_result(cuda, cudart.cudaRuntimeGetVersion(ctypes.byref(version)))
+ version = int(version.value)
+ major = version//1000
+ minor = (version-(major*1000))//10
+
+ return f'{major}{minor}'
+
+
+def get_cuda_lib_handle():
+ # 1. find libcuda.so library (GPU driver) (/usr/lib)
+ try:
+ cuda = ctypes.CDLL("libcuda.so")
+ except OSError:
+ # TODO: shouldn't we error or at least warn here?
+ print('ERROR: libcuda.so not found!')
+ return None
+ check_cuda_result(cuda, cuda.cuInit(0))
-def get_compute_capabilities():
+ return cuda
+
+
+def get_compute_capabilities(cuda):
"""
1. find libcuda.so library (GPU driver) (/usr/lib)
init_device -> init variables -> call function by reference
@@ -42,13 +72,6 @@ def get_compute_capabilities():
# bits taken from https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
"""
- # 1. find libcuda.so library (GPU driver) (/usr/lib)
- try:
- cuda = ctypes.CDLL("libcuda.so")
- except OSError:
- # TODO: shouldn't we error or at least warn here?
- print('ERROR: libcuda.so not found!')
- return None
nGpus = ctypes.c_int()
cc_major = ctypes.c_int()
@@ -57,8 +80,6 @@ def get_compute_capabilities():
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):
@@ -75,13 +96,13 @@ def get_compute_capabilities():
# def get_compute_capability()-> Union[List[str, ...], None]: # FIXME: error
-def get_compute_capability():
+def get_compute_capability(cuda):
"""
Extracts the highest compute capbility from all available GPUs, as compute
capabilities are downwards compatible. If no GPUs are detected, it returns
None.
"""
- ccs = get_compute_capabilities()
+ ccs = get_compute_capabilities(cuda)
if ccs is not None:
# TODO: handle different compute capabilities; for now, take the max
return ccs[-1]
@@ -89,10 +110,19 @@ def get_compute_capability():
def evaluate_cuda_setup():
- 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"
+ cudart_path = determine_cuda_runtime_lib_path()
+ if cudart_path is None:
+ print(
+ "WARNING: No libcudart.so found! Install CUDA or the cudatoolkit package (anaconda)!"
+ )
+ return binary_name
+
+ print(f"CUDA SETUP: CUDA path found: {cudart_path}")
+ cuda = get_cuda_lib_handle()
+ cc = get_compute_capability(cuda)
+ cuda_version_string = get_cuda_version(cuda, cudart_path)
+
if cc == '':
print(
@@ -107,15 +137,8 @@ def evaluate_cuda_setup():
# (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(".")
- )
- cuda_version_string = f"{major}{minor}"
print(f'CUDA_SETUP: Detected CUDA version {cuda_version_string}')
def get_binary_name():
diff --git a/bitsandbytes/cuda_setup/paths.py b/bitsandbytes/cuda_setup/paths.py
index c4a7465..b5e04f8 100644
--- a/bitsandbytes/cuda_setup/paths.py
+++ b/bitsandbytes/cuda_setup/paths.py
@@ -123,4 +123,4 @@ def determine_cuda_runtime_lib_path() -> Union[Path, None]:
warn_in_case_of_duplicates(cuda_runtime_libs)
- return next(iter(cuda_runtime_libs)) if cuda_runtime_libs else set()
+ return next(iter(cuda_runtime_libs)) if cuda_runtime_libs else None