cuda.pathfinder._find_nvidia_header_directory(): add support for CTK libs

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -1,13 +1,25 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+"""cuda.pathfinder public APIs"""
+
 from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError as DynamicLibNotFoundError
 from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL as LoadedDL
 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib
 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._headers.find_nvidia_headers import find_nvidia_header_directory as find_nvidia_header_directory
+from cuda.pathfinder._headers.supported_nvidia_headers import SUPPORTED_HEADERS_CTK as _SUPPORTED_HEADERS_CTK
 from cuda.pathfinder._version import __version__ as __version__
+
+# Indirection to help Sphinx find the docstring.
+#: Mapping from short CUDA Toolkit (CTK) library names to their canonical
+#: header basenames (used to validate a discovered include directory).
+#: Example: ``"cublas" → "cublas.h"``. The key set is platform-aware
+#: (e.g., ``"cufile"`` may be Linux-only).
+SUPPORTED_HEADERS_CTK = _SUPPORTED_HEADERS_CTK
+
+# Backward compatibility: _find_nvidia_header_directory was added in release 1.2.2.
+# It will be removed in release 1.2.4.
+_find_nvidia_header_directory = find_nvidia_header_directory

cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py

@@ -14,6 +14,7 @@
     SITE_PACKAGES_LIBDIRS_WINDOWS,
     is_suppressed_dll_file,
 )
+from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs, find_sub_dirs_all_sitepackages
 
 
@@ -79,15 +80,8 @@ def _find_dll_using_nvidia_bin_dirs(
     return None
 
 
-def _get_cuda_home() -> Optional[str]:
-    cuda_home = os.environ.get("CUDA_HOME")
-    if cuda_home is None:
-        cuda_home = os.environ.get("CUDA_PATH")
-    return cuda_home
-
-
 def _find_lib_dir_using_cuda_home(libname: str) -> Optional[str]:
-    cuda_home = _get_cuda_home()
+    cuda_home = get_cuda_home_or_path()
     if cuda_home is None:
         return None
     subdirs_list: tuple[tuple[str, ...], ...]

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -6,37 +6,145 @@
 import os
 from typing import Optional
 
-from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import IS_WINDOWS
+from cuda.pathfinder._headers import supported_nvidia_headers
+from cuda.pathfinder._headers.supported_nvidia_headers import IS_WINDOWS
+from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
 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=}")
+def _abs_norm(path: Optional[str]) -> Optional[str]:
+    if path:
+        return os.path.normpath(os.path.abspath(path))
+    return None
+
+
+def _joined_isfile(dirpath: str, basename: str) -> bool:
+    return os.path.isfile(os.path.join(dirpath, basename))
 
-    if libname == "nvshmem" and IS_WINDOWS:
+
+def _find_nvshmem_header_directory() -> Optional[str]:
+    if 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):
+        if _joined_isfile(hdr_dir, "nvshmem.h"):
             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")
-        nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
-        if os.path.isfile(nvshmem_h_path):
+        if _joined_isfile(hdr_dir, "nvshmem.h"):
             return hdr_dir
 
     for hdr_dir in sorted(glob.glob("/usr/include/nvshmem_*"), reverse=True):
-        nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
-        if os.path.isfile(nvshmem_h_path):
+        if _joined_isfile(hdr_dir, "nvshmem.h"):
             return hdr_dir
 
     return None
