Commit 230f0eb

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.

CUresult launch_host_function(

CUstream stream, CUhostFn fn, void *user_data, OpaqueHandle fn_owner, OpaqueHandle data_owner

) noexcept;

X(cuLaunchHostFunc, 10000) \

#include "py.hpp"

#include "api.hpp"

#include "driver_api.hpp"

#include "internal.hpp"

#include <new>

#include <utility>

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))

{

}

};

host_launch_trampoline(void *p) noexcept

{

auto *rec = static_cast<HostLaunchRecord *>(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<detail::DeferredCleanupItem *>(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

__all__ = ['launch', 'host_launch']

def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None:

"""Enqueue a host function on a stream in stream order.

from cuda.core._host_callback cimport _resolve_host_callback

from cuda.core._rt cimport as_cu, launch_host_function, OpaqueHandle

__all__ = ['launch', 'host_launch']

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