Initial version of cuda.pathfinder._find_nvidia_headers for nvshmem

.github/workflows/test-wheel-linux.yml

@@ -298,6 +298,7 @@ jobs:
       - name: Run cuda.pathfinder tests with see_what_works
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: see_what_works
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: see_what_works
         run: run-tests pathfinder
 
       - name: Run cuda.bindings tests
@@ -321,16 +322,16 @@ jobs:
             pip install $(ls cuda_python*.whl)[all]
           fi
 
-      - name: Install cuda.pathfinder nvidia_wheels_cu13
-        if: startsWith(matrix.CUDA_VER, '13.')
+      - name: Install cuda.pathfinder extra wheels for testing
         run: |
+          set -euo pipefail
           pushd cuda_pathfinder
-          pip install -v .[nvidia_wheels_cu13]
-          pip freeze
+          pip install -v ".[nvidia_wheels_cu${TEST_CUDA_MAJOR}]"
+          pip list
           popd
 
       - name: Run cuda.pathfinder tests with all_must_work
-        if: startsWith(matrix.CUDA_VER, '13.')
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: all_must_work
         run: run-tests pathfinder

.github/workflows/test-wheel-windows.yml

@@ -262,6 +262,7 @@ jobs:
       - name: Run cuda.pathfinder tests with see_what_works
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: see_what_works
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: see_what_works
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: run-tests pathfinder
 
@@ -288,18 +289,17 @@ jobs:
             pip install "$((Get-ChildItem -Filter cuda_python*.whl).FullName)[all]"
           }
 
-      - name: Install cuda.pathfinder nvidia_wheels_cu13
-        if: startsWith(matrix.CUDA_VER, '13.')
+      - name: Install cuda.pathfinder extra wheels for testing
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: |
           pushd cuda_pathfinder
-          pip install -v .[nvidia_wheels_cu13]
-          pip freeze
+          pip install -v ".[nvidia_wheels_cu${TEST_CUDA_MAJOR}]"
+          pip list
           popd
 
       - name: Run cuda.pathfinder tests with all_must_work
-        if: startsWith(matrix.CUDA_VER, '13.')
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: all_must_work
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: run-tests pathfinder

ci/tools/env-vars

@@ -69,6 +69,7 @@ elif [[ "${1}" == "test" ]]; then
   echo "SETUP_SANITIZER=${SETUP_SANITIZER}" >> $GITHUB_ENV
   echo "SKIP_CUDA_BINDINGS_TEST=${SKIP_CUDA_BINDINGS_TEST}" >> $GITHUB_ENV
   echo "SKIP_CYTHON_TEST=${SKIP_CYTHON_TEST}" >> $GITHUB_ENV
+  echo "TEST_CUDA_MAJOR=${TEST_CUDA_MAJOR}" >> $GITHUB_ENV
 fi
 
 echo "CUDA_BINDINGS_ARTIFACT_BASENAME=${CUDA_BINDINGS_ARTIFACT_BASENAME}" >> $GITHUB_ENV

ci/tools/run-tests

@@ -31,7 +31,9 @@ popd
 
 if [[ "${test_module}" == "pathfinder" ]]; then
   pushd ./cuda_pathfinder
-  echo "Running pathfinder tests with ${CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS}"
+  echo "Running pathfinder tests with " \
+      "LD:${CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS} " \
+      "FH:${CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS}"
   pwd
   pytest -ra -s -v tests/
   popd

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -7,4 +7,7 @@
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
     SUPPORTED_LIBNAMES as SUPPORTED_NVIDIA_LIBNAMES,  # noqa: F401
 )
