# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

"""GraphNode base class — factory, properties, and builder methods."""

from __future__ import annotations

from collections.abc import Iterable
from typing import TYPE_CHECKING

from libc.stddef cimport size_t
from libc.stdint cimport uintptr_t
from libc.string cimport memset as c_memset

from libcpp.vector cimport vector

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._subclasses cimport (
    AllocNode,
    ChildGraphNode,
    ConditionalNode,
    EmptyNode,
    EventRecordNode,
    EventWaitNode,
    FreeNode,
    HostCallbackNode,
    IfElseNode,
    IfNode,
    KernelNode,
    MemcpyNode,
    MemsetNode,
    SwitchNode,
    WhileNode,
)
from cuda.core._resource_handles cimport (
    GraphHandle,
    GraphNodeHandle,
    OpaqueHandle,
    PreparedAttachment,
    as_cu,
    as_intptr,
    as_py,
    create_child_graph_handle,
    create_graph_node_handle,
    graph_clone_attachments,
    graph_commit_attachment,
    graph_node_get_graph,
    graph_prepare_attachment,
    invalidate_child_graph_state,
    invalidate_graph_node,
    make_opaque_py,
)
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value

from cuda.core.graph._host_callback cimport (
    _resolve_host_callback,
)

import ctypes as ct
import weakref

from cuda.core.graph._adjacency_set_proxy import AdjacencySetProxy
from cuda.core._utils.validators import check_str_enum
from cuda.core._utils.cuda_utils import driver
from cuda.core.typing import GraphMemoryType

if TYPE_CHECKING:
    from cuda.core._device import Device

__all__ = ['GraphNode']

# See _cpp/REGISTRY_DESIGN.md (Level 2: Resource Handle -> Python Object)
_node_registry: weakref.WeakValueDictionary[int, GraphNode] = weakref.WeakValueDictionary()


cdef inline GraphNode _registered(GraphNode n):
    return _node_registry.setdefault(<uintptr_t>n._h_node.get(), n)


