Use path_finder.load_nvidia_dynamic_library("nvrtc") from cuda/bindings/_bindings/cynvrtc.pyx.in

Author: rwgk

cuda_bindings/cuda/bindings/_bindings/cynvrtc.pyx.in

@@ -9,13 +9,12 @@
 # This code was automatically generated with version 12.8.0. Do not modify it directly.
 {{if 'Windows' == platform.system()}}
 import os
-import site
-import struct
 import win32api
-from pywintypes import error
 {{else}}
 cimport cuda.bindings._lib.dlfcn as dlfcn
+from libc.stdint cimport uintptr_t
 {{endif}}
+from cuda.bindings import path_finder
 
 cdef bint __cuPythonInit = False
 {{if 'nvrtcGetErrorString' in found_functions}}cdef void *__nvrtcGetErrorString = NULL{{endif}}
@@ -47,74 +46,17 @@ cdef bint __cuPythonInit = False
 
 cdef int cuPythonInit() except -1 nogil:
     {{if 'Windows' != platform.system()}}
-    cdef char* err_msg
+    cdef void* handle = NULL
     {{endif}}
 
     global __cuPythonInit
     if __cuPythonInit:
         return 0
     __cuPythonInit = True
 
-    # Load library
-    {{if 'Windows' == platform.system()}}
-    with gil:
-        # First check if the DLL has been loaded by 3rd parties
-        try:
-            handle = win32api.GetModuleHandle("nvrtc64_120_0.dll")
-        except:
-            handle = None
-
-        # Else try default search
-        if not handle:
-            LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
-            try:
-                handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
-            except:
-                pass
-
-        # Final check if DLLs can be found within pip installations
-        if not handle:
-            site_packages = [site.getusersitepackages()] + site.getsitepackages()
-            for sp in site_packages:
-                mod_path = os.path.join(sp, "nvidia", "cuda_nvrtc", "bin")
-                if not os.path.isdir(mod_path):
-                    continue
-                os.add_dll_directory(mod_path)
-            LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
-            LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
-            try:
-                handle = win32api.LoadLibraryEx(
-                    # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                    os.path.join(mod_path, "nvrtc64_120_0.dll"),
-                    0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-
-                # Note: nvrtc64_120_0.dll calls into nvrtc-builtins64_*.dll which is
-                # located in the same mod_path.
-                # Update PATH environ so that the two dlls can find each other
-                os.environ["PATH"] = os.pathsep.join((os.environ.get("PATH", ""), mod_path))
-            except:
-                pass
-
-        if not handle:
-            raise RuntimeError('Failed to LoadLibraryEx nvrtc64_120_0.dll')
-    {{else}}
-    with gil:
-        print("\nLOOOK dlfcn.dlopen('libnvrtc.so.12', dlfcn.RTLD_NOW)", flush=True)
-    handle = dlfcn.dlopen('libnvrtc.so.12', dlfcn.RTLD_NOW)
-    if handle == NULL:
-        with gil:
-            err_msg = dlfcn.dlerror()
-            if err_msg == NULL:
-                err_msg_str = 'Unknown error'
-            else:
-                err_msg_str = err_msg.decode('utf-8', errors='backslashreplace')
-            raise RuntimeError(f'Failed to dlopen libnvrtc.so.12: {err_msg_str}')
-    {{endif}}
-
-
-    # Load function
     {{if 'Windows' == platform.system()}}
     with gil:
+        handle = path_finder.load_nvidia_dynamic_library("nvrtc")
         {{if 'nvrtcGetErrorString' in found_functions}}
         try:
             global __nvrtcGetErrorString
@@ -299,6 +241,8 @@ cdef int cuPythonInit() except -1 nogil:
         {{endif}}
 
     {{else}}
+    with gil:
+        handle = <void*><uintptr_t>path_finder.load_nvidia_dynamic_library("nvrtc")
     {{if 'nvrtcGetErrorString' in found_functions}}
     global __nvrtcGetErrorString
     __nvrtcGetErrorString = dlfcn.dlsym(handle, 'nvrtcGetErrorString')

cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py

@@ -45,6 +45,8 @@ def _windows_load_with_dll_basename(name: str) -> int:
 
     if name == "nvJitLink":
         dll_name = "nvJitLink_120_0.dll"
+    elif name == "nvrtc":
+        dll_name = "nvrtc64_120_0.dll"
     elif name == "nvvm":
         dll_name = "nvvm64_40_0.dll"