feat(cuda.core): add HOST_COPY_FN and pack_host_copy_args for GIL-free host copies

cuda_core/cuda/core/__init__.py

@@ -152,6 +152,8 @@ class _PatchedProperty(metaclass=_PatchedPropMeta):
 from cuda.core._graphics import __all__ as _graphics_all
 from cuda.core._host import *
 from cuda.core._host import __all__ as _host_all
+from cuda.core._host_callback import *
+from cuda.core._host_callback import __all__ as _host_callback_all
 from cuda.core._launch_config import *
 from cuda.core._launch_config import __all__ as _launch_config_all
 from cuda.core._launcher import *
@@ -180,6 +182,7 @@ class _PatchedProperty(metaclass=_PatchedPropMeta):
     *_event_all,
     *_graphics_all,
     *_host_all,
+    *_host_callback_all,
     *_launch_config_all,
     *_launcher_all,
     *_linker_all,

cuda_core/cuda/core/_host_callback.pyi

@@ -1,8 +1,20 @@
 # This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/_host_callback.pyx
 
 import sys
+from _ctypes import CFuncPtr
+from typing import TypedDict
 
+from cuda.core._memory._buffer import Buffer
+from cuda.core._utils.pycompat import BufferProtocol
+
+__all__ = ['HOST_COPY_FN', 'pack_host_copy_args']
 _CUHOSTFN_HINT = 'ctypes.CFUNCTYPE(None, ctypes.c_void_p)' if sys.platform != 'win32' else 'ctypes.CFUNCTYPE(None, ctypes.c_void_p) or ctypes.WINFUNCTYPE(None, ctypes.c_void_p)'
+HOST_COPY_FN: CFuncPtr = ...
+
+class _HostCopyArgs(TypedDict):
+    dst: int
+    src: int
+    nbytes: int
 
 def _cuhostfn_type_error(detail):
     """Build the rejection message for a non-conforming ctypes callback."""
