Add _inspect_environment() in find_nvidia_dynamic_library.py, call from nvjitlink_windows.pyx, nvvm_windows.pyx

Author: rwgk

cuda_bindings/cuda/bindings/_internal/nvjitlink_windows.pyx

@@ -10,6 +10,8 @@ from .utils cimport get_nvjitlink_dso_version_suffix
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
+from cuda.bindings import path_finder
+
 import os
 import site
 
@@ -88,6 +90,8 @@ cdef load_library(const int driver_ver):
     else:
         raise RuntimeError('Failed to load nvJitLink')
 
+    path_finder.find_nvidia_dynamic_library("nvJitLink", handle)
+
     assert handle != 0
     return handle
 

cuda_bindings/cuda/bindings/_internal/nvvm_windows.pyx

@@ -10,6 +10,8 @@ from .utils cimport get_nvvm_dso_version_suffix
 
 from .utils import FunctionNotFoundError, NotSupportedError
 
+from cuda.bindings import path_finder
+
 import os
 import site
 
@@ -86,6 +88,8 @@ cdef load_library(const int driver_ver):
     else:
         raise RuntimeError('Failed to load nvvm')
 
+    path_finder.find_nvidia_dynamic_library("nvvm", handle)
+
     assert handle != 0
     return handle
 

cuda_bindings/cuda/bindings/_path_finder_utils/find_nvidia_dynamic_library.py

@@ -5,8 +5,10 @@
 import functools
 import glob
 import os
+import sys
+import traceback
 
-from .cuda_paths import get_cuda_paths
+from .cuda_paths import IS_WIN32, get_cuda_paths
 from .find_nvidia_lib_dirs import find_nvidia_lib_dirs
 
 
@@ -70,8 +72,38 @@ def _find_using_cudalib_dir(so_basename, error_messages, attachments):
     return None
 
 
+def _inspect_environment(libbasename, handle):
+    if IS_WIN32:
+        import win32api
+
+        dll_path = win32api.GetModuleFileName(handle)
+        print(f"LOOOK {libbasename=} Loaded DLL path:", dll_path)
+    error_messages = []
+    lib_dir = _get_cuda_paths_info("cudalib_dir", error_messages)
+    if lib_dir is None:
+        print(f"LOOOK {libbasename=} {error_messages=}")
+    elif not os.path.isdir(lib_dir):
+        print(f"LOOOK {libbasename=} not isdir({lib_dir=})")
+    else:
+        print(f"LOOOK {libbasename=} cudalib_dir {lib_dir=}")
+        for node in sorted(os.listdir(lib_dir)):
+            print(f"LOOOK     {node}")
+    for lib_dir in find_nvidia_lib_dirs():
+        print(f"LOOOK {libbasename=} NVIDIA {lib_dir=}")
+        for node in sorted(os.listdir(lib_dir)):
+            print(f"LOOOK     {node}")
+
+
 @functools.cache
-def find_nvidia_dynamic_library(libbasename):
+def find_nvidia_dynamic_library(libbasename, handle=None):
+    if handle is not None:
+        try:
+            _inspect_environment(libbasename, handle)
+        except Exception as e:
+            print("LOOOK EXCEPTION:")
+            traceback.print_exception(type(e), e, e.__traceback__, file=sys.stdout)
+        if IS_WIN32:
+            return
     so_basename = f"lib{libbasename}.so"
     error_messages = []
     attachments = []

cuda_bindings/tests/path_finder.py

@@ -5,5 +5,5 @@
 for k, v in paths.items():
     print(f"{k}: {v}", flush=True)
 
-print(path_finder.find_nvidia_dynamic_library("nvvm"))
-print(path_finder.find_nvidia_dynamic_library("nvJitLink"))
+print(path_finder.find_nvidia_dynamic_library("nvvm", "TEST"))
+print(path_finder.find_nvidia_dynamic_library("nvJitLink", "TEST"))