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/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 a3da87071f4..a95a520bea2 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,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 `_) + 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 2e5547e2333..369190c6b64 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/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 diff --git a/cuda_bindings/tests/test_cudart.py b/cuda_bindings/tests/test_cudart.py index 82974586ef5..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,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(