Redesign PCH API: auto-retry on heap exhaustion, add pch_status property

Replace the low-level 1:1 wrappers (get_pch_create_status,
get_pch_heap_size_required) with a higher-level interface:

- compile() now automatically resizes the PCH heap and retries with a
  fresh NVRTC program when PCH creation fails due to heap exhaustion.
- program.pch_status property returns a clean string ("created",
  "not_attempted", "failed") or None.
- get_pch_heap_size() / set_pch_heap_size() retained for manual control.
- Factored compile+extract into _nvrtc_compile_and_extract() to support
  the retry path without duplicating code.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: cpcloud

cuda_core/cuda/core/_program.pxd

@@ -16,3 +16,5 @@ cdef class Program:
         object _compile_lock  # Per-instance lock for compile-time mutation
         bint _use_libdevice      # Flag for libdevice loading
         bint _libdevice_added
+        bytes _nvrtc_code       # Source code for NVRTC retry (PCH auto-resize)
+        str _pch_status         # PCH creation outcome after compile

cuda_core/cuda/core/_program.pyx

@@ -105,52 +105,24 @@ cdef class Program:
         """
         return Program_compile(self, target_type, name_expressions, logs)
 
-    def get_pch_create_status(self):
-        """Return NVRTC's PCH creation status for the latest compile.
-
-        Returns the ``nvrtcResult`` enum value describing the outcome of
-        the most recent PCH creation attempt. Possible values include
-        ``NVRTC_SUCCESS``, ``NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED``,
-        ``NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED``, and
-        ``NVRTC_ERROR_PCH_CREATE``.
-
-        This API is only available for ``code_type="c++"`` programs compiled
-        through NVRTC, and requires CUDA 12.8+ NVRTC bindings.
-
-        Returns
-        -------
-        :class:`~cuda.bindings.nvrtc.nvrtcResult`
-            The PCH creation status.
-
-        Raises
-        ------
-        NVRTCError
-            If the program handle is invalid.
-        """
-        return Program_get_pch_create_status(self)
-
-    def get_pch_heap_size_required(self):
-        """Return the PCH heap size required to compile this program.
-
-        This is only meaningful after a compile where
-        :meth:`get_pch_create_status` returned
-        ``NVRTC_SUCCESS`` or ``NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED``.
+    @property
+    def pch_status(self) -> str | None:
+        """PCH creation outcome from the most recent :meth:`compile` call.
 
-        This API is only available for ``code_type="c++"`` programs compiled
-        through NVRTC, and requires CUDA 12.8+ NVRTC bindings.
+        Possible values:
 
-        Returns
-        -------
-        int
-            Required PCH heap size in bytes.
+        * ``"created"`` — PCH file was written successfully.
+        * ``"not_attempted"`` — PCH creation was not attempted (e.g. the
+          compiler decided not to, or automatic PCH processing skipped it).
+        * ``"failed"`` — an error prevented PCH creation.
+        * ``None`` — PCH was not requested, or the program has not been
+          compiled yet, or the NVRTC bindings are too old to report status.
 
-        Raises
-        ------
-        NVRTCError
-            If the program handle is invalid or the query is not
-            applicable for the current PCH creation status.
+        When ``create_pch`` is set in :class:`ProgramOptions` and the PCH
+        heap is too small, :meth:`compile` automatically resizes the heap
+        and retries, so ``"created"`` should be the common outcome.
         """
-        return Program_get_pch_heap_size_required(self)
+        return self._pch_status
 
     @staticmethod
     def get_pch_heap_size():
@@ -543,11 +515,6 @@ def _find_libdevice_path():
     return find_bitcode_lib("device")
 
 
-cdef inline void _require_nvrtc_backend_for_pch(Program self) except *:
-    if self._backend != "NVRTC":
-        raise RuntimeError("PCH APIs are only available for Program instances using the NVRTC backend")
-
-
 cdef inline void _require_nvrtc_pch_api(str api_name) except *:
     if not hasattr(nvrtc, api_name):
         version = get_binding_version()
@@ -557,30 +524,6 @@ cdef inline void _require_nvrtc_pch_api(str api_name) except *:
         )
 
 
-cdef inline object Program_get_pch_create_status(Program self):
-    _require_nvrtc_backend_for_pch(self)
-    _require_nvrtc_pch_api("nvrtcGetPCHCreateStatus")
-    cdef cynvrtc.nvrtcProgram prog = as_cu(self._h_nvrtc)
-    cdef cynvrtc.nvrtcResult err
-    with nogil:
-        err = cynvrtc.nvrtcGetPCHCreateStatus(prog)
-    if err == cynvrtc.nvrtcResult.NVRTC_ERROR_INVALID_PROGRAM:
-        HANDLE_RETURN_NVRTC(prog, err)
-    return nvrtc.nvrtcResult(err)
-
-
-cdef inline size_t Program_get_pch_heap_size_required(Program self) except? 0:
-    _require_nvrtc_backend_for_pch(self)
-    _require_nvrtc_pch_api("nvrtcGetPCHHeapSizeRequired")
-    cdef cynvrtc.nvrtcProgram prog = as_cu(self._h_nvrtc)
-    cdef size_t size = 0
-    cdef cynvrtc.nvrtcResult err
-    with nogil:
-        err = cynvrtc.nvrtcGetPCHHeapSizeRequired(prog, &size)
-    HANDLE_RETURN_NVRTC(prog, err)
-    return size
-
-
 cdef inline size_t Program_get_pch_heap_size() except? 0:
     _require_nvrtc_pch_api("nvrtcGetPCHHeapSize")
     cdef size_t size = 0
