diff --git a/cuda_core/cuda/core/graph/_graph_node.pxd b/cuda_core/cuda/core/graph/_graph_node.pxd index 0a87b70ad62..3ef7440095e 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pxd +++ b/cuda_core/cuda/core/graph/_graph_node.pxd @@ -2,8 +2,10 @@ # # SPDX-License-Identifier: Apache-2.0 +from libc.stddef cimport size_t + from cuda.bindings cimport cydriver -from cuda.core._resource_handles cimport GraphHandle, GraphNodeHandle +from cuda.core._resource_handles cimport GraphHandle, GraphNodeHandle, OpaqueHandle cdef class GraphNode: @@ -13,3 +15,12 @@ cdef class GraphNode: @staticmethod cdef GraphNode _create(GraphHandle h_graph, cydriver.CUgraphNode node) + + +cdef OpaqueHandle _resolve_memcpy_operand( + object operand, object owner, str side, cydriver.CUdeviceptr* out_ptr) except * + +cdef void _init_memcpy_params( + cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t size, + cydriver.CUDA_MEMCPY3D* params, cydriver.CUmemorytype* dst_type, + cydriver.CUmemorytype* src_type) except * diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 9e1cab09e7b..0aa95814af1 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -889,7 +889,8 @@ cdef inline OpaqueHandle _buffer_attachment_owner(Buffer buf, str label): cdef inline OpaqueHandle _resolve_memcpy_operand( - object operand, object owner, str side, cydriver.CUdeviceptr* out_ptr): + object operand, object owner, str side, + cydriver.CUdeviceptr* out_ptr) except *: """Resolve an operand to a pointer and optional attachment owner. ``operand`` is a :class:`Buffer` or a raw integer address; its device @@ -969,9 +970,10 @@ cdef inline MemsetNode GN_memset( val, elem_size, width, height, pitch)) -cdef inline MemcpyNode GN_memcpy( - GraphNode self, cydriver.CUdeviceptr c_dst, OpaqueHandle dst_owner, - cydriver.CUdeviceptr c_src, OpaqueHandle src_owner, size_t size): +cdef void _init_memcpy_params( + cydriver.CUdeviceptr dst, cydriver.CUdeviceptr src, size_t size, + cydriver.CUDA_MEMCPY3D* params, cydriver.CUmemorytype* dst_type, + cydriver.CUmemorytype* src_type) except *: cdef unsigned int dst_mem_type = cydriver.CU_MEMORYTYPE_DEVICE cdef unsigned int src_mem_type = cydriver.CU_MEMORYTYPE_DEVICE cdef cydriver.CUresult ret @@ -979,36 +981,43 @@ cdef inline MemcpyNode GN_memcpy( ret = cydriver.cuPointerGetAttribute( &dst_mem_type, cydriver.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, - c_dst) + dst) if ret != cydriver.CUDA_SUCCESS and ret != cydriver.CUDA_ERROR_INVALID_VALUE: HANDLE_RETURN(ret) ret = cydriver.cuPointerGetAttribute( &src_mem_type, cydriver.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, - c_src) + src) if ret != cydriver.CUDA_SUCCESS and ret != cydriver.CUDA_ERROR_INVALID_VALUE: HANDLE_RETURN(ret) - cdef cydriver.CUmemorytype c_dst_type = dst_mem_type - cdef cydriver.CUmemorytype c_src_type = src_mem_type - - cdef cydriver.CUDA_MEMCPY3D params - c_memset(¶ms, 0, sizeof(params)) - - params.srcMemoryType = c_src_type - params.dstMemoryType = c_dst_type - if c_src_type == cydriver.CU_MEMORYTYPE_HOST: - params.srcHost = c_src + dst_type[0] = dst_mem_type + src_type[0] = src_mem_type + c_memset(params, 0, sizeof(params[0])) + params.srcMemoryType = src_type[0] + params.dstMemoryType = dst_type[0] + if src_type[0] == cydriver.CU_MEMORYTYPE_HOST: + params.srcHost = src else: - params.srcDevice = c_src - if c_dst_type == cydriver.CU_MEMORYTYPE_HOST: - params.dstHost = c_dst + params.srcDevice = src + if dst_type[0] == cydriver.CU_MEMORYTYPE_HOST: + params.dstHost = dst else: - params.dstDevice = c_dst + params.dstDevice = dst params.WidthInBytes = size params.Height = 1 params.Depth = 1 + +cdef inline MemcpyNode GN_memcpy( + GraphNode self, cydriver.CUdeviceptr c_dst, OpaqueHandle dst_owner, + cydriver.CUdeviceptr c_src, OpaqueHandle src_owner, size_t size): + cdef cydriver.CUDA_MEMCPY3D params + cdef cydriver.CUmemorytype c_dst_type + cdef cydriver.CUmemorytype c_src_type + _init_memcpy_params( + c_dst, c_src, size, ¶ms, &c_dst_type, &c_src_type) + cdef cydriver.CUgraphNode new_node = NULL cdef GraphHandle h_graph = graph_node_get_graph(self._h_node) cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node) diff --git a/cuda_core/cuda/core/graph/_subclasses.pyi b/cuda_core/cuda/core/graph/_subclasses.pyi index 345e6417c4d..4d144297b0f 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyi +++ b/cuda_core/cuda/core/graph/_subclasses.pyi @@ -5,6 +5,7 @@ from __future__ import annotations from cuda.core._event import Event from cuda.core._launch_config import LaunchConfig +from cuda.core._memory._buffer import Buffer from cuda.core._module import Kernel from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition from cuda.core.graph._graph_node import GraphNode @@ -139,6 +140,20 @@ class MemsetNode(GraphNode): def __repr__(self) -> str: ... + def update(self, dst: Buffer | int | None=None, value=None, width: int | None=None, height: int | None=None, pitch: int | None=None, *, dst_owner=None) -> None: + """Replace selected memset parameters. + + Omitted parameters preserve their current values. ``dst_owner`` may + only accompany a raw-address ``dst``. + + .. warning:: + + Use caution when a retained operand owner directly or indirectly + owns a graph. Any reference cycle involving the owner and a graph + that retains it cannot be broken by Python's cyclic garbage + collector. Use a weak reference to break such cycles. + """ + @property def dptr(self) -> int: """The destination device pointer.""" @@ -179,6 +194,20 @@ class MemcpyNode(GraphNode): def __repr__(self) -> str: ... + def update(self, dst: Buffer | int | None=None, src: Buffer | int | None=None, size: int | None=None, *, dst_owner=None, src_owner=None) -> None: + """Replace selected memcpy parameters. + + Omitted parameters preserve their current values. ``dst_owner`` and + ``src_owner`` may only accompany their corresponding raw addresses. + + .. warning:: + + Use caution when a retained operand owner directly or indirectly + owns a graph. Any reference cycle involving the owner and a graph + that retains it cannot be broken by Python's cyclic garbage + collector. Use a weak reference to break such cycles. + """ + @property def dst(self) -> int: """The destination pointer.""" @@ -219,6 +248,9 @@ class EventRecordNode(GraphNode): def __repr__(self) -> str: ... + def update(self, event: Event) -> None: + """Replace the event recorded by this node.""" + @property def event(self) -> Event: """The event being recorded.""" @@ -235,6 +267,9 @@ class EventWaitNode(GraphNode): def __repr__(self) -> str: ... + def update(self, event: Event) -> None: + """Replace the event waited on by this node.""" + @property def event(self) -> Event: """The event being waited on.""" @@ -251,6 +286,20 @@ class HostCallbackNode(GraphNode): def __repr__(self) -> str: ... + def update(self, fn, *, user_data=None) -> None: + """Replace the callback and user-data binding for this node. + + .. warning:: + + Callbacks must not call CUDA API functions. Doing so may + deadlock or corrupt driver state. + + Use caution when a Python callback retains an object that owns a + graph. Any reference cycle involving the callback and a graph that + retains it cannot be broken by Python's cyclic garbage collector. + Use a weak reference to break such cycles. + """ + @property def callback(self): """The Python callable, or None for ctypes function pointer callbacks.""" diff --git a/cuda_core/cuda/core/graph/_subclasses.pyx b/cuda_core/cuda/core/graph/_subclasses.pyx index 2fa08e2a6a1..db9141dd670 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyx +++ b/cuda_core/cuda/core/graph/_subclasses.pyx @@ -8,29 +8,44 @@ from __future__ import annotations from libc.stddef cimport size_t from libc.stdint cimport uintptr_t +from libc.string cimport memset as c_memset from cuda.bindings cimport cydriver from cuda.core._event cimport Event from cuda.core._launch_config cimport LaunchConfig +from cuda.core._memory._buffer cimport Buffer from cuda.core._module cimport Kernel from cuda.core.graph._graph_definition cimport GraphCondition, GraphDefinition -from cuda.core.graph._graph_node cimport GraphNode +from cuda.core.graph._graph_node cimport ( + GraphNode, + _init_memcpy_params, + _resolve_memcpy_operand, +) from cuda.core._resource_handles cimport ( EventHandle, GraphHandle, - KernelHandle, GraphNodeHandle, + KernelHandle, + OpaqueHandle, + PreparedAttachment, as_cu, as_intptr, - create_event_handle_ref, create_child_graph_handle, + create_event_handle_ref, create_kernel_handle_ref, + graph_commit_attachment, + graph_get_attachment, graph_node_get_graph, + graph_prepare_attachment, ) -from cuda.core._utils.cuda_utils cimport HANDLE_RETURN +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value +from cuda.core._utils.version cimport cy_binding_version, cy_driver_version -from cuda.core.graph._host_callback cimport _is_py_host_trampoline +from cuda.core.graph._host_callback cimport ( + _is_py_host_trampoline, + _resolve_host_callback, +) from cuda.core._utils.cuda_utils import driver, handle_return from cuda.core.typing import GraphConditionalType @@ -57,6 +72,36 @@ __all__ = [ cdef bint _has_cuGraphNodeGetParams = False cdef bint _version_checked = False + +cdef void _set_definition_node_params( + const GraphNodeHandle& h_node, + cydriver.CUgraphNodeParams* params, + OpaqueHandle owner0, + OpaqueHandle owner1=OpaqueHandle()) except *: + cdef tuple version = cy_driver_version() + if version < (12, 2, 0): + raise RuntimeError( + "Graph node mutation requires CUDA driver 12.2 or newer; " + f"using driver version {'.'.join(map(str, version))}" + ) + version = cy_binding_version() + if version < (12, 2, 0): + raise RuntimeError( + "Graph node mutation requires cuda.bindings 12.2 or newer; " + f"using cuda.bindings version {'.'.join(map(str, version))}" + ) + + cdef GraphHandle h_graph = graph_node_get_graph(h_node) + cdef cydriver.CUgraphNode node = as_cu(h_node) + cdef PreparedAttachment prepared + + HANDLE_RETURN(graph_prepare_attachment( + h_graph, owner0, owner1, &prepared)) + with nogil: + HANDLE_RETURN(cydriver.cuGraphNodeSetParams(node, params)) + HANDLE_RETURN(graph_commit_attachment(prepared, node)) + + cdef bint _check_node_get_params(): global _has_cuGraphNodeGetParams, _version_checked if not _version_checked: @@ -340,6 +385,77 @@ cdef class MemsetNode(GraphNode): return (f"") + def update( + self, + dst: Buffer | int | None = None, + value=None, + width: int | None = None, + height: int | None = None, + pitch: int | None = None, + *, + dst_owner=None, + ) -> None: + """Replace selected memset parameters. + + Omitted parameters preserve their current values. ``dst_owner`` may + only accompany a raw-address ``dst``. + + .. warning:: + + Use caution when a retained operand owner directly or indirectly + owns a graph. Any reference cycle involving the owner and a graph + that retains it cannot be broken by Python's cyclic garbage + collector. Use a weak reference to break such cycles. + """ + cdef cydriver.CUdeviceptr c_dst = self._dptr + cdef unsigned int c_value = self._value + cdef unsigned int c_element_size = self._element_size + cdef size_t c_width = self._width + cdef size_t c_height = self._height + cdef size_t c_pitch = self._pitch + cdef OpaqueHandle dst_attachment_owner + cdef GraphHandle h_graph + cdef cydriver.CUgraphNodeParams params + + if dst is None: + if dst_owner is not None: + raise ValueError("dst_owner requires dst") + h_graph = graph_node_get_graph(self._h_node) + HANDLE_RETURN(graph_get_attachment( + h_graph, as_cu(self._h_node), + &dst_attachment_owner, NULL)) + else: + dst_attachment_owner = _resolve_memcpy_operand( + dst, dst_owner, "dst", &c_dst) + + if value is not None: + c_value, c_element_size = _parse_fill_value(value) + if width is not None: + c_width = width + if height is not None: + c_height = height + if pitch is not None: + c_pitch = pitch + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMSET + params.memset.dst = c_dst + params.memset.value = c_value + params.memset.elementSize = c_element_size + params.memset.width = c_width + params.memset.height = c_height + params.memset.pitch = c_pitch + params.memset.ctx = NULL + + _set_definition_node_params( + self._h_node, ¶ms, dst_attachment_owner) + self._dptr = c_dst + self._value = c_value + self._element_size = c_element_size + self._width = c_width + self._height = c_height + self._pitch = c_pitch + @property def dptr(self) -> int: """The destination device pointer.""" @@ -428,6 +544,74 @@ cdef class MemcpyNode(GraphNode): return (f"") + def update( + self, + dst: Buffer | int | None = None, + src: Buffer | int | None = None, + size: int | None = None, + *, + dst_owner=None, + src_owner=None, + ) -> None: + """Replace selected memcpy parameters. + + Omitted parameters preserve their current values. ``dst_owner`` and + ``src_owner`` may only accompany their corresponding raw addresses. + + .. warning:: + + Use caution when a retained operand owner directly or indirectly + owns a graph. Any reference cycle involving the owner and a graph + that retains it cannot be broken by Python's cyclic garbage + collector. Use a weak reference to break such cycles. + """ + cdef cydriver.CUdeviceptr c_dst = self._dst + cdef cydriver.CUdeviceptr c_src = self._src + cdef size_t c_size = self._size + cdef OpaqueHandle dst_attachment_owner + cdef OpaqueHandle src_attachment_owner + cdef GraphHandle h_graph = graph_node_get_graph(self._h_node) + cdef cydriver.CUcontext ctx = NULL + cdef cydriver.CUgraphNodeParams params + cdef cydriver.CUmemorytype c_dst_type + cdef cydriver.CUmemorytype c_src_type + + HANDLE_RETURN(graph_get_attachment( + h_graph, as_cu(self._h_node), + &dst_attachment_owner, &src_attachment_owner)) + if dst is None: + if dst_owner is not None: + raise ValueError("dst_owner requires dst") + else: + dst_attachment_owner = _resolve_memcpy_operand( + dst, dst_owner, "dst", &c_dst) + if src is None: + if src_owner is not None: + raise ValueError("src_owner requires src") + else: + src_attachment_owner = _resolve_memcpy_operand( + src, src_owner, "src", &c_src) + if size is not None: + c_size = size + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMCPY + _init_memcpy_params( + c_dst, c_src, c_size, ¶ms.memcpy.copyParams, + &c_dst_type, &c_src_type) + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx)) + params.memcpy.copyCtx = ctx + + _set_definition_node_params( + self._h_node, ¶ms, + dst_attachment_owner, src_attachment_owner) + self._dst = c_dst + self._src = c_src + self._size = c_size + self._dst_type = c_dst_type + self._src_type = c_src_type + @property def dst(self) -> int: """The destination pointer.""" @@ -516,6 +700,19 @@ cdef class EventRecordNode(GraphNode): return (f"") + def update(self, event: Event) -> None: + """Replace the event recorded by this node.""" + cdef OpaqueHandle event_owner = event._h_event + cdef cydriver.CUgraphNodeParams params + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_EVENT_RECORD + params.eventRecord.event = as_cu(event._h_event) + + _set_definition_node_params( + self._h_node, ¶ms, event_owner) + self._h_event = event._h_event + @property def event(self) -> Event: """The event being recorded.""" @@ -554,6 +751,19 @@ cdef class EventWaitNode(GraphNode): return (f"") + def update(self, event: Event) -> None: + """Replace the event waited on by this node.""" + cdef OpaqueHandle event_owner = event._h_event + cdef cydriver.CUgraphNodeParams params + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_WAIT_EVENT + params.eventWait.event = as_cu(event._h_event) + + _set_definition_node_params( + self._h_node, ¶ms, event_owner) + self._h_event = event._h_event + @property def event(self) -> Event: """The event being waited on.""" @@ -604,6 +814,37 @@ cdef class HostCallbackNode(GraphNode): return (f"self._fn:x}>") + def update(self, fn, *, user_data=None) -> None: + """Replace the callback and user-data binding for this node. + + .. warning:: + + Callbacks must not call CUDA API functions. Doing so may + deadlock or corrupt driver state. + + Use caution when a Python callback retains an object that owns a + graph. Any reference cycle involving the callback and a graph that + retains it cannot be broken by Python's cyclic garbage collector. + Use a weak reference to break such cycles. + """ + cdef cydriver.CUhostFn c_fn + cdef void* c_user_data + cdef OpaqueHandle fn_owner, data_owner + cdef cydriver.CUgraphNodeParams params + + _resolve_host_callback( + fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner) + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_HOST + params.host.fn = c_fn + params.host.userData = c_user_data + + _set_definition_node_params( + self._h_node, ¶ms, fn_owner, data_owner) + self._callable = fn if _is_py_host_trampoline(c_fn) else None + self._fn = c_fn + self._user_data = c_user_data + @property def callback(self): """The Python callable, or None for ctypes function pointer callbacks.""" diff --git a/cuda_core/tests/graph/test_graph_node_update.py b/cuda_core/tests/graph/test_graph_node_update.py new file mode 100644 index 00000000000..3bbebcb95dc --- /dev/null +++ b/cuda_core/tests/graph/test_graph_node_update.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for updating individual graph node parameters.""" + +import ctypes +import threading +from dataclasses import dataclass +from typing import Callable + +import pytest + +from cuda.core import LegacyPinnedMemoryResource +from cuda.core._utils.cuda_utils import CUDAError +from cuda.core._utils.version import driver_version +from cuda.core.graph import GraphDefinition + + +@dataclass +class _DefinitionUpdateCase: + graph_def: GraphDefinition + node: object + original: object + replacement: object + update: Callable[[object], None] + assert_current: Callable[[object], None] + assert_exec_uses: Callable[[object, object], None] + invalid_update: Callable[[], None] | None + invalid_exception: type[BaseException] | None + invalid_argument_update: Callable[[], None] | None + cleanup: Callable[[], None] + + +def _assert_equal(actual, expected): + assert actual == expected + + +def _noop(): + pass + + +def _event_record_case(device): + """Keep the selected event pending to identify each exec's record target.""" + original = device.create_event() + replacement = device.create_event() + invalid_replacement = device.create_event() + invalid_replacement.close() + + callback_started = threading.Event() + callback_release = threading.Event() + + def blocking_callback(): + callback_started.set() + callback_release.wait(timeout=30) + + graph_def = GraphDefinition() + callback_node = graph_def.callback(blocking_callback) + node = callback_node.record(original) + + def assert_exec_uses(graph, expected): + callback_started.clear() + callback_release.clear() + stream = device.create_stream() + graph.launch(stream) + try: + assert callback_started.wait(timeout=5) + assert expected.is_done is False + unexpected = replacement if expected is original else original + assert unexpected.is_done is True + finally: + callback_release.set() + stream.sync() + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=node.update, + assert_current=lambda expected: _assert_equal(node.event, expected), + assert_exec_uses=assert_exec_uses, + invalid_update=lambda: node.update(invalid_replacement), + invalid_exception=CUDAError, + invalid_argument_update=lambda: node.update(object()), + cleanup=_noop, + ) + + +def _event_wait_case(device): + """Keep the selected event pending to identify each exec's wait target.""" + original = device.create_event() + replacement = device.create_event() + invalid_replacement = device.create_event() + invalid_replacement.close() + + callback_called = threading.Event() + graph_def = GraphDefinition() + node = graph_def.wait(original) + node.callback(callback_called.set) + + def assert_exec_uses(graph, expected): + producer_started = threading.Event() + producer_release = threading.Event() + + def blocking_callback(): + producer_started.set() + producer_release.wait(timeout=30) + + producer_def = GraphDefinition() + producer_def.callback(blocking_callback).record(expected) + producer_graph = producer_def.instantiate() + producer_stream = device.create_stream() + consumer_stream = device.create_stream() + + callback_called.clear() + producer_graph.launch(producer_stream) + try: + assert producer_started.wait(timeout=5) + graph.launch(consumer_stream) + assert not callback_called.wait(timeout=0.1) + finally: + producer_release.set() + producer_stream.sync() + consumer_stream.sync() + assert callback_called.is_set() + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=node.update, + assert_current=lambda expected: _assert_equal(node.event, expected), + assert_exec_uses=assert_exec_uses, + invalid_update=lambda: node.update(invalid_replacement), + invalid_exception=CUDAError, + invalid_argument_update=lambda: node.update(object()), + cleanup=_noop, + ) + + +def _host_callback_case(device): + """Use callbacks that report their identity to distinguish each exec.""" + called = [] + + def original(): + called.append(original) + + def replacement(): + called.append(replacement) + + graph_def = GraphDefinition() + node = graph_def.callback(original) + + def assert_exec_uses(graph, expected): + called.clear() + stream = device.create_stream() + graph.launch(stream) + stream.sync() + assert called == [expected] + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=node.update, + assert_current=lambda expected: _assert_equal(node.callback, expected), + assert_exec_uses=assert_exec_uses, + invalid_update=lambda: node.update(replacement, user_data=b"not valid for a Python callback"), + invalid_exception=ValueError, + invalid_argument_update=None, + cleanup=_noop, + ) + + +def _host_callback_ctypes_case(device): + """Use ctypes callbacks and copied payloads to distinguish each exec.""" + callback_type = ctypes.CFUNCTYPE(None, ctypes.c_void_p) + called = [] + + def read_byte(data): + return ctypes.cast(data, ctypes.POINTER(ctypes.c_uint8))[0] + + @callback_type + def original_fn(data): + called.append((original_fn, read_byte(data))) + + @callback_type + def replacement_fn(data): + called.append((replacement_fn, read_byte(data))) + + original = original_fn, bytes([0xA1]) + replacement = replacement_fn, bytes([0xB2]) + graph_def = GraphDefinition() + node = graph_def.callback(original_fn, user_data=original[1]) + + def update(value): + fn, user_data = value + node.update(fn, user_data=user_data) + + def assert_exec_uses(graph, expected): + called.clear() + stream = device.create_stream() + graph.launch(stream) + stream.sync() + assert called == [(expected[0], expected[1][0])] + + def invalid_update(): + node.update(lambda: None, user_data=b"not valid for a Python callback") + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=update, + assert_current=lambda _expected: _assert_equal(node.callback, None), + assert_exec_uses=assert_exec_uses, + invalid_update=invalid_update, + invalid_exception=ValueError, + invalid_argument_update=None, + cleanup=_noop, + ) + + +def _memset_case(device, *, replace_dst): + memory_resource = LegacyPinnedMemoryResource() + original_buffer = memory_resource.allocate(4) + replacement_buffer = memory_resource.allocate(4) if replace_dst else original_buffer + original = { + "dst": original_buffer, + "value": 0x11, + "element_size": 1, + "width": 4, + "height": 1, + "pitch": 0, + } + replacement = { + **original, + "dst": replacement_buffer, + "value": 0x22, + } + + graph_def = GraphDefinition() + node = graph_def.memset(original["dst"], original["value"], original["width"]) + + def update(expected): + if replace_dst: + node.update(dst=expected["dst"], value=expected["value"]) + else: + node.update(value=expected["value"]) + + def assert_current(expected): + assert node.dptr == int(expected["dst"].handle) + assert node.value == expected["value"] + assert node.element_size == expected["element_size"] + assert node.width == expected["width"] + assert node.height == expected["height"] + assert node.pitch == expected["pitch"] + + def as_bytes(buffer): + return (ctypes.c_uint8 * 4).from_address(int(buffer.handle)) + + def assert_exec_uses(graph, expected): + original_data = as_bytes(original_buffer) + replacement_data = as_bytes(replacement_buffer) + original_data[:] = [0] * 4 + replacement_data[:] = [0] * 4 + + stream = device.create_stream() + graph.launch(stream) + stream.sync() + + assert list(as_bytes(expected["dst"])) == [expected["value"]] * 4 + if replace_dst: + unexpected = replacement_buffer if expected["dst"] is original_buffer else original_buffer + assert list(as_bytes(unexpected)) == [0] * 4 + + def cleanup(): + node.destroy() + original_buffer.close() + if replacement_buffer is not original_buffer: + replacement_buffer.close() + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=update, + assert_current=assert_current, + assert_exec_uses=assert_exec_uses, + invalid_update=lambda: node.update(value=256), + invalid_exception=OverflowError, + invalid_argument_update=lambda: node.update(dst=object()), + cleanup=cleanup, + ) + + +def _memset_value_case(device): + """Change the fill value while preserving destination ownership.""" + return _memset_case(device, replace_dst=False) + + +def _memset_destination_case(device): + """Replace the destination and its retained allocation owner.""" + return _memset_case(device, replace_dst=True) + + +def _memcpy_case(device, *, replace_operand): + memory_resource = LegacyPinnedMemoryResource() + original_src = memory_resource.allocate(4) + original_dst = memory_resource.allocate(4) + replacement_src = memory_resource.allocate(4) if replace_operand == "src" else original_src + replacement_dst = memory_resource.allocate(4) if replace_operand == "dst" else original_dst + original = { + "dst": original_dst, + "src": original_src, + "size": 2 if replace_operand is None else 4, + } + replacement = { + "dst": replacement_dst, + "src": replacement_src, + "size": 4, + } + + graph_def = GraphDefinition() + node = graph_def.memcpy(original["dst"], original["src"], original["size"]) + + def update(expected): + if replace_operand == "src": + node.update(src=expected["src"]) + elif replace_operand == "dst": + node.update(dst=expected["dst"]) + else: + node.update(size=expected["size"]) + + def assert_current(expected): + assert node.dst == int(expected["dst"].handle) + assert node.src == int(expected["src"].handle) + assert node.size == expected["size"] + + def as_bytes(buffer): + return (ctypes.c_uint8 * 4).from_address(int(buffer.handle)) + + def assert_exec_uses(graph, expected): + as_bytes(original_src)[:] = [0x11] * 4 + as_bytes(original_dst)[:] = [0] * 4 + if replacement_src is not original_src: + as_bytes(replacement_src)[:] = [0x22] * 4 + if replacement_dst is not original_dst: + as_bytes(replacement_dst)[:] = [0] * 4 + + stream = device.create_stream() + graph.launch(stream) + stream.sync() + + source_value = 0x11 if expected["src"] is original_src else 0x22 + expected_data = [source_value] * expected["size"] + expected_data.extend([0] * (4 - expected["size"])) + assert list(as_bytes(expected["dst"])) == expected_data + if replacement_dst is not original_dst: + unexpected_dst = replacement_dst if expected["dst"] is original_dst else original_dst + assert list(as_bytes(unexpected_dst)) == [0] * 4 + + def cleanup(): + node.destroy() + original_src.close() + original_dst.close() + if replacement_src is not original_src: + replacement_src.close() + if replacement_dst is not original_dst: + replacement_dst.close() + + return _DefinitionUpdateCase( + graph_def=graph_def, + node=node, + original=original, + replacement=replacement, + update=update, + assert_current=assert_current, + assert_exec_uses=assert_exec_uses, + invalid_update=lambda: node.update(size=-1), + invalid_exception=OverflowError, + invalid_argument_update=lambda: node.update(src=object()), + cleanup=cleanup, + ) + + +def _memcpy_size_case(device): + """Change the copy size while preserving both operand owners.""" + return _memcpy_case(device, replace_operand=None) + + +def _memcpy_source_case(device): + """Replace the source while preserving destination ownership.""" + return _memcpy_case(device, replace_operand="src") + + +def _memcpy_destination_case(device): + """Replace the destination while preserving source ownership.""" + return _memcpy_case(device, replace_operand="dst") + + +@pytest.fixture( + params=[ + pytest.param(_event_record_case, id="event-record"), + pytest.param(_event_wait_case, id="event-wait"), + pytest.param(_host_callback_case, id="host-callback-python"), + pytest.param(_host_callback_ctypes_case, id="host-callback-ctypes"), + pytest.param(_memset_value_case, id="memset-value"), + pytest.param(_memset_destination_case, id="memset-destination"), + pytest.param(_memcpy_size_case, id="memcpy-size"), + pytest.param(_memcpy_source_case, id="memcpy-source"), + pytest.param(_memcpy_destination_case, id="memcpy-destination"), + ] +) +def definition_update_case(request, init_cuda): + if driver_version() < (12, 2, 0): + pytest.skip("individual graph node updates require CUDA 12.2+") + case = request.param(init_cuda) + yield case + case.cleanup() + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_definition_node_update_changes_future_instantiations( + definition_update_case, +): + case = definition_update_case + old_graph = case.graph_def.instantiate() + + case.update(case.replacement) + case.assert_current(case.replacement) + + new_graph = case.graph_def.instantiate() + case.assert_exec_uses(old_graph, case.original) + case.assert_exec_uses(new_graph, case.replacement) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_failed_definition_node_update_preserves_state( + definition_update_case, +): + case = definition_update_case + + assert case.invalid_update is not None + assert case.invalid_exception is not None + with pytest.raises(case.invalid_exception): + case.invalid_update() + + case.assert_current(case.original) + graph = case.graph_def.instantiate() + case.assert_exec_uses(graph, case.original) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_definition_node_update_rejects_wrong_type( + definition_update_case, +): + if definition_update_case.invalid_argument_update is None: + pytest.skip("update method has no typed positional argument") + with pytest.raises(TypeError): + definition_update_case.invalid_argument_update()