tjh.dev
1diff --git a/tinygrad/runtime/autogen/cuda.py b/tinygrad/runtime/autogen/cuda.py
2index 359083a9..3cd5f7be 100644
3--- a/tinygrad/runtime/autogen/cuda.py
4+++ b/tinygrad/runtime/autogen/cuda.py
5@@ -143,10 +143,25 @@ def char_pointer_cast(string, encoding='utf-8'):
6 return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
7
8
9+NAME_TO_PATHS = {
10+ "libcuda.so": ["@driverLink@/lib/libcuda.so"],
11+ "libnvrtc.so": ["@libnvrtc@"],
12+}
13+def _try_dlopen(name):
14+ try:
15+ return ctypes.CDLL(name)
16+ except OSError:
17+ pass
18+ for candidate in NAME_TO_PATHS.get(name, []):
19+ try:
20+ return ctypes.CDLL(candidate)
21+ except OSError:
22+ pass
23+ raise RuntimeError(f"{name} not found")
24
25 _libraries = {}
26-_libraries['libcuda.so'] = ctypes.CDLL(ctypes.util.find_library('cuda'))
27-_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc'))
28+_libraries['libcuda.so'] = _try_dlopen('libcuda.so')
29+_libraries['libnvrtc.so'] = _try_dlopen('libnvrtc.so')
30
31
32 cuuint32_t = ctypes.c_uint32