Generalize load_nvidia_dynamic_library.py to also work under Windows.

Author: rwgk

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -4,7 +4,7 @@
 #
 # This code was automatically generated across versions from 12.0.1 to 12.6.2. Do not modify it directly.
 
-from libc.stdint cimport intptr_t
+from libc.stdint cimport intptr_t, uintptr_t
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
@@ -14,7 +14,6 @@ import os
 import site
 
 import win32api
-import pywintypes
 
 
 ###############################################################################
@@ -44,25 +43,8 @@ cdef void* __nvJitLinkVersion = NULL
 
 
 cdef load_library(const int driver_ver):
-    cdef str dll_path = path_finder.find_nvidia_dynamic_library("nvJitLink")
-    cdef str dll_name = os.path.basename(dll_path)
-    cdef intptr_t handle = 0
-
-    # Check if already loaded
-    try:
-        handle = win32api.GetModuleHandle(dll_name)
-    except pywintypes.error:
-        pass
-    else:
-        return handle
-
-    # Not already loaded; load it
-    try:
-        handle = win32api.LoadLibrary(dll_path)
-    except pywintypes.error as e:
-        raise RuntimeError(f"Failed to load NVVM DLL at {dll_path}: {e}")
-
-    return handle
+    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_windows.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
 
@@ -14,7 +14,6 @@ import os
 import site
 
 import win32api
-import pywintypes
 
 
 ###############################################################################
@@ -42,25 +41,8 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef load_library(const int driver_ver):
-    cdef str dll_path = path_finder.find_nvidia_dynamic_library("nvvm")
-    cdef str dll_name = os.path.basename(dll_path)
-    cdef intptr_t handle = 0
-
-    # Check if already loaded
-    try:
-        handle = win32api.GetModuleHandle(dll_name)
-    except pywintypes.error:
-        pass
-    else:
-        return handle
-
-    # Not already loaded; load it
-    try:
-        handle = win32api.LoadLibrary(dll_path)
-    except pywintypes.error as e:
-        raise RuntimeError(f"Failed to load NVVM DLL at {dll_path}: {e}")
-
-    return handle
+    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

@@ -1,13 +1,25 @@
 import ctypes
+import functools
 import os
+import sys
 
 from .find_nvidia_dynamic_library import find_nvidia_dynamic_library
 
 
+@functools.cache
 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
+    dl_path = find_nvidia_dynamic_library(name)
+    if sys.platform == "win32":
+        try:
+            handle = ctypes.windll.kernel32.LoadLibraryW(dl_path)
+            if not handle:
+                raise ctypes.WinError(ctypes.get_last_error())
+        except Exception as e:
+            raise RuntimeError(f"Failed to load DLL at {dl_path}: {e}") from e
+        return handle
+    else:
+        try:
+            handle = ctypes.CDLL(dl_path, mode=os.RTLD_NOW | os.RTLD_GLOBAL)
+            return handle._handle  # Raw void* as int
+        except OSError as e:
+            raise RuntimeError(f"Failed to dlopen {dl_path}: {e}") from e