[12.9.x] Fix host-function callbacks in replayed graphs

cuda_bindings/cuda/bindings/driver.pyx

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 12.9.0. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=5883cb5633a3d97ee7d7582823371bd1928ee4a2388183d7433daf02032ece87
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=7bd998915de52155f2e27696c9a4d514071b9c721d94e05f0a00c7a07394a961
 from typing import Any, Optional
 import cython
 import ctypes
@@ -39995,19 +39995,6 @@ def cuLaunchCooperativeKernelMultiDevice(launchParamsList : Optional[tuple[CUDA_
         free(cylaunchParamsList)
     return (_CUresult(err),)
 
-ctypedef struct cuHostCallbackData_st:
-    cydriver.CUhostFn callback
-    void *userData
-
-ctypedef cuHostCallbackData_st cuHostCallbackData
-
-@cython.show_performance_hints(False)
-cdef void cuHostCallbackWrapper(void *data) nogil:
-    cdef cuHostCallbackData *cbData = <cuHostCallbackData *>data
-    with gil:
-        cbData.callback(cbData.userData)
-    free(cbData)
-
 @cython.embedsignature(True)
 def cuLaunchHostFunc(hStream, fn, userData):
     """ Enqueues a host function call in a stream.
@@ -40087,18 +40074,8 @@ def cuLaunchHostFunc(hStream, fn, userData):
     cyhStream = <cydriver.CUstream><void_ptr>phStream
     cdef _HelperInputVoidPtrStruct cyuserDataHelper
     cdef void* cyuserData = _helper_input_void_ptr(userData, &cyuserDataHelper)
-
-    cdef cuHostCallbackData *cbData = NULL
-    cbData = <cuHostCallbackData *>malloc(sizeof(cbData[0]))
-    if cbData == NULL:
-        return (CUresult.CUDA_ERROR_OUT_OF_MEMORY,)
-    cbData.callback = cyfn
-    cbData.userData = cyuserData
-
     with nogil:
-        err = cydriver.cuLaunchHostFunc(cyhStream, <cydriver.CUhostFn>cuHostCallbackWrapper, <void *>cbData)
-    if err != cydriver.CUDA_SUCCESS:
-        free(cbData)
+        err = cydriver.cuLaunchHostFunc(cyhStream, cyfn, cyuserData)
     _helper_input_void_ptr_free(&cyuserDataHelper)
     return (_CUresult(err),)
 

cuda_bindings/cuda/bindings/runtime.pyx.in

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 12.9.0. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=d6df3522c27bd5c4ff1ecf1bf0167f902aab2b381d9db0c473658a2922f5b6cd
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=e8e09d2b9dc5c50f5e7a6d85f94bd57c0c0776668393328be79deebdefce1c11
 from typing import Any, Optional
 import cython
 import ctypes
@@ -24776,19 +24776,6 @@ def cudaFuncSetAttribute(func, attr not None : cudaFuncAttribute, int value):
 
 {{if 'cudaLaunchHostFunc' in found_functions}}
 
-ctypedef struct cudaStreamHostCallbackData_st:
-    cyruntime.cudaHostFn_t callback
-    void *userData
-
-ctypedef cudaStreamHostCallbackData_st cudaStreamHostCallbackData
-
-@cython.show_performance_hints(False)
-cdef void cudaStreamRtHostCallbackWrapper(void *data) nogil:
-    cdef cudaStreamHostCallbackData *cbData = <cudaStreamHostCallbackData *>data
-    with gil:
-        cbData.callback(cbData.userData)
-    free(cbData)
-
 @cython.embedsignature(True)
 def cudaLaunchHostFunc(stream, fn, userData):
     """ Enqueues a host function call in a stream.
@@ -24869,20 +24856,9 @@ def cudaLaunchHostFunc(stream, fn, userData):
     cystream = <cyruntime.cudaStream_t><void_ptr>pstream
     cdef _HelperInputVoidPtrStruct cyuserDataHelper
     cdef void* cyuserData = _helper_input_void_ptr(userData, &cyuserDataHelper)
-
-    cdef cudaStreamHostCallbackData *cbData = NULL
-    cbData = <cudaStreamHostCallbackData *>malloc(sizeof(cbData[0]))
-    if cbData == NULL:
-        return (cudaError_t.cudaErrorMemoryAllocation,)
-    cbData.callback = cyfn
-    cbData.userData = cyuserData
-
     with nogil:
-        err = cyruntime.cudaLaunchHostFunc(cystream, <cyruntime.cudaHostFn_t>cudaStreamRtHostCallbackWrapper, <void *>cbData)
-    if err != cyruntime.cudaSuccess:
-        free(cbData)
+        err = cyruntime.cudaLaunchHostFunc(cystream, cyfn, cyuserData)
     _helper_input_void_ptr_free(&cyuserDataHelper)
-
     return (_cudaError_t(err),)
 {{endif}}
 

cuda_bindings/docs/source/overview.md

@@ -556,3 +556,22 @@ kernelTypes = (
 kernelArgs = (kernelValues, kernelTypes)
 ```
 
+
+(host-function-lifetime)=
+## Host functions
+
+{func}`cuda.bindings.driver.cuLaunchHostFunc` and
+{func}`cuda.bindings.runtime.cudaLaunchHostFunc` enqueue a host function on a
+stream. The caller owns the memory behind `fn` and `userData`; the bindings
+do not keep the corresponding Python objects alive.
+
+Keep strong references to the callback object (such as a `ctypes.CFUNCTYPE`
+instance) and the storage backing `userData` until the host function has
+completed, as confirmed by stream synchronization. A `CUhostFn` or
+`cudaHostFn_t` wrapper alone does not retain the underlying ctypes callback.
+
+When a host function is captured into a CUDA graph, it runs on every launch of
+that graph. Keep both references alive until the graph and all graph executables
+containing that host node have been destroyed and all pending launches have
+completed. Releasing either object after the first invocation can cause later
+graph launches to access freed memory.

cuda_bindings/docs/source/release/12.9.9-notes.rst

@@ -345,6 +345,15 @@ Bugfixes
   classes in ``cudla``, ``cufile``, and ``nvml``.
   (`PR #2875 <https://github.com/NVIDIA/cuda-python/pull/2875>`_)
 
+- ``driver.cuLaunchHostFunc`` and ``runtime.cudaLaunchHostFunc`` now pass the host
+  function and its user data to CUDA unchanged. They used to route the call through
+  an internal trampoline that freed its state after the first invocation, so a host
+  function captured into a CUDA graph crashed on the second launch of that graph.
+  (`Issue #790 <https://github.com/NVIDIA/cuda-python/issues/790>`_)
+  Callers must retain the callback object and user data storage for the lifetime
+  of the graph and its executables, including pending launches; see
+  :ref:`host-function-lifetime`.
+
 Known issues
 ------------
 

cuda_bindings/tests/pytest.ini

@@ -3,3 +3,7 @@
 
 [pytest]
 norecursedirs = cython
+markers =
+    agent_authored(model): agent-authored test not yet materially human-reviewed
+    human_reviewed: agent-authored test materially reviewed or rewritten by a human
+    human_authored: test authored primarily by a human

cuda_bindings/tests/test_cuda.py

@@ -779,6 +779,58 @@ def test_stream_capture():
     pass
 
 
+@pytest.mark.agent_authored(model="gpt-6")
+def test_cuda_launch_host_func_graph_replay():
+    (err,) = cuda.cuInit(0)
+    assert err == cuda.CUresult.CUDA_SUCCESS
+    err, device = cuda.cuDeviceGet(0)
+    assert err == cuda.CUresult.CUDA_SUCCESS
+    err, ctx = cuda.cuCtxCreate(0, device)
+    assert err == cuda.CUresult.CUDA_SUCCESS
+    try:
+        # A host function captured into a graph runs once per graph launch. The
+        # bindings must hand the function pointer and the user data to CUDA
+        # unchanged so that every launch reaches the same function (issue #790).
+        counter = ctypes.c_int(0)
+
+        @ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+        def host_fn(userData):
+            ctypes.c_int.from_address(userData).value += 1
+
+        callback = cuda.CUhostFn(_ptr=ctypes.addressof(host_fn))
+        err, stream = cuda.cuStreamCreate(0)
+        assert err == cuda.CUresult.CUDA_SUCCESS
+        try:
+            (err,) = cuda.cuStreamBeginCapture(stream, cuda.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL)
+            assert err == cuda.CUresult.CUDA_SUCCESS
+            (err,) = cuda.cuLaunchHostFunc(stream, callback, ctypes.addressof(counter))
+            assert err == cuda.CUresult.CUDA_SUCCESS
+            err, graph = cuda.cuStreamEndCapture(stream)
+            assert err == cuda.CUresult.CUDA_SUCCESS
+            try:
+                err, graphExec = cuda.cuGraphInstantiate(graph, 0)
+                assert err == cuda.CUresult.CUDA_SUCCESS
+                try:
+                    for launch in range(1, 6):
+                        (err,) = cuda.cuGraphLaunch(graphExec, stream)
+                        assert err == cuda.CUresult.CUDA_SUCCESS
+                        (err,) = cuda.cuStreamSynchronize(stream)
+                        assert err == cuda.CUresult.CUDA_SUCCESS
+                        assert counter.value == launch
+                finally:
+                    (err,) = cuda.cuGraphExecDestroy(graphExec)
+                    assert err == cuda.CUresult.CUDA_SUCCESS
+            finally:
+                (err,) = cuda.cuGraphDestroy(graph)
+                assert err == cuda.CUresult.CUDA_SUCCESS
+        finally:
+            (err,) = cuda.cuStreamDestroy(stream)
+            assert err == cuda.CUresult.CUDA_SUCCESS
+    finally:
+        (err,) = cuda.cuCtxDestroy(ctx)
+        assert err == cuda.CUresult.CUDA_SUCCESS
+
+
 def test_profiler():
     (err,) = cuda.cuInit(0)
     assert err == cuda.CUresult.CUDA_SUCCESS

cuda_bindings/tests/test_cudart.py

@@ -1291,33 +1291,32 @@ def test_cudart_eglFrame():
     assert int(frame.frame.pPitch[2].ysize) == 0
 
 
-def cudart_func_stream_callback(use_host_api):
+@pytest.mark.parametrize("use_host_api", [False, True], ids=["stream_callback", "host_function"])
+@pytest.mark.agent_authored(model="gpt-6")
+def test_cudart_func_callback(use_host_api):
     class testStruct(ctypes.Structure):
         _fields_ = [
             ("a", ctypes.c_int),
             ("b", ctypes.c_int),
             ("c", ctypes.c_int),
         ]
 
+    calls = []
+    stream_results = []
+
     def task_callback_host(userData):
         data = testStruct.from_address(userData)
-        assert data.a == 1
-        assert data.b == 2
-        assert data.c == 3
-        return 0
+        calls.append((userData, data.a, data.b, data.c))
 
     def task_callback_stream(stream, status, userData):
-        data = testStruct.from_address(userData)
-        assert data.a == 1
-        assert data.b == 2
-        assert data.c == 3
-        return 0
+        stream_results.append((stream, status))
+        task_callback_host(userData)
 
     if use_host_api:
-        callback_type = ctypes.PYFUNCTYPE(ctypes.c_int, ctypes.c_void_p)
+        callback_type = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
         target_task = task_callback_host
     else:
-        callback_type = ctypes.PYFUNCTYPE(ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p)
+        callback_type = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p)
         target_task = task_callback_stream
 
     # Construct ctype data
@@ -1333,19 +1332,64 @@ def task_callback_stream(stream, status, userData):
     # Run
     err, stream = cudart.cudaStreamCreate()
     assertSuccess(err)
-    if use_host_api:
-        (err,) = cudart.cudaLaunchHostFunc(stream, callback, ctypes.addressof(c_data))
+    try:
+        if use_host_api:
+            (err,) = cudart.cudaLaunchHostFunc(stream, callback, ctypes.addressof(c_data))
+        else:
+            (err,) = cudart.cudaStreamAddCallback(stream, callback, ctypes.addressof(c_data), 0)
         assertSuccess(err)
-    else:
-        (err,) = cudart.cudaStreamAddCallback(stream, callback, ctypes.addressof(c_data), 0)
+        (err,) = cudart.cudaStreamSynchronize(stream)
         assertSuccess(err)
-    (err,) = cudart.cudaDeviceSynchronize()
-    assertSuccess(err)
 
+        # ctypes swallows callback exceptions, so assert on the test thread.
+        assert calls == [(ctypes.addressof(c_data), 1, 2, 3)]
+        if not use_host_api:
+            assert stream_results == [(int(stream), int(cudart.cudaError_t.cudaSuccess))]
+    finally:
+        (err,) = cudart.cudaStreamDestroy(stream)
+        assertSuccess(err)
+
+
+@pytest.mark.agent_authored(model="claude-fable-5-1")
+def test_cudart_host_func_graph_replay():
+    # A host function captured into a graph runs once per graph launch. The
+    # bindings must hand the function pointer and the user data to CUDA
+    # unchanged so that every launch reaches the same function (issue #790).
+    counter = ctypes.c_int(0)
 
-def test_cudart_func_callback():
-    cudart_func_stream_callback(use_host_api=False)
-    cudart_func_stream_callback(use_host_api=True)
+    @ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+    def host_fn(userData):
+        ctypes.c_int.from_address(userData).value += 1
+
+    callback = cudart.cudaHostFn_t(_ptr=ctypes.addressof(host_fn))
+    err, stream = cudart.cudaStreamCreate()
+    assertSuccess(err)
+    try:
+        (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
+        assertSuccess(err)
+        (err,) = cudart.cudaLaunchHostFunc(stream, callback, ctypes.addressof(counter))
+        assertSuccess(err)
+        err, graph = cudart.cudaStreamEndCapture(stream)
+        assertSuccess(err)
+        try:
+            err, graphExec = cudart.cudaGraphInstantiate(graph, 0)
+            assertSuccess(err)
+            try:
+                for launch in range(1, 6):
+                    (err,) = cudart.cudaGraphLaunch(graphExec, stream)
+                    assertSuccess(err)
+                    (err,) = cudart.cudaStreamSynchronize(stream)
+                    assertSuccess(err)
+                    assert counter.value == launch
+            finally:
+                (err,) = cudart.cudaGraphExecDestroy(graphExec)
+                assertSuccess(err)
+        finally:
+            (err,) = cudart.cudaGraphDestroy(graph)
+            assertSuccess(err)
+    finally:
+        (err,) = cudart.cudaStreamDestroy(stream)
+        assertSuccess(err)
 
 
 @pytest.mark.skipif(