Skip to content

Commit 4c893de

Browse files
Andy-Jostclaude
andauthored
cuda.core: accept stream-captured memcpy nodes in MemcpyNode.update (#3029)
* cuda.core: accept stream-captured memcpy nodes in MemcpyNode.update The driver records both operands of a captured cuMemcpyAsync with the unified memory type. The descriptor check in MemcpyNode.update admitted only host and device operands, so every memcpy node produced by stream capture was rejected with a message about multidimensional copies. Accept the unified type for one-dimensional descriptors and read such operands from the device-pointer field. A replaced operand keeps the unified type, in the definition node and in the executable view, because the driver rejects a change of an operand's memory type in an executable update; Graph.update therefore applies the change as well. The node repr shows U for unified operands. Closes #2649 Co-Authored-By: Claude Fable 5.1 <[email protected]> * cuda.core: trim the memcpy node update docstrings The update() docstrings on MemcpyNode and ExecutableMemcpyNode now say only what a caller needs: nodes recorded by stream capture are supported. The driver's memory-type bookkeeping stays in code comments. Co-Authored-By: Claude Fable 5.1 <[email protected]> --------- Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent 7c8ce82 commit 4c893de

4 files changed

Lines changed: 206 additions & 20 deletions

File tree

‎cuda_core/cuda/core/graph/_subclasses.pyi‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ class MemcpyNode(GraphNode):
187187
Omitted parameters preserve their current values. ``dst_owner`` and
188188
``src_owner`` may only accompany their corresponding raw addresses.
189189
Multidimensional, pitched, offset, and array-backed memcpy nodes are
190-
not supported.
190+
not supported. Nodes recorded by stream capture are supported.
191191
192192
With drivers from CUDA 12.2 through 13.1, the node's intended CUDA
193193
context must be current when this method runs. With the CUDA 13 build
@@ -397,7 +397,10 @@ class ExecutableMemsetNode(ExecutableGraphNode):
397397
class ExecutableMemcpyNode(ExecutableGraphNode):
398398
"""An executable memcpy-node view."""
399399
def update(self, *, dst: Buffer | int, src: Buffer | int, size: int) -> None:
400-
"""Replace all one-dimensional memcpy parameters for future launches."""
400+
"""Replace all one-dimensional memcpy parameters for future launches.
401+
402+
Nodes recorded by stream capture are supported.
403+
"""
401404
@property
402405
def is_enabled(self) -> bool:
403406
"""Whether this node is enabled in the executable graph."""

‎cuda_core/cuda/core/graph/_subclasses.pyx‎

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -243,13 +243,33 @@ cdef void _reject_unsupported_kernel_node(
243243
"updating clustered or cooperative kernel nodes is not supported")
244244

245245

246+
cdef inline bint _is_linear_memcpy_memory_type(
247+
cydriver.CUmemorytype memory_type) noexcept nogil:
248+
# HOST and DEVICE name the location of a linear operand. UNIFIED is what
249+
# the driver records for a copy captured from cuMemcpyAsync: a generic
250+
# address that it resolves itself, stored in the device-pointer field.
251+
return (memory_type == cydriver.CU_MEMORYTYPE_HOST
252+
or memory_type == cydriver.CU_MEMORYTYPE_DEVICE
253+
or memory_type == cydriver.CU_MEMORYTYPE_UNIFIED)
254+
255+
256+
cdef str _memcpy_memory_type_tag(cydriver.CUmemorytype memory_type):
257+
if memory_type == cydriver.CU_MEMORYTYPE_HOST:
258+
return "H"
259+
if memory_type == cydriver.CU_MEMORYTYPE_DEVICE:
260+
return "D"
261+
if memory_type == cydriver.CU_MEMORYTYPE_UNIFIED:
262+
return "U"
263+
if memory_type == cydriver.CU_MEMORYTYPE_ARRAY:
264+
return "A"
265+
return str(int(memory_type))
266+
267+
246268
cdef bint _is_supported_memcpy_descriptor(
247269
cydriver.CUDA_MEMCPY3D* params) noexcept nogil:
248270
return (
249-
(params.srcMemoryType == cydriver.CU_MEMORYTYPE_HOST or
250-
params.srcMemoryType == cydriver.CU_MEMORYTYPE_DEVICE)
251-
and (params.dstMemoryType == cydriver.CU_MEMORYTYPE_HOST or
252-
params.dstMemoryType == cydriver.CU_MEMORYTYPE_DEVICE)
271+
_is_linear_memcpy_memory_type(params.srcMemoryType)
272+
and _is_linear_memcpy_memory_type(params.dstMemoryType)
253273
and params.srcXInBytes == 0
254274
and params.srcY == 0
255275
and params.srcZ == 0
@@ -819,8 +839,8 @@ cdef class MemcpyNode(GraphNode):
819839
params.dstMemoryType, params.srcMemoryType)
820840

821841
def __repr__(self) -> str:
822-
cdef str dt = "H" if self._dst_type == cydriver.CU_MEMORYTYPE_HOST else "D"
823-
cdef str st = "H" if self._src_type == cydriver.CU_MEMORYTYPE_HOST else "D"
842+
cdef str dt = _memcpy_memory_type_tag(self._dst_type)
843+
cdef str st = _memcpy_memory_type_tag(self._src_type)
824844
return (f"<MemcpyNode handle=0x{as_intptr(self._h_node):x}"
825845
f" dst=0x{self._dst:x}({dt}) src=0x{self._src:x}({st}) size={self._size}>")
826846

@@ -838,7 +858,7 @@ cdef class MemcpyNode(GraphNode):
838858
Omitted parameters preserve their current values. ``dst_owner`` and
839859
``src_owner`` may only accompany their corresponding raw addresses.
840860
Multidimensional, pitched, offset, and array-backed memcpy nodes are
841-
not supported.
861+
not supported. Nodes recorded by stream capture are supported.
842862

843863
With drivers from CUDA 12.2 through 13.1, the node's intended CUDA
844864
context must be current when this method runs. With the CUDA 13 build
@@ -894,32 +914,31 @@ cdef class MemcpyNode(GraphNode):
894914
"updating multidimensional, pitched, offset, or array-backed "
895915
"memcpy nodes is not supported")
896916

