Locate nvvm, libdevice, nvrtc, and cudart from nvidia-*-cu12 wheels

.github/workflows/pr.yaml

@@ -23,6 +23,7 @@ jobs:
       - build-wheels
       - test-wheels
       - test-wheels-pynvjitlink
+      - test-wheels-deps-wheels
       - build-docs
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.04
@@ -92,6 +93,14 @@ jobs:
       build_type: pull-request
       script: "ci/test_wheel_pynvjitlink.sh"
       matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "12.5.1" and .PY_VER == "3.12"))
+  test-wheels-deps-wheels:
+    needs:
+      - build-wheels
+    uses: ./.github/workflows/wheels-test.yaml
+    with:
+      build_type: pull-request
+      script: "ci/test_wheel_deps_wheels.sh"
+      matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "12.5.1" and .PY_VER == "3.12"))
   build-docs:
     needs:
       - build-conda

ci/test_wheel_deps_wheels.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+# Copyright (c) 2023-2024, NVIDIA CORPORATION
+
+set -euo pipefail
+
+rapids-logger "Install testing dependencies"
+# TODO: Replace with rapids-dependency-file-generator
+python -m pip install \
+    psutil \
+    cffi \
+    cuda-python \
+    nvidia-curand-cu12 \
+    nvidia-cuda-nvcc-cu12 \
+    pytest
+
+
+rapids-logger "Build tests"
+PY_SCRIPT="
+import numba_cuda
+root = numba_cuda.__file__.rstrip('__init__.py')
+test_dir = root + \"numba/cuda/tests/test_binary_generation/\"
+print(test_dir)
+"
+
+NUMBA_CUDA_TEST_BIN_DIR=$(python -c "$PY_SCRIPT")
+pushd $NUMBA_CUDA_TEST_BIN_DIR
+make
+popd
+
+rapids-logger "Install wheel"
+package=$(realpath wheel/numba_cuda*.whl)
+echo "Package path: $package"
+python -m pip install $package
+
+rapids-logger "Check GPU usage"
+nvidia-smi
+
+RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
+mkdir -p "${RAPIDS_TESTS_DIR}"
+pushd "${RAPIDS_TESTS_DIR}"
+
+rapids-logger "Show Numba system info"
+python -m numba --sysinfo
+
+# remove cuda-nvvm-12-5 leaving libnvvm.so from nvidia-cuda-nvcc-cu12 only 
+apt-get update
+apt remove --purge cuda-nvvm-12-5 -y
+
+rapids-logger "Run Tests"
+NUMBA_CUDA_TEST_BIN_DIR=$NUMBA_CUDA_TEST_BIN_DIR python -m numba.runtests numba.cuda.tests -v
+
+popd

numba_cuda/numba/cuda/cuda_paths.py

@@ -3,7 +3,8 @@
 import os
 from collections import namedtuple
 import platform
-
+import site
+from pathlib import Path
 from numba.core.config import IS_WIN32
 from numba.misc.findlib import find_lib, find_file
 from numba import config
@@ -29,9 +30,13 @@ def _get_libdevice_path_decision():
         ('Conda environment', get_conda_ctk()),
         ('Conda environment (NVIDIA package)', get_nvidia_libdevice_ctk()),
         ('CUDA_HOME', get_cuda_home('nvvm', 'libdevice')),
-        ('System', get_system_ctk('nvvm', 'libdevice')),
         ('Debian package', get_debian_pkg_libdevice()),
+        ('NVIDIA NVCC Wheel', get_libdevice_wheel()),
     ]
+    libdevice_ctk_dir = get_system_ctk('nvvm', 'libdevice')
+    if os.path.exists(libdevice_ctk_dir):
+        options.append(('System', libdevice_ctk_dir))
+
     by, libdir = _find_valid_path(options)
     return by, libdir
 
@@ -48,19 +53,58 @@ def _get_nvvm_path_decision():
         ('Conda environment', get_conda_ctk()),
         ('Conda environment (NVIDIA package)', get_nvidia_nvvm_ctk()),
         ('CUDA_HOME', get_cuda_home(*_nvvm_lib_dir())),
-        ('System', get_system_ctk(*_nvvm_lib_dir())),
+        ('NVIDIA NVCC Wheel', _get_nvvm_wheel()),
     ]
+    # need to ensure nvvm dir actually exists
+    nvvm_ctk_dir = get_system_ctk(*_nvvm_lib_dir())
+    if os.path.exists(nvvm_ctk_dir):
+        options.append(('System', nvvm_ctk_dir))
+
     by, path = _find_valid_path(options)
     return by, path
 
 
+def _get_nvvm_wheel():
+    site_paths = [
+        site.getusersitepackages()
+    ] + site.getsitepackages() + ["conda", None]
+    for sp in site_paths:
+        # The SONAME is taken based on public CTK 12.x releases
+        if sys.platform.startswith("linux"):
+            dso_dir = "lib64"
+            # Hack: libnvvm from Linux wheel
+            # does not have any soname (CUDAINST-3183)
+            dso_path = "libnvvm.so"
+        elif sys.platform.startswith("win32"):
+            dso_dir = "bin"
+            dso_path = "nvvm64_40_0.dll"
+        else:
+            raise AssertionError()
+
+        if sp is not None:
+            dso_dir = os.path.join(
+                sp,
+                "nvidia",
+                "cuda_nvcc",
+                "nvvm",
+                dso_dir
+            )
+            dso_path = os.path.join(dso_dir, dso_path)
+            if os.path.exists(dso_path):
+                return str(Path(dso_path).parent)
+
+
 def _get_libdevice_paths():
     by, libdir = _get_libdevice_path_decision()
-    # Search for pattern
-    pat = r'libdevice(\.\d+)*\.bc$'
-    candidates = find_file(re.compile(pat), libdir)
-    # Keep only the max (most recent version) of the bitcode files.
-    out = max(candidates, default=None)
+    if by == "NVIDIA NVCC Wheel":
+        # The NVVM path is a directory, not a file
+        out = os.path.join(libdir, "libdevice.10.bc")
+    else:
+        # Search for pattern
+        pat = r'libdevice(\.\d+)*\.bc$'
+        candidates = find_file(re.compile(pat), libdir)
+        # Keep only the max (most recent version) of the bitcode files.
+        out = max(candidates, default=None)
     return _env_path_tuple(by, out)
 
 
@@ -217,8 +261,12 @@ def get_cuda_home(*subdirs):
 
 def _get_nvvm_path():
     by, path = _get_nvvm_path_decision()
-    candidates = find_lib('nvvm', path)
-    path = max(candidates) if candidates else None
+    if by == "NVIDIA NVCC Wheel":
+        # The NVVM path is a directory, not a file
+        path = os.path.join(path, "libnvvm.so")
+    else:
+        candidates = find_lib('nvvm', path)
+        path = max(candidates) if candidates else None
     return _env_path_tuple(by, path)
 
 
@@ -261,6 +309,16 @@ def get_debian_pkg_libdevice():
     return pkg_libdevice_location
 
 
+def get_libdevice_wheel():
+    nvvm_path = _get_nvvm_wheel()
+    if nvvm_path is None:
+        return None
+    nvvm_path = Path(nvvm_path)
+    libdevice_path = nvvm_path.parent / "libdevice"
+
+    return str(libdevice_path)
+
+
 def get_current_cuda_target_name():
     """Determine conda's CTK target folder based on system and machine arch.