Factor out load_nvidia_dynamic_library() from _internal/nvjitlink_linux.pyx, nvvm_linux.pyx

Author: rwgk

cuda_bindings/cuda/bindings/_internal/nvjitlink_linux.pyx

@@ -6,7 +6,7 @@
 
 import os
 
-from libc.stdint cimport intptr_t
+from libc.stdint cimport intptr_t, uintptr_t
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
@@ -54,13 +54,9 @@ cdef void* __nvJitLinkGetInfoLog = NULL
 cdef void* __nvJitLinkVersion = NULL
 
 
-cdef void* load_library(const int driver_ver) except* with gil:
-    so_name = path_finder.find_nvidia_dynamic_library("nvJitLink")
-    cdef void* handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-    if handle != NULL:
-        return handle
-    err_msg = dlerror().decode(errors="backslashreplace")
-    raise RuntimeError(f"Failed to dlopen {so_name}: {err_msg}")
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink")
+    return <void*>handle
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvvm_linux.pyx

@@ -4,7 +4,7 @@
 #
 # This code was automatically generated across versions from 11.0.3 to 12.8.0. Do not modify it directly.
 
-from libc.stdint cimport intptr_t
+from libc.stdint cimport intptr_t, uintptr_t
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
@@ -51,12 +51,8 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    so_name = path_finder.find_nvidia_dynamic_library("nvvm")
-    cdef void* handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-    if handle != NULL:
-        return handle
-    err_msg = dlerror().decode(errors="backslashreplace")
-    raise RuntimeError(f"Failed to dlopen {so_name}: {err_msg}")
+    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm")
+    return <void*>handle
 
 
 cdef int _check_or_init_nvvm() except -1 nogil:

cuda_bindings/cuda/bindings/_path_finder_utils/load_nvidia_dynamic_library.py

@@ -0,0 +1,13 @@
+import ctypes
+import os
+
+from .find_nvidia_dynamic_library import find_nvidia_dynamic_library
+
+
+def load_nvidia_dynamic_library(name: str) -> int:
+    path = find_nvidia_dynamic_library(name)
+    try:
+        handle = ctypes.CDLL(path, mode=os.RTLD_NOW | os.RTLD_GLOBAL)
+        return handle._handle  # This is the actual `void*` value as an int
+    except OSError as e:
+        raise RuntimeError(f"Failed to dlopen {path}: {e}") from e

cuda_bindings/cuda/bindings/path_finder.py

@@ -17,9 +17,11 @@
     get_system_ctk,
 )
 from cuda.bindings._path_finder_utils.find_nvidia_dynamic_library import find_nvidia_dynamic_library
+from cuda.bindings._path_finder_utils.load_nvidia_dynamic_library import load_nvidia_dynamic_library
 
 __all__ = [
     "find_nvidia_dynamic_library",
+    "load_nvidia_dynamic_library",
     "get_conda_ctk",
     "get_conda_include_dir",
     "get_cuda_home",