Skip to content

Commit 858fb2f

Browse files
mdboommdboom
authored
Fix #2363: Support using bundled headers with NVRTC 13.3+. (#2753)
1 parent 4148644 commit 858fb2f

7 files changed

Lines changed: 175 additions & 54 deletions

File tree

‎cuda_core/cuda/core/_program.pyi‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,13 @@ class ProgramOptions:
218218
include_path : str | list[str], optional
219219
Add the directory or directories to the list of directories to be searched for headers.
220220
Default: None
221+
use_bundled_headers : bool, optional
222+
Use the CUDA and CCCL headers bundled with NVRTC, installed into a per-user cache
223+
directory, instead of requiring a full CUDA Toolkit installation. Implemented via NVRTC's
224+
``--use-bundled-headers=<dir>`` compiler option, which installs the headers into the cache
225+
directory (skipping installation if already present and up to date) and adds that
226+
directory to the include search path. NVRTC only.
227+
Default: False
221228
pre_include : str | list[str], optional
222229
Preinclude one or more headers during preprocessing. Can be either a string or a list of strings.
223230
Default: None
@@ -350,6 +357,7 @@ class ProgramOptions:
350357
define_macro: str | tuple[str, str] | list[str | tuple[str, str]] | tuple[str | tuple[str, str], ...] | None = None
351358
undefine_macro: str | list[str] | tuple[str] | None = None
352359
include_path: str | list[str] | tuple[str] | None = None
360+
use_bundled_headers: bool | None = None
353361
pre_include: str | list[str] | tuple[str] | None = None
354362
no_source_include: bool | None = None
355363
std: str | None = None

‎cuda_core/cuda/core/_program.pyx‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ from cuda.core._utils.cuda_utils import (
4747
is_sequence,
4848
)
4949
from cuda.core._utils.version import binding_version, driver_version
50+
from cuda.core.utils._cache_dir import _default_cache_dir
5051
from cuda.core.typing import ObjectCodeFormatType, CompilerBackendType, PCHStatusType, SourceCodeType
5152

