WIP+WIP: PR #604 + use path_finder in bindings (nvJitLink, nvrtc, nvvm)

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

@@ -9,13 +9,12 @@
 # This code was automatically generated with version 12.9.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
 
 from libc.stdint cimport intptr_t
 
@@ -48,65 +47,18 @@ cdef bint __cuPythonInit = False
 {{if 'nvrtcSetFlowCallback' in found_functions}}cdef void *__nvrtcSetFlowCallback = NULL{{endif}}
 
 cdef int cuPythonInit() except -1 nogil:
+    {{if 'Windows' != platform.system()}}
+    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
-
-        # Check if DLLs can be found within pip installations
-        if not handle:
-            LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
-            LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
-            site_packages = [site.getusersitepackages()] + site.getsitepackages()
-            for sp in site_packages:
-                mod_path = os.path.join(sp, "nvidia", "cuda_nvrtc", "bin")
-                if os.path.isdir(mod_path):
-                    os.add_dll_directory(mod_path)
-                    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
-                    else:
-                        break
-            else:
-                # Else try default search
-                # Only reached if DLL wasn't found in any site-package path
-                LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
-                try:
-                    handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
-                except:
-                    pass
-
-        if not handle:
-            raise RuntimeError('Failed to LoadLibraryEx nvrtc64_120_0.dll')
-    {{else}}
-    handle = dlfcn.dlopen('libnvrtc.so.12', dlfcn.RTLD_NOW)
-    if handle == NULL:
-        with gil:
-            raise RuntimeError('Failed to dlopen libnvrtc.so.12')
-    {{endif}}
-
-
-    # Load function
     {{if 'Windows' == platform.system()}}
     with gil:
+        handle = path_finder._load_nvidia_dynamic_library("nvrtc").handle
         {{if 'nvrtcGetErrorString' in found_functions}}
         try:
             global __nvrtcGetErrorString
@@ -291,6 +243,8 @@ cdef int cuPythonInit() except -1 nogil:
         {{endif}}
 
     {{else}}
+    with gil:
+        handle = <void*><uintptr_t>path_finder._load_nvidia_dynamic_library("nvrtc").handle
     {{if 'nvrtcGetErrorString' in found_functions}}
     global __nvrtcGetErrorString
     __nvrtcGetErrorString = dlfcn.dlsym(handle, 'nvrtcGetErrorString')

cuda_bindings/cuda/bindings/_internal/nvjitlink_linux.pyx

@@ -4,12 +4,12 @@
 #
 # This code was automatically generated across versions from 12.0.1 to 12.9.0. Do not modify it directly.
 
-from libc.stdint cimport intptr_t
-
-from .utils cimport get_nvjitlink_dso_version_suffix
+from libc.stdint cimport intptr_t, uintptr_t
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
+from cuda.bindings import path_finder
+
 ###############################################################################
 # Extern
 ###############################################################################
@@ -52,17 +52,9 @@ cdef void* __nvJitLinkGetInfoLog = NULL
 cdef void* __nvJitLinkVersion = NULL
 
 
-cdef void* load_library(const int driver_ver) except* with gil:
-    cdef void* handle
-    for suffix in get_nvjitlink_dso_version_suffix(driver_ver):
-        so_name = "libnvJitLink.so" + (f".{suffix}" if suffix else suffix)
-        handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-        if handle != NULL:
-            break
-    else:
-        err_msg = dlerror()
-        raise RuntimeError(f'Failed to dlopen libnvJitLink ({err_msg.decode()})')
-    return handle
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef uintptr_t handle = path_finder._load_nvidia_dynamic_library("nvJitLink").handle
+    return <void*>handle
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -6,12 +6,9 @@
 
 from libc.stdint cimport intptr_t
 
-from .utils cimport get_nvjitlink_dso_version_suffix
-
 from .utils import FunctionNotFoundError, NotSupportedError
 