@@ -14,3 +26,59 @@ def _validate_ctypes_host_callback(fn):
     keeps ctypes' defaults -- a ``c_int`` result and unspecified arguments --
     until the caller declares otherwise, so it must be declared to be accepted.
     """
+def pack_host_copy_args(dst: int | Buffer | BufferProtocol, src: int | Buffer | BufferProtocol, nbytes: int) -> bytes:
+    """Pack the arguments of :data:`HOST_COPY_FN` into ``user_data`` bytes.
+
+    The result is a record of three native-width fields in this order:
+    ``uintptr_t dst``, ``uintptr_t src``, ``size_t nbytes``, with no padding.
+    ``struct.pack("@PPN", dst, src, nbytes)`` builds the same record, and
+    ``struct.unpack("@PPN", data)`` reads one back. A caller that keeps such a
+    record in memory of its own can pass its address as an ``int``
+    ``user_data`` instead.
+
+    ``dst`` and ``src`` are not retained. The caller must keep both alive,
+    and must not resize them (a ``bytearray``, for example), until the copy
+    has run.
+
+    Parameters
+    ----------
+    dst : int | :class:`Buffer` | :obj:`collections.abc.Buffer`
+        The destination. Only an ``int`` is a host address. A buffer-protocol
+        object must be contiguous and writable.
+    src : int | :class:`Buffer` | :obj:`collections.abc.Buffer`
+        The source. Only an ``int`` is a host address. A buffer-protocol
+        object must be contiguous.
+    nbytes : int
+        The number of bytes to copy. Zero packs a copy that does nothing.
+
+    Returns
+    -------
+    bytes
+        The packed record, to pass as ``user_data``.
+
+    Raises
+    ------
+    TypeError
+        If ``dst`` or ``src`` is not an ``int``, a :class:`Buffer`, or a
+        buffer-protocol object, or is a ``bool``, a NumPy integer scalar, or a
+        ctypes scalar such as ``ctypes.c_void_p``. A buffer-protocol object
+        that cannot export a contiguous view, or a writable one for ``dst``,
+        raises its own error.
+    ValueError
+        If ``nbytes`` exceeds the size of ``dst`` or ``src``, if a
+        :class:`Buffer` comes from a memory resource that is not
+        host-accessible, if an address is null, or if the ranges overlap.
+    OverflowError
+        If ``nbytes`` or an ``int`` address is negative or does not fit in a
+        native-width field.
+    RuntimeError
+        If ``dst`` or ``src`` is a closed :class:`Buffer`.
+
+    Notes
+    -----
+    This function makes no CUDA calls, so a Python host callback can call it,
+    for example to fill a record that a later :data:`HOST_COPY_FN` reads
+    through an ``int`` ``user_data``. It checks host accessibility only for a
+    :class:`Buffer` with a memory resource. For an ``int`` or a :class:`Buffer`
+    without one, the caller must ensure that the memory is host-accessible.
+    """

cuda_core/cuda/core/_host_callback.pyx

@@ -2,12 +2,21 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from cpython.buffer cimport (
+    PyBUF_SIMPLE,
+    PyBUF_WRITABLE,
+    PyBuffer_Release,
+    PyObject_CheckBuffer,
+    PyObject_GetBuffer,
+    Py_buffer,
+)
 from libc.stdint cimport uintptr_t
 from libc.stdlib cimport malloc
 from libc.string cimport memcpy as c_memcpy
 
 from cuda.bindings cimport cydriver
 
+from cuda.core._memory._buffer cimport Buffer, Buffer_check_open
 from cuda.core._rt cimport (
     OpaqueHandle,
     make_opaque_malloc,
@@ -16,6 +25,13 @@ from cuda.core._rt cimport (
 
 import sys
 import ctypes as ct
+from _ctypes import CFuncPtr
+
+import numpy
+
+from cuda.core._utils.pycompat import BufferProtocol
+
+__all__ = ['HOST_COPY_FN', 'pack_host_copy_args']
 
 
 # CUhostFn is `void (CUDA_CB *)(void*)`. CUDA_CB is __stdcall on Windows and
@@ -63,6 +79,126 @@ cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil:
     return fn == <cydriver.CUhostFn>_py_host_trampoline
 
 
+ctypedef struct _HostCopyArgs:
+    uintptr_t dst
+    uintptr_t src
+    size_t nbytes
+
+
+cdef void _host_copy(void* data) noexcept nogil:
+    cdef _HostCopyArgs* args = <_HostCopyArgs*>data
+    if args.nbytes:
+        c_memcpy(<void*>args.dst, <const void*>args.src, args.nbytes)
+
+
+cdef object _host_copy_fn_pointer():
+    return ct.CFUNCTYPE(None, ct.c_void_p)(<uintptr_t><cydriver.CUhostFn>_host_copy)
+
+
+HOST_COPY_FN: CFuncPtr = _host_copy_fn_pointer()
+
+
+cdef uintptr_t _host_copy_address(object obj, size_t nbytes, bint writable, str name) except? 0:
+    cdef Buffer buf
+    cdef Py_buffer view
+    # These export the buffer protocol over their own storage, so a value
+    # meant as an address would otherwise pack the address of that storage.
+    if isinstance(obj, (bool, numpy.integer, ct._SimpleCData)):
+        raise TypeError(f"{name} must not be a {type(obj).__name__}; pass int(addr) for an address")
+    if isinstance(obj, int):
+        return <uintptr_t>obj
+    if isinstance(obj, Buffer):
+        buf = <Buffer>obj
+        Buffer_check_open(buf)
+        # A Buffer without a memory resource could only be checked with a
+        # driver query, which is not allowed on a host callback thread.
+        if buf._memory_resource is not None and not buf._memory_resource.is_host_accessible:
+            raise ValueError(f"{name} is a Buffer that is not host-accessible")
+        if nbytes > buf._size:
+            raise ValueError(f"nbytes ({nbytes}) exceeds the size of {name} ({buf._size})")
+        return <uintptr_t>buf.handle
+    if not PyObject_CheckBuffer(obj):
+        raise TypeError(
+            f"{name} must be an int, a Buffer, or a buffer-protocol object, got {type(obj).__name__}")
+    PyObject_GetBuffer(obj, &view, PyBUF_SIMPLE | PyBUF_WRITABLE if writable else PyBUF_SIMPLE)
+    try:
+        if nbytes > <size_t>view.len:
+            raise ValueError(f"nbytes ({nbytes}) exceeds the size of {name} ({view.len})")
+        return <uintptr_t>view.buf
+    finally:
+        PyBuffer_Release(&view)
+
+
+def pack_host_copy_args(
+        dst: int | Buffer | BufferProtocol, src: int | Buffer | BufferProtocol, nbytes: int) -> bytes:
+    """Pack the arguments of :data:`HOST_COPY_FN` into ``user_data`` bytes.
+
+    The result is a record of three native-width fields in this order:
+    ``uintptr_t dst``, ``uintptr_t src``, ``size_t nbytes``, with no padding.
+    ``struct.pack("@PPN", dst, src, nbytes)`` builds the same record, and
+    ``struct.unpack("@PPN", data)`` reads one back. A caller that keeps such a
+    record in memory of its own can pass its address as an ``int``
+    ``user_data`` instead.
+
+    ``dst`` and ``src`` are not retained. The caller must keep both alive,
+    and must not resize them (a ``bytearray``, for example), until the copy
+    has run.
+
+    Parameters
+    ----------
+    dst : int | :class:`Buffer` | :obj:`collections.abc.Buffer`
+        The destination. Only an ``int`` is a host address. A buffer-protocol
+        object must be contiguous and writable.
+    src : int | :class:`Buffer` | :obj:`collections.abc.Buffer`
+        The source. Only an ``int`` is a host address. A buffer-protocol
+        object must be contiguous.
+    nbytes : int
+        The number of bytes to copy. Zero packs a copy that does nothing.
+
+    Returns
+    -------
+    bytes
+        The packed record, to pass as ``user_data``.
+
+    Raises
+    ------
+    TypeError
+        If ``dst`` or ``src`` is not an ``int``, a :class:`Buffer`, or a
+        buffer-protocol object, or is a ``bool``, a NumPy integer scalar, or a
+        ctypes scalar such as ``ctypes.c_void_p``. A buffer-protocol object
+        that cannot export a contiguous view, or a writable one for ``dst``,
+        raises its own error.
+    ValueError
+        If ``nbytes`` exceeds the size of ``dst`` or ``src``, if a
+        :class:`Buffer` comes from a memory resource that is not
+        host-accessible, if an address is null, or if the ranges overlap.
+    OverflowError
+        If ``nbytes`` or an ``int`` address is negative or does not fit in a
+        native-width field.
+    RuntimeError
+        If ``dst`` or ``src`` is a closed :class:`Buffer`.
+
+    Notes
+    -----
+    This function makes no CUDA calls, so a Python host callback can call it,
+    for example to fill a record that a later :data:`HOST_COPY_FN` reads
+    through an ``int`` ``user_data``. It checks host accessibility only for a
+    :class:`Buffer` with a memory resource. For an ``int`` or a :class:`Buffer`
+    without one, the caller must ensure that the memory is host-accessible.
+    """
+    cdef size_t n = nbytes
+    cdef _HostCopyArgs args
+    args.dst = _host_copy_address(dst, n, True, "dst")
+    args.src = _host_copy_address(src, n, False, "src")
+    args.nbytes = n
+    if n:
+        if args.dst == 0 or args.src == 0:
+            raise ValueError("dst and src must not be null")
+        if (args.dst - args.src if args.dst > args.src else args.src - args.dst) < n:
+            raise ValueError("dst and src overlap")
+    return (<char*>&args)[:sizeof(args)]
+
+
 cdef void _resolve_host_callback(
         object fn, object user_data,
         cydriver.CUhostFn* out_fn, void** out_user_data,
@@ -78,15 +214,21 @@ cdef void _resolve_host_callback(
     ctypes callbacks are validated against the ``CUhostFn`` ABI before their
     address is passed to CUDA.
     """
