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

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

@@ -4,39 +4,114 @@
 import functools
 import glob
 import os
+import shutil
 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
+
 
-    if libname == "nvshmem" and IS_WINDOWS:
+def _joined_isfile(dirpath: str, basename: str) -> bool:
+    return os.path.isfile(os.path.join(dirpath, basename))
+
+
+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]:
+    if libname == "nvvm":
+        idir = os.path.join(anchor_point, "nvvm", "include")
+        if _joined_isfile(idir, h_basename):
+            return idir
+    else:
+        idir = os.path.join(anchor_point, "include")
+        if libname == "cccl":
+            cdir = os.path.join(idir, "cccl")
+            if _joined_isfile(cdir, h_basename):
+                return cdir
+        if _joined_isfile(idir, h_basename):
+            return idir
+
+    return None
+
+
+@functools.cache
+def _get_nvcc_home() -> Optional[str]:
+    # Fall back to the directory that owns nvcc (works for most local installs)
+    nvcc_path = shutil.which("nvcc")
+    if not nvcc_path:
+        return None
+    flds = nvcc_path.split(os.path.sep)
+    if len(flds) < 3:
+        return None
+    return os.path.sep.join(flds[:-2])
+
+
+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_ctk_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
+
+    nvcc_home = _get_nvcc_home()
+    if nvcc_home:  # noqa: SIM102
+        if result := _find_based_on_ctk_layout(libname, h_basename, nvcc_home):
+            return result
+
+    return None
+
+
+@functools.cache
+def find_nvidia_header_directory(libname: str) -> Optional[str]:
+    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,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+
+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": "cusolver_common.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 = SUPPORTED_HEADERS_CTK_WINDOWS if IS_WINDOWS else SUPPORTED_HEADERS_CTK_LINUX
+
+SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK = {
+    "cccl": (
+        "cuda/cccl/headers/include",  # cuda-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,12 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Optional
+
+
+def get_cuda_home_or_path() -> Optional[str]:
+    cuda_home = os.environ.get("CUDA_HOME")
+    if cuda_home is None:
+        cuda_home = os.environ.get("CUDA_PATH")
+    return cuda_home

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.3a0"

cuda_pathfinder/pyproject.toml

@@ -15,7 +15,7 @@ test = [
     "pytest>=6.2.4",
 ]
 nvidia_wheels_cu12 = [
-    "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg]==12.*",
+    "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl]==12.*",
     "cuda-toolkit[cufile]==12.*; sys_platform != 'win32'",
     "nvidia-cudss-cu12",
     "nvidia-cufftmp-cu12; sys_platform != 'win32'",
@@ -24,7 +24,7 @@ nvidia_wheels_cu12 = [
     "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[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,nvvm]==13.*",
     "cuda-toolkit[cufile]==13.*; sys_platform != 'win32'",
     "nvidia-nccl-cu13; sys_platform != 'win32'",
     "nvidia-nvshmem-cu13; sys_platform != 'win32'",

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -21,7 +21,12 @@
 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
+from cuda.pathfinder._headers.supported_nvidia_headers import (
+    IS_WINDOWS,
+    SUPPORTED_HEADERS_CTK,
+    SUPPORTED_HEADERS_CTK_ALL,
+    SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK,
+)
 
 STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS", "see_what_works")
 assert STRICTNESS in ("see_what_works", "all_must_work")
@@ -58,3 +63,19 @@ def test_find_libname_nvshmem(info_summary_append):
             assert hdr_dir.startswith(conda_prefix)
         else:
             assert hdr_dir.startswith("/usr/include/nvshmem_")
+
+
+def test_supported_headers_site_packages_ctk_consistency():
+    assert tuple(sorted(SUPPORTED_HEADERS_CTK_ALL)) == tuple(sorted(SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK.keys()))
+
+
+@pytest.mark.parametrize("libname", SUPPORTED_HEADERS_CTK.keys())
+def test_find_ctk_headers(info_summary_append, libname):
+    hdr_dir = find_nvidia_header_directory(libname)
+    info_summary_append(f"{hdr_dir=!r}")
+    if hdr_dir:
+        assert os.path.isdir(hdr_dir)
+        h_filename = SUPPORTED_HEADERS_CTK[libname]
+        assert os.path.isfile(os.path.join(hdr_dir, h_filename))
+    if STRICTNESS == "all_must_work":
+        assert hdr_dir is not None