@@ -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+
246268cdef 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(¶ms , 0, sizeof(params ))
15671592 params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMCPY
15681593 _init_memcpy_params(
15691594 c_dst , c_src , size , ¶ms.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
0 commit comments