diff --git a/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx b/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx index 54e5d51748..f9ae155ccd 100644 --- a/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx +++ b/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx @@ -406,7 +406,8 @@ cdef void* __nvmlDeviceSetRusdSettings_v1 = NULL cdef void* load_library() except* with gil: - return dlopen("libnvidia-ml.so.1", RTLD_NOW | RTLD_GLOBAL) + cdef uintptr_t handle = load_nvidia_dynamic_lib("nvml")._handle_uint + return handle cdef int _init_nvml() except -1 nogil: diff --git a/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx b/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx index 309b5a3039..50cc37b0d7 100644 --- a/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx +++ b/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx @@ -9,6 +9,8 @@ from libc.stdint cimport intptr_t import os import threading +from cuda.pathfinder import load_nvidia_dynamic_lib + from .utils import FunctionNotFoundError, NotSupportedError from libc.stddef cimport wchar_t @@ -422,35 +424,6 @@ cdef void* __nvmlDeviceReadPRMCounters_v1 = NULL cdef void* __nvmlDeviceSetRusdSettings_v1 = NULL -cdef uintptr_t load_library() except* with gil: - def do_load(path): - return LoadLibraryExW( - path, - 0, - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR - ) - - handle = do_load( - os.path.join( - os.getenv("WINDIR", "C:/Windows"), - "System32/nvml.dll" - ) - ) - if handle: - return handle - - handle = do_load( - os.path.join( - os.getenv("ProgramFiles", "C:/Program Files"), - "NVIDIA Corporation/NVSMI/nvml.dll" - ) - ) - if handle: - return handle - - return 0 - - cdef int _init_nvml() except -1 nogil: global __py_nvml_init @@ -458,7 +431,7 @@ cdef int _init_nvml() except -1 nogil: cdef uintptr_t handle with gil, __symbol_lock: - handle = load_library() + handle = load_nvidia_dynamic_lib("nvml")._handle_uint # Load function global __nvmlInit_v2