First version of cuda.bindings.path_finder

.github/actions/fetch_ctk/action.yml

@@ -123,4 +123,4 @@ runs:
         echo "CUDA_PATH=${CUDA_PATH}" >> $GITHUB_ENV
         echo "CUDA_HOME=${CUDA_PATH}" >> $GITHUB_ENV
         echo "${CUDA_PATH}/bin" >> $GITHUB_PATH
-        echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${CUDA_PATH}/lib:${CUDA_PATH}/nvvm/lib64" >> $GITHUB_ENV
+        echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${CUDA_PATH}/lib" >> $GITHUB_ENV

.github/workflows/test-wheel-windows.yml

@@ -164,13 +164,6 @@ jobs:
           method: 'network'
           sub-packages: ${{ env.MINI_CTK_DEPS }}
 
-      - name: Update PATH
-        if: ${{ inputs.local-ctk == '1' }}
-        run: |
-          # mimics actual CTK installation
-          echo $PATH
-          echo "$env:CUDA_PATH\nvvm\bin" >> $env:GITHUB_PATH
-
       - name: Run cuda.bindings tests
         if: ${{ env.SKIP_CUDA_BINDINGS_TEST == '0' }}
         run: |

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.8.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")
+    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,54 +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:
-            handle = win32api.GetModuleHandle(dll_name)
-        except:
-            pass
-        else:
-            break
-
-        # Next, check if DLLs are installed via pip
-        for sp in get_site_packages():
-            mod_path = os.path.join(sp, "nvidia", "nvJitLink", "bin")
-            if not os.path.isdir(mod_path):
-                continue
-            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, dll_name),
-                0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-        except:
-            pass
-        else:
-            break
-
-        # Finally, try default search
-        try:
-            handle = win32api.LoadLibrary(dll_name)
-        except:
-            pass
-        else:
-            break
-    else:
-        raise RuntimeError('Failed to load nvJitLink')
-
-    assert handle != 0
-    return handle
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvJitLink")
+    return <void*>handle
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:
@@ -98,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.8.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")
+    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,54 +37,9 @@ cdef void* __nvvmGetProgramLogSize = NULL
 cdef void* __nvvmGetProgramLog = 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_nvvm_dso_version_suffix(driver_ver):
-        if len(suffix) == 0:
-            continue
-        dll_name = "nvvm64_40_0"
-
-        # First check if the DLL has been loaded by 3rd parties
-        try:
-            handle = win32api.GetModuleHandle(dll_name)
-        except:
-            pass
-        else:
-            break
-
-        # Next, check if DLLs are installed via pip
-        for sp in get_site_packages():
-            mod_path = os.path.join(sp, "nvidia", "cuda_nvcc", "nvvm", "bin")
-            if not os.path.isdir(mod_path):
-                continue
-            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, dll_name),
-                0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-        except:
-            pass
-        else:
-            break
-
-        # Finally, try default search
-        try:
-            handle = win32api.LoadLibrary(dll_name)
-        except:
-            pass
-        else:
-            break
-    else:
-        raise RuntimeError('Failed to load nvvm')
-
-    assert handle != 0
-    return handle
+cdef void* load_library(int driver_ver) except* with gil:
+    cdef intptr_t handle = path_finder.load_nvidia_dynamic_library("nvvm")
+    return <void*>handle
 
 
 cdef int _check_or_init_nvvm() except -1 nogil:
@@ -96,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')