Skip to content

Commit d020a5d

Browse files
authored
fix(cuda.core): normalize PCH status after heap retry (#3132)
Convert the retry status to str before storing it in the Cython str field, matching the initial compilation path and preserving the public PCHStatusType result. Add a regression that forces PCH heap growth and restores the original heap size, and skip it when NVRTC_DISABLE_PCH is set. Signed-off-by: Rui Luo <[email protected]>
1 parent e957c5c commit d020a5d

3 files changed

Lines changed: 35 additions & 1 deletion

File tree

‎cuda_core/cuda/core/_program.pyx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1100,7 +1100,7 @@ cdef object Program_compile_nvrtc(Program self, str target_type, object name_exp
11001100
)
11011101

11021102
status = _read_pch_status(retry_prog)
1103-
self._pch_status = status if status is not None else str(PCHStatusType.FAILED)
1103+
self._pch_status = str(status) if status is not None else str(PCHStatusType.FAILED)
11041104
return result
11051105

11061106

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,9 @@ New features
104104
Fixes and enhancements
105105
----------------------
106106

107+
- Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH
108+
heap exhaustion.
109+
107110
- :attr:`system.Device.uuid_without_prefix` preserves UUIDs that NVML already
108111
reports without a prefix, such as on Orin. This also allows
109112
:meth:`system.Device.to_cuda_device` to match those devices correctly.

‎cuda_core/tests/test_program.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,37 @@ def test_cpp_program_pch_auto_creates(init_cuda, tmp_path):
298298
program.close()
299299

300300

301+
@nvrtc_pch_available
302+
@pytest.mark.skipif("NVRTC_DISABLE_PCH" in os.environ, reason="PCH is disabled by NVRTC_DISABLE_PCH")
303+
@pytest.mark.thread_unsafe(reason="resizes the process-global NVRTC PCH heap")
304+
@pytest.mark.agent_authored(model="gpt-6")
305+
def test_cpp_program_pch_heap_retry(tmp_path):
306+
"""PCH heap exhaustion retries compilation and returns a typed creation status."""
307+
source = '\n#pragma nv_hdrstop\nextern "C" __global__ void my_kernel() {}\n'
308+
pch_path = tmp_path / "retry.pch"
309+
arch = min(handle_return(nvrtc.nvrtcGetSupportedArchs()))
310+
options = ProgramOptions(arch=f"compute_{arch}", name="retry.cu", create_pch=str(pch_path))
311+
previous_heap = handle_return(nvrtc.nvrtcGetPCHHeapSize())
312+
try:
313+
# NVRTC rounds this up; compare against the actual initial heap size.
314+
handle_return(nvrtc.nvrtcSetPCHHeapSize(1))
315+
small_heap = handle_return(nvrtc.nvrtcGetPCHHeapSize())
316+
program = Program(source, "c++", options)
317+
try:
318+
result = program.compile("ptx")
319+
assert isinstance(result, ObjectCode)
320+
assert b"my_kernel" in result.code
321+
assert program.pch_status is PCHStatusType.CREATED
322+
assert pch_path.is_file()
323+
# Success alone could miss the retry path; require automatic growth.
324+
assert handle_return(nvrtc.nvrtcGetPCHHeapSize()) > small_heap
325+
finally:
326+
program.close()
327+
finally:
328+
handle_return(nvrtc.nvrtcSetPCHHeapSize(previous_heap))
329+
assert handle_return(nvrtc.nvrtcGetPCHHeapSize()) == previous_heap
330+
331+
301332
@bundled_headers_available
302333
@pytest.mark.agent_authored(model="claude-sonnet-5")
303334
def test_use_bundled_headers_installs_and_compiles(init_cuda, tmp_path, monkeypatch):

0 commit comments

Comments
 (0)