917+
# The descriptor check above admits HOST, DEVICE, and UNIFIED operands;
918+
# the last two keep their address in the device-pointer field.
897919
c_dst_type = params.memcpy.copyParams.dstMemoryType
898920
c_src_type = params.memcpy.copyParams.srcMemoryType
899921
if c_dst_type == cydriver.CU_MEMORYTYPE_HOST:
900922
c_dst = <cydriver.CUdeviceptr><uintptr_t>(
901923
params.memcpy.copyParams.dstHost)
902-
elif c_dst_type == cydriver.CU_MEMORYTYPE_DEVICE:
903-
c_dst = params.memcpy.copyParams.dstDevice
904924
else:
905-
raise NotImplementedError(
906-
f"unsupported destination memory type: {int(c_dst_type)}")
925+
c_dst = params.memcpy.copyParams.dstDevice
907926
if c_src_type == cydriver.CU_MEMORYTYPE_HOST:
908927
c_src = <cydriver.CUdeviceptr><uintptr_t>(
909928
params.memcpy.copyParams.srcHost)
910-
elif c_src_type == cydriver.CU_MEMORYTYPE_DEVICE:
911-
c_src = params.memcpy.copyParams.srcDevice
912929
else:
913-
raise NotImplementedError(
914-
f"unsupported source memory type: {int(c_src_type)}")
930+
c_src = params.memcpy.copyParams.srcDevice
915931

