2025-05-01 version of cuda.bindings.path_finder

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

@@ -9,12 +9,13 @@
 # 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}}
@@ -45,18 +46,65 @@ 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
@@ -241,8 +289,6 @@ 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,11 +4,11 @@
 #
 # This code was automatically generated across versions from 12.0.1 to 12.8.0. Do not modify it directly.
 
-from libc.stdint cimport intptr_t, uintptr_t
+from libc.stdint cimport intptr_t
 
-from .utils import FunctionNotFoundError, NotSupportedError
+from .utils cimport get_nvjitlink_dso_version_suffix
 
-from cuda.bindings import path_finder
+from .utils import FunctionNotFoundError, NotSupportedError
 
 ###############################################################################
 # Extern
@@ -52,9 +52,17 @@ cdef void* __nvJitLinkGetInfoLog = NULL
 cdef void* __nvJitLinkVersion = NULL
 
 
-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 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 int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -6,9 +6,12 @@
 
 from libc.stdint cimport intptr_t
 
+from .utils cimport get_nvjitlink_dso_version_suffix
+
 from .utils import FunctionNotFoundError, NotSupportedError
 
-from cuda.bindings import path_finder
+import os
+import site
 
 import win32api
 
@@ -39,9 +42,44 @@ cdef void* __nvJitLinkGetInfoLog = NULL
 cdef void* __nvJitLinkVersion = NULL
 
 
-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 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 int _check_or_init_nvjitlink() except -1 nogil:
@@ -50,24 +88,23 @@ 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:
-            nvcuda_handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
+            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(nvcuda_handle, 'cuDriverGetVersion')
+            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(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 = <intptr_t>load_library(driver_ver)
+        handle = load_library(driver_ver)
 
         # Load function
         global __nvJitLinkCreate

cuda_bindings/cuda/bindings/_internal/nvvm_linux.pyx

@@ -4,11 +4,11 @@
 #
 # 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, uintptr_t
+from libc.stdint cimport intptr_t
 
-from .utils import FunctionNotFoundError, NotSupportedError
+from .utils cimport get_nvvm_dso_version_suffix
 
-from cuda.bindings import path_finder
+from .utils import FunctionNotFoundError, NotSupportedError
 
 ###############################################################################
 # Extern
@@ -51,8 +51,16 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    cdef uintptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm").handle
-    return <void*>handle
+    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 int _check_or_init_nvvm() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx

@@ -6,9 +6,12 @@
 
 from libc.stdint cimport intptr_t
 
+from .utils cimport get_nvvm_dso_version_suffix
+
 from .utils import FunctionNotFoundError, NotSupportedError
 
-from cuda.bindings import path_finder
+import os
+import site
 
 import win32api
 
@@ -37,9 +40,52 @@ cdef void* __nvvmGetProgramLogSize = NULL
 cdef void* __nvvmGetProgramLog = NULL
 
 
-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 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 int _check_or_init_nvvm() except -1 nogil:
@@ -48,24 +94,23 @@ 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:
-            nvcuda_handle = win32api.LoadLibraryEx("nvcuda.dll", 0, LOAD_LIBRARY_SEARCH_SYSTEM32)
+            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(nvcuda_handle, 'cuDriverGetVersion')
+            __cuDriverGetVersion = <void*><intptr_t>win32api.GetProcAddress(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 = <intptr_t>load_library(driver_ver)
+        handle = load_library(driver_ver)
 
         # Load function
         global __nvvmVersion

cuda_bindings/cuda/bindings/_internal/utils.pxd

@@ -165,3 +165,6 @@ 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,3 +127,17 @@ 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')