+
+
+def _find_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str) -> Optional[str]:
+    parts = [anchor_point]
+    if libname == "nvvm":
+        parts.append(libname)
+    parts.append("include")
+    idir = os.path.join(*parts)
+    if libname == "cccl":
+        cdir = os.path.join(idir, "cccl")  # CTK 13
+        if _joined_isfile(cdir, h_basename):
+            return cdir
+    if _joined_isfile(idir, h_basename):
+        return idir
+    return None
+
+
+def _find_based_on_conda_layout(libname: str, h_basename: str, conda_prefix: str) -> Optional[str]:
+    if IS_WINDOWS:
+        anchor_point = os.path.join(conda_prefix, "Library")
+        if not os.path.isdir(anchor_point):
+            return None
+    else:
+        targets_include_path = glob.glob(os.path.join(conda_prefix, "targets", "*", "include"))
+        if not targets_include_path:
+            return None
+        if len(targets_include_path) != 1:
+            # Conda does not support multiple architectures.
+            # QUESTION(PR#956): Do we want to issue a warning?
+            return None
+        anchor_point = os.path.dirname(targets_include_path[0])
+    return _find_based_on_ctk_layout(libname, h_basename, anchor_point)
+
+
+def _find_ctk_header_directory(libname: str) -> Optional[str]:
+    h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_CTK[libname]
+    candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK[libname]
+
+    # Installed from a wheel
+    for cdir in candidate_dirs:
+        hdr_dir: str  # help mypy
+        for hdr_dir in find_sub_dirs_all_sitepackages(tuple(cdir.split("/"))):
+            if _joined_isfile(hdr_dir, h_basename):
+                return hdr_dir
+
+    conda_prefix = os.getenv("CONDA_PREFIX")
+    if conda_prefix:  # noqa: SIM102
+        if result := _find_based_on_conda_layout(libname, h_basename, conda_prefix):
+            return result
+
+    cuda_home = get_cuda_home_or_path()
+    if cuda_home:  # noqa: SIM102
+        if result := _find_based_on_ctk_layout(libname, h_basename, cuda_home):
+            return result
+
+    return None
+
+
+@functools.cache
+def find_nvidia_header_directory(libname: str) -> Optional[str]:
+    """Locate the header directory for a supported NVIDIA library.
+
+    Args:
+        libname (str): The short name of the library whose headers are needed
+            (e.g., ``"nvrtc"``, ``"cusolver"``, ``"nvshmem"``).
+
+    Returns:
+        str or None: Absolute path to the discovered header directory, or ``None``
+        if the headers cannot be found.
+
+    Raises:
+        RuntimeError: If ``libname`` is not in the supported set.
+
+    Search order:
+        1. **NVIDIA Python wheels**
+
+           - Scan installed distributions (``site-packages``) for header layouts
+             shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvrtc]``).
+
+        2. **Conda environments**
+
+           - Check Conda-style installation prefixes, which use platform-specific
+             include directory layouts.
+
+        3. **CUDA Toolkit environment variables**
+
+           - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
+
+    Notes:
+        - The ``SUPPORTED_HEADERS_CTK`` dictionary maps each supported CUDA Toolkit
+          (CTK) library to the name of its canonical header (e.g., ``"cublas" →
+          "cublas.h"``). This is used to verify that the located directory is valid.
+
+        - The only supported non-CTK library at present is ``nvshmem``.
+    """
+
+    if libname == "nvshmem":
+        return _abs_norm(_find_nvshmem_header_directory())
+
+    if libname in supported_nvidia_headers.SUPPORTED_HEADERS_CTK:
+        return _abs_norm(_find_ctk_header_directory(libname))
+
+    raise RuntimeError(f"UNKNOWN {libname=}")

cuda_pathfinder/cuda/pathfinder/_headers/supported_nvidia_headers.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from typing import Final
+
+IS_WINDOWS = sys.platform == "win32"
+
+SUPPORTED_HEADERS_CTK_COMMON = {
+    "cccl": "cuda/std/version",
+    "cublas": "cublas.h",
+    "cudart": "cuda_runtime.h",
+    "cufft": "cufft.h",
+    "curand": "curand.h",
+    "cusolver": "cusolverDn.h",
+    "cusparse": "cusparse.h",
+    "npp": "npp.h",
+    "nvcc": "fatbinary_section.h",
+    "nvfatbin": "nvFatbin.h",
+    "nvjitlink": "nvJitLink.h",
+    "nvjpeg": "nvjpeg.h",
+    "nvrtc": "nvrtc.h",
+    "nvvm": "nvvm.h",
+}
+
+SUPPORTED_HEADERS_CTK_LINUX_ONLY = {
+    "cufile": "cufile.h",
+}
+SUPPORTED_HEADERS_CTK_LINUX = SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_LINUX_ONLY
+
+SUPPORTED_HEADERS_CTK_WINDOWS_ONLY: dict[str, str] = {}
+SUPPORTED_HEADERS_CTK_WINDOWS = SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_WINDOWS_ONLY
+
+SUPPORTED_HEADERS_CTK_ALL = (
+    SUPPORTED_HEADERS_CTK_COMMON | SUPPORTED_HEADERS_CTK_LINUX_ONLY | SUPPORTED_HEADERS_CTK_WINDOWS_ONLY
+)
+SUPPORTED_HEADERS_CTK: Final[dict[str, str]] = (
+    SUPPORTED_HEADERS_CTK_WINDOWS if IS_WINDOWS else SUPPORTED_HEADERS_CTK_LINUX
+)
+
+SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK = {
+    "cccl": (
+        "nvidia/cu13/include/cccl",  # cuda-toolkit[cccl]==13.*
+        "nvidia/cuda_cccl/include",  # cuda-toolkit[cccl]==12.*
+    ),
+    "cublas": ("nvidia/cu13/include", "nvidia/cublas/include"),
+    "cudart": ("nvidia/cu13/include", "nvidia/cuda_runtime/include"),
+    "cufft": ("nvidia/cu13/include", "nvidia/cufft/include"),
+    "cufile": ("nvidia/cu13/include", "nvidia/cufile/include"),
+    "curand": ("nvidia/cu13/include", "nvidia/curand/include"),
+    "cusolver": ("nvidia/cu13/include", "nvidia/cusolver/include"),
+    "cusparse": ("nvidia/cu13/include", "nvidia/cusparse/include"),
+    "npp": ("nvidia/cu13/include", "nvidia/npp/include"),
+    "nvcc": ("nvidia/cu13/include", "nvidia/cuda_nvcc/include"),
+    "nvfatbin": ("nvidia/cu13/include", "nvidia/nvfatbin/include"),
+    "nvjitlink": ("nvidia/cu13/include", "nvidia/nvjitlink/include"),
+    "nvjpeg": ("nvidia/cu13/include", "nvidia/nvjpeg/include"),
+    "nvrtc": ("nvidia/cu13/include", "nvidia/cuda_nvrtc/include"),
+    "nvvm": ("nvidia/cu13/include", "nvidia/cuda_nvcc/nvvm/include"),
+}

