diff --git a/cuda_core/cuda/core/_cpp/rt/DESIGN.md b/cuda_core/cuda/core/_cpp/rt/DESIGN.md index adc8a0fdafc..1d550a4c03f 100644 --- a/cuda_core/cuda/core/_cpp/rt/DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/DESIGN.md @@ -124,6 +124,7 @@ cuda/core/ ├── context_scope.hpp # Scoped-context helpers (namespace detail) ├── internal.hpp # Registry, cleanup wrappers, deferred-cleanup item (namespace detail) ├── py_report.cpp, py_deferred_cleanup.cpp # Python-coupled bodies + ├── host_launch.cpp # The stream-side host function launch record └── context.cpp, stream.cpp, event.cpp, memory.cpp, program.cpp, graph.cpp, graph_exec.cpp, texture.cpp # One resource family per file ``` diff --git a/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md b/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md index fdaf77785b2..8363c4fcd3f 100644 --- a/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md +++ b/cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md @@ -127,6 +127,10 @@ deleters call CUDA or run Python finalizers. destructor callback only adds the attachment to the process-lifetime `DeferredCleanupQueue` and requests a `Py_AddPendingCall`. +The host launch record of `host_launch.cpp` is the one producer that is not a +graph attachment; +its trampoline enqueues the record after the user function returns. + One pending call drains all queued attachments from Python's main thread. The queue coalesces work because CPython's pending-call queue is bounded, and there could be many more deferred cleanup items than allowed pending calls. If diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index 082797e8d45..dabaac2a5c8 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -109,6 +109,14 @@ StreamHandle get_per_thread_stream(); // token is created. Returns an empty handle for an empty h_context. StreamHandle create_context_bound_legacy_stream(const ContextHandle &h_context); +// ============================================================================ +// Host function launch +// ============================================================================ + +CUresult launch_host_function( + CUstream stream, CUhostFn fn, void *user_data, OpaqueHandle fn_owner, OpaqueHandle data_owner +) noexcept; + // ============================================================================ // Event handle functions // ============================================================================ diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp index d759112385e..2b9affab5d7 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp @@ -139,6 +139,7 @@ namespace cuda_core::rt { /* cuda-bindings requests 7000 (PTDS) or 2000 (legacy) */ \ X(cuStreamSynchronize, 7000) \ X(cuStreamIsCapturing, 10000) \ + X(cuLaunchHostFunc, 10000) \ X(cuThreadExchangeStreamCaptureMode, 10010) #define CUDA_CORE_DECLARE_DRIVER_FN(name, introduced) extern decltype(&name) p_##name; diff --git a/cuda_core/cuda/core/_cpp/rt/host_launch.cpp b/cuda_core/cuda/core/_cpp/rt/host_launch.cpp new file mode 100644 index 00000000000..3e0789eb57e --- /dev/null +++ b/cuda_core/cuda/core/_cpp/rt/host_launch.cpp @@ -0,0 +1,57 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "py.hpp" +#include "api.hpp" +#include "driver_api.hpp" +#include "internal.hpp" +#include +#include + +namespace cuda_core::rt { + +namespace { + +struct HostLaunchRecord : detail::DeferredCleanupItem { + CUhostFn fn; + void *user_data; + OpaqueHandle fn_owner; + OpaqueHandle data_owner; + + HostLaunchRecord(CUhostFn fn, void *user_data, OpaqueHandle fn_owner, OpaqueHandle data_owner) noexcept + : fn(fn), user_data(user_data), fn_owner(std::move(fn_owner)), data_owner(std::move(data_owner)) + { + } +}; + +void +host_launch_trampoline(void *p) noexcept +{ + auto *rec = static_cast(p); + rec->fn(rec->user_data); + // Owners may run Python finalizers or CUDA calls, so release them on the main thread. + detail::enqueue_cleanup(static_cast(rec)); +} + +} // namespace + +CUresult +launch_host_function( + CUstream stream, CUhostFn fn, void *user_data, OpaqueHandle fn_owner, OpaqueHandle data_owner +) noexcept +{ + retry_deferred_cleanup(); + auto *record = new (std::nothrow) HostLaunchRecord(fn, user_data, std::move(fn_owner), std::move(data_owner)); + if (record == nullptr) { + return CUDA_ERROR_OUT_OF_MEMORY; + } + CUresult status = DRIVER_CALL(cuLaunchHostFunc, stream, &host_launch_trampoline, record); + if (status != CUDA_SUCCESS) { + delete record; + return status; + } + return CUDA_SUCCESS; +} + +} // namespace cuda_core::rt diff --git a/cuda_core/cuda/core/graph/_host_callback.pxd b/cuda_core/cuda/core/_host_callback.pxd similarity index 100% rename from cuda_core/cuda/core/graph/_host_callback.pxd rename to cuda_core/cuda/core/_host_callback.pxd diff --git a/cuda_core/cuda/core/graph/_host_callback.pyi b/cuda_core/cuda/core/_host_callback.pyi similarity index 95% rename from cuda_core/cuda/core/graph/_host_callback.pyi rename to cuda_core/cuda/core/_host_callback.pyi index 5f4cf3d2df9..3bbf6321be6 100644 --- a/cuda_core/cuda/core/graph/_host_callback.pyi +++ b/cuda_core/cuda/core/_host_callback.pyi @@ -1,4 +1,4 @@ -# This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/graph/_host_callback.pyx +# This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/_host_callback.pyx import sys diff --git a/cuda_core/cuda/core/graph/_host_callback.pyx b/cuda_core/cuda/core/_host_callback.pyx similarity index 100% rename from cuda_core/cuda/core/graph/_host_callback.pyx rename to cuda_core/cuda/core/_host_callback.pyx diff --git a/cuda_core/cuda/core/_launcher.pyi b/cuda_core/cuda/core/_launcher.pyi index 5d047fcbbfb..56f189fe72f 100644 --- a/cuda_core/cuda/core/_launcher.pyi +++ b/cuda_core/cuda/core/_launcher.pyi @@ -6,7 +6,7 @@ from cuda.core._stream import Stream from cuda.core.graph import GraphBuilder from cuda.core.typing import IsStreamType -__all__ = ['launch'] +__all__ = ['launch', 'host_launch'] def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, kernel: Kernel, *kernel_args) -> None: """Launches a :obj:`~_module.Kernel` @@ -27,3 +27,65 @@ def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, k launching kernel. """ +def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None: + """Enqueue a host function on a stream in stream order. + + The function runs on a driver-owned thread once all work enqueued on the + stream before this call has completed. Work enqueued after this call waits + until the function returns. This call returns when the function is enqueued; + observe completion with :meth:`Stream.sync` or an :class:`Event` recorded + after it. + + Passing a :class:`~graph.GraphBuilder` forwards to + :meth:`~graph.GraphBuilder.callback`. + + Parameters + ---------- + stream : Stream or graph.GraphBuilder or IsStreamType + The stream establishing execution order, or a graph builder capturing + the callback. + fn : callable or ctypes function pointer + A Python callable taking no arguments, which runs with the GIL, or a + ctypes function pointer declared as ``CUhostFn``, for example + ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``, which ``cuda.core`` + passes to the driver as is. A pointer to compiled code runs without + the GIL; a ctypes callback that wraps a Python function acquires the + GIL through ctypes and is subject to the same deadlock rule as a + Python callable. ``cuda.core`` keeps ``fn`` alive until the function + has run and releases it afterwards on the main thread. + user_data : int or bytes-like, optional + Only for ctypes function pointers. An ``int`` is a raw pointer that + the caller must keep valid until the function has run. Bytes-like data + is copied; ``cuda.core`` owns the copy until the function has run and + releases it afterwards on the main thread. + + Raises + ------ + TypeError + If ``stream`` is ``None`` or not a stream, or if ``fn`` is not + callable or its ctypes prototype is not ``CUhostFn``. + ValueError + If ``user_data`` is given for a Python callable. + RuntimeError + If ``stream`` is a :class:`Stream` that is capturing a graph, or a + :class:`~graph.GraphBuilder` that is closed or not building. + + Notes + ----- + An exception raised by a Python callable cannot propagate to the caller. + It is reported through :func:`sys.unraisablehook` and the stream continues. + + .. warning:: + + The function must not call CUDA, directly or through a library. The + driver does not permit such calls. Such a call may fail with an error, + or it may deadlock. + + A Python callable, or a ctypes callback that wraps a Python function, + acquires the GIL on the driver's thread. A thread + that holds the GIL while blocking inside a CUDA call waiting for this + stream therefore deadlocks with the callback. ``cuda.core`` and + ``cuda.bindings`` release the GIL around CUDA calls. A library that + calls CUDA with the GIL held must not be used to wait on a stream with + a pending Python callback. + """ diff --git a/cuda_core/cuda/core/_launcher.pyx b/cuda_core/cuda/core/_launcher.pyx index 1eef53beaf3..f5becd4b316 100644 --- a/cuda_core/cuda/core/_launcher.pyx +++ b/cuda_core/cuda/core/_launcher.pyx @@ -6,10 +6,11 @@ from libc.stdint cimport uintptr_t from cuda.bindings cimport cydriver +from cuda.core._host_callback cimport _resolve_host_callback from cuda.core._launch_config cimport LaunchConfig from cuda.core._kernel_arg_handler cimport ParamHolder from cuda.core._module cimport Kernel -from cuda.core._rt cimport as_cu +from cuda.core._rt cimport as_cu, launch_host_function, OpaqueHandle from cuda.core._stream cimport Stream_accept, Stream from cuda.core._utils.cuda_utils cimport ( check_or_create_options, @@ -24,7 +25,7 @@ if TYPE_CHECKING: from cuda.core.graph import GraphBuilder from cuda.core.typing import IsStreamType -__all__ = ['launch'] +__all__ = ['launch', 'host_launch'] def launch( @@ -69,6 +70,92 @@ def launch( HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL)) +def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None: + """Enqueue a host function on a stream in stream order. + + The function runs on a driver-owned thread once all work enqueued on the + stream before this call has completed. Work enqueued after this call waits + until the function returns. This call returns when the function is enqueued; + observe completion with :meth:`Stream.sync` or an :class:`Event` recorded + after it. + + Passing a :class:`~graph.GraphBuilder` forwards to + :meth:`~graph.GraphBuilder.callback`. + + Parameters + ---------- + stream : Stream or graph.GraphBuilder or IsStreamType + The stream establishing execution order, or a graph builder capturing + the callback. + fn : callable or ctypes function pointer + A Python callable taking no arguments, which runs with the GIL, or a + ctypes function pointer declared as ``CUhostFn``, for example + ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``, which ``cuda.core`` + passes to the driver as is. A pointer to compiled code runs without + the GIL; a ctypes callback that wraps a Python function acquires the + GIL through ctypes and is subject to the same deadlock rule as a + Python callable. ``cuda.core`` keeps ``fn`` alive until the function + has run and releases it afterwards on the main thread. + user_data : int or bytes-like, optional + Only for ctypes function pointers. An ``int`` is a raw pointer that + the caller must keep valid until the function has run. Bytes-like data + is copied; ``cuda.core`` owns the copy until the function has run and + releases it afterwards on the main thread. + + Raises + ------ + TypeError + If ``stream`` is ``None`` or not a stream, or if ``fn`` is not + callable or its ctypes prototype is not ``CUhostFn``. + ValueError + If ``user_data`` is given for a Python callable. + RuntimeError + If ``stream`` is a :class:`Stream` that is capturing a graph, or a + :class:`~graph.GraphBuilder` that is closed or not building. + + Notes + ----- + An exception raised by a Python callable cannot propagate to the caller. + It is reported through :func:`sys.unraisablehook` and the stream continues. + + .. warning:: + + The function must not call CUDA, directly or through a library. The + driver does not permit such calls. Such a call may fail with an error, + or it may deadlock. + + A Python callable, or a ctypes callback that wraps a Python function, + acquires the GIL on the driver's thread. A thread + that holds the GIL while blocking inside a CUDA call waiting for this + stream therefore deadlocks with the callback. ``cuda.core`` and + ``cuda.bindings`` release the GIL around CUDA calls. A library that + calls CUDA with the GIL held must not be used to wait on a stream with + a pending Python callback. + """ + from cuda.core.graph._graph_builder import GraphBuilder + + if isinstance(stream, GraphBuilder): + return stream.callback(fn, user_data=user_data) + + cdef Stream s = Stream_accept(stream, allow_stream_protocol=True) + cdef cydriver.CUstream c_stream = as_cu(s._h_stream) + cdef cydriver.CUstreamCaptureStatus status + with nogil: + HANDLE_RETURN(cydriver.cuStreamIsCapturing(c_stream, &status)) + if status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: + raise RuntimeError( + "host_launch cannot enqueue on a stream that is capturing a graph; " + "pass the GraphBuilder, or call GraphBuilder.callback") + + cdef cydriver.CUhostFn c_fn + cdef void* c_user_data = NULL + cdef OpaqueHandle fn_owner, data_owner + _resolve_host_callback(fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner) + with nogil: + HANDLE_RETURN(launch_host_function(c_stream, c_fn, c_user_data, fn_owner, data_owner)) + return None + + cdef _check_cooperative_launch(kernel: Kernel, config: LaunchConfig, stream: Stream): dev = stream.device num_sm = dev.properties.multiprocessor_count diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index 36fbd82acc4..4d5e63cb2bc 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -235,6 +235,9 @@ cdef StreamHandle create_stream_handle( cdef StreamHandle create_stream_handle_ref(cydriver.CUstream stream) except+ nogil cdef StreamHandle create_stream_handle_with_owner(cydriver.CUstream stream, object owner) except+ nogil cdef void retry_deferred_cleanup() noexcept +cdef cydriver.CUresult launch_host_function( + cydriver.CUstream stream, cydriver.CUhostFn fn, void* user_data, + OpaqueHandle fn_owner, OpaqueHandle data_owner) noexcept nogil cdef ContextHandle get_stream_context(const StreamHandle& h) noexcept nogil cdef StreamHandle get_legacy_stream() except+ nogil cdef StreamHandle get_per_thread_stream() except+ nogil diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index 1ebe3a1b460..49523bbce91 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -103,6 +103,9 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": cydriver.CUstream stream, object owner) except+ nogil void initialize_deferred_cleanup "cuda_core::rt::initialize_deferred_cleanup" () except+ void retry_deferred_cleanup "cuda_core::rt::retry_deferred_cleanup" () noexcept + cydriver.CUresult launch_host_function "cuda_core::rt::launch_host_function" ( + cydriver.CUstream stream, cydriver.CUhostFn fn, void* user_data, + OpaqueHandle fn_owner, OpaqueHandle data_owner) noexcept nogil ContextHandle get_stream_context "cuda_core::rt::get_stream_context" ( const StreamHandle& h) noexcept nogil StreamHandle get_legacy_stream "cuda_core::rt::get_legacy_stream" () except+ nogil diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyx b/cuda_core/cuda/core/graph/_graph_builder.pyx index ce45b4d272c..2ac92caacbc 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyx +++ b/cuda_core/cuda/core/graph/_graph_builder.pyx @@ -15,7 +15,7 @@ from cuda.core.graph._graph_definition cimport ( GD_check_valid, ) from cuda.core.graph._graph_node cimport GraphNode, GN_check_valid -from cuda.core.graph._host_callback cimport _resolve_host_callback +from cuda.core._host_callback cimport _resolve_host_callback from cuda.core.graph._subclasses cimport ( ChildGraphNode, EventRecordNode, diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 4c81860d39a..4abb1f7106c 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -66,7 +66,7 @@ from cuda.core._rt cimport ( ) from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value -from cuda.core.graph._host_callback cimport ( +from cuda.core._host_callback cimport ( _resolve_host_callback, ) diff --git a/cuda_core/cuda/core/graph/_subclasses.pyx b/cuda_core/cuda/core/graph/_subclasses.pyx index 50f09ca8f6b..747c9c9b452 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyx +++ b/cuda_core/cuda/core/graph/_subclasses.pyx @@ -62,7 +62,7 @@ from cuda.core._rt cimport ( from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value from cuda.core._utils.version cimport cy_driver_version -from cuda.core.graph._host_callback cimport ( +from cuda.core._host_callback cimport ( _is_py_host_trampoline, _resolve_host_callback, ) diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst index d9c4acf206c..a1be74a03df 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -21,6 +21,7 @@ Devices and execution Device Host launch + host_launch :template: autosummary/cyclass.rst diff --git a/cuda_core/docs/source/examples.rst b/cuda_core/docs/source/examples.rst index 05444b0bad3..e69ab5169bf 100644 --- a/cuda_core/docs/source/examples.rst +++ b/cuda_core/docs/source/examples.rst @@ -35,6 +35,9 @@ Kernels and launch configuration partitions SMs across kernels with green contexts. - :sample:`kernelNsysProfile ` annotates kernel launches with NVTX markers for Nsight Systems. +- :cuda-core-example:`host_launch.py ` + orders a host function between two kernel launches on a stream and reads a + pinned buffer from it. Memory management and interop ----------------------------- diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 5d6521d30a8..b6f7bb066a4 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -45,6 +45,14 @@ Breaking Changes New features ------------ +- Added :func:`host_launch`, which enqueues a host function on a stream in stream + order: a Python callable, or a ctypes function pointer with the ``CUhostFn`` + prototype, which ``cuda.core`` passes to the driver as is, so a pointer to + compiled code runs without the GIL. It is the stream-side counterpart of + :meth:`graph.GraphBuilder.callback` and accepts the same ``fn`` and + ``user_data`` forms. + (`#2058 `__) + - Added incremental linking to :class:`Linker` through ``LinkerOptions(incremental=True)`` on nvJitLink 13.2 or newer. A partial native result is returned as a CUBIN and can be passed directly to another diff --git a/cuda_core/examples/host_launch.py b/cuda_core/examples/host_launch.py new file mode 100644 index 00000000000..e70ec675b93 --- /dev/null +++ b/cuda_core/examples/host_launch.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# ################################################################################ +# +# This example reads a pinned host buffer in a stream-ordered host function. +# The callback observes the first kernel's result before the second overwrites it. +# +# ################################################################################ + +# /// script +# dependencies = ["cuda_bindings", "cuda_core"] +# /// + +import ctypes +from contextlib import closing + +from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, Program, host_launch, launch + + +def main(): + dev = Device(0) + dev.set_current() + size = 256 + code = r""" + extern "C" __global__ void fill(unsigned char *buffer, unsigned long long value) { + buffer[threadIdx.x] = static_cast(value); + } + """ + arch = "sm_" + "".join(str(part) for part in dev.compute_capability) + program = Program(code, code_type="c++", options={"arch": arch}) + module = program.compile("cubin") + kernel = module.get_kernel("fill") + config = LaunchConfig(grid=1, block=size) + checksums = [] + + with ( + closing(dev.create_stream()) as stream, + dev.allocate(size, stream=stream) as device_buffer, + LegacyPinnedMemoryResource().allocate(size) as host_buffer, + ): + + def record_checksum(): + data = ctypes.string_at(int(host_buffer.handle), size) + checksums.append(sum(data)) + + launch(stream, config, kernel, device_buffer, ctypes.c_uint64(0x11)) + device_buffer.copy_to(host_buffer, stream=stream) + host_launch(stream, record_checksum) + launch(stream, config, kernel, device_buffer, ctypes.c_uint64(0x22)) + device_buffer.copy_to(host_buffer, stream=stream) + stream.sync() + + assert checksums == [0x11 * size] + assert ctypes.string_at(int(host_buffer.handle), size) == b"\x22" * size + + print("Stream-ordered host function observed the first kernel's checksum.") + + +if __name__ == "__main__": + main() diff --git a/cuda_core/tests/test_host_launch.py b/cuda_core/tests/test_host_launch.py new file mode 100644 index 00000000000..cabbb94ffc4 --- /dev/null +++ b/cuda_core/tests/test_host_launch.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import ctypes +import sys +import threading +import time +import weakref +from contextlib import closing + +import pytest +from helpers.misc import StreamWrapper + +from cuda.core import Device, LegacyPinnedMemoryResource, host_launch + +N = 16 + + +def _as_bytes(buffer): + return bytes((ctypes.c_uint8 * N).from_address(int(buffer.handle))) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_orders_between_prior_and_later_work(init_cuda): + dev = Device(0) + dev.set_current() + seen = [] + with ( + closing(dev.create_stream()) as stream, + dev.allocate(N, stream=stream) as device_buffer, + LegacyPinnedMemoryResource().allocate(N) as host_buffer, + ): + device_buffer.fill(0x11, stream=stream) + device_buffer.copy_to(host_buffer, stream=stream) + result = host_launch(stream, lambda: seen.append(_as_bytes(host_buffer))) + device_buffer.fill(0x22, stream=stream) + device_buffer.copy_to(host_buffer, stream=stream) + stream.sync() + + assert result is None + assert seen == [b"\x11" * N] + assert _as_bytes(host_buffer) == b"\x22" * N + + +@pytest.mark.thread_unsafe(reason="the release is a main-thread pending call") +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_releases_callable_after_it_ran(init_cuda): + dev = Device(0) + dev.set_current() + ran = [] + + class Callback: + def __call__(self): + ran.append(True) + + with closing(dev.create_stream()) as stream: + callback = Callback() + ref = weakref.ref(callback) + host_launch(stream, callback) + del callback + stream.sync() + assert ran == [True] + + deadline = time.monotonic() + 5 + while ref() is not None and time.monotonic() < deadline: + time.sleep(0.01) + assert ref() is None, "host_launch did not release the callable after it ran" + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_ctypes_bytes_user_data_is_copied(init_cuda): + """The bytes are copied at launch; a later change to the caller's buffer is not seen.""" + dev = Device(0) + dev.set_current() + seen = [] + gate = threading.Event() + + @ctypes.CFUNCTYPE(None, ctypes.c_void_p) + def callback(data): + seen.append(ctypes.string_at(data, N)) + + with closing(dev.create_stream()) as stream: + data = bytearray(range(N)) + try: + # Hold the stream on the driver thread (Event.wait releases the GIL) while + # the caller overwrites its buffer; the copy made at launch is unaffected. + host_launch(stream, gate.wait) + host_launch(stream, callback, user_data=data) + data[:] = b"\xff" * N + finally: + gate.set() + stream.sync() + assert seen == [bytes(range(N))] + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_ctypes_int_user_data_passes_through(init_cuda): + dev = Device(0) + dev.set_current() + seen = [] + buf = ctypes.create_string_buffer(b"\x5a") + pointer = ctypes.addressof(buf) + + @ctypes.CFUNCTYPE(None, ctypes.c_void_p) + def callback(data): + seen.append((data, ctypes.c_uint8.from_address(data).value)) + + with closing(dev.create_stream()) as stream: + host_launch(stream, callback, user_data=pointer) + stream.sync() + assert seen == [(pointer, 0x5A)] + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_rejects_user_data_for_python_callable(init_cuda): + dev = Device(0) + dev.set_current() + with closing(dev.create_stream()) as stream, pytest.raises(ValueError): + host_launch(stream, lambda: None, user_data=b"data") + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_rejects_non_callable(init_cuda): + dev = Device(0) + dev.set_current() + with closing(dev.create_stream()) as stream, pytest.raises(TypeError): + host_launch(stream, 42) + + +@pytest.mark.parametrize("prototype", [ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p), ctypes.CFUNCTYPE(None)]) +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_rejects_wrong_ctypes_prototype(init_cuda, prototype): + dev = Device(0) + dev.set_current() + callback = prototype(lambda *_args: 0) + with closing(dev.create_stream()) as stream, pytest.raises(TypeError): + host_launch(stream, callback) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_forwards_graph_builder(init_cuda): + dev = Device(0) + dev.set_current() + ran = [] + with closing(dev.create_stream()) as stream: + gb = stream.create_graph_builder().begin_building() + host_launch(gb, lambda: ran.append(True)) + gb.end_building() + with closing(gb.complete()) as graph: + graph.upload(stream) + graph.launch(stream) + graph.launch(stream) + stream.sync() + assert ran == [True, True] + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_rejects_capturing_stream(init_cuda): + dev = Device(0) + dev.set_current() + ran = [] + with closing(dev.create_stream()) as stream: + gb = stream.create_graph_builder().begin_building() + try: + with pytest.raises(RuntimeError, match="capturing"): + host_launch(gb.stream, lambda: ran.append(True)) + finally: + gb.end_building() + graph = gb.complete() + graph.close() + assert ran == [] + + +@pytest.mark.thread_unsafe(reason="replaces the process-global sys.unraisablehook") +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_python_exception_is_reported_as_unraisable(init_cuda, monkeypatch): + dev = Device(0) + dev.set_current() + records = [] + monkeypatch.setattr(sys, "unraisablehook", records.append) + ran = [] + + def throws(): + raise RuntimeError("callback raised") + + with closing(dev.create_stream()) as stream: + host_launch(stream, throws) + stream.sync() + host_launch(stream, lambda: ran.append(True)) + stream.sync() + assert ran == [True] + assert [record.exc_type for record in records] == [RuntimeError] + assert "callback raised" in str(records[0].exc_value) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_two_streams_keep_their_own_data(init_cuda): + dev = Device(0) + dev.set_current() + seen1, seen2 = [], [] + with ( + closing(dev.create_stream()) as stream1, + closing(dev.create_stream()) as stream2, + dev.allocate(N, stream=stream1) as device_buffer1, + dev.allocate(N, stream=stream2) as device_buffer2, + LegacyPinnedMemoryResource().allocate(N) as host_buffer1, + LegacyPinnedMemoryResource().allocate(N) as host_buffer2, + ): + device_buffer1.fill(0x33, stream=stream1) + device_buffer1.copy_to(host_buffer1, stream=stream1) + host_launch(stream1, lambda: seen1.append(_as_bytes(host_buffer1))) + device_buffer2.fill(0x44, stream=stream2) + device_buffer2.copy_to(host_buffer2, stream=stream2) + host_launch(stream2, lambda: seen2.append(_as_bytes(host_buffer2))) + stream1.sync() + stream2.sync() + assert seen1 == [b"\x33" * N] + assert seen2 == [b"\x44" * N] + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_rejects_none_stream(init_cuda): + with pytest.raises(TypeError): + host_launch(None, lambda: None) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_host_launch_accepts_stream_protocol_with_deprecation(init_cuda): + dev = Device(0) + dev.set_current() + ran = [] + with closing(dev.create_stream()) as stream: + with pytest.warns(DeprecationWarning, match="stream protocol is deprecated"): + host_launch(StreamWrapper(stream), lambda: ran.append(True)) + stream.sync() + assert ran == [True]