From 22734193d57afa240fca4940a9338d2424d682d0 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 9 Oct 2026 08:00:13 -0700 Subject: [PATCH 1/2] cuda.bindings: pass host functions to cuLaunchHostFunc and cudaLaunchHostFunc unchanged The two wrappers routed the call through an internal trampoline that freed its (fn, userData) record after the first call. A host function captured into a CUDA graph is called again on every launch of that graph, so the second launch used freed memory and crashed. The wrappers now pass fn and userData to the driver unchanged, as the _v2 variants already do. The bindings change equals the output of the cybind generator with the matching fix. Two tests capture a host function into a graph and launch it repeatedly, through the driver and the runtime API. Fixes #790 Co-Authored-By: Claude Fable 5.1 --- cuda_bindings/cuda/bindings/driver.pyx | 27 +------------ cuda_bindings/cuda/bindings/runtime.pyx | 28 +------------ .../docs/source/release/13.5.0-notes.rst | 5 +++ cuda_bindings/tests/test_cuda.py | 39 +++++++++++++++++++ cuda_bindings/tests/test_cudart.py | 39 +++++++++++++++++++ 5 files changed, 87 insertions(+), 51 deletions(-) diff --git a/cuda_bindings/cuda/bindings/driver.pyx b/cuda_bindings/cuda/bindings/driver.pyx index a6570e88b39..0d7268422d9 100644 --- a/cuda_bindings/cuda/bindings/driver.pyx +++ b/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 = 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 = phStream cdef _HelperInputVoidPtrStruct cyuserDataHelper cdef void* cyuserData = _helper_input_void_ptr(userData, &cyuserDataHelper) - - cdef cuHostCallbackData *cbData = NULL - cbData = 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, cuHostCallbackWrapper, cbData) - if err != cydriver.CUDA_SUCCESS: - free(cbData) + err = cydriver.cuLaunchHostFunc(cyhStream, cyfn, cyuserData) _helper_input_void_ptr_free(&cyuserDataHelper) return (_CUresult(err),) diff --git a/cuda_bindings/cuda/bindings/runtime.pyx b/cuda_bindings/cuda/bindings/runtime.pyx index 2b61ac54135..4c6be9ee02d 100644 --- a/cuda_bindings/cuda/bindings/runtime.pyx +++ b/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 = 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 = pstream cdef _HelperInputVoidPtrStruct cyuserDataHelper cdef void* cyuserData = _helper_input_void_ptr(userData, &cyuserDataHelper) - - cdef cudaStreamHostCallbackData *cbData = NULL - cbData = malloc(sizeof(cbData[0])) - if cbData == NULL: - return (cudaError_t.cudaErrorMemoryAllocation,) - cbData.callback = cyfn - cbData.userData = cyuserData - with nogil: - err = cyruntime.cudaLaunchHostFunc(cystream, cudaStreamRtHostCallbackWrapper, 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) diff --git a/cuda_bindings/docs/source/release/13.5.0-notes.rst b/cuda_bindings/docs/source/release/13.5.0-notes.rst index 356a26a0c6f..2b20188dd00 100644 --- a/cuda_bindings/docs/source/release/13.5.0-notes.rst +++ b/cuda_bindings/docs/source/release/13.5.0-notes.rst @@ -318,6 +318,11 @@ 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 `_) Prerelease feature ------------------ diff --git a/cuda_bindings/tests/test_cuda.py b/cuda_bindings/tests/test_cuda.py index 33e9e7f6b5b..4eba446daa4 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/cuda_bindings/tests/test_cuda.py @@ -694,6 +694,45 @@ 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 + (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 + err, graphExec = cuda.cuGraphInstantiate(graph, 0) + assert err == cuda.CUresult.CUDA_SUCCESS + + launches = 5 + for _ in range(launches): + (err,) = cuda.cuGraphLaunch(graphExec, stream) + assert err == cuda.CUresult.CUDA_SUCCESS + (err,) = cuda.cuStreamSynchronize(stream) + assert err == cuda.CUresult.CUDA_SUCCESS + assert counter.value == launches + + (err,) = cuda.cuGraphExecDestroy(graphExec) + assert err == cuda.CUresult.CUDA_SUCCESS + (err,) = cuda.cuGraphDestroy(graph) + assert err == cuda.CUresult.CUDA_SUCCESS + (err,) = cuda.cuStreamDestroy(stream) + assert err == cuda.CUresult.CUDA_SUCCESS + + def test_profiler(): (err,) = cuda.cuProfilerStart() assert err == cuda.CUresult.CUDA_SUCCESS diff --git a/cuda_bindings/tests/test_cudart.py b/cuda_bindings/tests/test_cudart.py index b4ac2142f35..f582386032d 100644 --- a/cuda_bindings/tests/test_cudart.py +++ b/cuda_bindings/tests/test_cudart.py @@ -1454,6 +1454,45 @@ def test_cudart_func_callback(): cudart_func_stream_callback(use_host_api=True) +@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 + + callback = cudart.cudaHostFn_t(_ptr=ctypes.addressof(host_fn)) + err, stream = cudart.cudaStreamCreate() + assertSuccess(err) + (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) + err, graphExec = cudart.cudaGraphInstantiate(graph, 0) + assertSuccess(err) + + launches = 5 + for _ in range(launches): + (err,) = cudart.cudaGraphLaunch(graphExec, stream) + assertSuccess(err) + (err,) = cudart.cudaStreamSynchronize(stream) + assertSuccess(err) + assert counter.value == launches + + (err,) = cudart.cudaGraphExecDestroy(graphExec) + assertSuccess(err) + (err,) = cudart.cudaGraphDestroy(graph) + assertSuccess(err) + (err,) = cudart.cudaStreamDestroy(stream) + assertSuccess(err) + + @pytest.mark.skipif( driver_version_less_than(12030) or not supportsCudaAPI("cudaGraphConditionalHandleCreate"), reason="Conditional graph APIs required", From 82b5f82cc2fb533518d02b415c877c10f4c8a0c0 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 13 Oct 2026 17:09:39 -0700 Subject: [PATCH 2/2] test(cuda.bindings): validate host callbacks after synchronization Keep coverage of cudaStreamAddCallback and cudaLaunchHostFunc while checking callback results on the test thread. Check every captured graph replay and clean up CUDA resources on failure. Document callback and user data lifetimes for captured graphs. Signed-off-by: Andy Jost --- cuda_bindings/docs/source/overview.rst | 22 +++++ .../docs/source/release/13.5.0-notes.rst | 3 + cuda_bindings/tests/test_cuda.py | 45 ++++----- cuda_bindings/tests/test_cudart.py | 91 ++++++++++--------- 4 files changed, 97 insertions(+), 64 deletions(-) diff --git a/cuda_bindings/docs/source/overview.rst b/cuda_bindings/docs/source/overview.rst index aefe7819f51..2888c06f70e 100644 --- a/cuda_bindings/docs/source/overview.rst +++ b/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. diff --git a/cuda_bindings/docs/source/release/13.5.0-notes.rst b/cuda_bindings/docs/source/release/13.5.0-notes.rst index 2b20188dd00..7c0d55ab97f 100644 --- a/cuda_bindings/docs/source/release/13.5.0-notes.rst +++ b/cuda_bindings/docs/source/release/13.5.0-notes.rst @@ -323,6 +323,9 @@ Bugfixes 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 `_) + 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 ------------------ diff --git a/cuda_bindings/tests/test_cuda.py b/cuda_bindings/tests/test_cuda.py index d40ef730b5b..369190c6b64 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/cuda_bindings/tests/test_cuda.py @@ -705,29 +705,32 @@ def host_fn(userData): callback = cuda.CUhostFn(_ptr=ctypes.addressof(host_fn)) err, stream = cuda.cuStreamCreate(0) assert err == cuda.CUresult.CUDA_SUCCESS - (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 - err, graphExec = cuda.cuGraphInstantiate(graph, 0) - assert err == cuda.CUresult.CUDA_SUCCESS - - launches = 5 - for _ in range(launches): - (err,) = cuda.cuGraphLaunch(graphExec, stream) + try: + (err,) = cuda.cuStreamBeginCapture(stream, cuda.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL) assert err == cuda.CUresult.CUDA_SUCCESS - (err,) = cuda.cuStreamSynchronize(stream) + (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 - assert counter.value == launches - - (err,) = cuda.cuGraphExecDestroy(graphExec) - assert err == cuda.CUresult.CUDA_SUCCESS - (err,) = cuda.cuGraphDestroy(graph) - assert err == cuda.CUresult.CUDA_SUCCESS - (err,) = cuda.cuStreamDestroy(stream) - assert err == cuda.CUresult.CUDA_SUCCESS def test_profiler(): diff --git a/cuda_bindings/tests/test_cudart.py b/cuda_bindings/tests/test_cudart.py index c420426659d..d0be4166c0d 100644 --- a/cuda_bindings/tests/test_cudart.py +++ b/cuda_bindings/tests/test_cudart.py @@ -1395,7 +1395,9 @@ 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), @@ -1403,25 +1405,22 @@ class testStruct(ctypes.Structure): ("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,22 @@ 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) - -def test_cudart_func_callback(): - cudart_func_stream_callback(use_host_api=False) - cudart_func_stream_callback(use_host_api=True) + # 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") @@ -1466,29 +1468,32 @@ def host_fn(userData): callback = cudart.cudaHostFn_t(_ptr=ctypes.addressof(host_fn)) err, stream = cudart.cudaStreamCreate() assertSuccess(err) - (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) - err, graphExec = cudart.cudaGraphInstantiate(graph, 0) - assertSuccess(err) - - launches = 5 - for _ in range(launches): - (err,) = cudart.cudaGraphLaunch(graphExec, stream) + try: + (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) assertSuccess(err) - (err,) = cudart.cudaStreamSynchronize(stream) + (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) - assert counter.value == launches - - (err,) = cudart.cudaGraphExecDestroy(graphExec) - assertSuccess(err) - (err,) = cudart.cudaGraphDestroy(graph) - assertSuccess(err) - (err,) = cudart.cudaStreamDestroy(stream) - assertSuccess(err) @pytest.mark.skipif(