updates

Author: brandon-b-miller

numba_cuda/numba/cuda/cuda_paths.py

@@ -10,6 +10,8 @@
 from numba.cuda.misc.findlib import find_lib
 from numba.cuda import config
 import warnings
+from cuda import pathfinder
+import pathlib
 
 _env_path_tuple = namedtuple("_env_path_tuple", ["by", "info"])
 
@@ -516,3 +518,25 @@ def _get_include_dir():
     ]
     by, include_dir = _find_valid_path(options)
     return _env_path_tuple(by, include_dir)
+
+
+def _get_nvvm():
+    try:
+        nvvm = pathfinder.load_nvidia_dynamic_lib("nvvm")
+        return nvvm
+    except pathfinder.DynamicLibNotFoundError:
+        nvrtc = _get_nvrtc()
+        path = pathlib.Path(nvrtc.abs_path)
+        nvvm = path.parents[1] / "nvvm" / "lib64" / "libnvvm.so"
+
+        if nvvm.exists():
+            dl = pathfinder._dynamic_libs.load_nvidia_dynamic_lib.load_with_abs_path(
+                "nvvm", path, "system-search"
+            )
+            return dl
+        else:
+            raise pathfinder.DynamicLibNotFoundError("nvvm not found")
+
+
+def _get_nvrtc():
+    return pathfinder.load_nvidia_dynamic_lib("nvrtc")

numba_cuda/numba/cuda/cudadrv/libs.py

@@ -17,7 +17,7 @@
 import ctypes
 
 from numba.cuda.misc.findlib import find_lib
-from numba.cuda.cuda_paths import get_cuda_paths
+from numba.cuda.cuda_paths import get_cuda_paths, _get_nvvm, _get_nvrtc
 from numba.cuda.cudadrv.driver import locate_driver_and_loader, load_driver
 from numba.cuda.cudadrv.error import CudaSupportError
 from numba.core import config
@@ -54,8 +54,10 @@ def get_cudalib(lib, static=False):
     'libnvvm.so' for 'nvvm') so that we may attempt to load it using the system
     loader's search mechanism.
     """
-    if lib in {"nvrtc", "nvvm"}:
-        return pathfinder.load_nvidia_dynamic_lib(lib).abs_path
+    if lib == "nvrtc":
+        return _get_nvrtc().abs_path
+    elif lib == "nvvm":
+        return _get_nvvm().abs_path
     else:
         # cudart, cudadevrt
         dir_type = "static_cudalib_dir" if static else "cudalib_dir"