Address comments from PR

mdboom · mdboom · commit e0e868a78a9f · 2025-09-05T13:01:13.000-04:00
diff --git a/cuda_bindings/cuda/bindings/_internal/cufile_linux.pyx b/cuda_bindings/cuda/bindings/_internal/cufile_linux.pyx
@@ -18,6 +18,8 @@ import cython
 # Extern
 ###############################################################################
 
+from .utils import NotSupportedError
+
 cdef extern from "<dlfcn.h>" nogil:
     void* dlopen(const char*, int)
     char* dlerror()
@@ -32,14 +34,31 @@ cdef extern from "<dlfcn.h>" nogil:
 
     const void* RTLD_DEFAULT 'RTLD_DEFAULT'
 
+cdef int get_cuda_version():
+    cdef void* handle = NULL
+    cdef int err, driver_ver = 0
+
+    # Load driver to check version
+    handle = dlopen('libcuda.so.1', RTLD_NOW | RTLD_GLOBAL)
+    if handle == NULL:
+        err_msg = dlerror()
+        raise NotSupportedError(f'CUDA driver is not found ({err_msg.decode()})')
+    cuDriverGetVersion = dlsym(handle, "cuDriverGetVersion")
+    if cuDriverGetVersion == NULL:
+        raise RuntimeError('something went wrong')
+    err = (<int (*)(int*) noexcept nogil>cuDriverGetVersion)(&driver_ver)
+    if err != 0:
+        raise RuntimeError('something went wrong')
+
+    return driver_ver
+
 
 ###############################################################################
 # Wrapper init
 ###############################################################################
 
 cdef object __symbol_lock = threading.Lock()
 cdef bint __py_cufile_init = False
-cdef void* __cuDriverGetVersion = NULL
 
 cdef void* __cuFileHandleRegister = NULL
 cdef void* __cuFileHandleDeregister = NULL
@@ -97,24 +116,9 @@ cdef int _check_or_init_cufile() except -1 nogil:
         return 0
 
     cdef void* handle = NULL
-    cdef int err, driver_ver = 0
 
     with gil, __symbol_lock:
-        # Load driver to check version
-        handle = dlopen('libcuda.so.1', RTLD_NOW | RTLD_GLOBAL)
-        if handle == NULL:
-            err_msg = dlerror()
-            raise NotSupportedError(f'CUDA driver is not found ({err_msg.decode()})')
-        global __cuDriverGetVersion
-        if __cuDriverGetVersion == NULL:
-            __cuDriverGetVersion = dlsym(handle, "cuDriverGetVersion")
-        if __cuDriverGetVersion == NULL:
-            raise RuntimeError('something went wrong')
-        err = (<int (*)(int*) noexcept nogil>__cuDriverGetVersion)(&driver_ver)
-        if err != 0:
-            raise RuntimeError('something went wrong')
-        #dlclose(handle)
-        handle = NULL
+        driver_ver = get_cuda_version()
 
         # Load function
         global __cuFileHandleRegister
diff --git a/cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx b/cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx
@@ -13,7 +13,7 @@ from cuda.pathfinder import load_nvidia_dynamic_lib
 
 from libc.stddef cimport wchar_t
 from libc.stdint cimport uintptr_t
-from cpython cimport PyUnicode_AsWideCharString
+from cpython cimport PyUnicode_AsWideCharString, PyMem_Free
 
 from .utils import NotSupportedError
 
@@ -23,6 +23,7 @@ cdef extern from "windows.h" nogil:
     ctypedef void* FARPROC
     ctypedef unsigned long DWORD
     ctypedef const wchar_t *LPCWSTR
+    ctypedef const char *LPCSTR
 
     cdef DWORD LOAD_LIBRARY_SEARCH_SYSTEM32 = 0x00000800
     cdef DWORD LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
@@ -36,19 +37,23 @@ cdef extern from "windows.h" nogil:
 
     HMODULE _LoadLibraryW "LoadLibraryW"(LPCWSTR lpLibFileName)
 
-    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, const char* lpProcName)
+    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, LPCSTR lpProcName)
 
     HMODULE _GetModuleHandleW "GetModuleHandleW"(LPCWSTR lpModuleName)
 
-cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags) nogil:
-    cdef wchar_t* wpath
-    with gil:
-        wpath = PyUnicode_AsWideCharString(path, NULL)
-    return <uintptr_t>_LoadLibraryExW(
-        wpath,
-        hFile,
-        dwFlags
-    )
+cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags):
+    cdef uintptr_t result
+    cdef wchar_t* wpath = PyUnicode_AsWideCharString(path, NULL)
+    if wpath == NULL:
+        raise
+    with nogil:
+        result = <uintptr_t>_LoadLibraryExW(
+            wpath,
+            hFile,
+            dwFlags
+        )
+    PyMem_Free(wpath)
+    return result
 
 cdef inline uintptr_t LoadLibraryW(str path) nogil:
     cdef wchar_t* wpath
diff --git a/cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx b/cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx
@@ -13,7 +13,7 @@ from cuda.pathfinder import load_nvidia_dynamic_lib
 
 from libc.stddef cimport wchar_t
 from libc.stdint cimport uintptr_t
-from cpython cimport PyUnicode_AsWideCharString
+from cpython cimport PyUnicode_AsWideCharString, PyMem_Free
 
 from .utils import NotSupportedError
 
@@ -23,6 +23,7 @@ cdef extern from "windows.h" nogil:
     ctypedef void* FARPROC
     ctypedef unsigned long DWORD
     ctypedef const wchar_t *LPCWSTR
+    ctypedef const char *LPCSTR
 
     cdef DWORD LOAD_LIBRARY_SEARCH_SYSTEM32 = 0x00000800
     cdef DWORD LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
@@ -36,19 +37,23 @@ cdef extern from "windows.h" nogil:
 
     HMODULE _LoadLibraryW "LoadLibraryW"(LPCWSTR lpLibFileName)
 
-    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, const char* lpProcName)
+    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, LPCSTR lpProcName)
 
     HMODULE _GetModuleHandleW "GetModuleHandleW"(LPCWSTR lpModuleName)
 
-cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags) nogil:
-    cdef wchar_t* wpath
-    with gil:
-        wpath = PyUnicode_AsWideCharString(path, NULL)
-    return <uintptr_t>_LoadLibraryExW(
-        wpath,
-        hFile,
-        dwFlags
-    )
+cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags):
+    cdef uintptr_t result
+    cdef wchar_t* wpath = PyUnicode_AsWideCharString(path, NULL)
+    if wpath == NULL:
+        raise
+    with nogil:
+        result = <uintptr_t>_LoadLibraryExW(
+            wpath,
+            hFile,
+            dwFlags
+        )
+    PyMem_Free(wpath)
+    return result
 
 cdef inline uintptr_t LoadLibraryW(str path) nogil:
     cdef wchar_t* wpath
diff --git a/cuda_bindings/cuda/bindings/_lib/windll.pxd b/cuda_bindings/cuda/bindings/_lib/windll.pxd
@@ -3,34 +3,39 @@
 
 from libc.stddef cimport wchar_t
 from libc.stdint cimport uintptr_t
-from cpython cimport PyUnicode_AsWideCharString
+from cpython cimport PyUnicode_AsWideCharString, PyMem_Free
 
-cdef extern from "windows.h":
+cdef extern from "windows.h" nogil:
     ctypedef void* HMODULE
     ctypedef void* HANDLE
     ctypedef void* FARPROC
     ctypedef unsigned long DWORD
     ctypedef const wchar_t *LPCWSTR
+    ctypedef const char *LPCSTR
 
     cdef DWORD LOAD_LIBRARY_SEARCH_SYSTEM32 = 0x00000800
 
     HMODULE _LoadLibraryExW "LoadLibraryExW"(
         LPCWSTR lpLibFileName,
         HANDLE hFile,
         DWORD dwFlags
-    ) nogil
+    )
 
-    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, const char* lpProcName) nogil
+    FARPROC _GetProcAddress "GetProcAddress"(HMODULE hModule, LPCSTR lpProcName)
 
-cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags) nogil:
-    cdef wchar_t* wpath
-    with gil:
-        wpath = PyUnicode_AsWideCharString(path, NULL)
-    return <uintptr_t>_LoadLibraryExW(
-        wpath,
-        hFile,
-        dwFlags
-    )
+cdef inline uintptr_t LoadLibraryExW(str path, HANDLE hFile, DWORD dwFlags):
+    cdef uintptr_t result
+    cdef wchar_t* wpath = PyUnicode_AsWideCharString(path, NULL)
+    if wpath is NULL:
+        raise
+    with nogil:
+        result = <uintptr_t>_LoadLibraryExW(
+            wpath,
+            hFile,
+            dwFlags
+        )
+    PyMem_Free(wpath)
+    return result
 
 cdef inline FARPROC GetProcAddress(uintptr_t hModule, const char* lpProcName) nogil:
     return _GetProcAddress(<HMODULE>hModule, lpProcName)
