forked from mrq/bitsandbytes-rocm
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
|
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)
|