cuda.bindings: pass host functions to cuLaunchHostFunc and cudaLaunchHostFunc unchanged

cuda_bindings/cuda/bindings/driver.pyx

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 13.4.1. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=175f5b5d51fbe64e7d8f5eca92e8b557cab5dea406091f66e10a4d19400dcefa
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=548b03ca3d3021957ab650d666fcd33a16d9af80a5004254c2733f807836679b
 from typing import Any, Optional
 import cython
 import ctypes
@@ -43815,19 +43815,6 @@ def cuLaunchCooperativeKernel(f, unsigned int gridDimX, unsigned int gridDimY, u
         err = cydriver.cuLaunchCooperativeKernel(cyf, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY, blockDimZ, sharedMemBytes, cyhStream, cykernelParams_ptr)
     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.
@@ -43907,18 +43894,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

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 13.4.1. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=4c952a0aa0dcb7a2948769d3eee64d0ef217ca1a8bd72873d99d94e9a34f2a65
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=ea95eb12334cbd4c1bf3fc8e890fb6f9493063b12faea853bb3d8c6288934182
 from typing import Any, Optional
 import cython
 import ctypes
@@ -25014,19 +25014,6 @@ def cudaFuncGetParamCount(func):
         return (_cudaError_t(err), None)
     return (_cudaError_t_SUCCESS, paramCount)
 
-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.
@@ -25107,20 +25094,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),)
 
 @cython.embedsignature(True)

cuda_bindings/docs/source/overview.rst

@@ -571,3 +571,25 @@ For ctypes, we leverage the special handling of ``None`` type since each Python
        None,
    )
    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. The same lifetime requirements apply to
+the ``_v2`` variants of these APIs.

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

@@ -323,6 +323,14 @@ Bugfixes
 - A number of incorrect docstrings have been corrected.
 - Setting a ``cuda_bindings`` class instance to itself (``obj = cuda_bindings.driver.Class(); obj[0] = obj``) will no longer crash or use garbage values.
 - Fix a potential use-after-free error in the destructors of classes in ``driver`` and ``runtime``.
+- ``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`.
 
 Prerelease feature
 ------------------

cuda_bindings/tests/test_cuda.py

@@ -691,6 +691,48 @@ def test_stream_capture():
     pass
 
 
+@pytest.mark.agent_authored(model="claude-fable-5-1")
+def test_cuda_launch_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)
+
+    @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
+
+
 def test_profiler():
     (err,) = cuda.cuProfilerStart()
     assert err == cuda.CUresult.CUDA_SUCCESS

cuda_bindings/tests/test_cudart.py

@@ -1395,33 +1395,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
@@ -1437,19 +1436,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)
+
+    @ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+    def host_fn(userData):
+        ctypes.c_int.from_address(userData).value += 1
 
-def test_cudart_func_callback():
-    cudart_func_stream_callback(use_host_api=False)
-    cudart_func_stream_callback(use_host_api=True)
+    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(