-import os
-import site
+from cuda.bindings import path_finder
 
 import win32api
 
@@ -42,44 +39,9 @@ cdef void* __nvJitLinkGetInfoLog = NULL
 cdef void* __nvJitLinkVersion = NULL
 
 
-cdef inline list get_site_packages():
-    return [site.getusersitepackages()] + site.getsitepackages()
-
-
-cdef load_library(const int driver_ver):
-    handle = 0
-
-    for suffix in get_nvjitlink_dso_version_suffix(driver_ver):
-        if len(suffix) == 0:
-            continue
-        dll_name = f"nvJitLink_{suffix}0_0.dll"
-
-        # First check if the DLL has been loaded by 3rd parties
-        try:
-            return win32api.GetModuleHandle(dll_name)
-        except:
-            pass
-
-        # Next, check if DLLs are installed via pip
-        for sp in get_site_packages():
-            mod_path = os.path.join(sp, "nvidia", "nvJitLink", "bin")
-            if os.path.isdir(mod_path):
-                os.add_dll_directory(mod_path)
-                try:
-                    return win32api.LoadLibraryEx(
-                        # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                        os.path.join(mod_path, dll_name),
-                        0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-                except:
-                    pass
-        # Finally, try default search
-        # Only reached if DLL wasn't found in any site-package path
-        try:
-            return win32api.LoadLibrary(dll_name)
-        except:
-            pass
-
-    raise RuntimeError('Failed to load nvJitLink')
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef intptr_t handle = path_finder._load_nvidia_dynamic_library("nvJitLink").handle
+    return <void*>handle
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:
@@ -88,23 +50,24 @@ cdef int _check_or_init_nvjitlink() except -1 nogil:
         return 0
 
     cdef int err, driver_ver
+    cdef intptr_t handle
     with gil:
         # Load driver to check version
         try:
-            handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
+            nvcuda_handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
         except Exception as e:
             raise NotSupportedError(f'CUDA driver is not found ({e})')
         global __cuDriverGetVersion
         if __cuDriverGetVersion == NULL:
-            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(handle, 'cuDriverGetVersion')
+            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(nvcuda_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')
 
         # Load library
-        handle = load_library(driver_ver)
+        handle = <intptr_t>load_library(driver_ver)
 
         # Load function
         global __nvJitLinkCreate

cuda_bindings/cuda/bindings/_internal/nvvm_linux.pyx

@@ -4,12 +4,12 @@
 #
 # This code was automatically generated across versions from 11.0.3 to 12.9.0. Do not modify it directly.
 
-from libc.stdint cimport intptr_t
-
-from .utils cimport get_nvvm_dso_version_suffix
+from libc.stdint cimport intptr_t, uintptr_t
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
+from cuda.bindings import path_finder
+
 ###############################################################################
 # Extern
 ###############################################################################
@@ -51,16 +51,8 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    cdef void* handle
-    for suffix in get_nvvm_dso_version_suffix(driver_ver):
-        so_name = "libnvvm.so" + (f".{suffix}" if suffix else suffix)
-        handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-        if handle != NULL:
-            break
-    else:
-        err_msg = dlerror()
-        raise RuntimeError(f'Failed to dlopen libnvvm ({err_msg.decode()})')
-    return handle
+    cdef uintptr_t handle = path_finder._load_nvidia_dynamic_library("nvvm").handle
+    return <void*>handle
 
 
 cdef int _check_or_init_nvvm() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx

@@ -6,12 +6,9 @@
 
 from libc.stdint cimport intptr_t
 
-from .utils cimport get_nvvm_dso_version_suffix
-
 from .utils import FunctionNotFoundError, NotSupportedError
 
-import os
-import site
+from cuda.bindings import path_finder
 
 import win32api
 
@@ -40,52 +37,9 @@ cdef void* __nvvmGetProgramLogSize = NULL
 cdef void* __nvvmGetProgramLog = NULL
 
 
-cdef inline list get_site_packages():
-    return [site.getusersitepackages()] + site.getsitepackages() + ["conda"]
-
-
-cdef load_library(const int driver_ver):
-    handle = 0
-
-    for suffix in get_nvvm_dso_version_suffix(driver_ver):
-        if len(suffix) == 0:
-            continue
-        dll_name = "nvvm64_40_0.dll"
-
-        # First check if the DLL has been loaded by 3rd parties
-        try:
-            return win32api.GetModuleHandle(dll_name)
-        except:
-            pass
-
-        # Next, check if DLLs are installed via pip or conda
-        for sp in get_site_packages():
-            if sp == "conda":
-                # nvvm is not under $CONDA_PREFIX/lib, so it's not in the default search path
-                conda_prefix = os.environ.get("CONDA_PREFIX")
-                if conda_prefix is None:
-                    continue
-                mod_path = os.path.join(conda_prefix, "Library", "nvvm", "bin")
-            else:
-                mod_path = os.path.join(sp, "nvidia", "cuda_nvcc", "nvvm", "bin")
-            if os.path.isdir(mod_path):
-                os.add_dll_directory(mod_path)
-                try:
-                    return win32api.LoadLibraryEx(
-                        # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                        os.path.join(mod_path, dll_name),
-                        0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-                except:
-                    pass
-
-        # Finally, try default search
-        # Only reached if DLL wasn't found in any site-package path
-        try:
-            return win32api.LoadLibrary(dll_name)
-        except:
-            pass
-
-    raise RuntimeError('Failed to load nvvm')
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef intptr_t handle = path_finder._load_nvidia_dynamic_library("nvvm").handle
+    return <void*>handle
 
 
 cdef int _check_or_init_nvvm() except -1 nogil:
@@ -94,23 +48,24 @@ cdef int _check_or_init_nvvm() except -1 nogil:
         return 0
 
     cdef int err, driver_ver
+    cdef intptr_t handle
     with gil:
         # Load driver to check version
         try:
-            handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
+            nvcuda_handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
         except Exception as e:
             raise NotSupportedError(f'CUDA driver is not found ({e})')
         global __cuDriverGetVersion
         if __cuDriverGetVersion == NULL:
-            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(handle, 'cuDriverGetVersion')
+            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(nvcuda_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')
 
         # Load library
-        handle = load_library(driver_ver)
+        handle = <intptr_t>load_library(driver_ver)
 
         # Load function
         global __nvvmVersion

cuda_bindings/cuda/bindings/_internal/utils.pxd

@@ -165,6 +165,3 @@ cdef int get_nested_resource_ptr(nested_resource[ResT] &in_out_ptr, object obj,
 
 cdef bint is_nested_sequence(data)
 cdef void* get_buffer_pointer(buf, Py_ssize_t size, readonly=*) except*
-
-cdef tuple get_nvjitlink_dso_version_suffix(int driver_ver)
-cdef tuple get_nvvm_dso_version_suffix(int driver_ver)

cuda_bindings/cuda/bindings/_internal/utils.pyx

@@ -127,17 +127,3 @@ cdef int get_nested_resource_ptr(nested_resource[ResT] &in_out_ptr, object obj,
 class FunctionNotFoundError(RuntimeError): pass
 
 class NotSupportedError(RuntimeError): pass
-
-
-cdef tuple get_nvjitlink_dso_version_suffix(int driver_ver):
-    if 12000 <= driver_ver < 13000:
-        return ('12', '')
-    raise NotSupportedError(f'CUDA driver version {driver_ver} is not supported')
-
-
-cdef tuple get_nvvm_dso_version_suffix(int driver_ver):
-    if 11000 <= driver_ver < 11020:
-        return ('3', '')
-    if 11020 <= driver_ver < 13000:
-        return ('4', '')
-    raise NotSupportedError(f'CUDA driver version {driver_ver} is not supported')