diff --git a/cuda_core/cuda/core/__init__.py b/cuda_core/cuda/core/__init__.py index 06f7ad79bfe..4b87a57529e 100644 --- a/cuda_core/cuda/core/__init__.py +++ b/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, diff --git a/cuda_core/cuda/core/_host_callback.pyi b/cuda_core/cuda/core/_host_callback.pyi index 3bbf6321be6..6fd4c52030f 100644 --- a/cuda_core/cuda/core/_host_callback.pyi +++ b/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. + """ diff --git a/cuda_core/cuda/core/_host_callback.pyx b/cuda_core/cuda/core/_host_callback.pyx index de974ab257f..82e1e59fafa 100644 --- a/cuda_core/cuda/core/_host_callback.pyx +++ b/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 == _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(args.dst, args.src, args.nbytes) + + +cdef object _host_copy_fn_pointer(): + return ct.CFUNCTYPE(None, ct.c_void_p)(_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 obj + if isinstance(obj, Buffer): + buf = 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 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 > view.len: + raise ValueError(f"nbytes ({nbytes}) exceeds the size of {name} ({view.len})") + return 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 (&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] = ct.cast(fn, ct.c_void_p).value + is_host_copy = out_fn[0] == _host_copy if user_data is None: out_user_data[0] = NULL elif isinstance(user_data, int): out_user_data[0] = 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( diff --git a/cuda_core/cuda/core/_launcher.pyi b/cuda_core/cuda/core/_launcher.pyi index 56f189fe72f..80b43563e4f 100644 --- a/cuda_core/cuda/core/_launcher.pyi +++ b/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. diff --git a/cuda_core/cuda/core/_launcher.pyx b/cuda_core/cuda/core/_launcher.pyx index f5becd4b316..86f028367ed 100644 --- a/cuda_core/cuda/core/_launcher.pyx +++ b/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. diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyi b/cuda_core/cuda/core/graph/_graph_builder.pyi index bf20e0cb3bf..10ae6032228 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyi +++ b/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: diff --git a/cuda_core/cuda/core/graph/_graph_builder.pyx b/cuda_core/cuda/core/graph/_graph_builder.pyx index 2ac92caacbc..bb19db6e688 100644 --- a/cuda_core/cuda/core/graph/_graph_builder.pyx +++ b/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) diff --git a/cuda_core/cuda/core/graph/_graph_node.pyi b/cuda_core/cuda/core/graph/_graph_node.pyi index b22f3d64dcd..82b838663de 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyi +++ b/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. diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 4abb1f7106c..562733d046c 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/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) diff --git a/cuda_core/cuda/core/graph/_subclasses.pyi b/cuda_core/cuda/core/graph/_subclasses.pyi index 97330bdecfb..ef1599a882c 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyi +++ b/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:: diff --git a/cuda_core/cuda/core/graph/_subclasses.pyx b/cuda_core/cuda/core/graph/_subclasses.pyx index 747c9c9b452..ac92f037df5 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyx +++ b/cuda_core/cuda/core/graph/_subclasses.pyx @@ -1220,7 +1220,9 @@ cdef 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:: @@ -1688,7 +1690,9 @@ cdef 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:: diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst index a1be74a03df..af6dc186043 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -22,6 +22,7 @@ Devices and execution Host launch host_launch + pack_host_copy_args :template: autosummary/cyclass.rst @@ -52,6 +53,21 @@ Devices and execution default stream, and work launched on it can execute concurrently with work on other non-blocking streams. +.. data:: HOST_COPY_FN + + A ctypes function pointer with the ``CUhostFn`` prototype that copies host + memory. Pass it as ``fn`` to :func:`host_launch`, + :meth:`graph.GraphBuilder.callback`, or the other host callback methods, + with ``user_data`` from :func:`pack_host_copy_args`. It is compiled code, so + it runs without the GIL. It copies ``nbytes`` bytes from ``src`` to ``dst`` + with ``memcpy``. + + The driver may run all host functions of a context on one thread, in order, + whatever their stream. A long copy then delays every other host function in + that context, including Python callbacks. Split a large copy into chunks + that each take no longer than the other host functions can wait. Chunks of + a few MiB keep each copy well under a millisecond on current hosts. + Memory management ----------------- diff --git a/cuda_core/docs/source/examples.rst b/cuda_core/docs/source/examples.rst index e69ab5169bf..7a890944133 100644 --- a/cuda_core/docs/source/examples.rst +++ b/cuda_core/docs/source/examples.rst @@ -38,6 +38,9 @@ Kernels and launch configuration - :cuda-core-example:`host_launch.py ` orders a host function between two kernel launches on a stream and reads a pinned buffer from it. +- :cuda-core-example:`host_copy.py ` + copies a device buffer into a NumPy array in chunks through two pinned + staging slots, with the host copies done by ``HOST_COPY_FN``. 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 a9c24a420ed..503826ad9a6 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -53,6 +53,13 @@ New features ``user_data`` forms. (`#2058 `__) +- Added :data:`HOST_COPY_FN`, a host function that copies host memory without + the GIL, and :func:`pack_host_copy_args`, which packs its ``user_data``. Pass + both to :func:`host_launch` or a graph host callback method for a host copy + in stream order, for example from a pinned staging buffer into a NumPy array, + without a compiled extension of your own. + (`#3256 `__) + - 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_copy.py b/cuda_core/examples/host_copy.py new file mode 100644 index 00000000000..6e2cae374f1 --- /dev/null +++ b/cuda_core/examples/host_copy.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# ################################################################################ +# +# This example copies a device buffer into a NumPy array in chunks through two +# pinned staging slots. Each chunk is a device-to-pinned copy followed by +# HOST_COPY_FN, which copies the slot into the array without the GIL. Each slot +# has its own stream, so the copy into one slot overlaps the host copy out of +# the other, and stream order keeps a slot from being reused too early. +# +# ################################################################################ + +# /// script +# dependencies = ["cuda_bindings", "cuda_core", "numpy"] +# /// + +from contextlib import closing + +import numpy as np + +from cuda.core import HOST_COPY_FN, Buffer, Device, LegacyPinnedMemoryResource, host_launch, pack_host_copy_args + + +def main(): + dev = Device(0) + dev.set_current() + chunk = 4 << 20 + total = 16 * chunk + result = np.empty(total, dtype=np.uint8) + pinned_mr = LegacyPinnedMemoryResource() + + with ( + closing(dev.create_stream()) as stream0, + closing(dev.create_stream()) as stream1, + dev.allocate(total, stream=stream0) as device_buffer, + pinned_mr.allocate(chunk) as slot0, + pinned_mr.allocate(chunk) as slot1, + ): + device_buffer.fill(0x5A, stream=stream0) + stream1.wait(stream0) + + streams = (stream0, stream1) + slots = (slot0, slot1) + for i, offset in enumerate(range(0, total, chunk)): + stream, slot = streams[i % 2], slots[i % 2] + src = Buffer.from_handle(int(device_buffer.handle) + offset, chunk, owner=device_buffer) + src.copy_to(slot, stream=stream) + dst = memoryview(result)[offset : offset + chunk] + host_launch(stream, HOST_COPY_FN, user_data=pack_host_copy_args(dst, slot, chunk)) + + stream0.sync() + stream1.sync() + + assert (result == 0x5A).all() + print(f"Copied {total >> 20} MiB into a NumPy array through two pinned slots.") + + +if __name__ == "__main__": + main() diff --git a/cuda_core/tests/test_host_copy.py b/cuda_core/tests/test_host_copy.py new file mode 100644 index 00000000000..4abd3cf25f4 --- /dev/null +++ b/cuda_core/tests/test_host_copy.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +import ctypes +import struct +from contextlib import closing + +import numpy as np +import pytest + +from cuda.core import HOST_COPY_FN, Buffer, Device, LegacyPinnedMemoryResource, host_launch, pack_host_copy_args + +N = 4096 + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_copy_fn_has_cuhostfn_prototype(): + assert HOST_COPY_FN.restype is None + assert tuple(HOST_COPY_FN.argtypes) == (ctypes.c_void_p,) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_copy_orders_between_prior_and_later_work(init_cuda): + dev = Device(0) + dev.set_current() + dst = np.zeros(N, dtype=np.uint8) + with ( + closing(dev.create_stream()) as stream, + dev.allocate(N, stream=stream) as device_buffer, + LegacyPinnedMemoryResource().allocate(N) as pinned, + ): + device_buffer.fill(0x11, stream=stream) + device_buffer.copy_to(pinned, stream=stream) + host_launch(stream, HOST_COPY_FN, user_data=pack_host_copy_args(dst, pinned, N)) + device_buffer.fill(0x22, stream=stream) + device_buffer.copy_to(pinned, stream=stream) + stream.sync() + assert (dst == 0x11).all() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_copy_int_user_data_points_to_a_packed_record(init_cuda): + dev = Device(0) + dev.set_current() + src = np.arange(N, dtype=np.uint8) + dst = np.zeros(N, dtype=np.uint8) + record = ctypes.create_string_buffer(struct.pack("@PPN", dst.ctypes.data, src.ctypes.data, N), 24) + with closing(dev.create_stream()) as stream: + host_launch(stream, HOST_COPY_FN, user_data=ctypes.addressof(record)) + stream.sync() + assert (dst == src).all() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_copy_replays_in_graph(init_cuda): + dev = Device(0) + dev.set_current() + src = np.full(N, 0x5A, dtype=np.uint8) + dst = np.zeros(N, dtype=np.uint8) + with closing(dev.create_stream()) as stream: + gb = stream.create_graph_builder().begin_building() + host_launch(gb, HOST_COPY_FN, user_data=pack_host_copy_args(dst, src, N)) + gb.end_building() + with closing(gb.complete()) as graph: + graph.launch(stream) + stream.sync() + assert (dst == 0x5A).all() + src[:] = 0x6B + graph.launch(stream) + stream.sync() + assert (dst == 0x6B).all() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_copy_zero_bytes_does_nothing(init_cuda): + dev = Device(0) + dev.set_current() + with closing(dev.create_stream()) as stream: + host_launch(stream, HOST_COPY_FN, user_data=pack_host_copy_args(0, 0, 0)) + stream.sync() + + +@pytest.mark.parametrize("user_data", [None, 0, b"", b"\x00" * 23]) +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_host_launch_rejects_host_copy_without_a_record(init_cuda, user_data): + dev = Device(0) + dev.set_current() + with closing(dev.create_stream()) as stream, pytest.raises(ValueError, match="HOST_COPY_FN"): + host_launch(stream, HOST_COPY_FN, user_data=user_data) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_layout(): + data = pack_host_copy_args(0x1000, 0x2000, 16) + assert data == struct.pack("@PPN", 0x1000, 0x2000, 16) + assert struct.unpack("@PPN", data) == (0x1000, 0x2000, 16) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_takes_addresses_from_objects(init_cuda): + src = np.arange(N, dtype=np.uint8) + dst = bytearray(N) + with LegacyPinnedMemoryResource().allocate(N) as pinned: + assert struct.unpack("@PPN", pack_host_copy_args(pinned, src, N)) == (int(pinned.handle), src.ctypes.data, N) + dst_address = ctypes.addressof((ctypes.c_char * N).from_buffer(dst)) + assert struct.unpack("@PPN", pack_host_copy_args(dst, pinned, N)) == (dst_address, int(pinned.handle), N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_accepts_buffer_without_memory_resource(): + src = np.arange(N, dtype=np.uint8) + dst = np.zeros(N, dtype=np.uint8) + buffer = Buffer.from_handle(dst.ctypes.data, N, owner=dst) + assert struct.unpack("@PPN", pack_host_copy_args(buffer, src, N)) == (dst.ctypes.data, src.ctypes.data, N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_checks_bounds(init_cuda): + small = np.zeros(N - 1, dtype=np.uint8) + large = np.zeros(N, dtype=np.uint8) + with pytest.raises(ValueError, match="exceeds the size of dst"): + pack_host_copy_args(small, large, N) + with pytest.raises(ValueError, match="exceeds the size of src"): + pack_host_copy_args(large, small, N) + with LegacyPinnedMemoryResource().allocate(N - 1) as pinned, pytest.raises(ValueError, match="exceeds the size"): + pack_host_copy_args(large, pinned, N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_device_buffer(init_cuda): + dev = Device(0) + dev.set_current() + host = np.zeros(N, dtype=np.uint8) + with closing(dev.create_stream()) as stream, dev.allocate(N, stream=stream) as device_buffer: + with pytest.raises(ValueError, match="not host-accessible"): + pack_host_copy_args(device_buffer, host, N) + with pytest.raises(ValueError, match="not host-accessible"): + pack_host_copy_args(host, device_buffer, N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_closed_buffer(init_cuda): + pinned = LegacyPinnedMemoryResource().allocate(N) + pinned.close() + with pytest.raises(RuntimeError, match="closed"): + pack_host_copy_args(np.zeros(N, dtype=np.uint8), pinned, N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_read_only_dst(): + with pytest.raises(BufferError): + pack_host_copy_args(bytes(N), bytearray(N), N) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_unsupported_type(): + with pytest.raises(TypeError, match="dst must be"): + pack_host_copy_args(1.0, bytearray(N), N) + + +@pytest.mark.parametrize("address", [True, np.uintp(0x1000), ctypes.c_void_p(0x1000)]) +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_scalar_address(address): + with pytest.raises(TypeError, match=r"int\(addr\)"): + pack_host_copy_args(bytearray(8), address, 1) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_overlap(): + data = bytearray(N) + view = memoryview(data) + with pytest.raises(ValueError, match="overlap"): + pack_host_copy_args(view[1:], view, N - 1) + with pytest.raises(ValueError, match="overlap"): + pack_host_copy_args(view, view[1:], N - 1) + pack_host_copy_args(view[N // 2 :], view, N // 2) + pack_host_copy_args(view, view[N // 2 :], N // 2) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_null_address(): + with pytest.raises(ValueError, match="null"): + pack_host_copy_args(0, bytearray(N), N) + + +@pytest.mark.parametrize( + ("dst", "nbytes"), + [(0x1000, -1), (0x1000, 1 << 64), (-1, 1), (1 << 64, 1)], +) +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_pack_host_copy_args_rejects_out_of_range_int(dst, nbytes): + with pytest.raises(OverflowError): + pack_host_copy_args(dst, 0x2000, nbytes)