cuda_pathfinder/cuda/pathfinder/_utils/env_vars.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import warnings
+from typing import Optional
+
+
+def _paths_differ(a: str, b: str) -> bool:
+    """
+    Return True if paths are observably different.
+
+    Strategy:
+    1) Compare os.path.normcase(os.path.normpath(...)) for quick, robust textual equality.
+       - Handles trailing slashes and case-insensitivity on Windows.
+    2) If still different AND both exist, use os.path.samefile to resolve symlinks/junctions.
+    3) Otherwise (nonexistent paths or samefile unavailable), treat as different.
+    """
+    norm_a = os.path.normcase(os.path.normpath(a))
+    norm_b = os.path.normcase(os.path.normpath(b))
+    if norm_a == norm_b:
+        return False
+
+    try:
+        if os.path.exists(a) and os.path.exists(b):
+            # samefile raises on non-existent paths; only call when both exist.
+            return not os.path.samefile(a, b)
+    except OSError:
+        # Fall through to "different" if samefile isn't applicable/available.
+        pass
+
+    # If normalized strings differ and we couldn't prove they're the same entry, treat as different.
+    return True
+
+
+def get_cuda_home_or_path() -> Optional[str]:
+    cuda_home = os.environ.get("CUDA_HOME")
+    cuda_path = os.environ.get("CUDA_PATH")
+
+    if cuda_home and cuda_path and _paths_differ(cuda_home, cuda_path):
+        warnings.warn(
+            "Both CUDA_HOME and CUDA_PATH are set but differ:\n"
+            f"  CUDA_HOME={cuda_home}\n"
+            f"  CUDA_PATH={cuda_path}\n"
+            "Using CUDA_HOME (higher priority).",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    if cuda_home is not None:
+        return cuda_home
+    return cuda_path

cuda_pathfinder/cuda/pathfinder/_version.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "1.2.2"
+__version__ = "1.2.3"

cuda_pathfinder/docs/nv-versions.json

@@ -3,6 +3,10 @@
         "version": "latest",
         "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/latest/"
     },
+    {
+        "version": "1.2.3",
+        "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/1.2.3/"
+    },
     {
         "version": "1.2.2",
         "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/1.2.2/"

cuda_pathfinder/docs/source/api.rst

@@ -6,10 +6,8 @@
 ``cuda.pathfinder`` API Reference
 =================================
 
-The ``cuda.pathfinder`` module provides utilities for loading NVIDIA dynamic libraries.
-
-Public API
------------
+The ``cuda.pathfinder`` module provides utilities for loading NVIDIA dynamic libraries,
+and experimental APIs for locating NVIDIA C/C++ header directories.
 
 .. autosummary::
    :toctree: generated/
@@ -18,3 +16,6 @@ Public API
    load_nvidia_dynamic_lib
    LoadedDL
    DynamicLibNotFoundError
+
+   SUPPORTED_HEADERS_CTK
+   find_nvidia_header_directory

cuda_pathfinder/docs/source/release.rst

@@ -7,6 +7,7 @@ Release Notes
 .. toctree::
    :maxdepth: 3
 
+   1.2.3 <release/1.2.3-notes>
    1.2.2 <release/1.2.2-notes>
    1.2.1 <release/1.2.1-notes>
    1.2.0 <release/1.2.0-notes>

cuda_pathfinder/docs/source/release/1.2.3-notes.rst

@@ -0,0 +1,17 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: Apache-2.0
+
+.. module:: cuda.pathfinder
+
+``cuda-pathfinder`` 1.2.3 Release notes
+=======================================
+
+Released on Sep 17, 2025
+
+
+Highlights
+----------
+
+* Extend experimental ``cuda.pathfinder._find_nvidia_headers`` API
+  to support CTK library headers
+  (`PR #956 <https://github.com/NVIDIA/cuda-python/pull/956>`_)