cuda.core: accept stream-captured memcpy nodes in MemcpyNode.update

[Andy-Jost]

Summary

MemcpyNode.update() raised NotImplementedError for every memcpy node produced by stream capture. The driver records both operands of a captured cuMemcpyAsync (Buffer.copy_from, Buffer.copy_to) with the unified memory type, and the descriptor check admitted only host and device operands. The error message blamed multidimensional, pitched, or array-backed copies, which did not describe the rejected node.

Changes

• _is_supported_memcpy_descriptor accepts CU_MEMORYTYPE_UNIFIED for one-dimensional, unpitched, unoffset descriptors. The pitch and offset checks are unchanged, so multidimensional and array-backed nodes are still rejected with the existing message.
• MemcpyNode.update() reads a unified operand from the device-pointer field. A replaced operand keeps the unified type and its new address goes into that field. The driver rejects a change of an operand's memory type in an executable update (cuGraphExecMemcpyNodeSetParams and cuGraphExecUpdate both compare the type of each operand), so keeping the type lets Graph.update() apply the change to an instantiated graph. Operands recorded as host or device memory are classified as before.
• ExecutableMemcpyNode.update() reads the recorded descriptor of the node and keeps a unified operand unified for the same reason. Without this, the executable view rejected every node from stream capture with CUDA_ERROR_INVALID_VALUE.
• MemcpyNode.__repr__ shows U for unified operands and A for array operands instead of D for every non-host type.
• Tests: a captured copy accepts a size update; a source or destination replacement keeps the recorded types, and both a fresh instantiation and Graph.update() on an earlier one copy what the updated node describes; the executable view of a captured node accepts new operands and a new size. The existing rejection test for pitched descriptors is unchanged.
• Release note under "Fixes and enhancements".

Related Work

Closes #2649. Multidimensional, pitched, and array-backed descriptors remain out of scope here and are tracked in #2420.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Andy-Jost]

/ok to test fb71159

[github-actions]

Doc Preview CI
Preview removed because the pull request was closed or merged.

[juenglin]

/ok to test 46ca8ac

[juenglin]

Reviewed at 46ca8ac; built it locally and ran tests/graph (531 passed, 23 skipped on an RTX 6000 Ada). The core logic looks right to me. Three small follow-ups, in order of importance:

1. The release note overstates what Graph.update() accepts.

The note says a replaced operand stays unified because the driver rejects a memory-type change, "so Graph.update() applies the change as well." That holds when the replacement is a device buffer. It does not hold when the replacement is a host buffer. I captured a device-to-device copy and replaced src with a pinned host buffer, then with a pageable host pointer (ctypes array):

Step	Pinned host	Pageable host
MemcpyNode.update(src=...), then a fresh instantiate()	works, both types stay U, copy is correct	works, both types stay U, copy is correct
Graph.update() on an earlier instantiation	fails, CU_GRAPH_EXEC_UPDATE_ERROR_PARAMETERS_CHANGED	same
ExecutableMemcpyNode.update(src=...)	fails, CUDA_ERROR_INVALID_VALUE	same

This is not caused by this PR. A control with an ordinary GraphDefinition().memcpy() node (recorded as DEVICE, no capture) fails the same way when src is swapped from a device buffer to a pinned host buffer. That matches the driver docs for executable updates: the operand's device allocation or mapping cannot change, and changing the memory type is unsupported. Suggest narrowing the sentence so it doesn't imply every replacement can be applied to an instantiated graph, for example "replacing a device operand with another device buffer".

2. (Optional) Add a test that documents that boundary.

No test replaces a captured node's operand with a host buffer. A short one would capture a device-to-device copy, call update(src=<pinned host buffer>), check the copy after a fresh instantiate(), and check that Graph.update() and the executable view raise. It would pin down the behaviour in item 1 and show it is a driver restriction. This is new coverage for the feature, not a regression guard, because captured nodes raised NotImplementedError before this PR.

3. Add a comment in ExecutableMemcpyNode.update().

It reads the recorded memory types from the definition node (cuGraphMemcpyNodeGetParams on self._h_node), not from the instantiated graph. That is correct only while MemcpyNode.update() preserves unified types, which this PR guarantees. A one-line comment would record the dependency, so a later change to the definition-side update doesn't silently break the executable view.

Minor nits, no action needed: the three new tests unpack builder but never use it (it only keeps the captured graph alive, so a comment or del would show that), and the driver docs disallow CU_MEMORYTYPE_UNIFIED for managed-memory operands when any device reports CONCURRENT_MANAGED_ACCESS == 0. I did not test that case.