diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index b3d2e3fe373..2102d36f22f 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -361,6 +361,15 @@ class HandleRegistry { map_.erase(key); } + void register_handles(const std::vector& handles) { + std::lock_guard lock(mutex_); + for (const Handle& h : handles) { + if (h) { + map_[*h] = h; + } + } + } + Handle lookup(const Key& key) { std::lock_guard lock(mutex_); auto it = map_.find(key); @@ -1341,7 +1350,8 @@ struct GraphHierarchy { }; // See REGISTRY_DESIGN.md (Level 1: Driver Handle -> Resource Handle) -static HandleRegistry graph_registry; +using GraphRegistry = HandleRegistry; +static GraphRegistry graph_registry; // Immutable resource owners for one version of a graph node's parameters. // Inheriting DeferredCleanupItem lets CUDA's user-object destructor enqueue @@ -1404,52 +1414,71 @@ CUresult rekey_attachments( return CUDA_SUCCESS; } -// Recursively copy and rekey attachments for a cloned graph hierarchy. -// The caller must release the GIL before calling this function. -CUresult copy_attachments( +struct StagedGraphMetadata { + const GraphBox* source; + GraphBox* clone; + GraphAttachmentMap* attachments; +}; +using StagedGraphMetadataList = std::vector; + +// Copy a source hierarchy into detached metadata before CUDA mutation. +void stage_graph_metadata( const GraphBox& source, GraphBox& clone, GraphAttachmentMap& attachments, - std::list& subgraphs) { - if (!p_cuGraphNodeFindInClone || !p_cuGraphChildGraphNodeGetGraph) { - return CUDA_ERROR_NOT_SUPPORTED; - } - + std::list& subgraphs, + StagedGraphMetadataList& staged) { attachments = source.attachments; - CUresult status = rekey_attachments(attachments, clone.resource); - if (status != CUDA_SUCCESS) { - return status; - } + staged.push_back({&source, &clone, &attachments}); for (const GraphBox& source_child : source.hierarchy->graphs) { if (source_child.parent != &source || !source_child.resource) { continue; } - - CUgraphNode cloned_owner = nullptr; - status = p_cuGraphNodeFindInClone( - &cloned_owner, source_child.owner_node, clone.resource); - if (status != CUDA_SUCCESS) { - return status; - } - - CUgraph cloned_graph = nullptr; - status = p_cuGraphChildGraphNodeGetGraph( - cloned_owner, &cloned_graph); - if (status != CUDA_SUCCESS) { - return status; - } - GraphBox& cloned_child = subgraphs.emplace_back( - cloned_graph, + nullptr, clone.hierarchy, &clone, - cloned_owner); - status = copy_attachments( + nullptr); + stage_graph_metadata( source_child, cloned_child, cloned_child.attachments, - subgraphs); + subgraphs, + staged); + } +} + +// Bind staged metadata to a CUDA-cloned hierarchy. The root clone resource +// must be populated before entry. The caller must release the GIL. +CUresult rekey_graph_metadata( + StagedGraphMetadataList& staged) { + if (!p_cuGraphNodeFindInClone || !p_cuGraphChildGraphNodeGetGraph) { + return CUDA_ERROR_NOT_SUPPORTED; + } + + CUresult status; + for (size_t i = 0; i < staged.size(); ++i) { + const GraphBox& source = *staged[i].source; + GraphBox& clone = *staged[i].clone; + if (i != 0) { + CUgraphNode cloned_owner = nullptr; + status = p_cuGraphNodeFindInClone( + &cloned_owner, + source.owner_node, + clone.parent->resource); + if (status == CUDA_SUCCESS) { + status = p_cuGraphChildGraphNodeGetGraph( + cloned_owner, &clone.resource); + } + if (status != CUDA_SUCCESS) { + return status; + } + clone.owner_node = cloned_owner; + } + + status = rekey_attachments( + *staged[i].attachments, clone.resource); if (status != CUDA_SUCCESS) { return status; } @@ -1497,6 +1526,29 @@ void rollback_prepared_attachment( delete state; } +// Detached metadata for a replacement embedded graph hierarchy. Preparation +// copies every attachment map and allocates every GraphBox before CUDA destroys +// the old embedded graph. Commit only rekeys and publishes it. +struct PreparedChildGraphUpdateState { + GraphHandle h_parent; + GraphHandle h_source; + GraphBox* old_root = nullptr; + CUgraphNode owner_node = nullptr; + std::list replacement; + StagedGraphMetadataList staged; + std::vector handles; + + PreparedChildGraphUpdateState( + GraphHandle h_parent_, + GraphHandle h_source_, + GraphBox* old_root_, + CUgraphNode owner_node_) + : h_parent(std::move(h_parent_)), + h_source(std::move(h_source_)), + old_root(old_root_), + owner_node(owner_node_) {} +}; + GraphHandle create_graph_handle(CUgraph graph) { if (!graph) { return {}; @@ -1543,15 +1595,112 @@ GraphHandle create_child_graph_handle( child_graph, hierarchy, parent, owner_node); GraphHandle h_child(h_parent, &child.resource); - try { - graph_registry.register_handle(child_graph, h_child); - } catch (...) { - hierarchy->graphs.pop_back(); - throw; - } + graph_registry.register_handle(child_graph, h_child); return h_child; } +CUresult graph_prepare_child_graph_update( + const GraphHandle& h_parent, + const GraphHandle& h_old_child, + CUgraphNode owner_node, + const GraphHandle& h_source, + PreparedChildGraphUpdate* out_prepared) { + if (!h_parent || !h_old_child || !owner_node || + !h_source || !out_prepared) { + return CUDA_ERROR_INVALID_VALUE; + } + out_prepared->reset(); + + GraphBox* parent = get_box(h_parent); + GraphBox* old_root = get_box(h_old_child); + GraphBox* source = get_box(h_source); + // A source from the destination hierarchy can include the old embedded + // subtree whose raw node keys CUDA destroys during replacement. + if (!parent->resource || !old_root->resource || !source->resource || + old_root->parent != parent || + old_root->owner_node != owner_node || + source->hierarchy == parent->hierarchy) { + return CUDA_ERROR_INVALID_VALUE; + } + + PreparedChildGraphUpdate prepared = + std::make_shared( + h_parent, h_source, old_root, owner_node); + + GraphBox& replacement_root = + prepared->replacement.emplace_back( + nullptr, parent->hierarchy, parent, owner_node); + stage_graph_metadata( + *source, + replacement_root, + replacement_root.attachments, + prepared->replacement, + prepared->staged); + + const size_t graph_count = prepared->staged.size(); + prepared->handles.reserve(graph_count); + for (const StagedGraphMetadata& graph : prepared->staged) { + prepared->handles.emplace_back( + h_parent, &graph.clone->resource); + } + + *out_prepared = std::move(prepared); + return CUDA_SUCCESS; +} + +void publish_child_graph_update( + PreparedChildGraphUpdateState& state, + GraphHandle* out_child) { + GraphBox* parent = get_box(state.h_parent); + parent->hierarchy->graphs.splice( + parent->hierarchy->graphs.end(), state.replacement); + *out_child = state.handles.front(); + graph_registry.register_handles(state.handles); +} + +CUresult graph_commit_child_graph_update( + PreparedChildGraphUpdate& prepared, + GraphHandle* out_child) { + if (!prepared || !out_child) { + return CUDA_ERROR_INVALID_VALUE; + } + out_child->reset(); + + PreparedChildGraphUpdateState& state = *prepared; + GraphBox* parent = get_box(state.h_parent); + if (!parent->resource || !state.old_root->resource) { + prepared.reset(); + return CUDA_ERROR_INVALID_VALUE; + } + + CUresult status = CUDA_ERROR_NOT_SUPPORTED; + CUgraph cloned_root = nullptr; + if (p_cuGraphChildGraphNodeGetGraph) { + GILReleaseGuard gil; + status = p_cuGraphChildGraphNodeGetGraph( + state.owner_node, &cloned_root); + if (status == CUDA_SUCCESS) { + state.staged.front().clone->resource = cloned_root; + status = rekey_graph_metadata(state.staged); + } + } + + // CUDA has already destroyed the old embedded graph. No replacement + // metadata is visible yet, so this selects only the old generation. + invalidate_child_graph_state( + state.h_parent, state.owner_node); + + if (status != CUDA_SUCCESS) { + prepared.reset(); + throw std::runtime_error( + "failed to update graph metadata after child graph replacement"); + } + + publish_child_graph_update(state, out_child); + prepared.reset(); + return status; +} + CUresult graph_get_attachment( const GraphHandle& h_graph, CUgraphNode node, OpaqueHandle* owner0, OpaqueHandle* owner1) { @@ -1727,13 +1876,22 @@ CUresult graph_clone_attachments( // Build and rekey the clone metadata off-hierarchy so a CUDA mapping error // cannot partially publish it. - GraphAttachmentMap attachments = source->attachments; + GraphAttachmentMap attachments; std::list subgraphs; + StagedGraphMetadataList staged; + stage_graph_metadata( + *source, *clone, attachments, subgraphs, staged); + + std::vector handles; + handles.reserve(subgraphs.size()); + for (GraphBox& graph : subgraphs) { + handles.emplace_back(h_clone, &graph.resource); + } + CUresult status; { GILReleaseGuard gil; - status = copy_attachments( - *source, *clone, attachments, subgraphs); + status = rekey_graph_metadata(staged); } if (status != CUDA_SUCCESS) { return status; @@ -1744,13 +1902,9 @@ CUresult graph_clone_attachments( return CUDA_SUCCESS; } - auto first = subgraphs.begin(); clone->hierarchy->graphs.splice( clone->hierarchy->graphs.end(), subgraphs); - for (auto it = first; it != clone->hierarchy->graphs.end(); ++it) { - GraphHandle h_graph(h_clone, &it->resource); - graph_registry.register_handle(it->resource, h_graph); - } + graph_registry.register_handles(handles); return CUDA_SUCCESS; } diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index 3a9d2d75cff..60a9f3c53a9 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -516,6 +516,12 @@ struct PreparedAttachmentDeleter { using PreparedAttachment = std::unique_ptr; +struct PreparedChildGraphUpdateState; +// Opaque unpublished hierarchy transaction; releasing it discards staged +// metadata unless graph_commit_child_graph_update publishes the replacement. +using PreparedChildGraphUpdate = + std::shared_ptr; + // Copy requested owners from node's current attachment. Pass nullptr to ignore // either owner; a missing attachment produces empty handles. CUresult graph_get_attachment( @@ -543,6 +549,21 @@ CUresult graph_clone_attachments( const GraphHandle& h_clone, const GraphHandle& h_source); +// Stage a complete metadata replacement before CUDA replaces an embedded +// graph. Dropping the prepared state leaves the current hierarchy unchanged. +CUresult graph_prepare_child_graph_update( + const GraphHandle& h_parent, + const GraphHandle& h_old_child, + CUgraphNode owner_node, + const GraphHandle& h_source, + PreparedChildGraphUpdate* out_prepared); + +// Rekey staged metadata to CUDA's replacement clone, retire the old embedded +// hierarchy, and publish the replacement handle. +CUresult graph_commit_child_graph_update( + PreparedChildGraphUpdate& prepared, + GraphHandle* out_child); + // Invalidate cuda.core state for child graphs CUDA destroyed with owner_node. void invalidate_child_graph_state( const GraphHandle& h_parent, diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index 2f481fee4f8..b0ae65d1666 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -71,6 +71,12 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": PreparedAttachmentState, PreparedAttachmentDeleter ] PreparedAttachment + cppclass PreparedChildGraphUpdateState: + pass + ctypedef shared_ptr[ + PreparedChildGraphUpdateState + ] PreparedChildGraphUpdate + # as_cu() - extract the raw CUDA handle (inline C++) cydriver.CUcontext as_cu(ContextHandle h) noexcept nogil cydriver.CUgreenCtx as_cu(GreenCtxHandle h) noexcept nogil @@ -253,6 +259,12 @@ cdef cydriver.CUresult graph_commit_attachment( PreparedAttachment& prepared, cydriver.CUgraphNode node) except+ cdef cydriver.CUresult graph_clone_attachments( const GraphHandle& h_clone, const GraphHandle& h_source) except+ +cdef cydriver.CUresult graph_prepare_child_graph_update( + const GraphHandle& h_parent, const GraphHandle& h_old_child, + cydriver.CUgraphNode owner_node, const GraphHandle& h_source, + PreparedChildGraphUpdate* out_prepared) except+ +cdef cydriver.CUresult graph_commit_child_graph_update( + PreparedChildGraphUpdate& prepared, GraphHandle* out_child) except+ cdef void invalidate_child_graph_state( const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept diff --git a/cuda_core/cuda/core/_resource_handles.pyi b/cuda_core/cuda/core/_resource_handles.pyi index f11f6f08e00..c1bbb3983c9 100644 --- a/cuda_core/cuda/core/_resource_handles.pyi +++ b/cuda_core/cuda/core/_resource_handles.pyi @@ -26,4 +26,5 @@ MipmappedArrayHandle = shared_ptr TexObjectHandle = shared_ptr SurfObjectHandle = shared_ptr OpaqueHandle = shared_ptr -PreparedAttachment = unique_ptr \ No newline at end of file +PreparedAttachment = unique_ptr +PreparedChildGraphUpdate = shared_ptr \ No newline at end of file diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index f6fb6ac4e20..a1b0d912a71 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -167,6 +167,12 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": PreparedAttachment& prepared, cydriver.CUgraphNode node) except+ cydriver.CUresult graph_clone_attachments "cuda_core::graph_clone_attachments" ( const GraphHandle& h_clone, const GraphHandle& h_source) except+ + cydriver.CUresult graph_prepare_child_graph_update "cuda_core::graph_prepare_child_graph_update" ( + const GraphHandle& h_parent, const GraphHandle& h_old_child, + cydriver.CUgraphNode owner_node, const GraphHandle& h_source, + PreparedChildGraphUpdate* out_prepared) except+ + cydriver.CUresult graph_commit_child_graph_update "cuda_core::graph_commit_child_graph_update" ( + PreparedChildGraphUpdate& prepared, GraphHandle* out_child) except+ void invalidate_child_graph_state "cuda_core::invalidate_child_graph_state" ( const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept diff --git a/cuda_core/cuda/core/graph/_graph_node.pxd b/cuda_core/cuda/core/graph/_graph_node.pxd index 0a87b70ad62..2ad4851a54c 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,15 @@ 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 cydriver.CUmemorytype _get_memcpy_memory_type( + cydriver.CUdeviceptr 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.pyi b/cuda_core/cuda/core/graph/_graph_node.pyi index 23bcbf191a3..ab503183f63 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyi +++ b/cuda_core/cuda/core/graph/_graph_node.pyi @@ -94,6 +94,9 @@ class GraphNode: def launch(self, config: LaunchConfig, kernel: Kernel, *args) -> KernelNode: """Add a kernel launch node depending on this node. + Clustered and cooperative launch configurations are not currently + supported for graph kernel nodes. + .. warning:: Use caution when a retained kernel argument directly or indirectly diff --git a/cuda_core/cuda/core/graph/_graph_node.pyx b/cuda_core/cuda/core/graph/_graph_node.pyx index 9e1cab09e7b..a2a830f2ede 100644 --- a/cuda_core/cuda/core/graph/_graph_node.pyx +++ b/cuda_core/cuda/core/graph/_graph_node.pyx @@ -209,6 +209,9 @@ cdef class GraphNode: def launch(self, config: LaunchConfig, kernel: Kernel, *args) -> KernelNode: """Add a kernel launch node depending on this node. + Clustered and cooperative launch configurations are not currently + supported for graph kernel nodes. + .. warning:: Use caution when a retained kernel argument directly or indirectly @@ -722,6 +725,10 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker, cdef OpaqueHandle kernel_owner, args_owner cdef PreparedAttachment prepared + if conf.cluster is not None or conf.is_cooperative: + raise NotImplementedError( + "clustered or cooperative graph kernel nodes are not supported") + if pred_node != NULL: deps = &pred_node num_deps = 1 @@ -889,7 +896,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,46 +977,52 @@ 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 unsigned int dst_mem_type = cydriver.CU_MEMORYTYPE_DEVICE - cdef unsigned int src_mem_type = cydriver.CU_MEMORYTYPE_DEVICE +cdef cydriver.CUmemorytype _get_memcpy_memory_type( + cydriver.CUdeviceptr ptr) except *: + cdef unsigned int memory_type = cydriver.CU_MEMORYTYPE_DEVICE cdef cydriver.CUresult ret with nogil: ret = cydriver.cuPointerGetAttribute( - &dst_mem_type, + &memory_type, cydriver.CU_POINTER_ATTRIBUTE_MEMORY_TYPE, - c_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) - 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 + ptr) + if ret != cydriver.CUDA_SUCCESS and ret != cydriver.CUDA_ERROR_INVALID_VALUE: + HANDLE_RETURN(ret) + return memory_type + + +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 *: + dst_type[0] = _get_memcpy_memory_type(dst) + src_type[0] = _get_memcpy_memory_type(src) + + 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/_host_callback.pyx b/cuda_core/cuda/core/graph/_host_callback.pyx index 27f251abeae..5fd71f8653f 100644 --- a/cuda_core/cuda/core/graph/_host_callback.pyx +++ b/cuda_core/cuda/core/graph/_host_callback.pyx @@ -54,6 +54,9 @@ cdef void _resolve_host_callback( else: out_user_data[0] = NULL else: + if not callable(fn): + raise TypeError( + f"callback must be callable, got {type(fn).__name__}") if user_data is not None: raise ValueError( "user_data is only supported with ctypes function pointers") diff --git a/cuda_core/cuda/core/graph/_subclasses.pyi b/cuda_core/cuda/core/graph/_subclasses.pyi index 345e6417c4d..2e6d4dd9529 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 @@ -37,6 +38,21 @@ class KernelNode(GraphNode): def __repr__(self) -> str: ... + def update(self, *, config: LaunchConfig | None=None, kernel: Kernel | None=None, args=None) -> None: + """Replace selected kernel launch parameters. + + Omitted parameters preserve their current values. Changing ``kernel`` + requires ``args``, including ``args=()`` for a no-argument kernel. + Clustered and cooperative kernel nodes are not supported. + + .. warning:: + + Use caution when a retained kernel argument directly or indirectly + owns a graph. Any reference cycle involving the argument 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 grid(self) -> tuple[int, int, int]: """Grid dimensions as a 3-tuple (gridDimX, gridDimY, gridDimZ).""" @@ -139,6 +155,24 @@ 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``. + + With CUDA 12.2 through 13.1, the node's intended CUDA context must be + current when this method is called. CUDA driver and ``cuda.bindings`` + versions 13.2 and newer preserve the recorded context automatically. + + .. 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 +213,26 @@ 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. + Multidimensional, pitched, offset, and array-backed memcpy nodes are + not supported. + + With CUDA 12.2 through 13.1, the node's intended CUDA context must be + current when this method is called. CUDA driver and ``cuda.bindings`` + versions 13.2 and newer preserve the recorded context automatically. + + .. 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.""" @@ -203,6 +257,12 @@ class ChildGraphNode(GraphNode): def __repr__(self) -> str: ... + def update(self, child: GraphDefinition) -> None: + """Replace the embedded graph with a clone of ``child``. + + ``child`` must belong to an independent graph hierarchy. + """ + @property def child_graph(self) -> GraphDefinition: """The embedded graph definition (non-owning wrapper).""" @@ -219,6 +279,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 +298,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 +317,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..f896b865a8d 100644 --- a/cuda_core/cuda/core/graph/_subclasses.pyx +++ b/cuda_core/cuda/core/graph/_subclasses.pyx @@ -8,29 +8,49 @@ 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._kernel_arg_handler cimport ParamHolder 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, + _get_memcpy_memory_type, + _resolve_memcpy_operand, +) from cuda.core._resource_handles cimport ( EventHandle, GraphHandle, - KernelHandle, GraphNodeHandle, + KernelHandle, + OpaqueHandle, + PreparedAttachment, + PreparedChildGraphUpdate, as_cu, as_intptr, - create_event_handle_ref, create_child_graph_handle, + create_event_handle_ref, create_kernel_handle_ref, + graph_commit_attachment, + graph_commit_child_graph_update, + graph_get_attachment, graph_node_get_graph, + graph_prepare_attachment, + graph_prepare_child_graph_update, + make_opaque_py, ) -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 +77,54 @@ __all__ = [ cdef bint _has_cuGraphNodeGetParams = False cdef bint _version_checked = False + +cdef void _require_graph_node_update_support() 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 void _set_definition_node_params( + const GraphNodeHandle& h_node, + cydriver.CUgraphNodeParams* params, + OpaqueHandle owner0, + OpaqueHandle owner1=OpaqueHandle(), + cydriver.CUcontext update_ctx=NULL) except *: + _require_graph_node_update_support() + + cdef GraphHandle h_graph = graph_node_get_graph(h_node) + cdef cydriver.CUgraphNode node = as_cu(h_node) + cdef cydriver.CUcontext previous_ctx = NULL + cdef bint restore_ctx = False + cdef PreparedAttachment prepared + + HANDLE_RETURN(graph_prepare_attachment( + h_graph, owner0, owner1, &prepared)) + if update_ctx != NULL: + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&previous_ctx)) + if previous_ctx != update_ctx: + HANDLE_RETURN(cydriver.cuCtxSetCurrent(update_ctx)) + restore_ctx = True + try: + with nogil: + HANDLE_RETURN(cydriver.cuGraphNodeSetParams(node, params)) + finally: + if restore_ctx: + with nogil: + HANDLE_RETURN(cydriver.cuCtxSetCurrent(previous_ctx)) + HANDLE_RETURN(graph_commit_attachment(prepared, node)) + + cdef bint _check_node_get_params(): global _has_cuGraphNodeGetParams, _version_checked if not _version_checked: @@ -68,6 +136,54 @@ cdef bint _check_node_get_params(): return _has_cuGraphNodeGetParams +cdef void _reject_unsupported_kernel_node( + cydriver.CUgraphNode node) except *: + cdef cydriver.CUkernelNodeAttrValue cluster + cdef cydriver.CUkernelNodeAttrValue cooperative + + c_memset(&cluster, 0, sizeof(cluster)) + c_memset(&cooperative, 0, sizeof(cooperative)) + with nogil: + HANDLE_RETURN(cydriver.cuGraphKernelNodeGetAttribute( + node, ( + cydriver.CU_KERNEL_NODE_ATTRIBUTE_CLUSTER_DIMENSION), + &cluster)) + HANDLE_RETURN(cydriver.cuGraphKernelNodeGetAttribute( + node, ( + cydriver.CU_KERNEL_NODE_ATTRIBUTE_COOPERATIVE), + &cooperative)) + if (cluster.clusterDim.x != 0 or cluster.clusterDim.y != 0 or + cluster.clusterDim.z != 0 or cooperative.cooperative != 0): + raise NotImplementedError( + "updating clustered or cooperative kernel nodes is not supported") + + +cdef bint _is_supported_memcpy_descriptor( + cydriver.CUDA_MEMCPY3D* params) noexcept nogil: + return ( + (params.srcMemoryType == cydriver.CU_MEMORYTYPE_HOST or + params.srcMemoryType == cydriver.CU_MEMORYTYPE_DEVICE) + and (params.dstMemoryType == cydriver.CU_MEMORYTYPE_HOST or + params.dstMemoryType == cydriver.CU_MEMORYTYPE_DEVICE) + and params.srcXInBytes == 0 + and params.srcY == 0 + and params.srcZ == 0 + and params.srcLOD == 0 + and params.srcPitch == 0 + and params.srcHeight == 0 + and params.dstXInBytes == 0 + and params.dstY == 0 + and params.dstZ == 0 + and params.dstLOD == 0 + and params.dstPitch == 0 + and params.dstHeight == 0 + and params.Height == 1 + and params.Depth == 1 + and params.reserved0 == NULL + and params.reserved1 == NULL + ) + + cdef class EmptyNode(GraphNode): """An empty (synchronization) node.""" @@ -130,6 +246,99 @@ cdef class KernelNode(GraphNode): return (f"") + def update( + self, + *, + config: LaunchConfig | None = None, + kernel: Kernel | None = None, + args=None, + ) -> None: + """Replace selected kernel launch parameters. + + Omitted parameters preserve their current values. Changing ``kernel`` + requires ``args``, including ``args=()`` for a no-argument kernel. + Clustered and cooperative kernel nodes are not supported. + + .. warning:: + + Use caution when a retained kernel argument directly or indirectly + owns a graph. Any reference cycle involving the argument and a + graph that retains it cannot be broken by Python's cyclic garbage + collector. Use a weak reference to break such cycles. + """ + cdef LaunchConfig c_config + cdef Kernel c_kernel + cdef ParamHolder arg_holder + cdef object kernel_args + cdef KernelHandle h_kernel = self._h_kernel + cdef OpaqueHandle kernel_owner + cdef OpaqueHandle args_owner + cdef GraphHandle h_graph = graph_node_get_graph(self._h_node) + cdef cydriver.CUgraphNode node = as_cu(self._h_node) + cdef cydriver.CUgraphNodeParams params + + if config is not None: + c_config = config + if (c_config.cluster is not None or + c_config.is_cooperative): + raise NotImplementedError( + "updating clustered or cooperative kernel nodes is not " + "supported") + _require_graph_node_update_support() + _reject_unsupported_kernel_node(node) + if kernel is not None: + if args is None: + raise ValueError("changing kernel requires args") + c_kernel = kernel + h_kernel = c_kernel._h_kernel + if args is not None: + arg_holder = ParamHolder(args) + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_KERNEL + with nogil: + HANDLE_RETURN(cydriver.cuGraphKernelNodeGetParams( + node, ¶ms.kernel)) + HANDLE_RETURN(graph_get_attachment( + h_graph, node, &kernel_owner, &args_owner)) + + if config is not None: + params.kernel.gridDimX = c_config.grid[0] + params.kernel.gridDimY = c_config.grid[1] + params.kernel.gridDimZ = c_config.grid[2] + params.kernel.blockDimX = c_config.block[0] + params.kernel.blockDimY = c_config.block[1] + params.kernel.blockDimZ = c_config.block[2] + params.kernel.sharedMemBytes = c_config.shmem_size + if kernel is not None: + params.kernel.kern = as_cu(h_kernel) + params.kernel.func = NULL + params.kernel.ctx = NULL + kernel_owner = h_kernel + if args is not None: + params.kernel.kernelParams = arg_holder.ptr + params.kernel.extra = NULL + kernel_args = arg_holder.kernel_args + if kernel_args is None: + args_owner = OpaqueHandle() + else: + args_owner = make_opaque_py(kernel_args) + + _set_definition_node_params( + self._h_node, ¶ms, kernel_owner, args_owner) + self._grid = ( + params.kernel.gridDimX, + params.kernel.gridDimY, + params.kernel.gridDimZ, + ) + self._block = ( + params.kernel.blockDimX, + params.kernel.blockDimY, + params.kernel.blockDimZ, + ) + self._shmem_size = params.kernel.sharedMemBytes + self._h_kernel = h_kernel + @property def grid(self) -> tuple[int, int, int]: """Grid dimensions as a 3-tuple (gridDimX, gridDimY, gridDimZ).""" @@ -340,6 +549,114 @@ 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``. + + With CUDA 12.2 through 13.1, the node's intended CUDA context must be + current when this method is called. CUDA driver and ``cuda.bindings`` + versions 13.2 and newer preserve the recorded context automatically. + + .. 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 + cdef unsigned int c_value + cdef unsigned int c_element_size + cdef size_t c_width + cdef size_t c_height + cdef size_t c_pitch + cdef OpaqueHandle dst_attachment_owner + cdef GraphHandle h_graph + cdef cydriver.CUgraphNode node = as_cu(self._h_node) + cdef cydriver.CUcontext ctx = NULL + cdef cydriver.CUDA_MEMSET_NODE_PARAMS current + cdef cydriver.CUgraphNodeParams params + cdef object queried + + if dst is None and dst_owner is not None: + raise ValueError("dst_owner requires dst") + if (dst is None and value is None and width is None and + height is None and pitch is None): + return + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMSET + with nogil: + HANDLE_RETURN(cydriver.cuGraphMemsetNodeGetParams( + node, ¤t)) + if _check_node_get_params(): + queried = handle_return(driver.cuGraphNodeGetParams( + node)) + ctx = int(queried.memset.ctx) + else: + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx)) + params.memset.dst = current.dst + params.memset.value = current.value + params.memset.elementSize = current.elementSize + params.memset.width = current.width + params.memset.height = current.height + params.memset.pitch = current.pitch + params.memset.ctx = ctx + + c_dst = params.memset.dst + c_value = params.memset.value + c_element_size = params.memset.elementSize + c_width = params.memset.width + c_height = params.memset.height + c_pitch = params.memset.pitch + + if dst is None: + h_graph = graph_node_get_graph(self._h_node) + HANDLE_RETURN(graph_get_attachment( + h_graph, 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 + + 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 + + _set_definition_node_params( + self._h_node, ¶ms, dst_attachment_owner, + OpaqueHandle(), params.memset.ctx) + 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 +745,127 @@ 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. + Multidimensional, pitched, offset, and array-backed memcpy nodes are + not supported. + + With CUDA 12.2 through 13.1, the node's intended CUDA context must be + current when this method is called. CUDA driver and ``cuda.bindings`` + versions 13.2 and newer preserve the recorded context automatically. + + .. 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 OpaqueHandle dst_attachment_owner + cdef OpaqueHandle src_attachment_owner + cdef GraphHandle h_graph = graph_node_get_graph(self._h_node) + cdef cydriver.CUgraphNode node = as_cu(self._h_node) + cdef cydriver.CUcontext ctx = NULL + cdef cydriver.CUgraphNodeParams params + cdef cydriver.CUmemorytype c_dst_type + cdef cydriver.CUmemorytype c_src_type + cdef object queried + + if dst is None and dst_owner is not None: + raise ValueError("dst_owner requires dst") + if src is None and src_owner is not None: + raise ValueError("src_owner requires src") + if dst is None and src is None and size is None: + return + + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_MEMCPY + with nogil: + HANDLE_RETURN(cydriver.cuGraphMemcpyNodeGetParams( + node, ¶ms.memcpy.copyParams)) + if _check_node_get_params(): + queried = handle_return(driver.cuGraphNodeGetParams( + node)) + ctx = int( + queried.memcpy.copyCtx) + else: + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx)) + params.memcpy.copyCtx = ctx + + if not _is_supported_memcpy_descriptor(¶ms.memcpy.copyParams): + raise NotImplementedError( + "updating multidimensional, pitched, offset, or array-backed " + "memcpy nodes is not supported") + + c_dst_type = params.memcpy.copyParams.dstMemoryType + c_src_type = params.memcpy.copyParams.srcMemoryType + if c_dst_type == cydriver.CU_MEMORYTYPE_HOST: + c_dst = ( + params.memcpy.copyParams.dstHost) + elif c_dst_type != cydriver.CU_MEMORYTYPE_ARRAY: + c_dst = params.memcpy.copyParams.dstDevice + if c_src_type == cydriver.CU_MEMORYTYPE_HOST: + c_src = ( + params.memcpy.copyParams.srcHost) + elif c_src_type != cydriver.CU_MEMORYTYPE_ARRAY: + c_src = params.memcpy.copyParams.srcDevice + + HANDLE_RETURN(graph_get_attachment( + h_graph, node, + &dst_attachment_owner, &src_attachment_owner)) + if dst is not None: + dst_attachment_owner = _resolve_memcpy_operand( + dst, dst_owner, "dst", &c_dst) + c_dst_type = _get_memcpy_memory_type(c_dst) + params.memcpy.copyParams.dstMemoryType = c_dst_type + params.memcpy.copyParams.dstHost = NULL + params.memcpy.copyParams.dstDevice = 0 + params.memcpy.copyParams.dstArray = NULL + params.memcpy.copyParams.reserved1 = NULL + if c_dst_type == cydriver.CU_MEMORYTYPE_HOST: + params.memcpy.copyParams.dstHost = c_dst + else: + params.memcpy.copyParams.dstDevice = c_dst + if src is not None: + src_attachment_owner = _resolve_memcpy_operand( + src, src_owner, "src", &c_src) + c_src_type = _get_memcpy_memory_type(c_src) + params.memcpy.copyParams.srcMemoryType = c_src_type + params.memcpy.copyParams.srcHost = NULL + params.memcpy.copyParams.srcDevice = 0 + params.memcpy.copyParams.srcArray = NULL + params.memcpy.copyParams.reserved0 = NULL + if c_src_type == cydriver.CU_MEMORYTYPE_HOST: + params.memcpy.copyParams.srcHost = c_src + else: + params.memcpy.copyParams.srcDevice = c_src + if size is not None: + params.memcpy.copyParams.WidthInBytes = size + + _set_definition_node_params( + self._h_node, ¶ms, + dst_attachment_owner, src_attachment_owner, + params.memcpy.copyCtx) + self._dst = c_dst + self._src = c_src + self._size = params.memcpy.copyParams.WidthInBytes + self._dst_type = c_dst_type + self._src_type = c_src_type + @property def dst(self) -> int: """The destination pointer.""" @@ -478,6 +916,37 @@ cdef class ChildGraphNode(GraphNode): return (f"") + def update(self, child: GraphDefinition) -> None: + """Replace the embedded graph with a clone of ``child``. + + ``child`` must belong to an independent graph hierarchy. + """ + cdef GraphHandle h_parent = graph_node_get_graph(self._h_node) + cdef GraphHandle h_replacement + cdef cydriver.CUgraphNode node = as_cu(self._h_node) + cdef cydriver.CUgraphNodeParams params + cdef cydriver.CUresult commit_status + cdef PreparedChildGraphUpdate prepared + + _require_graph_node_update_support() + c_memset(¶ms, 0, sizeof(params)) + params.type = cydriver.CU_GRAPH_NODE_TYPE_GRAPH + params.graph.graph = as_cu(child._h_graph) + + HANDLE_RETURN(graph_prepare_child_graph_update( + h_parent, self._h_child_graph, node, + child._h_graph, &prepared)) + with nogil: + HANDLE_RETURN(cydriver.cuGraphNodeSetParams( + node, ¶ms)) + try: + commit_status = graph_commit_child_graph_update( + prepared, &h_replacement) + finally: + if h_replacement: + self._h_child_graph = h_replacement + HANDLE_RETURN(commit_status) + @property def child_graph(self) -> GraphDefinition: """The embedded graph definition (non-owning wrapper).""" @@ -516,6 +985,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 +1036,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 +1099,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/docs/source/api.rst b/cuda_core/docs/source/api.rst index 089e68576c9..da18ee27c80 100644 --- a/cuda_core/docs/source/api.rst +++ b/cuda_core/docs/source/api.rst @@ -154,6 +154,22 @@ Every graph node is a subclass of :class:`~graph.GraphNode`, which provides the common interface (dependencies, successors, destruction). Each subclass exposes attributes unique to its operation type. +Parameter-bearing definition nodes expose subclass-specific ``update()`` +methods: :class:`~graph.KernelNode`, :class:`~graph.MemcpyNode`, +:class:`~graph.MemsetNode`, :class:`~graph.ChildGraphNode`, +:class:`~graph.EventRecordNode`, :class:`~graph.EventWaitNode`, and +:class:`~graph.HostCallbackNode`. These methods require CUDA driver and +``cuda.bindings`` versions 12.2 or newer. Updates affect future graph +instantiations; executable graphs that were already instantiated continue +using their previous parameters and retained resources. Omitted optional +arguments preserve their current values where supported. +On CUDA 12.2 through 13.1, the intended CUDA context must be current when +updating memcpy or memset nodes. CUDA driver and ``cuda.bindings`` versions +13.2 and newer preserve the recorded context automatically. +Multidimensional or array-backed memcpy nodes and clustered or cooperative +kernel nodes cannot currently be updated. Clustered and cooperative kernel +nodes also cannot currently be constructed explicitly. + .. autosummary:: :toctree: generated/ diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 6a047c9cfe8..de255d01cbb 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -18,6 +18,13 @@ Fixes and enhancements (`#2357 `__, `#2371 `__) +- Added ``update()`` methods to kernel, memcpy, memset, child-graph, event + record, event wait, and host-callback graph definition nodes. Updates change + parameters used by future graph instantiations without affecting existing + executable graphs. This feature requires CUDA driver and ``cuda.bindings`` + versions 12.2 or newer. + (`#2352 `__) + Deprecation Notices ------------------- diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index 804cccc1923..1364f437ea6 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -77,6 +77,8 @@ def _wait_until(predicate, timeout=None, interval=0.02): from cuda.core import Device, DeviceMemoryResource, EventOptions, Kernel, LaunchConfig +from cuda.core._utils.cuda_utils import CUDAError +from cuda.core._utils.version import driver_version from cuda.core.graph import ( ChildGraphNode, ConditionalNode, @@ -424,6 +426,97 @@ def test_destroying_child_node_invalidates_embedded_handles(init_cuda): assert not embedded_callback.is_valid +@pytest.mark.agent_authored(model="gpt-5.6") +def test_updating_child_node_replaces_embedded_handles(init_cuda): + """A successful replacement invalidates only the old embedded hierarchy.""" + if driver_version() < (12, 2, 0): + pytest.skip("individual graph node updates require CUDA 12.2+") + + old_inner = GraphDefinition() + old_inner.callback(lambda: None) + old_middle = GraphDefinition() + old_middle.embed(old_inner) + parent = GraphDefinition() + child_node = parent.embed(old_middle) + + embedded_middle = child_node.child_graph + embedded_child = next(node for node in embedded_middle.nodes() if isinstance(node, ChildGraphNode)) + embedded_inner = embedded_child.child_graph + embedded_callback = next(node for node in embedded_inner.nodes() if isinstance(node, HostCallbackNode)) + + # Sources from the destination hierarchy may contain handles CUDA destroys + # during replacement, so cuda-core rejects them before mutation. + with pytest.raises(CUDAError): + child_node.update(embedded_middle) + with pytest.raises(CUDAError): + child_node.update(parent) + assert int(embedded_middle.handle) != 0 + assert int(embedded_inner.handle) != 0 + assert embedded_child.is_valid + assert embedded_callback.is_valid + + replacement_inner = GraphDefinition() + replacement_inner.callback(lambda: None) + replacement_middle = GraphDefinition() + replacement_middle.embed(replacement_inner) + child_node.update(replacement_middle) + + assert child_node.is_valid + assert int(embedded_middle.handle) == 0 + assert int(embedded_inner.handle) == 0 + assert not embedded_child.is_valid + assert not embedded_callback.is_valid + + new_middle = child_node.child_graph + new_child = next(node for node in new_middle.nodes() if isinstance(node, ChildGraphNode)) + assert int(new_middle.handle) != 0 + assert int(new_child.child_graph.handle) != 0 + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_child_update_replaces_nested_attachments(init_cuda): + """Replacement drops old owners and imports nested replacement owners.""" + if driver_version() < (12, 2, 0): + pytest.skip("individual graph node updates require CUDA 12.2+") + + def old_callback(): + pass + + old_callback_weak = weakref.ref(old_callback) + old_child = GraphDefinition() + old_child.callback(old_callback) + parent = GraphDefinition() + child_node = parent.embed(old_child) + + del old_callback, old_child + gc.collect() + assert old_callback_weak() is not None + + def replacement_callback(): + pass + + replacement_callback_weak = weakref.ref(replacement_callback) + replacement_inner = GraphDefinition() + replacement_inner.callback(replacement_callback) + replacement = GraphDefinition() + replacement.embed(replacement_inner) + child_node.update(replacement) + + _wait_until(lambda: old_callback_weak() is None) + del replacement_callback, replacement_inner, replacement + gc.collect() + assert replacement_callback_weak() is not None + + embedded = child_node.child_graph + embedded_child = next(node for node in embedded.nodes() if isinstance(node, ChildGraphNode)) + embedded_callback = next(node for node in embedded_child.child_graph.nodes() if isinstance(node, HostCallbackNode)) + assert embedded_callback.callback is replacement_callback_weak() + + del embedded_callback, embedded_child, embedded + child_node.destroy() + _wait_until(lambda: replacement_callback_weak() is None) + + @pytest.mark.agent_authored(model="gpt-5.6") def test_builder_embedded_clone_releases_attachment_on_node_destroy(init_cuda): """GraphBuilder.embed imports metadata from the captured child graph.""" 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..baf7c770ac1 --- /dev/null +++ b/cuda_core/tests/graph/test_graph_node_update.py @@ -0,0 +1,728 @@ +# 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 helpers.graph_kernels import compile_common_kernels + +from cuda.core import LaunchConfig, LegacyPinnedMemoryResource +from cuda.core._utils.cuda_utils import CUDAError, driver, handle_return +from cuda.core._utils.version import driver_version +from cuda.core.graph import GraphDefinition, HostCallbackNode + + +@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=lambda: node.update(object()), + 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") + + +def _kernel_case(device, *, replace): + module = compile_common_kernels() + add_one = module.get_kernel("add_one") + empty_kernel = module.get_kernel("empty_kernel") + write_launch_dims = module.get_kernel("write_launch_dims") + memory_resource = LegacyPinnedMemoryResource() + original_buffer = memory_resource.allocate(ctypes.sizeof(ctypes.c_int)) + replacement_buffer = memory_resource.allocate(ctypes.sizeof(ctypes.c_int)) if replace == "args" else original_buffer + + original_config = LaunchConfig(grid=1, block=1) + replacement_config = LaunchConfig(grid=2, block=3) if replace == "config" else original_config + original_kernel = write_launch_dims if replace == "config" else add_one + replacement_kernel = empty_kernel if replace == "kernel" else original_kernel + original_args = (original_buffer,) + if replace == "kernel": + replacement_args = () + elif replace == "args": + replacement_args = (replacement_buffer,) + else: + replacement_args = original_args + + original = { + "config": original_config, + "kernel": original_kernel, + "args": original_args, + "output": original_buffer, + "expected": 1001 if replace == "config" else 1, + } + replacement = { + "config": replacement_config, + "kernel": replacement_kernel, + "args": replacement_args, + "output": replacement_buffer, + "expected": 2003 if replace == "config" else int(replace != "kernel"), + } + + graph_def = GraphDefinition() + node = graph_def.launch(original["config"], original["kernel"], *original["args"]) + + def update(expected): + if replace == "config": + node.update(config=expected["config"]) + elif replace == "args": + node.update(args=expected["args"]) + else: + node.update(kernel=expected["kernel"], args=expected["args"]) + + def assert_current(expected): + assert node.config == expected["config"] + assert int(node.kernel.handle) == int(expected["kernel"].handle) + + def as_int(buffer): + return ctypes.c_int.from_address(int(buffer.handle)) + + def assert_exec_uses(graph, expected): + as_int(original_buffer).value = 0 + as_int(replacement_buffer).value = 0 + + stream = device.create_stream() + graph.launch(stream) + stream.sync() + + assert as_int(expected["output"]).value == expected["expected"] + if replacement_buffer is not original_buffer: + unexpected = replacement_buffer if expected["output"] is original_buffer else original_buffer + assert as_int(unexpected).value == 0 + + def invalid_update(): + if replace == "kernel": + node.update(kernel=replacement_kernel) + elif replace == "args": + node.update(args=(object(),)) + else: + node.update(config=object()) + + invalid_exception = ValueError if replace == "kernel" else TypeError + + 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=invalid_update, + invalid_exception=invalid_exception, + invalid_argument_update=lambda: node.update(config=object()), + cleanup=cleanup, + ) + + +def _kernel_config_case(device): + """Replace launch dimensions while preserving the kernel and arguments.""" + return _kernel_case(device, replace="config") + + +def _kernel_args_case(device): + """Replace arguments while preserving the kernel and configuration.""" + return _kernel_case(device, replace="args") + + +def _kernel_function_case(device): + """Replace a kernel and explicitly supply its coupled arguments.""" + return _kernel_case(device, replace="kernel") + + +def _child_graph_case(device): + """Replace the embedded clone while preserving existing executables.""" + called = [] + + def original_callback(): + called.append(original_callback) + + def replacement_callback(): + called.append(replacement_callback) + + original_child = GraphDefinition() + original_child.callback(original_callback) + replacement_child = GraphDefinition() + replacement_child.callback(replacement_callback) + original = { + "child": original_child, + "callback": original_callback, + } + replacement = { + "child": replacement_child, + "callback": replacement_callback, + } + + graph_def = GraphDefinition() + node = graph_def.embed(original_child) + invalid_child = node.child_graph + + def update(expected): + node.update(expected["child"]) + + def assert_current(expected): + callback_node = next( + child_node for child_node in node.child_graph.nodes() if isinstance(child_node, HostCallbackNode) + ) + assert callback_node.callback is expected["callback"] + + def assert_exec_uses(graph, expected): + called.clear() + stream = device.create_stream() + graph.launch(stream) + stream.sync() + assert called == [expected["callback"]] + + 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(invalid_child), + invalid_exception=CUDAError, + invalid_argument_update=lambda: node.update(object()), + cleanup=node.destroy, + ) + + +@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"), + pytest.param(_kernel_config_case, id="kernel-config"), + pytest.param(_kernel_args_case, id="kernel-args"), + pytest.param(_kernel_function_case, id="kernel-function"), + pytest.param(_child_graph_case, id="child-graph"), + ] +) +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_memcpy_update_rejects_unsupported_descriptor(init_cuda): + if driver_version() < (12, 2, 0): + pytest.skip("individual graph node updates require CUDA 12.2+") + + memory_resource = LegacyPinnedMemoryResource() + src = memory_resource.allocate(8) + dst = memory_resource.allocate(8) + graph_def = GraphDefinition() + node = graph_def.memcpy(dst, src, 4) + + # cuda.core cannot construct this descriptor, but imported graphs can + # contain one; use cuda.bindings to exercise that rejection path. + params = driver.CUDA_MEMCPY3D() + params.srcXInBytes = 1 + params.srcMemoryType = driver.CUmemorytype.CU_MEMORYTYPE_HOST + params.srcHost = int(src.handle) + params.srcPitch = 4 + params.srcHeight = 2 + params.dstMemoryType = driver.CUmemorytype.CU_MEMORYTYPE_HOST + params.dstHost = int(dst.handle) + params.dstPitch = 4 + params.dstHeight = 2 + params.WidthInBytes = 2 + params.Height = 2 + params.Depth = 1 + handle_return(driver.cuGraphMemcpyNodeSetParams(node.handle, params)) + + with pytest.raises(NotImplementedError, match="multidimensional"): + node.update(size=3) + + unchanged = handle_return(driver.cuGraphMemcpyNodeGetParams(node.handle)) + assert unchanged.srcXInBytes == 1 + assert unchanged.WidthInBytes == 2 + assert unchanged.Height == 2 + + node.destroy() + src.close() + dst.close() + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_kernel_update_rejects_unsupported_config(init_cuda): + if driver_version() < (12, 2, 0): + pytest.skip("individual graph node updates require CUDA 12.2+") + + kernel = compile_common_kernels().get_kernel("empty_kernel") + graph_def = GraphDefinition() + node = graph_def.launch(LaunchConfig(grid=1, block=1), kernel) + + clustered = LaunchConfig(grid=1, block=1) + clustered.cluster = (1, 1, 1) + with pytest.raises(NotImplementedError, match="clustered or cooperative"): + node.update(config=clustered) + with pytest.raises(NotImplementedError, match="clustered or cooperative"): + graph_def.launch(clustered, kernel) + + cooperative = LaunchConfig(grid=1, block=1) + cooperative.is_cooperative = True + with pytest.raises(NotImplementedError, match="clustered or cooperative"): + node.update(config=cooperative) + with pytest.raises(NotImplementedError, match="clustered or cooperative"): + graph_def.launch(cooperative, kernel) + + node.destroy() + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_partial_memory_updates_are_keyword_only(init_cuda): + memory_resource = LegacyPinnedMemoryResource() + src = memory_resource.allocate(4) + dst = memory_resource.allocate(4) + graph_def = GraphDefinition() + memset_node = graph_def.memset(dst, 0, 4) + memcpy_node = graph_def.memcpy(dst, src, 4) + + with pytest.raises(TypeError): + memset_node.update(dst) + with pytest.raises(TypeError): + memcpy_node.update(dst) + + memset_node.destroy() + memcpy_node.destroy() + src.close() + dst.close() + + +@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() diff --git a/cuda_core/tests/helpers/graph_kernels.py b/cuda_core/tests/helpers/graph_kernels.py index 54caedd165c..d08837585fe 100644 --- a/cuda_core/tests/helpers/graph_kernels.py +++ b/cuda_core/tests/helpers/graph_kernels.py @@ -19,15 +19,24 @@ def compile_common_kernels(): Returns a module with: - empty_kernel: does nothing - add_one: increments an int pointer by 1 + - write_launch_dims: encodes the launch dimensions in an int """ code = """ __global__ void empty_kernel() {} __global__ void add_one(int *a) { *a += 1; } + __global__ void write_launch_dims(int *a) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + *a = gridDim.x * 1000 + blockDim.x; + } + } """ arch = "".join(f"{i}" for i in Device().compute_capability) program_options = ProgramOptions(std="c++17", arch=f"sm_{arch}") prog = Program(code, code_type="c++", options=program_options) - mod = prog.compile("cubin", name_expressions=("empty_kernel", "add_one")) + mod = prog.compile( + "cubin", + name_expressions=("empty_kernel", "add_one", "write_launch_dims"), + ) return mod diff --git a/cuda_core/tests/test_green_context.py b/cuda_core/tests/test_green_context.py index 693dffacdc7..a0fd8eabf98 100644 --- a/cuda_core/tests/test_green_context.py +++ b/cuda_core/tests/test_green_context.py @@ -21,7 +21,9 @@ WorkqueueResourceOptions, launch, ) -from cuda.core._utils.cuda_utils import CUDAError +from cuda.core._utils.cuda_utils import CUDAError, driver, handle_return +from cuda.core._utils.version import binding_version, driver_version +from cuda.core.graph import GraphDefinition from cuda.core.typing import WorkqueueSharingScopeType # --------------------------------------------------------------------------- @@ -160,6 +162,38 @@ def _use_green_ctx(dev, ctx): dev.set_current(prev) +@pytest.mark.agent_authored(model="gpt-5.6") +def test_memory_node_updates_preserve_green_context( + init_cuda, + green_ctx, +): + if driver_version() < (13, 2, 0) or binding_version() < (13, 2, 0): + pytest.skip("generic graph node parameter queries require CUDA 13.2+") + + memory_resource = LegacyPinnedMemoryResource() + src = memory_resource.allocate(4) + dst = memory_resource.allocate(4) + with _use_green_ctx(init_cuda, green_ctx): + graph_def = GraphDefinition() + memset_node = graph_def.memset(dst, 0, 4) + memcpy_node = graph_def.memcpy(dst, src, 4) + original_memset = handle_return(driver.cuGraphNodeGetParams(memset_node.handle)) + original_memcpy = handle_return(driver.cuGraphNodeGetParams(memcpy_node.handle)) + + memset_node.update(value=1) + memcpy_node.update(size=2) + updated_memset = handle_return(driver.cuGraphNodeGetParams(memset_node.handle)) + updated_memcpy = handle_return(driver.cuGraphNodeGetParams(memcpy_node.handle)) + + assert int(updated_memset.memset.ctx) == int(original_memset.memset.ctx) + assert int(updated_memcpy.memcpy.copyCtx) == int(original_memcpy.memcpy.copyCtx) + + memset_node.destroy() + memcpy_node.destroy() + src.close() + dst.close() + + # --------------------------------------------------------------------------- # Construction / type tests # ---------------------------------------------------------------------------