@@ -666,6 +609,8 @@ cdef inline int Program_init(Program self, object code, str code_type, object op
     self._use_libdevice = False
     self._libdevice_added = False
 
+    self._pch_status = None
+
     if code_type == "c++":
         assert_type(code, str)
         if options.extra_sources is not None:
@@ -680,6 +625,7 @@ cdef inline int Program_init(Program self, object code, str code_type, object op
             HANDLE_RETURN_NVRTC(NULL, cynvrtc.nvrtcCreateProgram(
                 &nvrtc_prog, code_ptr, name_ptr, 0, NULL, NULL))
         self._h_nvrtc = create_nvrtc_program_handle(nvrtc_prog)
+        self._nvrtc_code = code_bytes
         self._backend = "NVRTC"
         self._linker = None
 
@@ -767,9 +713,15 @@ cdef inline int Program_init(Program self, object code, str code_type, object op
     return 0
 
 
-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 object _nvrtc_compile_and_extract(
+    cynvrtc.nvrtcProgram prog, str target_type, object name_expressions,
+    object logs, list options_list, str name,
+):
+    """Run nvrtcCompileProgram on *prog* and extract the output.
+
+    This is the inner compile+extract loop, factored out so the PCH
+    auto-retry path can call it on a fresh program handle.
+    """
     cdef size_t output_size = 0
     cdef size_t logsize = 0
     cdef vector[const char*] options_vec
@@ -787,7 +739,6 @@ cdef object Program_compile_nvrtc(Program self, str target_type, object name_exp
             HANDLE_RETURN_NVRTC(prog, cynvrtc.nvrtcAddNameExpression(prog, name_ptr))
 
     # Build options array
-    options_list = self._options.as_bytes("nvrtc", target_type)
     options_vec.resize(len(options_list))
     for i in range(len(options_list)):
         options_vec[i] = <const char*>(<bytes>options_list[i])
@@ -834,7 +785,72 @@ cdef object Program_compile_nvrtc(Program self, str target_type, object name_exp
                 HANDLE_RETURN_NVRTC(prog, cynvrtc.nvrtcGetProgramLog(prog, data_ptr))
             logs.write(log.decode("utf-8", errors="backslashreplace"))
 
-    return ObjectCode._init(bytes(data), target_type, symbol_mapping=symbol_mapping, name=self._options.name)
+    return ObjectCode._init(bytes(data), target_type, symbol_mapping=symbol_mapping, name=name)
+
+
+cdef bint _has_nvrtc_pch_apis():
+    return hasattr(nvrtc, "nvrtcGetPCHCreateStatus")
+
+
+cdef str _PCH_STATUS_CREATED = "created"
+cdef str _PCH_STATUS_NOT_ATTEMPTED = "not_attempted"
+cdef str _PCH_STATUS_FAILED = "failed"
+
+
+cdef str _read_pch_status(cynvrtc.nvrtcProgram prog):
+    """Query nvrtcGetPCHCreateStatus and translate to a high-level string."""
+    cdef cynvrtc.nvrtcResult err
+    with nogil:
+        err = cynvrtc.nvrtcGetPCHCreateStatus(prog)
+    if err == cynvrtc.nvrtcResult.NVRTC_SUCCESS:
+        return _PCH_STATUS_CREATED
+    if err == cynvrtc.nvrtcResult.NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED:
+        return None  # sentinel: caller should auto-retry
+    if err == cynvrtc.nvrtcResult.NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED:
+        return _PCH_STATUS_NOT_ATTEMPTED
+    return _PCH_STATUS_FAILED
+
+
+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)
+
+    result = _nvrtc_compile_and_extract(
+        prog, target_type, name_expressions, logs, options_list, self._options.name,
+    )
+
+    if not self._options.create_pch or not _has_nvrtc_pch_apis():
+        self._pch_status = None
+        return result
+
+    # PCH was requested — check creation status
+    cdef str status = _read_pch_status(prog)
+    if status is not None:
+        self._pch_status = status
+        return result
+
+    # Heap exhausted — auto-resize and retry with a fresh program
+    cdef size_t required = 0
+    with nogil:
+        cynvrtc.nvrtcGetPCHHeapSizeRequired(prog, &required)
+        cynvrtc.nvrtcSetPCHHeapSize(required)
+
+    cdef cynvrtc.nvrtcProgram retry_prog
+    cdef const char* code_ptr = <const char*>self._nvrtc_code
+    cdef const char* name_ptr = <const char*>self._options._name
+    with nogil:
+        HANDLE_RETURN_NVRTC(NULL, cynvrtc.nvrtcCreateProgram(
+            &retry_prog, code_ptr, name_ptr, 0, NULL, NULL))
+    self._h_nvrtc = create_nvrtc_program_handle(retry_prog)
+
+    result = _nvrtc_compile_and_extract(
+        retry_prog, target_type, name_expressions, logs, options_list, self._options.name,
+    )
+
+    status = _read_pch_status(retry_prog)
+    self._pch_status = status if status is not None else _PCH_STATUS_FAILED
+    return result
 
 
 cdef object Program_compile_nvvm(Program self, str target_type, object logs):

cuda_core/docs/source/release/0.6.0-notes.rst

@@ -54,9 +54,12 @@ New features
 - Added CUDA version compatibility check at import time to detect mismatches between
   ``cuda.core`` and the installed ``cuda-bindings`` version.
 
-- Added NVRTC PCH runtime APIs on ``Program``: ``get_pch_create_status()``,
-  ``get_pch_heap_size_required()``, ``get_pch_heap_size()``, and
-  ``set_pch_heap_size()``.
+- ``Program.compile()`` now automatically resizes the NVRTC PCH heap and
+  retries when precompiled header creation fails due to heap exhaustion.
+  The ``pch_status`` property reports the PCH creation outcome
+  (``"created"``, ``"not_attempted"``, ``"failed"``, or ``None``).
+  ``Program.get_pch_heap_size()`` and ``Program.set_pch_heap_size()``
+  are available for manual heap management.
 
 
 Fixes and enhancements

cuda_core/tests/test_program.py

@@ -64,10 +64,7 @@ def _has_nvrtc_pch_apis_for_tests():
         "nvrtcGetPCHCreateStatus",
         "nvrtcGetPCHHeapSizeRequired",
     )
-    try:
-        return all(hasattr(nvrtc, name) for name in required)
-    except Exception:
-        return False
+    return all(hasattr(nvrtc, name) for name in required)
 
 
 nvrtc_pch_available = pytest.mark.skipif(
@@ -336,28 +333,13 @@ def test_cpp_program_with_pch_options(init_cuda, tmp_path):
 
 
 @nvrtc_pch_available
-def test_cpp_program_pch_runtime_apis(init_cuda, tmp_path):
+def test_cpp_program_pch_auto_creates(init_cuda, tmp_path):
     code = 'extern "C" __global__ void my_kernel() {}'
-    options = ProgramOptions(create_pch=str(tmp_path / "test.pch"))
-    program = Program(code, "c++", options)
-    assert program.backend == "NVRTC"
+    pch_path = str(tmp_path / "test.pch")
+    program = Program(code, "c++", ProgramOptions(create_pch=pch_path))
+    assert program.pch_status is None  # not compiled yet
     program.compile("ptx")
-
-    status = program.get_pch_create_status()
-    valid_status_names = (
-        "NVRTC_ERROR_NO_PCH_CREATE_ATTEMPTED",
-        "NVRTC_ERROR_PCH_CREATE",
-        "NVRTC_ERROR_PCH_CREATE_HEAP_EXHAUSTED",
-    )
-    valid_statuses = {nvrtc.nvrtcResult.NVRTC_SUCCESS}
-    valid_statuses.update(
-        getattr(nvrtc.nvrtcResult, name) for name in valid_status_names if hasattr(nvrtc.nvrtcResult, name)
-    )
-    assert status in valid_statuses
-
-    required_heap_size = program.get_pch_heap_size_required()
-    assert isinstance(required_heap_size, int)
-    assert required_heap_size > 0
+    assert program.pch_status in ("created", "not_attempted", "failed")
     program.close()
 
 
@@ -378,12 +360,11 @@ def test_cpp_program_pch_set_heap_size_rejects_negative():
         Program.set_pch_heap_size(-1)
 
 
-def test_cpp_program_pch_runtime_apis_require_nvrtc_backend(init_cuda, ptx_code_object):
-    program = Program(ptx_code_object.code.decode(), "ptx")
-    with pytest.raises(RuntimeError, match="only available for Program instances using the NVRTC backend"):
-        program.get_pch_create_status()
-    with pytest.raises(RuntimeError, match="only available for Program instances using the NVRTC backend"):
-        program.get_pch_heap_size_required()
+def test_cpp_program_pch_status_none_without_pch(init_cuda):
+    code = 'extern "C" __global__ void my_kernel() {}'
+    program = Program(code, "c++")
+    program.compile("ptx")
+    assert program.pch_status is None
     program.close()