diff --git a/ci/test_thirdparty_awkward.sh b/ci/test_thirdparty_awkward.sh index a52c31503..6349ac1c2 100755 --- a/ci/test_thirdparty_awkward.sh +++ b/ci/test_thirdparty_awkward.sh @@ -5,6 +5,7 @@ set -euo pipefail CUDA_VER_MAJOR_MINOR=${CUDA_VER%.*} +CUDA_VER_MAJOR=${CUDA_VER%.*.*} AWKWARD_VERSION="2.8.10" rapids-logger "Install awkward and related libraries" @@ -15,10 +16,7 @@ rapids-logger "Install wheel with test dependencies" package=$(realpath "${NUMBA_CUDA_ARTIFACTS_DIR}"/*.whl) echo "Package path: ${package}" python -m pip install \ - "${package}" \ - "cuda-python==${CUDA_VER_MAJOR_MINOR%.*}.*" \ - "cuda-core" \ - "nvidia-nvjitlink-cu12" \ + "${package}[cu${CUDA_VER_MAJOR}]" \ --group test diff --git a/numba_cuda/numba/cuda/cuda_paths.py b/numba_cuda/numba/cuda/cuda_paths.py index 4f86064f7..0b7f031d3 100644 --- a/numba_cuda/numba/cuda/cuda_paths.py +++ b/numba_cuda/numba/cuda/cuda_paths.py @@ -9,6 +9,9 @@ from numba.cuda.core.config import IS_WIN32 from numba.cuda.misc.findlib import find_lib from numba.cuda import config +from cuda import pathfinder +import pathlib +from contextlib import contextmanager _env_path_tuple = namedtuple("_env_path_tuple", ["by", "info"]) @@ -20,6 +23,20 @@ ] +@contextmanager +def temporary_env_var(key, value): + """Context manager to temporarily set an environment variable.""" + old_value = os.environ.get(key) + os.environ[key] = value + try: + yield + finally: + if old_value is None: + os.environ.pop(key, None) + else: + os.environ[key] = old_value + + def _get_distribution(distribution_name): """Get the distribution path using importlib.metadata, returning None if not found.""" try: @@ -83,112 +100,6 @@ def _get_libdevice_path_decision(): return _find_first_valid_lazy(options) -def _get_nvvm_path_decision(): - options = _build_options( - [ - ("Conda environment", _get_nvvm_conda_path), - ("NVIDIA NVCC Wheel", _get_nvvm_wheel_path), - ("CUDA_HOME", _get_nvvm_cuda_home_path), - ("System", _get_nvvm_system_path), - ] - ) - return _find_first_valid_lazy(options) - - -def _get_nvrtc_path_decision(): - options = _build_options( - [ - ("Conda environment", get_conda_ctk_libdir), - ("NVIDIA NVCC Wheel", _get_nvrtc_wheel_libdir), - ("CUDA_HOME", get_cuda_home_libdir), - ("System", get_system_ctk_libdir), - ] - ) - return _find_first_valid_lazy(options) - - -def _get_nvvm_wheel_path(): - dso_path = None - # CUDA 12 - nvcc_distribution = _get_distribution("nvidia-cuda-nvcc-cu12") - if nvcc_distribution is not None: - site_packages_path = nvcc_distribution.locate_file("") - nvvm_lib_dir = os.path.join( - site_packages_path, - "nvidia", - "cuda_nvcc", - "nvvm", - "bin" if IS_WIN32 else "lib64", - ) - dso_path = os.path.join( - nvvm_lib_dir, "nvvm64_40_0.dll" if IS_WIN32 else "libnvvm.so" - ) - - # CUDA 13 - if dso_path is None: - nvcc_distribution = _get_distribution("nvidia-nvvm") - if ( - nvcc_distribution is not None - and nvcc_distribution.version.startswith("13.") - ): - site_packages_path = nvcc_distribution.locate_file("") - nvvm_lib_dir = os.path.join( - site_packages_path, - "nvidia", - "cu13", - "bin" if IS_WIN32 else "lib", - "x86_64" if IS_WIN32 else "", - ) - dso_path = os.path.join( - nvvm_lib_dir, "nvvm64_40_0.dll" if IS_WIN32 else "libnvvm.so.4" - ) - - if dso_path and os.path.isfile(dso_path): - return dso_path - return None - - -def _get_nvrtc_wheel_libdir(): - dso_path = None - # CUDA 12 - nvrtc_distribution = _get_distribution("nvidia-cuda-nvrtc-cu12") - if nvrtc_distribution is not None: - site_packages_path = nvrtc_distribution.locate_file("") - nvrtc_lib_dir = os.path.join( - site_packages_path, - "nvidia", - "cuda_nvrtc", - "bin" if IS_WIN32 else "lib", - ) - dso_path = os.path.join( - nvrtc_lib_dir, "nvrtc64_120_0.dll" if IS_WIN32 else "libnvrtc.so.12" - ) - - # CUDA 13 - if dso_path is None: - nvrtc_distribution = _get_distribution("nvidia-cuda-nvrtc") - if ( - nvrtc_distribution is not None - and nvrtc_distribution.version.startswith("13.") - ): - site_packages_path = nvrtc_distribution.locate_file("") - nvrtc_lib_dir = os.path.join( - site_packages_path, - "nvidia", - "cu13", - "bin" if IS_WIN32 else "lib", - "x86_64" if IS_WIN32 else "", - ) - dso_path = os.path.join( - nvrtc_lib_dir, - "nvrtc64_130_0.dll" if IS_WIN32 else "libnvrtc.so.13", - ) - - if dso_path and os.path.isfile(dso_path): - return os.path.dirname(dso_path) - return None - - def _get_libdevice_path(): by, out = _get_libdevice_path_decision() if not out: @@ -321,22 +232,6 @@ def get_system_ctk_include(): return None -def _get_nvvm_system_path(): - nvvm_lib_dir = get_system_ctk("nvvm") - if nvvm_lib_dir is None: - return None - nvvm_lib_dir = os.path.join(nvvm_lib_dir, "bin" if IS_WIN32 else "lib64") - if IS_WIN32 and os.path.isdir(os.path.join(nvvm_lib_dir, "x64")): - nvvm_lib_dir = os.path.join(nvvm_lib_dir, "x64") - - nvvm_path = os.path.join( - nvvm_lib_dir, "nvvm64_40_0.dll" if IS_WIN32 else "libnvvm.so.4" - ) - # if os.path.isfile(nvvm_path): - # return nvvm_path - return nvvm_path - - def get_conda_ctk_libdir(): """Return path to directory containing the shared libraries of cudatoolkit.""" is_conda_env = os.path.isdir(os.path.join(sys.prefix, "conda-meta")) @@ -378,29 +273,6 @@ def get_libdevice_conda_path(): return None -def _get_nvvm_conda_path(): - """Return path to directory containing the nvvm library.""" - is_conda_env = os.path.isdir(os.path.join(sys.prefix, "conda-meta")) - if not is_conda_env: - return None - nvvm_dir = os.path.join( - sys.prefix, - "Library" if IS_WIN32 else "", - "nvvm", - "bin" if IS_WIN32 else "lib64", - ) - # Windows CUDA 13.0.0 puts in "bin\x64" directory but 13.0.1+ just uses "bin" directory - if IS_WIN32 and os.path.isdir(os.path.join(nvvm_dir, "x64")): - nvvm_dir = os.path.join(nvvm_dir, "x64") - - nvvm_path = os.path.join( - nvvm_dir, "nvvm64_40_0.dll" if IS_WIN32 else "libnvvm.so.4" - ) - if os.path.isfile(nvvm_path): - return nvvm_path - return None - - def get_wheel_static_libdir(): cuda_module_static_lib_dir = None # CUDA 12 @@ -528,43 +400,13 @@ def get_cuda_home_include(): return None -def _get_nvvm_cuda_home_path(): - nvvm_lib_dir = get_cuda_home("nvvm") - if nvvm_lib_dir is None: - return - nvvm_lib_dir = os.path.join(nvvm_lib_dir, "bin" if IS_WIN32 else "lib64") - if IS_WIN32 and os.path.isdir(os.path.join(nvvm_lib_dir, "x64")): - nvvm_lib_dir = os.path.join(nvvm_lib_dir, "x64") - - nvvm_path = os.path.join( - nvvm_lib_dir, "nvvm64_40_0.dll" if IS_WIN32 else "libnvvm.so.4" - ) - # if os.path.isfile(nvvm_path): - # return nvvm_path - return nvvm_path - - -def _get_nvvm_path(): - by, out = _get_nvvm_path_decision() - if not out: - return _env_path_tuple(by, None) - return _env_path_tuple(by, out) - - -def _get_nvrtc_path(): - by, path = _get_nvrtc_path_decision() - candidates = find_lib("nvrtc", libdir=path) - path = max(candidates) if candidates else None - return _env_path_tuple(by, path) - - def get_cuda_paths(): """Returns a dictionary mapping component names to a 2-tuple of (source_variable, info). The returned dictionary will have the following keys and infos: - - "nvvm": file_path - "nvrtc": file_path + - "nvvm": file_path - "libdevice": file_path - "cudalib_dir": directory_path - "static_cudalib_dir": directory_path @@ -578,8 +420,8 @@ def get_cuda_paths(): else: # Not in cache d = { - "nvvm": _get_nvvm_path(), "nvrtc": _get_nvrtc_path(), + "nvvm": _get_nvvm_path(), "libdevice": _get_libdevice_path(), "cudalib_dir": _get_cudalib_dir(), "static_cudalib_dir": _get_static_cudalib_dir(), @@ -689,3 +531,92 @@ def _get_include_dir(): ] by, include_dir = _find_valid_path(options) return _env_path_tuple(by, include_dir) + + +def _find_cuda_home_from_lib_path(lib_path): + """ + Walk up from a library path to find a directory containing 'nvvm' subdirectory. + + For example, given /usr/local/cuda/lib64/libnvrtc.so.12, + this would find /usr/local/cuda (which contains nvvm/). + + Returns the path if found, None otherwise. + """ + current = pathlib.Path(lib_path).resolve() + + # Walk up the directory tree + for parent in current.parents: + nvvm_subdir = parent / "nvvm" + if nvvm_subdir.is_dir(): + return str(parent) + + return None + + +def _get_nvvm(): + # Strategy: + # 1. Try pathfinder directly + # 2. If CUDA_HOME/CUDA_PATH are set, pathfinder would have found it - give up + # 3. Use nvrtc's location to infer CUDA installation root + # 4. Temporarily set CUDA_HOME and retry pathfinder + # First, try pathfinder directly + try: + return pathfinder.load_nvidia_dynamic_lib("nvvm") + except pathfinder.DynamicLibNotFoundError as e: + nvvm_exc = e + + def _raise_original(reason: str) -> None: + raise pathfinder.DynamicLibNotFoundError( + f"{reason}; original nvvm error: {nvvm_exc}" + ) from nvvm_exc + + # If CUDA_HOME or CUDA_PATH is set, pathfinder would have found libnvvm + # based on the environment variable(s) - nothing more we can do + if os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH"): + _raise_original("nvvm not found and CUDA_HOME/CUDA_PATH is set") + # Try to locate nvrtc - this library is almost certainly needed if nvvm is needed (in the context of numba-cuda) + try: + loaded_nvrtc = _get_nvrtc() + except Exception as nvrtc_exc: + raise pathfinder.DynamicLibNotFoundError( + f"nvrtc load failed while inferring CUDA_HOME; original nvvm error: {nvvm_exc}" + ) from nvrtc_exc + # If nvrtc was not found via system-search, we can't reliably determine + # the CUDA installation structure + if loaded_nvrtc.found_via != "system-search": + _raise_original( + f"nvrtc found via {loaded_nvrtc.found_via}, cannot infer CUDA_HOME" + ) + # Search backward from nvrtc's location to find a directory with "nvvm" subdirectory + cuda_home = _find_cuda_home_from_lib_path(loaded_nvrtc.abs_path) + if cuda_home is None: + _raise_original( + f"nvrtc path did not map to CUDA_HOME ({loaded_nvrtc.abs_path})" + ) + # Temporarily set CUDA_HOME and retry pathfinder + with temporary_env_var("CUDA_HOME", cuda_home): + try: + library = pathfinder.load_nvidia_dynamic_lib("nvvm") + except pathfinder.DynamicLibNotFoundError as exc: + raise pathfinder.DynamicLibNotFoundError( + f"nvvm not found after inferring CUDA_HOME={cuda_home}; " + f"original nvvm error: {nvvm_exc}" + ) from exc + library.found_via = "system-search" + return library + + +def _get_nvrtc(): + return pathfinder.load_nvidia_dynamic_lib("nvrtc") + + +def _get_nvrtc_path(): + # the pathfinder API will either find the library or raise + nvrtc = _get_nvrtc() + return _env_path_tuple(nvrtc.found_via, nvrtc.abs_path) + + +def _get_nvvm_path(): + # the pathfinder API will either find the library or raise + nvvm = _get_nvvm() + return _env_path_tuple(nvvm.found_via, nvvm.abs_path) diff --git a/numba_cuda/numba/cuda/cudadrv/libs.py b/numba_cuda/numba/cuda/cudadrv/libs.py index a4a6a1764..f99eb2b71 100644 --- a/numba_cuda/numba/cuda/cudadrv/libs.py +++ b/numba_cuda/numba/cuda/cudadrv/libs.py @@ -20,7 +20,7 @@ from numba.cuda.cuda_paths import get_cuda_paths from numba.cuda.cudadrv.driver import locate_driver_and_loader, load_driver from numba.cuda.cudadrv.error import CudaSupportError -from numba.cuda.core import config +from numba.cuda import config if sys.platform == "win32": @@ -53,7 +53,9 @@ def get_cudalib(lib, static=False): loader's search mechanism. """ if lib in {"nvrtc", "nvvm"}: - return get_cuda_paths()[lib].info or _dllnamepattern % lib + # System search either invoked inside cuda-pathfinder + # or, for nvvm, using custom logic inside cuda-paths + return get_cuda_paths()[lib].info dir_type = "static_cudalib_dir" if static else "cudalib_dir" libdir = get_cuda_paths()[dir_type].info @@ -69,7 +71,6 @@ def get_cuda_include_dir(): Note that this does not list the `CUDA_INCLUDE_PATH` entry in user configuration. """ - return get_cuda_paths()["include_dir"].info diff --git a/numba_cuda/numba/cuda/tests/nocuda/test_library_lookup.py b/numba_cuda/numba/cuda/tests/nocuda/test_library_lookup.py index 351f2d548..83f9de5cf 100644 --- a/numba_cuda/numba/cuda/tests/nocuda/test_library_lookup.py +++ b/numba_cuda/numba/cuda/tests/nocuda/test_library_lookup.py @@ -5,7 +5,7 @@ import os import multiprocessing as mp import warnings - +import pathlib from numba.cuda.core.config import IS_WIN32 from numba.cuda.core.errors import NumbaWarning @@ -17,10 +17,10 @@ ) from numba.cuda.cuda_paths import ( _get_libdevice_path_decision, - _get_nvvm_path_decision, _get_cudalib_dir_path_decision, get_system_ctk, get_system_ctk_libdir, + _find_cuda_home_from_lib_path, ) @@ -134,58 +134,6 @@ def do_set_cuda_home(): return True, _get_libdevice_path_decision() -@skip_on_cudasim("Library detection unsupported in the simulator") -@unittest.skipUnless(has_mp_get_context, "mp.get_context not available") -@skip_unless_conda_cudatoolkit("test assumes conda installed cudatoolkit") -class TestNvvmLookUp(LibraryLookupBase): - def test_nvvm_path_decision(self): - # Check that the default is using conda environment - by, info, warns = self.remote_do(self.do_clear_envs) - if has_cuda: - self.assertEqual(by, "Conda environment") - else: - self.assertEqual(by, "") - self.assertIsNone(info) - self.assertFalse(warns) - # Check that CUDA_HOME works by removing conda-env - by, info, warns = self.remote_do(self.do_set_cuda_home) - self.assertEqual(by, "CUDA_HOME") - self.assertFalse(warns) - if IS_WIN32: - self.assertEqual( - os.path.dirname(info), os.path.join("mycudahome", "nvvm", "bin") - ) - else: - self.assertEqual( - os.path.dirname(info), - os.path.join("mycudahome", "nvvm", "lib64"), - ) - - if get_system_ctk("nvvm") is None: - # Fake remove conda environment so no cudatoolkit is available - by, info, warns = self.remote_do(self.do_clear_envs) - self.assertEqual(by, "") - self.assertIsNone(info) - self.assertFalse(warns) - else: - # Use system available cudatoolkit - by, info, warns = self.remote_do(self.do_clear_envs) - self.assertEqual(by, "System") - self.assertFalse(warns) - - @staticmethod - def do_clear_envs(): - remove_env("CUDA_HOME") - remove_env("CUDA_PATH") - return True, _get_nvvm_path_decision() - - @staticmethod - def do_set_cuda_home(): - os.environ["CUDA_HOME"] = os.path.join("mycudahome") - _fake_non_conda_env() - return True, _get_nvvm_path_decision() - - @skip_on_cudasim("Library detection unsupported in the simulator") @unittest.skipUnless(has_mp_get_context, "mp.get_context not available") @skip_unless_conda_cudatoolkit("test assumes conda installed cudatoolkit") @@ -248,5 +196,40 @@ def _fake_non_conda_env(): sys.prefix = "" +class TestCudaHomeDetection(unittest.TestCase): + def test_find_cuda_home(self): + """Test the directory walking logic.""" + import tempfile + + # Create a mock CUDA installation structure + with tempfile.TemporaryDirectory() as tmpdir: + cuda_root = pathlib.Path(tmpdir) / "cuda" + lib64 = cuda_root / "lib64" + nvvm = cuda_root / "nvvm" + nvvm_lib64 = nvvm / "lib64" + + lib64.mkdir(parents=True) + nvvm_lib64.mkdir(parents=True) + + # Create mock library files + nvrtc_path = lib64 / "libnvrtc.so.12" + nvrtc_path.touch() + + nvvm_lib = nvvm_lib64 / "libnvvm.so.4" + nvvm_lib.touch() + + # Test finding CUDA_HOME from nvrtc path + found_cuda_home = _find_cuda_home_from_lib_path(str(nvrtc_path)) + + # Compare resolved paths to handle Windows short path names (e.g., RGROSS~1) + expected = str(cuda_root.resolve()) + assert found_cuda_home == expected, ( + f"Expected {expected}, got {found_cuda_home}" + ) + + # Test that the nvvm directory exists at the found location + assert (pathlib.Path(found_cuda_home) / "nvvm").is_dir() + + if __name__ == "__main__": unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 0c951689e..0f66fbb98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,12 +27,14 @@ dependencies = ["numba>=0.60.0", "cuda-bindings>=12.9.1,<14.0.0", "cuda-core>=0. cu12 = [ "cuda-bindings>=12.9.1,<13.0.0", "cuda-core>=0.3.0,<1.0.0", + "cuda-pathfinder>=1.3.1,<2.0.0", # install nvcc for libNVVM "cuda-toolkit[cudart,nvcc,nvrtc,nvjitlink,cccl]==12.*", ] cu13 = [ "cuda-bindings==13.*", "cuda-core>=0.3.2,<1.0.0", + "cuda-pathfinder>=1.3.1,<2.0.0", "cuda-toolkit[cudart,nvvm,nvrtc,nvjitlink,cccl]==13.*", ]