cdef class GraphNode:
    """A node in a graph definition.

    Nodes are created by calling builder methods on GraphDefinition (for
    entry-point nodes with no dependencies) or on other Nodes (for
    nodes that depend on a predecessor).
    """

    @staticmethod
    cdef GraphNode _create(GraphHandle h_graph, cydriver.CUgraphNode node):
        """Factory: dispatch to the right subclass based on node type."""
        return GN_create(h_graph, node)

    def __repr__(self) -> str:
        cdef cydriver.CUgraphNode node = as_cu(self._h_node)
        if node == NULL:
            return "<GraphNode entry>"
        return f"<GraphNode handle=0x{<uintptr_t>node:x}>"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, GraphNode):
            return NotImplemented
        cdef GraphNode o = <GraphNode>other
        cdef GraphHandle self_graph = graph_node_get_graph(self._h_node)
        cdef GraphHandle other_graph = graph_node_get_graph(o._h_node)
        return (as_intptr(self._h_node) == as_intptr(o._h_node)
                and as_intptr(self_graph) == as_intptr(other_graph))

    def __hash__(self) -> int:
        cdef GraphHandle g = graph_node_get_graph(self._h_node)
        return hash((as_intptr(self._h_node), as_intptr(g)))

    @property
    def type(self) -> driver.CUgraphNodeType | None:
        """Return the CUDA graph node type.

        Returns
        -------
        CUgraphNodeType or None
            The node type enum value, or None for the entry node.
        """
        cdef cydriver.CUgraphNode node = as_cu(self._h_node)
        if node == NULL:
            return None
        cdef cydriver.CUgraphNodeType node_type
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphNodeGetType(node, &node_type))
        return driver.CUgraphNodeType(<int>node_type)

    @property
    def graph(self) -> GraphDefinition:
        """Return the GraphDefinition this node belongs to."""
        return GraphDefinition._from_handle(graph_node_get_graph(self._h_node))

    @property
    def handle(self) -> driver.CUgraphNode:
        """Return the underlying driver CUgraphNode handle.

        Returns None for the entry node.
        """
        return as_py(self._h_node)

    @property
    def is_valid(self) -> bool:
        """Whether this node is valid (not destroyed).

        Returns ``False`` after :meth:`destroy` has been called.
        """
        return as_intptr(self._h_node) != 0

    def destroy(self) -> None:
        """Destroy this node and remove all its edges from the parent graph.

        After this call, :attr:`is_valid` returns ``False`` and the node
        cannot be re-added to any graph.  Safe to call on an
        already-destroyed node (no-op).
        """
        cdef cydriver.CUgraphNode node = as_cu(self._h_node)
        cdef PreparedAttachment prepared
        if node == NULL:
            return

        cdef GraphHandle h_graph = graph_node_get_graph(self._h_node)
        # Allocate the cleanup transaction before asking CUDA to destroy the
        # node. A failed CUDA call leaves metadata and wrappers unchanged.
        HANDLE_RETURN(graph_prepare_attachment(
            h_graph, OpaqueHandle(), OpaqueHandle(), &prepared))
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphDestroyNode(node))

        # Publish attachment removal before invalidating graph and node aliases.
        cdef cydriver.CUresult cleanup_status = graph_commit_attachment(prepared, node)
        invalidate_child_graph_state(h_graph, node)
        _node_registry.pop(<uintptr_t>self._h_node.get(), None)
        invalidate_graph_node(self._h_node)
        HANDLE_RETURN(cleanup_status)

    @property
    def pred(self) -> AdjacencySetProxy:
        """A mutable set-like view of this node's predecessors."""
        return AdjacencySetProxy(self, False)

    @pred.setter
    def pred(self, value: Iterable[GraphNode]) -> None:
        p = AdjacencySetProxy(self, False)
        p.clear()
        p.update(value)

    @property
    def succ(self) -> AdjacencySetProxy:
        """A mutable set-like view of this node's successors."""
        return AdjacencySetProxy(self, True)

    @succ.setter
    def succ(self, value: Iterable[GraphNode]) -> None:
        s = AdjacencySetProxy(self, True)
        s.clear()
        s.update(value)

    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
            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.

        Parameters
        ----------
        config : LaunchConfig
            Launch configuration (grid, block, shared memory, etc.)
        kernel : Kernel
            The kernel to launch.
        *args
            Kernel arguments.

        Returns
        -------
        KernelNode
            A new KernelNode representing the kernel launch.
        """
        return GN_launch(self, config, <Kernel>kernel, ParamHolder(args))

    def join(self, *nodes: GraphNode) -> EmptyNode:
        """Create an empty node that depends on this node and all given nodes.

        This is used to synchronize multiple branches of execution.

        Parameters
        ----------
        *nodes : GraphNode
            Additional nodes to depend on.

        Returns
        -------
        EmptyNode
            A new EmptyNode that depends on all input nodes.
        """
        return GN_join(self, nodes)

    def allocate(self, size_t size, *, device: Device | int | None = None,
                 memory_type: GraphMemoryType = GraphMemoryType.DEVICE,
                 peer_access: list[Device | int] | None = None) -> AllocNode:
        """Add a memory allocation node depending on this node.

        Parameters
        ----------
        size : int
            Number of bytes to allocate.
        device : int or Device, optional
            The device on which to allocate memory. If None (default),
            uses the current CUDA context's device.
        memory_type : GraphMemoryType or str, optional
            Type of memory to allocate. One of:

            - ``GraphMemoryType.DEVICE`` (default): Pinned device memory,
              optimal for GPU kernels.
            - ``GraphMemoryType.HOST``: Pinned host memory, accessible from
              both host and device. Useful for graphs containing host
              callback nodes. Note: may not be supported on all
              systems/drivers.
            - ``GraphMemoryType.MANAGED``: Managed/unified memory that
              automatically migrates between host and device. Useful for
              mixed host/device access patterns.

        peer_access : list of int or Device, optional
            List of devices that should have read-write access to the
            allocated memory. If None (default), only the allocating
            device has access.

        Returns
        -------
        AllocNode
            A new AllocNode representing the allocation. Access the allocated
            device pointer via the dptr property.

        Notes
        -----
        IPC (inter-process communication) is not supported for graph
        memory allocation nodes per CUDA documentation.
        """
        return GN_alloc(self, size, device, memory_type, peer_access)

    def deallocate(self, dptr: int) -> FreeNode:
        """Add a memory free node depending on this node.

        Parameters
        ----------
        dptr : int
            Device pointer to free (typically from AllocNode.dptr).

        Returns
        -------
        FreeNode
            A new FreeNode representing the free operation.
        """
        return GN_free(self, <cydriver.CUdeviceptr>dptr)

    def memset(
        self,
        dst: Buffer | int,
        value,
        size_t width,
        size_t height=1,
        size_t pitch=0,
        *,
        dst_owner=None,
    ) -> MemsetNode:
        """Add a memset node depending on this node.

        .. 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.

        Parameters
        ----------
        dst : Buffer or int
            Destination. When ``dst`` is a :class:`Buffer`, the underlying
            allocation is retained for the graph's lifetime. A raw pointer
            (``int``) is used as-is; the caller must keep the underlying memory
            alive, or supply ``dst_owner`` to have the graph retain it.
        value : int or buffer-protocol object
            Fill value. int for 1-byte fill (range [0, 256)),
            or buffer-protocol object of 1, 2, or 4 bytes.
        width : int
            Width of the row in elements.
        height : int, optional
            Number of rows (default 1).
        pitch : int, optional
            Pitch of destination in bytes (default 0, unused if height is 1).
        dst_owner : object, optional
            Object retained for the graph's lifetime when ``dst`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation,
            not the wrapper. Must not be passed when ``dst`` is a :class:`Buffer`.

        Returns
        -------
        MemsetNode
            A new MemsetNode representing the memset operation.

        Raises
        ------
        ValueError
            If ``dst_owner`` is given together with a :class:`Buffer` ``dst``.
        """
        cdef cydriver.CUdeviceptr c_dst
        cdef unsigned int val
        cdef unsigned int elem_size
        cdef OpaqueHandle dst_attachment_owner = _resolve_memcpy_operand(
            dst, dst_owner, "dst", &c_dst)
        val, elem_size = _parse_fill_value(value)
        return GN_memset(
            self, c_dst, dst_attachment_owner,
            val, elem_size, width, height, pitch)

    def memcpy(
        self,
        dst: Buffer | int,
        src: Buffer | int,
        size_t size,
        *,
        dst_owner=None,
        src_owner=None,
    ) -> MemcpyNode:
        """Add a memcpy node depending on this node.

        Copies ``size`` bytes from ``src`` to ``dst``. Memory types are
        auto-detected via the driver, so both device and pinned host
        pointers are supported.

        .. 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.

        Parameters
        ----------
        dst : Buffer or int
            Destination (device or pinned host). When a :class:`Buffer` is given,
            the underlying allocation is retained for the graph's lifetime. A raw
            pointer (``int``) is used as-is; the caller must keep the underlying
            memory alive, or supply ``dst_owner`` to have the graph retain it.
        src : Buffer or int
            Source (device or pinned host). Same retention rules as ``dst``;
            use ``src_owner`` for a raw pointer.
        size : int
            Number of bytes to copy.
        dst_owner : object, optional
            Object retained for the graph's lifetime when ``dst`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation.
            Must not be passed when ``dst`` is a :class:`Buffer`.
        src_owner : object, optional
            Object retained for the graph's lifetime when ``src`` is a raw
            pointer. A :class:`Buffer` owner retains its underlying allocation.
            Must not be passed when ``src`` is a :class:`Buffer`.

        Returns
        -------
        MemcpyNode
            A new MemcpyNode representing the copy operation.

        Raises
        ------
        ValueError
            If ``dst_owner`` or ``src_owner`` is given together with a
            :class:`Buffer` ``dst`` or ``src`` respectively.
        """
        cdef cydriver.CUdeviceptr c_dst
        cdef cydriver.CUdeviceptr c_src
        cdef OpaqueHandle dst_attachment_owner = _resolve_memcpy_operand(
            dst, dst_owner, "dst", &c_dst)
        cdef OpaqueHandle src_attachment_owner = _resolve_memcpy_operand(
            src, src_owner, "src", &c_src)
        return GN_memcpy(
            self, c_dst, dst_attachment_owner,
            c_src, src_attachment_owner, size)

    def embed(self, child: GraphDefinition) -> ChildGraphNode:
        """Add a child graph node depending on this node.

        Embeds a clone of the given graph definition as a sub-graph node.
        The child graph must not contain allocation, free, or conditional
        nodes.

        Parameters
        ----------
        child : GraphDefinition
            The graph definition to embed (will be cloned).

        Returns
        -------
        ChildGraphNode
            A new ChildGraphNode representing the embedded sub-graph.
        """
        return GN_embed(self, <GraphDefinition>child)

    def record(self, event: Event) -> EventRecordNode:
        """Add an event record node depending on this node.

        Parameters
        ----------
        event : Event
            The event to record.

        Returns
        -------
        EventRecordNode
            A new EventRecordNode representing the event record operation.
        """
        return GN_record_event(self, <Event>event)

    def wait(self, event: Event) -> EventWaitNode:
        """Add an event wait node depending on this node.

        Parameters
        ----------
        event : Event
            The event to wait for.

        Returns
        -------
        EventWaitNode
            A new EventWaitNode representing the event wait operation.
        """
        return GN_wait_event(self, <Event>event)

    def callback(self, fn, *, user_data=None) -> object:
        """Add a host callback node depending on this node.

        The callback runs on the host CPU when the graph reaches this node.
        Two modes are supported:

        - **Python callable**: Pass any callable. The GIL is acquired
          automatically. The callable must take no arguments; use closures
          or ``functools.partial`` to bind state.
        - **ctypes function pointer**: The function receives a single
          ``void*`` argument (the ``user_data``), and the caller must keep
          the ctypes wrapper alive for the lifetime of the graph. Its
          declared prototype must match the driver's ``CUhostFn``
          (``void (*)(void*)``): ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``,
          or ``ctypes.WINFUNCTYPE(None, ctypes.c_void_p)`` on Windows.

        .. 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.

        Parameters
        ----------
        fn : callable or ctypes function pointer
            The callback function.
        user_data : int or bytes-like, optional
            Only for ctypes function pointers. If ``int``, passed as a raw
            pointer (caller manages lifetime). If bytes-like, the data is
            copied and its lifetime is tied to the graph.

        Returns
        -------
        HostCallbackNode
            A new HostCallbackNode representing the callback.

        Raises
        ------
        TypeError
            If ``fn`` is a ctypes function pointer whose declared prototype
            does not match ``CUhostFn``.
        ValueError
            If ``user_data`` is given for a Python callable.
        """
        return GN_callback(self, fn, user_data)

    def if_then(self, condition: GraphCondition) -> IfNode:
        """Add an if-conditional node depending on this node.

        The body graph executes only when the condition evaluates to
        a non-zero value at runtime.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        IfNode
            A new IfNode with one branch accessible via ``.then``.
        """
        return _make_conditional_node(
            self, condition,
            cydriver.CU_GRAPH_COND_TYPE_IF, 1, IfNode)

    def if_else(self, condition: GraphCondition) -> IfElseNode:
        """Add an if-else conditional node depending on this node.

        Two body graphs: the first executes when the condition is
        non-zero, the second when it is zero.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        IfElseNode
            A new IfElseNode with branches accessible via
            ``.then`` and ``.else_``.
        """
        return _make_conditional_node(
            self, condition,
            cydriver.CU_GRAPH_COND_TYPE_IF, 2, IfElseNode)

    def while_loop(self, condition: GraphCondition) -> WhileNode:
        """Add a while-loop conditional node depending on this node.

        The body graph executes repeatedly while the condition
        evaluates to a non-zero value.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.

        Returns
        -------
        WhileNode
            A new WhileNode with body accessible via ``.body``.
        """
        return _make_conditional_node(
            self, condition,
            cydriver.CU_GRAPH_COND_TYPE_WHILE, 1, WhileNode)

    def switch(self, condition: GraphCondition, unsigned int count) -> SwitchNode:
        """Add a switch conditional node depending on this node.

        The condition value selects which branch to execute. If the
        value is out of range, no branch executes.

        Parameters
        ----------
        condition : GraphCondition
            GraphCondition from :meth:`GraphDefinition.create_condition`.
        count : int
            Number of switch cases (branches).

        Returns
        -------
        SwitchNode
            A new SwitchNode with branches accessible via ``.branches``.
        """
        return _make_conditional_node(
            self, condition,
            cydriver.CU_GRAPH_COND_TYPE_SWITCH, count, SwitchNode)


cdef inline ConditionalNode _make_conditional_node(
        GraphNode pred,
        GraphCondition condition,
        cydriver.CUgraphConditionalNodeType cond_type,
        unsigned int size,
        type node_cls):
    if not isinstance(condition, GraphCondition):
        raise TypeError(
            f"condition must be a GraphCondition object (from "
            f"GraphDefinition.create_condition()), got {type(condition).__name__}")
    cdef cydriver.CUgraphNodeParams params
    cdef cydriver.CUgraphNode new_node = NULL

    c_memset(&params, 0, sizeof(params))
    params.type = cydriver.CU_GRAPH_NODE_TYPE_CONDITIONAL
    params.conditional.handle = condition._c_handle
    params.conditional.type = cond_type
    params.conditional.size = size

    cdef cydriver.CUcontext ctx = NULL
    cdef GraphHandle h_graph = graph_node_get_graph(pred._h_node)
    cdef cydriver.CUgraphNode pred_node = as_cu(pred._h_node)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    with nogil:
        HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
    params.conditional.ctx = ctx

    with nogil:
        IF CUDA_CORE_BUILD_MAJOR >= 13:
            HANDLE_RETURN(cydriver.cuGraphAddNode(
                &new_node, as_cu(h_graph), deps, NULL, num_deps, &params))
        ELSE:
            HANDLE_RETURN(cydriver.cuGraphAddNode(
                &new_node, as_cu(h_graph), deps, num_deps, &params))

    cdef list branch_list = []
    cdef unsigned int i
    cdef cydriver.CUgraph bg
    cdef GraphHandle h_branch
    for i in range(size):
        bg = params.conditional.phGraph_out[i]
        h_branch = create_child_graph_handle(bg, h_graph, new_node)
        branch_list.append(GraphDefinition._from_handle(h_branch))
    cdef tuple branches = tuple(branch_list)

    cdef ConditionalNode n = node_cls.__new__(node_cls)
    n._h_node = create_graph_node_handle(new_node, h_graph)
    n._condition = condition
    n._cond_type = cond_type
    n._branches = branches

    return _registered(n)

cdef inline GraphNode GN_create(GraphHandle h_graph, cydriver.CUgraphNode node):
    cdef GraphNodeHandle h_node = create_graph_node_handle(node, h_graph)

    # Sentinel: virtual node to represent the graph entry point.
    if node == NULL:
        n = GraphNode.__new__(GraphNode)
        (<GraphNode>n)._h_node = h_node
        return n

    # Return a registered object or create and register a new one.
    registered = _node_registry.get(<uintptr_t>h_node.get())
    if registered is not None:
        return <GraphNode>registered
    else:
        return _registered(GN_create_impl(h_node))


cdef inline GraphNode GN_create_impl(GraphNodeHandle h_node):
    cdef cydriver.CUgraphNodeType node_type
    with nogil:
        HANDLE_RETURN(cydriver.cuGraphNodeGetType(as_cu(h_node), &node_type))

    if node_type == cydriver.CU_GRAPH_NODE_TYPE_EMPTY:
        return EmptyNode._create_impl(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_KERNEL:
        return KernelNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_MEM_ALLOC:
        return AllocNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_MEM_FREE:
        return FreeNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_MEMSET:
        return MemsetNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_MEMCPY:
        return MemcpyNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_GRAPH:
        return ChildGraphNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_EVENT_RECORD:
        return EventRecordNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_WAIT_EVENT:
        return EventWaitNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_HOST:
        return HostCallbackNode._create_from_driver(h_node)
    elif node_type == cydriver.CU_GRAPH_NODE_TYPE_CONDITIONAL:
        return ConditionalNode._create_from_driver(h_node)
    else:
        n = GraphNode.__new__(GraphNode)
        (<GraphNode>n)._h_node = h_node
        return n


cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker, ParamHolder ker_args):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef OpaqueHandle 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

    cdef cydriver.CUDA_KERNEL_NODE_PARAMS node_params = cydriver.CUDA_KERNEL_NODE_PARAMS(
        kern=as_cu(ker._h_kernel),
        func=<cydriver.CUfunction>NULL,
        gridDimX=conf.grid[0],
        gridDimY=conf.grid[1],
        gridDimZ=conf.grid[2],
        blockDimX=conf.block[0],
        blockDimY=conf.block[1],
        blockDimZ=conf.block[2],
        sharedMemBytes=conf.shmem_size,
        kernelParams=<void**><uintptr_t>(ker_args.ptr),
        extra=NULL,
        ctx=<cydriver.CUcontext>NULL,
    )

    # Keep the kernel and argument objects alive because CUDA copies argument
    # values but does not retain the resources they reference.
    cdef OpaqueHandle kernel_owner = ker._h_kernel
    kernel_args = ker_args.kernel_args
    if kernel_args is not None:
        args_owner = make_opaque_py(kernel_args)
    HANDLE_RETURN(graph_prepare_attachment(
        h_graph, kernel_owner, args_owner, &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddKernelNode(
            &new_node, as_cu(h_graph), deps, num_deps, &node_params))

    HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    return _registered(KernelNode._create_with_params(
        create_graph_node_handle(new_node, h_graph),
        conf.grid, conf.block, conf.shmem_size,
        ker._h_kernel))


cdef inline EmptyNode GN_join(GraphNode self, tuple nodes):
    cdef vector[cydriver.CUgraphNode] deps
    cdef cydriver.CUgraphNode new_node = NULL
    cdef GraphHandle h_graph = graph_node_get_graph(self._h_node)
    cdef GraphNode other
    cdef cydriver.CUgraphNode* deps_ptr = NULL
    cdef size_t num_deps = 0
    cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node)

    if pred_node != NULL:
        deps.push_back(pred_node)
    for other in nodes:
        if as_cu((<GraphNode>other)._h_node) != NULL:
            deps.push_back(as_cu((<GraphNode>other)._h_node))

    num_deps = deps.size()
    if num_deps > 0:
        deps_ptr = deps.data()

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddEmptyNode(
            &new_node, as_cu(h_graph), deps_ptr, num_deps))

    return _registered(EmptyNode._create_impl(create_graph_node_handle(new_node, h_graph)))


cdef inline AllocNode GN_alloc(GraphNode self, size_t size, object device,
                               object memory_type, object peer_access):
    cdef int device_id
    cdef cydriver.CUdevice dev

    if device is None:
        with nogil:
            HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
        device_id = <int>dev
    else:
        device_id = getattr(device, 'device_id', device)

    cdef cydriver.CUDA_MEM_ALLOC_NODE_PARAMS alloc_params
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    cdef vector[cydriver.CUmemAccessDesc] access_descs
    cdef int peer_id
    cdef list peer_ids = []

    if peer_access is not None:
        for peer_dev in peer_access:
            peer_id = getattr(peer_dev, 'device_id', peer_dev)
            peer_ids.append(peer_id)
            access_descs.push_back(cydriver.CUmemAccessDesc_st(
                cydriver.CUmemLocation_st(
                    cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE,
                    peer_id
                ),
                cydriver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
            ))

    cdef str memory_type_str = "device" if memory_type is None else str(memory_type)

    c_memset(&alloc_params, 0, sizeof(alloc_params))
    alloc_params.poolProps.handleTypes = cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE
    alloc_params.bytesize = size

    check_str_enum(memory_type_str, GraphMemoryType)
    if memory_type_str == "device":
        alloc_params.poolProps.allocType = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
        alloc_params.poolProps.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
        alloc_params.poolProps.location.id = device_id
    elif memory_type_str == "host":
        alloc_params.poolProps.allocType = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
        alloc_params.poolProps.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
        alloc_params.poolProps.location.id = 0
    elif memory_type_str == "managed":
        IF CUDA_CORE_BUILD_MAJOR >= 13:
            alloc_params.poolProps.allocType = cydriver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_MANAGED
            alloc_params.poolProps.location.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
            alloc_params.poolProps.location.id = device_id
        ELSE:
            raise ValueError("memory_type='managed' requires CUDA 13.0 or later")

    if access_descs.size() > 0:
        alloc_params.accessDescs = access_descs.data()
        alloc_params.accessDescCount = access_descs.size()

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddMemAllocNode(
            &new_node, as_cu(h_graph), deps, num_deps, &alloc_params))

    return _registered(AllocNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), alloc_params.dptr, size,
        device_id, memory_type_str, tuple(peer_ids)))


cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddMemFreeNode(
            &new_node, as_cu(h_graph), deps, num_deps, c_dptr))

    return _registered(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))


cdef inline OpaqueHandle _buffer_attachment_owner(Buffer buf, str label):
    """Copy a Buffer's device-pointer handle into an attachment owner."""
    if not buf._h_ptr:
        raise ValueError(f"{label} Buffer has no active allocation")
    # The local is required: Cython permits the DevicePtrHandle -> OpaqueHandle
    # conversion on assignment, but not directly in a return statement.
    cdef OpaqueHandle attachment_owner = buf._h_ptr
    return attachment_owner