916932
HANDLE_RETURN(graph_get_attachment(
917933
h_graph, node,
918934
&dst_attachment_owner, &src_attachment_owner))
935+
# A unified operand stays unified: any address is valid for that type,
936+
# and an executable update rejects a change of memory type.
919937
if dst is not None:
920938
dst_attachment_owner = _resolve_memcpy_operand(
921939
dst, dst_owner, "dst", &c_dst)
922-
c_dst_type = _get_memcpy_memory_type(c_dst)
940+
if c_dst_type != cydriver.CU_MEMORYTYPE_UNIFIED:
941+
c_dst_type = _get_memcpy_memory_type(c_dst)
923942
params.memcpy.copyParams.dstMemoryType = c_dst_type
924943
params.memcpy.copyParams.dstHost = NULL
925944
params.memcpy.copyParams.dstDevice = 0
@@ -932,7 +951,8 @@ cdef class MemcpyNode(GraphNode):
932951
if src is not None:
933952
src_attachment_owner = _resolve_memcpy_operand(
934953
src, src_owner, "src", &c_src)
935-
c_src_type = _get_memcpy_memory_type(c_src)
954+
if c_src_type != cydriver.CU_MEMORYTYPE_UNIFIED:
955+
c_src_type = _get_memcpy_memory_type(c_src)
936956
params.memcpy.copyParams.srcMemoryType = c_src_type
937957
params.memcpy.copyParams.srcHost = NULL
938958
params.memcpy.copyParams.srcDevice = 0
@@ -1551,7 +1571,10 @@ cdef class ExecutableMemcpyNode(ExecutableGraphNode):
15511571
src: Buffer | int,
15521572
size_t size,
15531573
) -> None:
1554-
"""Replace all one-dimensional memcpy parameters for future launches."""
1574+
"""Replace all one-dimensional memcpy parameters for future launches.
1575+
1576+
Nodes recorded by stream capture are supported.
1577+
"""
15551578
cdef cydriver.CUdeviceptr c_dst
15561579
cdef cydriver.CUdeviceptr c_src
15571580
cdef OpaqueHandle dst_owner = _resolve_memcpy_operand(
@@ -1562,12 +1585,29 @@ cdef class ExecutableMemcpyNode(ExecutableGraphNode):
15621585
cdef cydriver.CUmemorytype src_type
15631586
cdef cydriver.CUcontext ctx = NULL
15641587
cdef cydriver.CUgraphNodeParams params
1588+
cdef cydriver.CUDA_MEMCPY3D recorded
1589+
cdef cydriver.CUgraphNode node = as_cu(self._h_node)
15651590

15661591
c_memset(&params, 0, sizeof(params))
15671592
params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMCPY
15681593
_init_memcpy_params(
15691594
c_dst, c_src, size, &params.memcpy.copyParams,
15701595
&dst_type, &src_type)
1596+
# A destroyed node is reported by _set_executable_node_params below.
1597+
if node != NULL:
1598+
with nogil:
1599+
HANDLE_RETURN(cydriver.cuGraphMemcpyNodeGetParams(
1600+
node, &recorded))
1601+
if recorded.dstMemoryType == cydriver.CU_MEMORYTYPE_UNIFIED:
1602+
params.memcpy.copyParams.dstMemoryType = (
1603+
cydriver.CU_MEMORYTYPE_UNIFIED)
1604+
params.memcpy.copyParams.dstHost = NULL
1605+
params.memcpy.copyParams.dstDevice = c_dst
1606+
if recorded.srcMemoryType == cydriver.CU_MEMORYTYPE_UNIFIED:
1607+
params.memcpy.copyParams.srcMemoryType = (
1608+
cydriver.CU_MEMORYTYPE_UNIFIED)
1609+
params.memcpy.copyParams.srcHost = NULL
1610+
params.memcpy.copyParams.srcDevice = c_src
15711611
with nogil:
15721612
HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
15731613
params.memcpy.copyCtx = ctx

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,3 +257,13 @@ Fixes and enhancements
257257
- The ``StridedMemoryView`` DLPack C exchange API now preserves the underlying
258258
Python exception when an export or import fails, as required by the DLPack
259259
contract, instead of returning ``-1`` with no exception set.
260+
261+
- :meth:`~graph.MemcpyNode.update` and the executable view of a memcpy node
262+
now accept nodes recorded by stream capture. The driver stores the operands
263+
of a captured :meth:`Buffer.copy_from` or :meth:`Buffer.copy_to` as unified
264+
addresses, and ``update`` rejected that descriptor with a message about
265+
multidimensional copies. Such an operand stays unified when it is replaced,
266+
because the driver does not accept a change of an operand's memory type in
267+
an instantiated graph, so :meth:`~graph.Graph.update` applies the change as
268+
well. The node's ``repr`` now shows ``U`` for unified operands.
269+
(`#2649 <https://github.com/NVIDIA/cuda-python/issues/2649>`__)

‎cuda_core/tests/graph/test_graph_node_update.py‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -745,6 +745,139 @@ def test_memcpy_update_between_host_and_device(init_cuda, device_operand):
745745
assert list(host_dst_bytes) == [0x5A] * 4
746746

747747

748+
def _capture_device_memcpy(device, stream, size=64):
749+
"""Capture ``dst.copy_from(src)`` between two device buffers.
750+
751+
The driver records both operands of a captured ``cuMemcpyAsync`` as
752+
unified addresses rather than as device memory (#2649). The builder is
753+
returned so that the captured graph stays alive with the definition view.
754+
"""
755+
src = device.memory_resource.allocate(size, stream=stream)
756+
dst = device.memory_resource.allocate(size, stream=stream)
757+
src.fill(0x5A, stream=stream)
758+
dst.fill(0, stream=stream)
759+
stream.sync()
760+
761+
builder = device.create_graph_builder().begin_building()
762+
dst.copy_from(src, stream=builder)
763+
builder.end_building()
764+
graph_def = builder.graph_definition
765+
node = next(n for n in graph_def.nodes() if isinstance(n, MemcpyNode))
766+
return builder, graph_def, node, dst, src
767+
768+
769+
def _read_device_bytes(buffer, stream):
770+
host = LegacyPinnedMemoryResource().allocate(buffer.size)
771+
buffer.copy_to(host, stream=stream)
772+
stream.sync()
773+
return list((ctypes.c_uint8 * buffer.size).from_address(int(host.handle)))
774+
775+
776+
_MEMORY_TYPE_TAGS = {
777+
driver.CUmemorytype.CU_MEMORYTYPE_HOST: "H",
778+
driver.CUmemorytype.CU_MEMORYTYPE_DEVICE: "D",
779+
driver.CUmemorytype.CU_MEMORYTYPE_UNIFIED: "U",
780+
}
781+
782+
783+
@pytest.mark.agent_authored(model="claude-fable-5-1")
784+
def test_memcpy_update_captured_node_size(init_cuda):
785+
"""A memcpy node recorded by stream capture accepts a size update (#2649)."""
786+
if driver_version() < (12, 2, 0):
787+
pytest.skip("individual graph node updates require CUDA 12.2+")
788+
789+
stream = init_cuda.create_stream()
790+
builder, graph_def, node, dst, src = _capture_device_memcpy(init_cuda, stream)
791+
recorded = handle_return(driver.cuGraphMemcpyNodeGetParams(node.handle))
792+
assert node.size == 64
793+
assert node.dst == int(dst.handle)
794+
assert node.src == int(src.handle)
795+
assert f"({_MEMORY_TYPE_TAGS[recorded.dstMemoryType]})" in repr(node)
796+
assert f"({_MEMORY_TYPE_TAGS[recorded.srcMemoryType]})" in repr(node)
797+
798+
node.update(size=32)
799+
800+
assert node.size == 32
801+
assert node.dst == int(dst.handle)
802+
assert node.src == int(src.handle)
803+
updated = handle_return(driver.cuGraphMemcpyNodeGetParams(node.handle))
804+
assert updated.WidthInBytes == 32
805+
assert updated.dstMemoryType == recorded.dstMemoryType
806+
assert updated.srcMemoryType == recorded.srcMemoryType
807+
808+
graph = graph_def.instantiate()
809+
graph.launch(stream)
810+
assert _read_device_bytes(dst, stream) == [0x5A] * 32 + [0] * 32
811+
812+
813+
@pytest.mark.parametrize("operand", ["src", "dst"])
814+
@pytest.mark.agent_authored(model="claude-fable-5-1")
815+
def test_memcpy_update_captured_node_operand(init_cuda, operand):
816+
"""Replacing one operand of a captured memcpy node keeps the other as recorded (#2649)."""
817+
if driver_version() < (12, 2, 0):
818+
pytest.skip("individual graph node updates require CUDA 12.2+")
819+
820+
stream = init_cuda.create_stream()
821+
builder, graph_def, node, dst, src = _capture_device_memcpy(init_cuda, stream)
822+
recorded = handle_return(driver.cuGraphMemcpyNodeGetParams(node.handle))
823+
instantiated_before = graph_def.instantiate()
824+
replacement = init_cuda.memory_resource.allocate(64, stream=stream)
825+
replacement.fill(0xA5 if operand == "src" else 0, stream=stream)
826+
stream.sync()
827+
828+
if operand == "src":
829+
node.update(src=replacement)
830+
assert node.src == int(replacement.handle)
831+
assert node.dst == int(dst.handle)
832+
copied_into, expected = dst, [0xA5] * 64
833+
else:
834+
node.update(dst=replacement)
835+
assert node.dst == int(replacement.handle)
836+
assert node.src == int(src.handle)
837+
copied_into, expected = replacement, [0x5A] * 64
838+
assert node.size == 64
839+
840+
# The driver rejects a change of memory type in an executable update, so
841+
# the replaced operand keeps the recorded type.
842+
updated = handle_return(driver.cuGraphMemcpyNodeGetParams(node.handle))
843+
assert updated.dstMemoryType == recorded.dstMemoryType
844+
assert updated.srcMemoryType == recorded.srcMemoryType
845+
846+
fresh = graph_def.instantiate()
847+
fresh.launch(stream)
848+
assert _read_device_bytes(copied_into, stream) == expected
849+
850+
copied_into.fill(0, stream=stream)
851+
stream.sync()
852+
instantiated_before.update(graph_def)
853+
instantiated_before.launch(stream)
854+
assert _read_device_bytes(copied_into, stream) == expected
855+
if operand == "dst":
856+
assert _read_device_bytes(dst, stream) == [0] * 64
857+
858+
859+
@pytest.mark.agent_authored(model="claude-fable-5-1")
860+
def test_executable_memcpy_update_on_captured_node(init_cuda):
861+
"""The executable view of a captured memcpy node accepts new operands (#2649)."""
862+
if driver_version() < (12, 2, 0):
863+
pytest.skip("individual graph node updates require CUDA 12.2+")
864+
865+
stream = init_cuda.create_stream()
866+
builder, graph_def, node, dst, src = _capture_device_memcpy(init_cuda, stream)
867+
new_src = init_cuda.memory_resource.allocate(64, stream=stream)
868+
new_dst = init_cuda.memory_resource.allocate(64, stream=stream)
869+
new_src.fill(0xA5, stream=stream)
870+
new_dst.fill(0, stream=stream)
871+
stream.sync()
872+
873+
graph = graph_def.instantiate()
874+
graph[node].update(dst=new_dst, src=new_src, size=32)
875+
graph.launch(stream)
876+
877+
assert _read_device_bytes(new_dst, stream) == [0xA5] * 32 + [0] * 32
878+
assert _read_device_bytes(dst, stream) == [0] * 64
879+
880+
748881
@pytest.mark.agent_authored(model="gpt-5.6")
749882
def test_definition_node_update_changes_future_instantiations(
750883
definition_update_case,

0 commit comments

Comments
 (0)