1- # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22# SPDX-License-Identifier: Apache-2.0
33
44from __future__ import annotations
77import glob
88import os
99from dataclasses import dataclass
10+ from typing import TYPE_CHECKING
1011
1112from cuda .pathfinder ._dynamic_libs .load_nvidia_dynamic_lib import (
1213 _resolve_system_loaded_abs_path_in_subprocess ,
1314)
1415from cuda .pathfinder ._dynamic_libs .search_steps import derive_ctk_root
15- from cuda .pathfinder ._headers import supported_nvidia_headers
16+ from cuda .pathfinder ._headers . header_descriptor import HEADER_DESCRIPTORS
1617from cuda .pathfinder ._utils .env_vars import get_cuda_home_or_path
1718from cuda .pathfinder ._utils .find_sub_dirs import find_sub_dirs_all_sitepackages
1819from cuda .pathfinder ._utils .platform_aware import IS_WINDOWS
1920
21+ if TYPE_CHECKING :
22+ from cuda .pathfinder ._headers .header_descriptor import HeaderDescriptor
23+
2024
2125@dataclass
2226class LocatedHeaderDir :
@@ -38,37 +42,33 @@ def _joined_isfile(dirpath: str, basename: str) -> bool:
3842
3943
4044def _locate_under_site_packages (sub_dir : str , h_basename : str ) -> LocatedHeaderDir | None :
41- # Installed from a wheel
4245 hdr_dir : str # help mypy
4346 for hdr_dir in find_sub_dirs_all_sitepackages (tuple (sub_dir .split ("/" ))):
4447 if _joined_isfile (hdr_dir , h_basename ):
4548 return LocatedHeaderDir (abs_path = hdr_dir , found_via = "site-packages" )
4649 return None
4750
4851
49- def _locate_based_on_ctk_layout (libname : str , h_basename : str , anchor_point : str ) -> str | None :
50- parts = [anchor_point ]
51- if libname == "nvvm" :
52- parts .append (libname )
53- parts .append ("include" )
54- idir = os .path .join (* parts )
55- if libname == "cccl" :
52+ def _locate_in_anchor_layout (desc : HeaderDescriptor , anchor_point : str ) -> str | None :
53+ """Search for a header under *anchor_point* using the descriptor's layout fields."""
54+ h_basename = desc .header_basename
55+ for rel_dir in desc .anchor_include_rel_dirs :
56+ idir = os .path .join (anchor_point , rel_dir )
5657 if IS_WINDOWS :
57- cdir_ctk12 = os .path .join (idir , "targets" , "x64" ) # conda has this anomaly
58- cdir_ctk13 = os .path .join (cdir_ctk12 , "cccl" )
59- if _joined_isfile (cdir_ctk13 , h_basename ):
60- return cdir_ctk13
61- if _joined_isfile (cdir_ctk12 , h_basename ):
62- return cdir_ctk12
63- cdir = os .path .join (idir , "cccl" ) # CTK 13
64- if _joined_isfile (cdir , h_basename ):
65- return cdir
66- if _joined_isfile (idir , h_basename ):
67- return idir
58+ for subdir in desc .include_subdirs_windows :
59+ cdir = os .path .join (idir , subdir )
60+ if _joined_isfile (cdir , h_basename ):
61+ return cdir
62+ for subdir in desc .include_subdirs :
63+ cdir = os .path .join (idir , subdir )
64+ if _joined_isfile (cdir , h_basename ):
65+ return cdir
66+ if _joined_isfile (idir , h_basename ):
67+ return idir
6868 return None
6969
7070
71- def _find_based_on_conda_layout ( libname : str , h_basename : str , ctk_layout : bool ) -> LocatedHeaderDir | None :
71+ def _find_in_conda ( desc : HeaderDescriptor ) -> LocatedHeaderDir | None :
7272 conda_prefix = os .environ .get ("CONDA_PREFIX" )
7373 if not conda_prefix :
7474 return None
@@ -77,25 +77,23 @@ def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool)
7777 if not os .path .isdir (anchor_point ):
7878 return None
7979 else :
80- if ctk_layout :
80+ if desc . packaged_with == "ctk" :
8181 targets_include_path = glob .glob (os .path .join (conda_prefix , "targets" , "*" , "include" ))
8282 if not targets_include_path :
8383 return None
8484 if len (targets_include_path ) != 1 :
85- # Conda does not support multiple architectures.
86- # QUESTION(PR#956): Do we want to issue a warning?
8785 return None
8886 include_path = targets_include_path [0 ]
8987 else :
9088 include_path = os .path .join (conda_prefix , "include" )
9189 anchor_point = os .path .dirname (include_path )
92- found_header_path = _locate_based_on_ctk_layout ( libname , h_basename , anchor_point )
90+ found_header_path = _locate_in_anchor_layout ( desc , anchor_point )
9391 if found_header_path :
9492 return LocatedHeaderDir (abs_path = found_header_path , found_via = "conda" )
9593 return None
9694
9795
98- def _find_ctk_header_directory_via_canary ( libname : str , h_basename : str ) -> str | None :
96+ def _find_via_ctk_root_canary ( desc : HeaderDescriptor ) -> str | None :
9997 """Try CTK header lookup via CTK-root canary probing.
10098
10199 Uses the same canary as dynamic-library CTK-root discovery: system-load
@@ -109,26 +107,43 @@ def _find_ctk_header_directory_via_canary(libname: str, h_basename: str) -> str
109107 ctk_root = derive_ctk_root (canary_abs_path )
110108 if ctk_root is None :
111109 return None
112- return _locate_based_on_ctk_layout (libname , h_basename , ctk_root )
110+ return _locate_in_anchor_layout (desc , ctk_root )
111+
112+
113+ def _find_in_system_install_dirs (desc : HeaderDescriptor ) -> LocatedHeaderDir | None :
114+ """Search system install directories (glob patterns) for non-CTK headers."""
115+ for pattern in desc .system_install_dirs :
116+ for hdr_dir in sorted (glob .glob (pattern ), reverse = True ):
117+ if _joined_isfile (hdr_dir , desc .header_basename ):
118+ return LocatedHeaderDir (abs_path = hdr_dir , found_via = "supported_install_dir" )
119+ return None
113120
114121
115- def _find_ctk_header_directory ( libname : str ) -> LocatedHeaderDir | None :
116- h_basename = supported_nvidia_headers . SUPPORTED_HEADERS_CTK [ libname ]
117- candidate_dirs = supported_nvidia_headers . SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK [ libname ]
122+ def _locate_header_directory ( desc : HeaderDescriptor ) -> LocatedHeaderDir | None :
123+ """Unified header search using a descriptor."""
124+ h_basename = desc . header_basename
118125
119- for cdir in candidate_dirs :
126+ # 1. Site-packages (pip wheels)
127+ for cdir in desc .site_packages_dirs :
120128 if hdr_dir := _locate_under_site_packages (cdir , h_basename ):
121129 return hdr_dir
122130
123- if hdr_dir := _find_based_on_conda_layout (libname , h_basename , True ):
131+ # 2. Conda
132+ if hdr_dir := _find_in_conda (desc ):
124133 return hdr_dir
125134
126- cuda_home = get_cuda_home_or_path ()
127- if cuda_home and (result := _locate_based_on_ctk_layout (libname , h_basename , cuda_home )):
128- return LocatedHeaderDir (abs_path = result , found_via = "CUDA_HOME" )
135+ # CTK-only steps: CUDA_HOME and canary probe
136+ if desc .packaged_with == "ctk" :
137+ cuda_home = get_cuda_home_or_path ()
138+ if cuda_home and (result := _locate_in_anchor_layout (desc , cuda_home )):
139+ return LocatedHeaderDir (abs_path = result , found_via = "CUDA_HOME" )
129140
130- if result := _find_ctk_header_directory_via_canary (libname , h_basename ):
131- return LocatedHeaderDir (abs_path = result , found_via = "system-ctk-root" )
141+ if result := _find_via_ctk_root_canary (desc ):
142+ return LocatedHeaderDir (abs_path = result , found_via = "system-ctk-root" )
143+
144+ # Non-CTK fallback: system install directories
145+ if desc .system_install_dirs and (hdr_dir := _find_in_system_install_dirs (desc )):
146+ return hdr_dir
132147
133148 return None
134149
@@ -160,41 +175,24 @@ def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None:
160175 - Check Conda-style installation prefixes, which use platform-specific
161176 include directory layouts.
162177
163- 3. **CUDA Toolkit environment variables**
178+ 3. **CUDA Toolkit environment variables** *(CTK headers only)*
164179
165180 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
166181
167- 4. **CTK root canary probe**
182+ 4. **CTK root canary probe** *(CTK headers only)*
168183
169184 - Probe a system-loaded ``cudart`` in a spawned child process,
170185 derive the CTK root from the resolved library path, then search
171186 CTK include layout under that root.
172- """
173187
174- if libname in supported_nvidia_headers .SUPPORTED_HEADERS_CTK :
175- return _find_ctk_header_directory (libname )
188+ 5. **System install directories** *(non-CTK headers only)*
176189
177- h_basename = supported_nvidia_headers .SUPPORTED_HEADERS_NON_CTK .get (libname )
178- if h_basename is None :
190+ - Search well-known system install paths (e.g., ``/usr/include/nvshmem_*``).
191+ """
192+ desc = HEADER_DESCRIPTORS .get (libname )
193+ if desc is None :
179194 raise RuntimeError (f"UNKNOWN { libname = } " )
180-
181- candidate_dirs = supported_nvidia_headers .SUPPORTED_SITE_PACKAGE_HEADER_DIRS_NON_CTK .get (libname , [])
182-
183- for cdir in candidate_dirs :
184- if found_hdr := _locate_under_site_packages (cdir , h_basename ):
185- return found_hdr
186-
187- if found_hdr := _find_based_on_conda_layout (libname , h_basename , False ):
188- return found_hdr
189-
190- # Fall back to system install directories
191- candidate_dirs = supported_nvidia_headers .SUPPORTED_INSTALL_DIRS_NON_CTK .get (libname , [])
192- for cdir in candidate_dirs :
193- for hdr_dir in sorted (glob .glob (cdir ), reverse = True ):
194- if _joined_isfile (hdr_dir , h_basename ):
195- # For system installs, we don't have a clear found_via, so use "system"
196- return LocatedHeaderDir (abs_path = hdr_dir , found_via = "supported_install_dir" )
197- return None
195+ return _locate_header_directory (desc )
198196
199197
200198def find_nvidia_header_directory (libname : str ) -> str | None :
@@ -222,15 +220,19 @@ def find_nvidia_header_directory(libname: str) -> str | None:
222220 - Check Conda-style installation prefixes, which use platform-specific
223221 include directory layouts.
224222
225- 3. **CUDA Toolkit environment variables**
223+ 3. **CUDA Toolkit environment variables** *(CTK headers only)*
226224
227225 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
228226
229- 4. **CTK root canary probe**
227+ 4. **CTK root canary probe** *(CTK headers only)*
230228
231229 - Probe a system-loaded ``cudart`` in a spawned child process,
232230 derive the CTK root from the resolved library path, then search
233231 CTK include layout under that root.
232+
233+ 5. **System install directories** *(non-CTK headers only)*
234+
235+ - Search well-known system install paths (e.g., ``/usr/include/nvshmem_*``).
234236 """
235237 found = locate_nvidia_header_directory (libname )
236238 return found .abs_path if found else None
0 commit comments