5253
__all__ = ["Program", "ProgramOptions"]
@@ -430,6 +431,13 @@ class ProgramOptions:
430431
include_path : str | list[str], optional
431432
Add the directory or directories to the list of directories to be searched for headers.
432433
Default: None
434+
use_bundled_headers : bool, optional
435+
Use the CUDA and CCCL headers bundled with NVRTC, installed into a per-user cache
436+
directory, instead of requiring a full CUDA Toolkit installation. Implemented via NVRTC's
437+
``--use-bundled-headers=<dir>`` compiler option, which installs the headers into the cache
438+
directory (skipping installation if already present and up to date) and adds that
439+
directory to the include search path. NVRTC only.
440+
Default: False
433441
pre_include : str | list[str], optional
434442
Preinclude one or more headers during preprocessing. Can be either a string or a list of strings.
435443
Default: None
@@ -563,6 +571,7 @@ class ProgramOptions:
563571
define_macro: str | tuple[str, str] | list[str | tuple[str, str]] | tuple[str | tuple[str, str], ...] | None = None
564572
undefine_macro: str | list[str] | tuple[str] | None = None
565573
include_path: str | list[str] | tuple[str] | None = None
574+
use_bundled_headers: bool | None = None
566575
pre_include: str | list[str] | tuple[str] | None = None
567576
no_source_include: bool | None = None
568577
std: str | None = None
@@ -606,6 +615,16 @@ class ProgramOptions:
606615
# Set arch to default if not provided
607616
if self.arch is None:
608617
self.arch = f"sm_{Device().arch}"
618+
if self.use_bundled_headers:
619+
# --use-bundled-headers (and the bundled CUDA/CCCL headers themselves) were
620+
# introduced in NVRTC 13.3.
621+
nvrtc_major, nvrtc_minor = handle_return(nvrtc.nvrtcVersion())
622+
if (nvrtc_major, nvrtc_minor) < (13, 3):
623+
raise RuntimeError(
624+
"use_bundled_headers requires NVRTC >= 13.3, but found "
625+
f"{nvrtc_major}.{nvrtc_minor}. Upgrade the CUDA Toolkit / driver providing "
626+
"libnvrtc, or set use_bundled_headers=False and supply include_path manually."
627+
)
609628
if self.extra_sources is not None:
610629
if not is_sequence(self.extra_sources):
611630
raise TypeError(
@@ -765,8 +784,6 @@ def _find_libdevice_path() -> object:
765784
return find_bitcode_lib("device")
766785

767786

768-
769-
770787
cdef inline bint _process_define_macro_inner(list options, object macro) except? -1:
771788
"""Process a single define macro, returning True if successful."""
772789
if isinstance(macro, str):
@@ -1234,6 +1251,8 @@ cdef inline list _prepare_nvrtc_options_impl(object opts):
12341251
elif is_sequence(opts.undefine_macro):
12351252
for macro in opts.undefine_macro:
12361253
options.append(f"--undefine-macro={macro}")
1254+
if opts.use_bundled_headers:
1255+
options.append(f"--use-bundled-headers={_default_cache_dir() / 'nvrtc-bundled-headers'}")
12371256
if opts.include_path is not None:
12381257
if isinstance(opts.include_path, str):
12391258
options.append(f"--include-path={opts.include_path}")
@@ -1398,6 +1417,8 @@ cdef inline object _prepare_nvvm_options_impl(object opts, bint as_bytes):
13981417
unsupported.append("undefine_macro")
13991418
if opts.include_path is not None:
14001419
unsupported.append("include_path")
1420+
if opts.use_bundled_headers:
1421+
unsupported.append("use_bundled_headers")
14011422
if opts.pre_include is not None:
14021423
unsupported.append("pre_include")
14031424
if opts.no_source_include is not None and opts.no_source_include:

‎cuda_core/cuda/core/utils/_cache_dir.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
"""Shared user-cache-root resolution for cuda.core's on-disk caches."""
6+
7+
from __future__ import annotations
8+
9+
import os
10+
from pathlib import Path
11+
12+
# Exposed as a module-level flag so tests can toggle it without monkeypatching
13+
# ``os.name`` itself (pathlib reads ``os.name`` at instantiation time).
14+
_IS_WINDOWS = os.name == "nt"
15+
16+
17+
def _default_cache_dir() -> Path:
18+
"""OS-conventional root for cuda.core's on-disk caches.
19+
20+
Resolves to the user-cache root for the calling user, with a
21+
``cuda-python`` vendor leaf so callers can each place their own cache
22+
under a stable, shared root:
23+
24+
* Linux: ``$XDG_CACHE_HOME/cuda-python``
25+
(default ``~/.cache/cuda-python`` per the XDG Base Directory spec).
26+
* Windows: ``%LOCALAPPDATA%\\cuda-python``
27+
(Windows uses local AppData -- caches don't roam; falls back to
28+
``~/AppData/Local`` if the env var is unset).
29+
30+
CUDA does not support macOS, so no macOS branch is provided.
31+
32+
Callers append their own leaf directory, e.g. ``program-cache`` or
33+
``nvrtc-headers``.
34+
"""
35+
if _IS_WINDOWS:
36+
local_app_data = os.environ.get("LOCALAPPDATA")
37+
root = Path(local_app_data) if local_app_data else Path.home() / "AppData" / "Local"
38+
else:
39+
xdg = os.environ.get("XDG_CACHE_HOME")
40+
root = Path(xdg) if xdg else Path.home() / ".cache"
41+
return root / "cuda-python"

‎cuda_core/cuda/core/utils/_program_cache/_file_stream.py‎

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from typing import Any, Callable, Iterable
2424

2525
from cuda.core._module import ObjectCode
26+
from cuda.core.utils._cache_dir import _default_cache_dir as _user_cache_dir
2627

2728
from ._abc import ProgramCacheResource, _as_key_bytes, _extract_bytes
2829

@@ -57,28 +58,10 @@ def _stat_key(st: os.stat_result) -> tuple[int, int, int]:
5758

5859

5960
def _default_cache_dir() -> Path:
60-
"""OS-conventional default location for the file-stream cache.
61-
62-
Resolves to the user-cache root for the calling user, with a
63-
``program-cache`` leaf so future tooling can place sibling caches
64-
under the same ``cuda-python`` vendor directory:
65-
66-
* Linux: ``$XDG_CACHE_HOME/cuda-python/program-cache``
67-
(default ``~/.cache/cuda-python/program-cache`` per the XDG Base
68-
Directory spec).
69-
* Windows: ``%LOCALAPPDATA%\\cuda-python\\program-cache``
70-
(Windows uses local AppData -- caches don't roam; falls back to
71-
``~/AppData/Local`` if the env var is unset).
72-
73-
CUDA does not support macOS, so no macOS branch is provided.
61+
"""Default location for the file-stream cache: the ``program-cache`` leaf under the shared
62+
``cuda-python`` user-cache root (see :func:`cuda.core.utils._cache_dir._default_cache_dir`).
7463
"""
75-
if _IS_WINDOWS:
76-
local_app_data = os.environ.get("LOCALAPPDATA")
77-
root = Path(local_app_data) if local_app_data else Path.home() / "AppData" / "Local"
78-
else:
79-
xdg = os.environ.get("XDG_CACHE_HOME")
80-
root = Path(xdg) if xdg else Path.home() / ".cache"
81-
return root / "cuda-python" / "program-cache"
64+
return _user_cache_dir() / "program-cache"
8265

8366

8467
def _with_sharing_retry(

‎cuda_core/tests/test_cache_dir.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
import pytest
6+
7+
8+
@pytest.mark.agent_authored(model="claude-sonnet-5")
9+
def test_user_cache_dir_lives_under_platform_root(monkeypatch, tmp_path):
10+
"""The shared user-cache root (``cuda.core.utils._cache_dir``) is platform-specific:
11+
12+
* Linux: ``$XDG_CACHE_HOME`` or ``~/.cache``.
13+
* Windows: ``%LOCALAPPDATA%`` or ``~/AppData/Local``.
14+
15+
Both branches must end in ``cuda-python``; that suffix is what guarantees a
16+
stable on-disk layout across releases, and callers (e.g. the file-stream
17+
cache, NVRTC's bundled-headers cache) each append their own leaf under it.
18+
"""
19+
from pathlib import Path
20+
21+
from cuda.core.utils import _cache_dir
22+
from cuda.core.utils._cache_dir import _default_cache_dir
23+
24+
# Path must end with cuda-python regardless of platform.
25+
assert _default_cache_dir().parts[-1] == "cuda-python"
26+
27+
# Linux branch: XDG_CACHE_HOME wins when set.
28+
monkeypatch.setattr(_cache_dir, "_IS_WINDOWS", False)
29+
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
30+
assert _default_cache_dir() == tmp_path / "xdg" / "cuda-python"
31+
32+
# Linux branch: falls back to ``~/.cache`` when XDG_CACHE_HOME is unset.
33+
monkeypatch.delenv("XDG_CACHE_HOME", raising=False)
34+
monkeypatch.setattr(Path, "home", classmethod(lambda _cls: tmp_path / "home"))
35+
assert _default_cache_dir() == tmp_path / "home" / ".cache" / "cuda-python"
36+
37+
# Windows branch: LOCALAPPDATA wins when set.
38+
monkeypatch.setattr(_cache_dir, "_IS_WINDOWS", True)
39+
monkeypatch.setenv("LOCALAPPDATA", str(tmp_path / "appdata"))
40+
assert _default_cache_dir() == tmp_path / "appdata" / "cuda-python"
41+
42+
# Windows branch: falls back to ``~/AppData/Local`` when LOCALAPPDATA is unset.
43+
monkeypatch.delenv("LOCALAPPDATA", raising=False)
44+
assert _default_cache_dir() == tmp_path / "home" / "AppData" / "Local" / "cuda-python"

‎cuda_core/tests/test_program.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,11 @@ def _has_nvrtc_pch_apis_for_tests():
7272
reason="PCH runtime APIs require NVRTC >= 12.8 bindings",
7373
)
7474

75+
bundled_headers_available = pytest.mark.skipif(
76+
(_get_nvrtc_version_for_tests() or 0) < 13300,
77+
reason="use_bundled_headers requires NVRTC >= 13.3",
78+
)
79+
7580

7681
def _has_check_nvvm_compiler_options():
7782
try:
@@ -302,6 +307,48 @@ def test_cpp_program_pch_auto_creates(init_cuda, tmp_path):
302307
program.close()
303308

304309

310+
@bundled_headers_available
311+
@pytest.mark.agent_authored(model="claude-sonnet-5")
312+
def test_use_bundled_headers_installs_and_compiles(init_cuda, tmp_path, monkeypatch):
313+
"""``use_bundled_headers`` should install NVRTC's bundled CUDA/CCCL headers into the
314+
(monkeypatched) cache directory and make them available on the include path, without
315+
a CUDA Toolkit or any user-supplied ``include_path``."""
316+
import cuda.core._program as _program_module
317+
318+
cache_root = tmp_path / "cache-root"
319+
monkeypatch.setattr(_program_module, "_default_cache_dir", lambda: cache_root)
320+
321+
code = """
322+
#include <cuda/std/type_traits>
323+
extern "C" __global__ void my_kernel(int *out) {
324+
*out = cuda::std::is_integral<int>::value;
325+
}
326+
"""
327+
headers_dir = cache_root / "nvrtc-bundled-headers"
328+
assert not headers_dir.exists()
329+
330+
# Sanity check: without use_bundled_headers, the CCCL header isn't found (proves the
331+
# option -- not some ambient CUDA Toolkit install -- is what makes the compile below work).
332+
program = Program(code, "c++")
333+
try:
334+
with pytest.raises(CUDAError, match="could not open source file"):
335+
program.compile("ptx")
336+
finally:
337+
program.close()
338+
339+
program = Program(code, "c++", ProgramOptions(use_bundled_headers=True))
340+
try:
341+
object_code = program.compile("ptx")
342+
finally:
343+
program.close()
344+
assert isinstance(object_code, ObjectCode)
345+
346+
assert headers_dir.is_dir()
347+
assert (headers_dir / ".nvrtc_headers_version").is_file()
348+
assert (headers_dir / "cccl").is_dir()
349+
assert (headers_dir / "cccl" / "cuda" / "std" / "type_traits").is_file()
350+
351+
305352
def test_cpp_program_pch_status_none_without_pch(init_cuda):
306353
code = 'extern "C" __global__ void my_kernel() {}'
307354
program = Program(code, "c++")

‎cuda_core/tests/test_program_cache.py‎

Lines changed: 8 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1526,41 +1526,18 @@ def test_filestream_cache_rejects_non_positive_size_cap(tmp_path, bad):
15261526
FileStreamProgramCache(tmp_path / "fc", max_size_bytes=bad)
15271527

15281528

1529-
def test_default_cache_dir_lives_under_user_cache_root(monkeypatch, tmp_path):
1530-
"""The cache root is platform-specific:
1531-
1532-
* Linux: ``$XDG_CACHE_HOME`` or ``~/.cache``.
1533-
* Windows: ``%LOCALAPPDATA%`` or ``~/AppData/Local``.
1534-
1535-
Both branches must end in ``cuda-python/program-cache``; that suffix
1536-
is what guarantees a stable on-disk layout across releases.
1537-
"""
1538-
from pathlib import Path
1539-
1529+
@pytest.mark.agent_authored(model="claude-sonnet-5")
1530+
def test_file_stream_default_cache_dir_appends_program_cache_leaf(monkeypatch, tmp_path):
1531+
"""The file-stream cache's default dir is the shared user-cache root plus a
1532+
``program-cache`` leaf, so it can live alongside sibling caches (e.g. NVRTC's
1533+
bundled-headers cache) under the same ``cuda-python`` vendor directory."""
15401534
from cuda.core.utils import _program_cache
1541-
from cuda.core.utils._program_cache._file_stream import _default_cache_dir
15421535

1543-
# Path must end with cuda-python/program-cache regardless of platform.
1544-
assert _default_cache_dir().parts[-2:] == ("cuda-python", "program-cache")
1536+
monkeypatch.setattr(_program_cache._file_stream, "_user_cache_dir", lambda: tmp_path / "root")
15451537

1546-
# Linux branch: XDG_CACHE_HOME wins when set.
1547-
monkeypatch.setattr(_program_cache._file_stream, "_IS_WINDOWS", False)
1548-
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg"))
1549-
assert _default_cache_dir() == tmp_path / "xdg" / "cuda-python" / "program-cache"
1550-
1551-
# Linux branch: falls back to ``~/.cache`` when XDG_CACHE_HOME is unset.
1552-
monkeypatch.delenv("XDG_CACHE_HOME", raising=False)
1553-
monkeypatch.setattr(Path, "home", classmethod(lambda _cls: tmp_path / "home"))
1554-
assert _default_cache_dir() == tmp_path / "home" / ".cache" / "cuda-python" / "program-cache"
1555-
1556-
# Windows branch: LOCALAPPDATA wins when set.
1557-
monkeypatch.setattr(_program_cache._file_stream, "_IS_WINDOWS", True)
1558-
monkeypatch.setenv("LOCALAPPDATA", str(tmp_path / "appdata"))
1559-
assert _default_cache_dir() == tmp_path / "appdata" / "cuda-python" / "program-cache"
1538+
from cuda.core.utils._program_cache._file_stream import _default_cache_dir
15601539

1561-
# Windows branch: falls back to ``~/AppData/Local`` when LOCALAPPDATA is unset.
1562-
monkeypatch.delenv("LOCALAPPDATA", raising=False)
1563-
assert _default_cache_dir() == tmp_path / "home" / "AppData" / "Local" / "cuda-python" / "program-cache"
1540+
assert _default_cache_dir() == tmp_path / "root" / "program-cache"
15641541

15651542

15661543
def test_filestream_cache_uses_default_dir_when_path_omitted(tmp_path, monkeypatch):

0 commit comments

Comments
 (0)