Repository navigation
feat(cuda.core): add HOST_COPY_FN and pack_host_copy_args for GIL-free host copies - #3269
srimanachanta wants to merge 8 commits into
Conversation
Add cuda.core.host_launch(stream, fn, *, user_data=None), the stream-side counterpart of GraphBuilder.callback. It enqueues a Python callable or a ctypes function pointer with the CUhostFn prototype on a stream through cuLaunchHostFunc; a GraphBuilder argument forwards to its callback. The callable and a copied user_data stay alive until the function has run and are released afterwards on the main thread through the deferred-cleanup queue, never on the driver thread. A plain stream under active capture is rejected, because the owners could not be attached to the captured node. The shared callback helper moves from cuda/core/graph/_host_callback to cuda/core/_host_callback, since a core module must not import the graph package at module level. The launch record and trampoline live in _cpp/rt/host_launch.cpp, and cuLaunchHostFunc joins the driver table. Written by Codex (gpt-6.1-sol); reviewed by Claude. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
State that a ctypes callback that wraps a Python function acquires the GIL through ctypes, so the deadlock rule applies to it; complete the Raises section; say the driver may reject CUDA calls from the callback rather than that it does. Upcast the launch record to the cleanup base before the void* parameter, as the graph producers do. Prove the user_data copy by changing the caller's buffer while the stream is held, pin the unraisable type and message, and link the example from the examples page. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
HOST_COPY_FN is a CUhostFn that memcpys without the GIL. pack_host_copy_args packs (dst, src, nbytes) and checks bounds, host accessibility of Buffer memory resources, null addresses, and overlap without calling the driver. The host callback resolver rejects HOST_COPY_FN without a packed record.
Covers the shared per-context host function thread and chunk sizing.
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Runtime host-function launch path cuda_core/cuda/core/_cpp/rt/api.hpp, cuda_core/cuda/core/_cpp/rt/driver_api.hpp, cuda_core/cuda/core/_cpp/rt/host_launch.cpp, cuda_core/cuda/core/_rt.pxd, cuda_core/cuda/core/_rt.pyx, cuda_core/cuda/core/_cpp/rt/DESIGN.md, cuda_core/cuda/core/_cpp/rt/GRAPH_ATTACHMENTS.md |
Adds the driver entry and runtime launch function. The launch record retains callback and data owners, invokes the callback, and queues the record for deferred cleanup. |
Callback resolution and copy arguments cuda_core/cuda/core/_host_callback.pyx, cuda_core/cuda/core/_host_callback.pyi |
Adds HOST_COPY_FN, pack_host_copy_args, callback validation, and callback/data ownership handling. The packer validates addresses and buffers, then returns a native-width argument record. |
Python launch and graph callback integration cuda_core/cuda/core/__init__.py, cuda_core/cuda/core/_launcher.pyx, cuda_core/cuda/core/_launcher.pyi, cuda_core/cuda/core/graph/* |
Exports the host callback symbols and adds host_launch for streams and GraphBuilder inputs. Graph callback modules import the shared resolver from cuda.core._host_callback and document the packed-record requirement for HOST_COPY_FN. |
Examples, documentation, and tests cuda_core/tests/test_host_copy.py, cuda_core/tests/test_host_launch.py, cuda_core/examples/host_copy.py, cuda_core/examples/host_launch.py, cuda_core/docs/source/api.rst, cuda_core/docs/source/examples.rst, cuda_core/docs/source/release/1.3.0-notes.rst |
Adds tests for stream ordering, graph replay, callback arguments, and validation. Adds examples for stream callbacks and chunked host copies, plus API, example, and release documentation. |
Assessment against linked issues
| Objective | Addressed | Explanation |
|---|---|---|
Provide a GIL-free CUhostFn host copy callback and a helper to pack (dst, src, nbytes) for use with host launch and graph callbacks [#3256] |
✅ |
Out-of-scope changes
| Code Change | Explanation |
|---|---|
Adds a two-stream pinned staging-slot transfer example (cuda_core/examples/host_copy.py:1-61) |
Issue #3256 explicitly excludes a staged device-to-host copy utility and ring policy. This example demonstrates chunking and slot reuse rather than only the callback primitive. |
Possibly related PRs
- NVIDIA/cuda-python#3258: Adds the stream-side host launch path used by this PR.
- NVIDIA/cuda-python#3034: Replaces the hand-written release page with generated release notes; this PR adds host-function content to that artifact.
- NVIDIA/cuda-python#2925: Covers callback contracts reused by this PR and the relocated graph callback resolver.
Suggested labels: feature
Suggested reviewers: andy-jost
Priority:
Change: Feature
Merge Risk: 
Minimal · up to 27942
-
The change is ready to merge after normal checks.
Comment @coderabbitai help to get the list of available commands.
|
/ok to test b11ec0f |
|
Andy-Jost
left a comment
There was a problem hiding this comment.
Thanks, this implements #3256 as agreed. I reviewed only the three commits on top of #3258. _host_copy is noexcept nogil, reads the documented 24-byte @PPN record, and skips memcpy for zero bytes; _resolve_host_callback rejects None, 0, empty, and wrong-length records before it allocates, and the packed bytes ride the existing user_data ownership on the stream, builder, definition, and update() paths. pack_host_copy_args makes no driver call on any built-in path: it reads _memory_resource.is_host_accessible directly and so avoids the cuPointerGetAttributes that Buffer.is_host_accessible falls back to. The example orders its pinned slots correctly. Since the branch carries #3258, it merges after #3258.
Three changes before merge:
- The Notes of
pack_host_copy_argssay a Python host callback can use it "to enqueue the next copy". Enqueueing is a CUDA call, whichhost_launchforbids inside a host function. Say instead that a host callback can call it, for example to fill a record that a laterHOST_COPY_FNreads through anintuser_data. - A negative or oversized
nbytesorintaddress raisesOverflowError(and a test pins it), but Raises lists onlyTypeError,ValueError, andRuntimeError. Document it, in the.pyxand the.pyi, or convert it toValueError. dstandsrchave no annotations; the public-API rule wantsint | Buffer | BufferProtocol, asBuffer.filluses.
Smaller points:
- Only a Python
intis an address.Truepacks as address 1, and a NumPy integer scalar (np.uintp(addr)) or actypes.c_void_pexports the buffer protocol, so it packs the address of its own 8 bytes rather than the address it holds. Rejectbooland these scalars with aTypeErrorthat says to passint(addr), or say in the docstring that only anintis an address. - "The caller must keep both alive" is not quite enough for a
bytearray: it must also not be resized until the copy has run. HostCallbackNode.updateandExecutableHostCallbackNode.updateraise the sameHOST_COPY_FNrecord error but their docstrings do not say so.- Tests: a
Buffer.from_handle(ptr, n)with no memory resource would pin the no-query path (GPU-free), and the overlap test covers onlydst > src; add thesrc > dstcase and its adjacent boundary.
True packed as address 1, and a NumPy integer scalar or a ctypes scalar exports the buffer protocol, so it packed the address of its own storage instead of the address it holds. Raise TypeError and point to int(addr).
Annotate dst and src as Buffer.fill does. List OverflowError, say a bytearray must not be resized before the copy runs, and replace the Notes claim that a host callback can enqueue work, which host_launch forbids. Document the HOST_COPY_FN record error on the host callback update methods.
Pin the GPU-free path for a Buffer without a memory resource, cover the src > dst overlap and its adjacent boundary, and check OverflowError for out-of-range int addresses as well as nbytes.
Closes #3256. Built on #3258; the first two commits are #3258 rebased onto
main.Summary
A pure-Python package cannot give
host_launcha GIL-free host function: a ctypes callback built from Python takes the GIL, and libcmemcpydoes not matchCUhostFn. This PR adds one incuda.core, as agreed in #3256.cuda.core.HOST_COPY_FN: a ctypesCUhostFnpointer to anoexcept nogilfunction that copiesnbytesfromsrctodst. It works asfninhost_launch, the graph callback methods, andupdate().cuda.core.pack_host_copy_args(dst, src, nbytes) -> bytes: packs the record foruser_data. It accepts ints,Buffer, and buffer-protocol objects, and checks bounds, overlap, null addresses, and the host accessibility of aBuffer's memory resource. It makes no driver calls, so a Python host callback can use it.dstandsrcare not retained.The record layout is public:
uintptr_t dst,uintptr_t src,size_t nbytes, the same asstruct.pack("@PPN", ...). A caller can pre-pack records and pass an address as an intuser_data. The host callback resolver rejectsHOST_COPY_FNwith nouser_data, a zero address, or bytes of the wrong length, so a bad call raises instead of crashing the driver thread.The
HOST_COPY_FNdocs note that the driver may run all host functions of a context on one thread, so a long copy delays other host functions, including Python callbacks. They recommend chunks of a few MiB.Changes
cuda_core/cuda/core/_host_callback.pyx: the copy function,HOST_COPY_FN,pack_host_copy_args, and the resolver check. Exported fromcuda.core.cuda_core/docs/source/api.rst:HOST_COPY_FNandpack_host_copy_args. Thehost_launchand graph callback docstrings point to them. A release-note entry for 1.3.0.cuda_core/examples/host_copy.py: a chunked device-to-NumPy copy through two pinned slots, one stream per slot.cuda_core/tests/test_host_copy.py: stream ordering, intuser_data, graph replay, zero-byte copies, the packed layout, and each validation error.