+from cuda.pathfinder._headers.find_nvidia_headers import (
+    find_nvidia_header_directory as _find_nvidia_header_directory,  # noqa: F401
+)
 from cuda.pathfinder._version import __version__ as __version__

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import glob
+import os
+from typing import Optional
+
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import IS_WINDOWS
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+
+
+@functools.cache
+def find_nvidia_header_directory(libname: str) -> Optional[str]:
+    if libname != "nvshmem":
+        raise RuntimeError(f"UNKNOWN {libname=}")
+
+    if libname == "nvshmem" and IS_WINDOWS:
+        # nvshmem has no Windows support.
+        return None
+
+    # Installed from a wheel
+    nvidia_sub_dirs = ("nvidia", "nvshmem", "include")
+    hdr_dir: str  # help mypy
+    for hdr_dir in find_sub_dirs_all_sitepackages(nvidia_sub_dirs):
+        nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+        if os.path.isfile(nvshmem_h_path):
+            return hdr_dir
+
+    conda_prefix = os.environ.get("CONDA_PREFIX")
+    if conda_prefix and os.path.isdir(conda_prefix):
+        hdr_dir = os.path.join(conda_prefix, "include")
+        if os.path.isdir(hdr_dir):
+            nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+            if os.path.isfile(nvshmem_h_path):
+                return hdr_dir
+
+    for hdr_dir in sorted(glob.glob("/usr/include/nvshmem_*"), reverse=True):
+        if os.path.isdir(hdr_dir):
+            nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+            if os.path.isfile(nvshmem_h_path):
+                return hdr_dir
+
+    return None

cuda_pathfinder/pyproject.toml

@@ -17,10 +17,12 @@ test = [
 nvidia_wheels_cu12 = [
     "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg]==12.*",
     "cuda-toolkit[cufile]==12.*; sys_platform != 'win32'",
+    "nvidia-nvshmem-cu12; sys_platform != 'win32'",
 ]
 nvidia_wheels_cu13 = [
     "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,nvvm]==13.*",
     "cuda-toolkit[cufile]==13.*; sys_platform != 'win32'",
+    "nvidia-nvshmem-cu13; sys_platform != 'win32'",
 ]
 
 [project.urls]

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Currently these installations are only manually tested:
+
+# conda create -y -n nvshmem python=3.12
+# conda activate nvshmem
+# conda install -y conda-forge::libnvshmem3 conda-forge::libnvshmem-dev
+
+# wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+# sudo dpkg -i cuda-keyring_1.1-1_all.deb
+# sudo apt update
+# sudo apt install libnvshmem3-cuda-12 libnvshmem3-dev-cuda-12
+# sudo apt install libnvshmem3-cuda-13 libnvshmem3-dev-cuda-13
+
+import functools
+import importlib.metadata
+import os
+import re
+
+import pytest
+
+from cuda.pathfinder import _find_nvidia_header_directory as find_nvidia_header_directory
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import IS_WINDOWS
+
+STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS", "see_what_works")
+assert STRICTNESS in ("see_what_works", "all_must_work")
+
+
+@functools.cache
+def have_nvidia_nvshmem_package() -> bool:
+    pattern = re.compile(r"^nvidia-nvshmem-.*$")
+    return any(
+        pattern.match(dist.metadata["Name"]) for dist in importlib.metadata.distributions() if "Name" in dist.metadata
+    )
+
+
+def test_unknown_libname():
+    with pytest.raises(RuntimeError, match=r"^UNKNOWN libname='unknown-libname'$"):
+        find_nvidia_header_directory("unknown-libname")
+
+
+def test_find_libname_nvshmem(info_summary_append):
+    hdr_dir = find_nvidia_header_directory("nvshmem")
+    info_summary_append(f"{hdr_dir=!r}")
+    if IS_WINDOWS:
+        assert hdr_dir is None
+        pytest.skip("nvshmem has no Windows support.")
+    if STRICTNESS == "all_must_work" or have_nvidia_nvshmem_package():
+        assert hdr_dir is not None
+        if have_nvidia_nvshmem_package():
+            hdr_dir_parts = hdr_dir.split(os.path.sep)
+            assert any(
+                sub_dir in hdr_dir_parts
+                for sub_dir in (
+                    "site-packages",  # pip install
+                    "dist-packages",  # apt install
+                )
+            ), hdr_dir