diff --git a/cuda_core/cuda/core/_jit_source.py b/cuda_core/cuda/core/_jit_source.py new file mode 100644 index 00000000000..acf0ed9ba78 --- /dev/null +++ b/cuda_core/cuda/core/_jit_source.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import contextlib +import hashlib +import os +import tempfile +import threading +from pathlib import Path + +_DIGEST_CHARS = 32 + + +_lock = threading.Lock() +_source_dir: tempfile.TemporaryDirectory[str] | None = None + + +def _ensure_source_dir() -> Path: + """Return the store directory, creating it if needed. Caller holds ``_lock``.""" + global _source_dir + if _source_dir is None: + _source_dir = tempfile.TemporaryDirectory(prefix="cuda-core-jit-") + root = Path(_source_dir.name) + # Re-created rather than assumed: a /tmp reaper can delete the tree out from + # under a long-running process. + root.mkdir(parents=True, exist_ok=True) + return root + + +def source_dir() -> Path: + """The process-scoped directory holding materialized JIT source.""" + with _lock: + return _ensure_source_dir() + + +def materialize(code: bytes, suffix: str = ".cu") -> str | None: + + digest = hashlib.sha256(code).hexdigest()[:_DIGEST_CHARS] + try: + with _lock: + target = _ensure_source_dir() / f"{digest}{suffix}" + # Writers are serialized and the directory is ours alone, so an + # entry that exists is one a previous caller finished writing. + if not target.exists(): + try: + target.write_bytes(code) + except BaseException: + # Otherwise a half-written entry survives to be mistaken + # for a complete one. + with contextlib.suppress(OSError): + target.unlink() + raise + return os.fspath(target) + except OSError: + # A read only or full filesystem, or a sandbox that forbids the temp + # dir should pass + return None diff --git a/cuda_core/cuda/core/_program.pxd b/cuda_core/cuda/core/_program.pxd index cea430c3f20..aa49bb5ad5e 100644 --- a/cuda_core/cuda/core/_program.pxd +++ b/cuda_core/cuda/core/_program.pxd @@ -20,3 +20,5 @@ cdef class Program: bytes _code # Source code as bytes: used for key derivation and NVRTC PCH retry str _code_type # Normalised code_type ("c++", "ptx", "nvvm") str _pch_status # PCH creation outcome after compile + bytes _source_name # Name handed to NVRTC/NVVM as the DWARF source path + list _extra_options # Compiler options Program adds on top of ProgramOptions diff --git a/cuda_core/cuda/core/_program.pyx b/cuda_core/cuda/core/_program.pyx index 1d07b88bbf4..3e3144ecc3b 100644 --- a/cuda_core/cuda/core/_program.pyx +++ b/cuda_core/cuda/core/_program.pyx @@ -10,6 +10,7 @@ This module provides :class:`Program` for compiling source code into from __future__ import annotations from dataclasses import dataclass +import os import threading from typing import TYPE_CHECKING from warnings import warn @@ -30,6 +31,7 @@ from ._resource_handles cimport ( ) from cuda.bindings cimport cynvrtc, cynvvm from cuda.core._utils.cuda_utils cimport HANDLE_RETURN_NVRTC, HANDLE_RETURN_NVVM +from cuda.core import _jit_source from cuda.core._device import Device from cuda.core._linker import Linker, LinkerHandleT, LinkerOptions from cuda.core._module import ObjectCode @@ -768,6 +770,31 @@ cdef inline object _translate_program_options(object options): ) +cdef inline int _program_setup_debug_source(Program self, object options, bytes code_bytes) except -1: + """Point the NVRTC source name at an on-disk copy of the source + """ + if not (options.debug or options.lineinfo): + return 0 + + # Resolved before anything is written, so a failure here leaves the original + # name in place rather than a redirected name with broken includes. + try: + include_dir = os.path.dirname(os.path.abspath(options.name)) + except OSError: + return 0 + + source_path = _jit_source.materialize(code_bytes) + if source_path is None: + return 0 + + self._source_name = source_path.encode() + + # NVRTC searches the dir of the name it's given for quoted includes, so + # hand the original one back or every #include "foo.h" stops resolving. + self._extra_options = [b"--include-path=" + include_dir.encode()] + return 0 + + cdef inline int Program_init(Program self, object code, str code_type, object options) except -1: """Initialize a Program instance.""" cdef cynvrtc.nvrtcProgram nvrtc_prog @@ -788,6 +815,10 @@ cdef inline int Program_init(Program self, object code, str code_type, object op self._libdevice_added = False self._pch_status = None + # ProgramOptions may be shared across Programs and is never mutated here, + # so any name override or added compiler option is held per-Program. + self._source_name = options._name + self._extra_options = [] if code_type == "c++": assert_type(code, str) @@ -796,8 +827,9 @@ cdef inline int Program_init(Program self, object code, str code_type, object op # TODO: support pre-loaded headers & include names code_bytes = code.encode() + _program_setup_debug_source(self, options, code_bytes) code_ptr = code_bytes - name_ptr = options._name + name_ptr = self._source_name with nogil: HANDLE_RETURN_NVRTC(NULL, cynvrtc.nvrtcCreateProgram( @@ -829,7 +861,7 @@ cdef inline int Program_init(Program self, object code, str code_type, object op # Use self._code (strictly bytes) for the C pointer so a bytearray # input doesn't trip the `code` cast at runtime. code_ptr = self._code - name_ptr = options._name + name_ptr = self._source_name code_len = len(self._code) with nogil: @@ -966,7 +998,7 @@ cdef object _read_pch_status(cynvrtc.nvrtcProgram prog): cdef object Program_compile_nvrtc(Program self, str target_type, object name_expressions, object logs): """Compile using NVRTC backend and return ObjectCode.""" cdef cynvrtc.nvrtcProgram prog = as_cu(self._h_nvrtc) - cdef list options_list = self._options.as_bytes("nvrtc", target_type) + cdef list options_list = self._options.as_bytes("nvrtc", target_type) + self._extra_options result = _nvrtc_compile_and_extract( prog, target_type, name_expressions, logs, options_list, self._options.name, @@ -997,7 +1029,7 @@ cdef object Program_compile_nvrtc(Program self, str target_type, object name_exp cdef cynvrtc.nvrtcProgram retry_prog cdef const char* code_ptr = self._code - cdef const char* name_ptr = self._options._name + cdef const char* name_ptr = self._source_name with nogil: HANDLE_RETURN_NVRTC(NULL, cynvrtc.nvrtcCreateProgram( &retry_prog, code_ptr, name_ptr, 0, NULL, NULL)) diff --git a/cuda_core/tests/test_program.py b/cuda_core/tests/test_program.py index 72f8b8f5942..f19b0d34da5 100644 --- a/cuda_core/tests/test_program.py +++ b/cuda_core/tests/test_program.py @@ -3,7 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +import os import re +import shutil +import subprocess import warnings import pytest @@ -367,6 +370,49 @@ def test_program_options_name_accepts_none(name): assert options._name == expected.encode() +@pytest.mark.human_reviewed +@pytest.mark.parametrize("name", [None, "my_program"]) +def test_program_string_source_debug(name, tmp_path): + # when a file doesn't exist on disk, nvrtc + # still ends up referencing a file it thinks is there + # in the dwarf table. As a WAR, we put the file there + # ourselves. This test verifies its there and that its + # contents match the passed in source code. + code = 'extern "C" __global__ void my_kernel() {}' + program = Program(code, "c++", options={"name": name, "debug": True, "lineinfo": True}) + cubin = program.compile("cubin") + + # read dwarf table + nvdisasm = shutil.which("nvdisasm") + + cubin_path = tmp_path / "program.cubin" + cubin_path.write_bytes(bytes(cubin.code)) + + result = subprocess.run( # noqa: S603 + [nvdisasm, "-g", str(cubin_path)], + capture_output=True, + text=True, + errors="replace", + ) + if result.returncode != 0: + pytest.fail(f"nvdisasm -g failed with exit code {result.returncode}\n{result.stderr}", pytrace=False) + + # -g annotates each instruction with the line-table entry that covers it: + # //## File "/abs/dir/name.cu", line 12 + # The directory comes from the line table's directory table, so this is the + # fully resolved path cuda-gdb will try to open. + paths = set(re.findall(r'//## File "([^"]+)", line \d+', result.stdout)) + assert len(paths) == 1, f"expected exactly one source file in the line table, got {sorted(paths)}" + dwarf_path = paths.pop() + + # cuda-gdb opens this path literally, so the source has to actually be + # sitting there for source-level debugging to work. + assert os.path.isfile(dwarf_path) + + with open(dwarf_path, encoding="utf-8") as source_file: + assert source_file.read().splitlines() == code.splitlines() + + # This is tested against the current device's arch def test_program_compile_valid_target_type(init_cuda): code = 'extern "C" __global__ void my_kernel() {}'