summaryrefslogtreecommitdiff
path: root/bitsandbytes/cuda_setup
diff options
context:
space:
mode:
Diffstat (limited to 'bitsandbytes/cuda_setup')
-rw-r--r--bitsandbytes/cuda_setup/compute_capability.py56
1 files changed, 35 insertions, 21 deletions
diff --git a/bitsandbytes/cuda_setup/compute_capability.py b/bitsandbytes/cuda_setup/compute_capability.py
index 19ceb3b..7a3f463 100644
--- a/bitsandbytes/cuda_setup/compute_capability.py
+++ b/bitsandbytes/cuda_setup/compute_capability.py
@@ -2,27 +2,28 @@ import ctypes
from dataclasses import dataclass, field
-CUDA_SUCCESS = 0
-
@dataclass
class CudaLibVals:
- # code bits taken from
+ # 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()
+ nGpus: ctypes.c_int = field(default=ctypes.c_int())
+ cc_major: ctypes.c_int = field(default=ctypes.c_int())
+ cc_minor: ctypes.c_int = field(default=ctypes.c_int())
+ device: ctypes.c_int = field(default=ctypes.c_int())
+ error_str: ctypes.c_char_p = field(default=ctypes.c_char_p())
cuda: ctypes.CDLL = field(init=False, repr=False)
ccs: List[str, ...] = field(init=False)
- def load_cuda_lib(self):
+ def _initialize_driver_API(self):
+ self.check_cuda_result(self.cuda.cuInit(0))
+
+ 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")
+ libnames = "libcuda.so"
for libname in libnames:
try:
self.cuda = ctypes.CDLL(libname)
@@ -33,32 +34,45 @@ class CudaLibVals:
else:
raise OSError("could not load any of: " + " ".join(libnames))
- def check_cuda_result(self, result_val):
+ def call_cuda_func(self, function_obj, **kwargs):
+ CUDA_SUCCESS = 0 # constant taken from cuda.h
+ pass
+ # if (CUDA_SUCCESS := function_obj(
+
+ def _error_handle(cuda_lib_call_return_value):
"""
- 2. call extern C function to determine CC
+ 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__)
+ CUDA_SUCCESS = 0 # constant taken from cuda.h
- if result_val != 0:
- self.cuda.cuGetErrorString(result_val, ctypes.byref(self.error_str))
+ if cuda_lib_call_return_value != CUDA_SUCCESS:
+ self.cuda.cuGetErrorString(
+ cuda_lib_call_return_value,
+ ctypes.byref(self.error_str),
+ )
print("Count not initialize CUDA - failure!")
raise Exception("CUDA exception!")
- return result_val
+ return cuda_lib_call_return_value
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)))
+ self._load_cuda_lib()
+ self._initialize_driver_API()
+ 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)
+ 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
+ ctypes.byref(self.cc_major),
+ ctypes.byref(self.cc_minor),
+ self.device,
),
)
tmp_ccs.append(f"{self.cc_major.value}.{self.cc_minor.value}")