+    cdef bint is_host_copy
     if isinstance(fn, ct._CFuncPtr):
         _validate_ctypes_host_callback(fn)
         out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(fn, ct.c_void_p).value
+        is_host_copy = out_fn[0] == <cydriver.CUhostFn>_host_copy
         if user_data is None:
             out_user_data[0] = NULL
         elif isinstance(user_data, int):
             out_user_data[0] = <void*><uintptr_t>user_data
         else:
             buf = bytes(user_data)
+            if is_host_copy and len(buf) != sizeof(_HostCopyArgs):
+                raise ValueError(
+                    f"HOST_COPY_FN user_data must be {sizeof(_HostCopyArgs)} bytes "
+                    f"from pack_host_copy_args, got {len(buf)}")
             if len(buf):
                 out_user_data[0] = malloc(len(buf))
                 if out_user_data[0] == NULL:
@@ -95,6 +237,8 @@ cdef void _resolve_host_callback(
                 out_data_owner[0] = make_opaque_malloc(out_user_data[0])
             else:
                 out_user_data[0] = NULL
+        if is_host_copy and out_user_data[0] == NULL:
+            raise ValueError("HOST_COPY_FN requires user_data from pack_host_copy_args or a nonzero address")
     else:
         if not callable(fn):
             raise TypeError(

cuda_core/cuda/core/_launcher.pyi

@@ -51,8 +51,9 @@ def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=N
         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.
+        Python callable. :data:`HOST_COPY_FN` is such a pointer, for a host
+        copy that runs without the GIL. ``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
@@ -65,7 +66,8 @@ def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=N
         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.
+        If ``user_data`` is given for a Python callable, or if ``fn`` is
+        :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
     RuntimeError
         If ``stream`` is a :class:`Stream` that is capturing a graph, or a
         :class:`~graph.GraphBuilder` that is closed or not building.

cuda_core/cuda/core/_launcher.pyx

@@ -94,8 +94,9 @@ def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=N
         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.
+        Python callable. :data:`HOST_COPY_FN` is such a pointer, for a host
+        copy that runs without the GIL. ``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
@@ -108,7 +109,8 @@ def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=N
         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.
+        If ``user_data`` is given for a Python callable, or if ``fn`` is
+        :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
     RuntimeError
         If ``stream`` is a :class:`Stream` that is capturing a graph, or a
         :class:`~graph.GraphBuilder` that is closed or not building.

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -449,7 +449,8 @@ class GraphBuilder:
             If ``fn`` is a ctypes function pointer whose declared prototype
             does not match ``CUhostFn``.
         ValueError
-            If ``user_data`` is given for a Python callable.
+            If ``user_data`` is given for a Python callable, or if ``fn`` is
+            :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
         """
 
 class Graph:

cuda_core/cuda/core/graph/_graph_builder.pyx

@@ -988,7 +988,8 @@ cdef class GraphBuilder:
             If ``fn`` is a ctypes function pointer whose declared prototype
             does not match ``CUhostFn``.
         ValueError
-            If ``user_data`` is given for a Python callable.
+            If ``user_data`` is given for a Python callable, or if ``fn`` is
+            :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
         """
         GB_callback(self, fn, user_data, False)
 

cuda_core/cuda/core/graph/_graph_node.pyi

@@ -343,7 +343,8 @@ class GraphNode:
             If ``fn`` is a ctypes function pointer whose declared prototype
             does not match ``CUhostFn``.
         ValueError
-            If ``user_data`` is given for a Python callable.
+            If ``user_data`` is given for a Python callable, or if ``fn`` is
+            :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
         """
     def if_then(self, condition: GraphCondition) -> IfNode:
         """Add an if-conditional node depending on this node.

cuda_core/cuda/core/graph/_graph_node.pyx

@@ -541,7 +541,8 @@ cdef class GraphNode:
             If ``fn`` is a ctypes function pointer whose declared prototype
             does not match ``CUhostFn``.
         ValueError
-            If ``user_data`` is given for a Python callable.
+            If ``user_data`` is given for a Python callable, or if ``fn`` is
+            :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
         """
         return GN_callback(self, fn, user_data)
 

cuda_core/cuda/core/graph/_subclasses.pyi

@@ -274,7 +274,9 @@ class HostCallbackNode(GraphNode):
         ``fn`` accepts the same forms as :meth:`~graph.GraphNode.callback`: a
         Python callable, or a ctypes function pointer whose declared prototype
         matches ``CUhostFn`` (``void (*)(void*)``). A mismatched ctypes
-        prototype raises ``TypeError``.
+        prototype raises ``TypeError``. ``user_data`` with a Python callable,
+        or :data:`HOST_COPY_FN` without a packed record, raises
+        ``ValueError``.
 
         .. warning::
 
@@ -431,7 +433,9 @@ class ExecutableHostCallbackNode(ExecutableGraphNode):
 
         ``fn`` may be a Python callable, or a ctypes function pointer whose
         declared prototype matches ``CUhostFn`` (``void (*)(void*)``); a
-        mismatched prototype raises ``TypeError``.
+        mismatched prototype raises ``TypeError``. ``user_data`` with a Python
+        callable, or :data:`HOST_COPY_FN` without a packed record, raises
+        ``ValueError``.
 
         .. warning::