cdef inline OpaqueHandle _resolve_memcpy_operand(
        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
    pointer is written to ``out_ptr``. For a :class:`Buffer` operand, returns an
    owner that retains the underlying allocation (not the wrapper). For a raw
    pointer, returns an owner built from ``owner`` (or an empty handle when
    ``owner`` is ``None``).

    Raises
    ------
    ValueError
        If ``operand`` is a :class:`Buffer` and ``owner`` is not ``None``.
        If a :class:`Buffer` operand or ``*_owner`` has no active allocation.
    """
    cdef Buffer buf
    cdef OpaqueHandle attachment_owner

    if isinstance(operand, Buffer):
        if owner is not None:
            raise ValueError(
                f"{side}_owner cannot be used when {side} is a Buffer"
            )
        buf = operand
        attachment_owner = _buffer_attachment_owner(buf, side)
        out_ptr[0] = as_cu(buf._h_ptr)
        return attachment_owner
    out_ptr[0] = <cydriver.CUdeviceptr><uintptr_t>operand
    if owner is None:
        return OpaqueHandle()
    if isinstance(owner, Buffer):
        return _buffer_attachment_owner(owner, f"{side}_owner")
    return make_opaque_py(owner)


cdef inline MemsetNode GN_memset(
        GraphNode self, cydriver.CUdeviceptr c_dst, OpaqueHandle dst_owner,
        unsigned int val, unsigned int elem_size,
        size_t width, size_t height, size_t pitch):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef PreparedAttachment prepared

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    cdef cydriver.CUcontext ctx = NULL
    with nogil:
        HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))

    cdef cydriver.CUDA_MEMSET_NODE_PARAMS memset_params = cydriver.CUDA_MEMSET_NODE_PARAMS(
        dst=c_dst,
        pitch=pitch,
        value=val,
        elementSize=elem_size,
        width=width,
        height=height,
    )

    if dst_owner:
        HANDLE_RETURN(graph_prepare_attachment(
            h_graph, dst_owner, OpaqueHandle(), &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddMemsetNode(
            &new_node, as_cu(h_graph), deps, num_deps,
            &memset_params, ctx))

    if dst_owner:
        HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    return _registered(MemsetNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), c_dst,
        val, elem_size, width, height, pitch))


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(
            &memory_type,
            cydriver.CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
            ptr)
    if ret != cydriver.CUDA_SUCCESS and ret != cydriver.CUDA_ERROR_INVALID_VALUE:
        HANDLE_RETURN(ret)
    return <cydriver.CUmemorytype>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 = <void*><uintptr_t>src
    else:
        params.srcDevice = src
    if dst_type[0] == cydriver.CU_MEMORYTYPE_HOST:
        params.dstHost = <void*><uintptr_t>dst
    else:
        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, &params, &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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef PreparedAttachment prepared

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    cdef cydriver.CUcontext ctx = NULL
    if dst_owner or src_owner:
        HANDLE_RETURN(graph_prepare_attachment(
            h_graph, dst_owner, src_owner, &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuCtxGetCurrent(&ctx))
        HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode(
            &new_node, as_cu(h_graph), deps, num_deps, &params, ctx))

    if dst_owner or src_owner:
        HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    return _registered(MemcpyNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), c_dst, c_src, size,
        c_dst_type, c_src_type))


cdef inline ChildGraphNode GN_embed(GraphNode self, GraphDefinition child_def):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef cydriver.CUresult rollback_status

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddChildGraphNode(
            &new_node, as_cu(h_graph), deps, num_deps, as_cu(child_def._h_graph)))

    cdef cydriver.CUgraph embedded_graph = NULL
    cdef GraphHandle h_embedded
    try:
        with nogil:
            HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph(
                new_node, &embedded_graph))
        h_embedded = create_child_graph_handle(
            embedded_graph, h_graph, new_node)
        HANDLE_RETURN(graph_clone_attachments(
            h_embedded, child_def._h_graph))
    except:
        with nogil:
            rollback_status = cydriver.cuGraphDestroyNode(new_node)
        if rollback_status == cydriver.CUDA_SUCCESS:
            invalidate_child_graph_state(h_graph, new_node)
        raise

    return _registered(ChildGraphNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), h_embedded))


cdef inline EventRecordNode GN_record_event(GraphNode self, Event ev):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef PreparedAttachment prepared

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    cdef OpaqueHandle owner = ev._h_event
    HANDLE_RETURN(graph_prepare_attachment(
        h_graph, owner, OpaqueHandle(), &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddEventRecordNode(
            &new_node, as_cu(h_graph), deps, num_deps, as_cu(ev._h_event)))

    HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    return _registered(EventRecordNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), ev._h_event))


cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef PreparedAttachment prepared

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    cdef OpaqueHandle owner = ev._h_event
    HANDLE_RETURN(graph_prepare_attachment(
        h_graph, owner, OpaqueHandle(), &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddEventWaitNode(
            &new_node, as_cu(h_graph), deps, num_deps, as_cu(ev._h_event)))

    HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    return _registered(EventWaitNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), ev._h_event))


cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_data):
    cdef cydriver.CUDA_HOST_NODE_PARAMS node_params
    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)
    cdef cydriver.CUgraphNode* deps = NULL
    cdef size_t num_deps = 0
    cdef OpaqueHandle fn_owner, data_owner
    cdef PreparedAttachment prepared

    if pred_node != NULL:
        deps = &pred_node
        num_deps = 1

    _resolve_host_callback(
        fn, user_data, &node_params.fn, &node_params.userData,
        &fn_owner, &data_owner)
    HANDLE_RETURN(graph_prepare_attachment(
        h_graph, fn_owner, data_owner, &prepared))

    with nogil:
        HANDLE_RETURN(cydriver.cuGraphAddHostNode(
            &new_node, as_cu(h_graph), deps, num_deps, &node_params))

    HANDLE_RETURN(graph_commit_attachment(prepared, new_node))

    cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
    return _registered(HostCallbackNode._create_with_params(
        create_graph_node_handle(new_node, h_graph), callable_obj,
        node_params.fn, node_params.userData))
