cuda.core: add explicit dst/src_owner for graph memcpy/memset

Keyword-only *_owner args retain arbitrary objects for raw pointer
operands; Buffer+owner combinations are rejected. Strengthen owner tests
with weakref retention checks and add src_owner rejection test.

Author: Andy-Jost

cuda_core/cuda/core/graph/_graph_definition.pyx

@@ -33,6 +33,7 @@ if TYPE_CHECKING:
     from cuda.core._device import Device
     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_builder import (
         Graph,
@@ -156,17 +157,21 @@ cdef class GraphDefinition:
 
     def memset(
         self,
-        dst: int,
+        dst: Buffer | int,
         value,
         size_t width,
+        *,
         size_t height=1,
-        size_t pitch=0
+        size_t pitch=0,
+        dst_owner=None,
     ) -> MemsetNode:
         """Add an entry-point memset node (no dependencies).
 
         See :meth:`GraphNode.memset` for full documentation.
         """
-        return self._entry.memset(dst, value, width, height, pitch)
+        return self._entry.memset(
+            dst, value, width, height=height, pitch=pitch, dst_owner=dst_owner
+        )
 
     def launch(self, config: LaunchConfig, kernel: Kernel, *args) -> KernelNode:
         """Add an entry-point kernel launch node (no dependencies).
@@ -200,12 +205,12 @@ cdef class GraphDefinition:
         """
         return self._entry.join(*nodes)
 
-    def memcpy(self, dst: int, src: int, size_t size) -> MemcpyNode:
+    def memcpy(self, dst: Buffer | int, src: Buffer | int, size_t size, *, dst_owner=None, src_owner=None) -> MemcpyNode:
         """Add an entry-point memcpy node (no dependencies).
 
         See :meth:`GraphNode.memcpy` for full documentation.
         """
-        return self._entry.memcpy(dst, src, size)
+        return self._entry.memcpy(dst, src, size, dst_owner=dst_owner, src_owner=src_owner)
 
     def embed(self, child: GraphDefinition) -> ChildGraphNode:
         """Add an entry-point child graph node (no dependencies).

cuda_core/cuda/core/graph/_graph_node.pyx

@@ -284,15 +284,25 @@ cdef class GraphNode:
         """
         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) -> MemsetNode:
+    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.
 
         Parameters
         ----------
         dst : Buffer or int
-            Destination. If a :class:`Buffer` is given it is retained for the
-            graph's lifetime; if a raw device pointer (int) is given it is used
-            as-is and the caller must keep the underlying memory alive.
+            Destination. A :class:`Buffer` 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.
@@ -302,20 +312,37 @@ cdef class GraphNode:
             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. Must not be passed when ``dst`` is a :class:`Buffer`, which
+            is retained automatically.
 
         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
-        dst_owner = _resolve_ptr(dst, &c_dst)
+        dst_keepalive = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst)
         val, elem_size = _parse_fill_value(value)
-        return GN_memset(self, c_dst, dst_owner, val, elem_size, width, height, pitch)
-
-    def memcpy(self, dst: Buffer | int, src: Buffer | int, size_t size) -> MemcpyNode:
+        return GN_memset(self, c_dst, dst_keepalive, 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
@@ -325,24 +352,40 @@ cdef class GraphNode:
         Parameters
         ----------
         dst : Buffer or int
-            Destination (device or pinned host). If a :class:`Buffer` is given
-            it is retained for the graph's lifetime; a raw pointer (int) is used
-            as-is and the caller must keep the underlying memory alive.
+            Destination (device or pinned host). A :class:`Buffer` 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). Retained like ``dst`` when a Buffer.
+            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. Must not be passed when ``dst`` is a :class:`Buffer`, which
+            is retained automatically.
+        src_owner : object, optional
+            Object retained for the graph's lifetime when ``src`` is a raw
+            pointer. Must not be passed when ``src`` is a :class:`Buffer`, which
+            is retained automatically.
 
         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
-        dst_owner = _resolve_ptr(dst, &c_dst)
-        src_owner = _resolve_ptr(src, &c_src)
-        return GN_memcpy(self, c_dst, dst_owner, c_src, src_owner, size)
+        dst_keepalive = _resolve_memcpy_operand(dst, dst_owner, "dst", &c_dst)
+        src_keepalive = _resolve_memcpy_operand(src, src_owner, "src", &c_src)
+        return GN_memcpy(self, c_dst, dst_keepalive, c_src, src_keepalive, size)
 
     def embed(self, child: GraphDefinition) -> ChildGraphNode:
         """Add a child graph node depending on this node.
@@ -789,20 +832,28 @@ cdef inline FreeNode GN_free(GraphNode self, cydriver.CUdeviceptr c_dptr):
     return _registered(FreeNode._create_with_params(create_graph_node_handle(new_node, h_graph), c_dptr))
 
 
-cdef inline object _resolve_ptr(object value, cydriver.CUdeviceptr* out_ptr):
-    """Resolve a memcpy/memset operand into a device pointer.
+cdef inline object _resolve_memcpy_operand(
+        object operand, object owner, str side, cydriver.CUdeviceptr* out_ptr):
+    """Resolve a memcpy/memset operand to a device pointer and an owner.
+
+    ``operand`` is a :class:`Buffer` or a raw integer address; its device
+    pointer is written to ``out_ptr``. Returns the object to retain on the
+    graph's slot table for the graph's lifetime: the Buffer itself, or
+    ``owner`` (possibly ``None``) for a raw pointer. ``side`` is ``"dst"`` or
+    ``"src"`` and is used only to compose the error message.
 
-    ``value`` is a :class:`Buffer` or a raw integer address. For a Buffer the
-    device pointer is read from the buffer and the buffer is returned so the
-    caller can retain it on the node's slot table for the graph's lifetime. For
-    a raw integer no owner is returned and the caller is responsible for keeping
-    the underlying memory alive.
+    Raises ValueError if ``operand`` is a Buffer and ``owner`` is not None,
+    since a Buffer is already retained automatically.
     """
-    if isinstance(value, Buffer):
-        out_ptr[0] = as_cu((<Buffer>value)._h_ptr)
-        return value
-    out_ptr[0] = <cydriver.CUdeviceptr><uintptr_t>value
-    return None
+    if isinstance(operand, Buffer):
+        if owner is not None:
+            raise ValueError(
+                f"{side}_owner cannot be used when {side} is a Buffer"
+            )
+        out_ptr[0] = as_cu((<Buffer>operand)._h_ptr)
+        return operand
+    out_ptr[0] = <cydriver.CUdeviceptr><uintptr_t>operand
+    return owner
 
 
 cdef inline MemsetNode GN_memset(
@@ -837,7 +888,7 @@ cdef inline MemsetNode GN_memset(
             &new_node, as_cu(h_graph), deps, num_deps,
             &memset_params, ctx))
 
-    # Retain a Buffer destination for the graph's lifetime (slot 0).
+    # Retain the destination owner for the graph's lifetime (slot 0).
     if dst_owner is not None:
         HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, make_opaque_py(dst_owner)))
 
@@ -902,7 +953,7 @@ cdef inline MemcpyNode GN_memcpy(
         HANDLE_RETURN(cydriver.cuGraphAddMemcpyNode(
             &new_node, as_cu(h_graph), deps, num_deps, &params, ctx))
 
-    # Retain Buffer operands for the graph's lifetime (dst -> slot 0, src -> slot 1).
+    # Retain operand owners for the graph's lifetime (dst -> slot 0, src -> slot 1).
     if dst_owner is not None:
         HANDLE_RETURN(graph_set_slot(h_graph, new_node, 0, make_opaque_py(dst_owner)))
     if src_owner is not None:

cuda_core/tests/graph/test_graph_definition_lifetime.py

@@ -695,3 +695,135 @@ def test_memcpy_buffers_survive_graph_clone(init_cuda):
     out = (ctypes.c_uint8 * 4)(0)
     handle_return(driver.cuMemcpyDtoH(out, dst_dptr, 4))
     assert list(out) == [0xCD] * 4
+
+
+# =============================================================================
+# Explicit dst_owner / src_owner for raw pointer operands
+# =============================================================================
+
+
+def test_memset_raw_ptr_with_dst_owner(init_cuda):
+    """Raw dst address plus dst_owner: the graph retains the owner until destroyed."""
+    from cuda.core._utils.cuda_utils import driver, handle_return
+
+    _skip_if_no_mempool()
+    dev = Device()
+    mr = DeviceMemoryResource(dev)
+    buf = mr.allocate(4, stream=dev.default_stream)
+    dev.default_stream.sync()
+    buf_weak = weakref.ref(buf)
+    dptr = int(buf.handle)
+
+    g = GraphDefinition()
+    g.memset(dptr, 0xAB, 4, dst_owner=buf)
+
+    del buf
+    gc.collect()
+    assert buf_weak() is not None  # graph retains the explicit owner
+
+    stream = dev.create_stream()
+    g.instantiate().launch(stream)
+    stream.sync()
+
+    out = (ctypes.c_uint8 * 4)(0)
+    handle_return(driver.cuMemcpyDtoH(out, dptr, 4))
+    assert list(out) == [0xAB] * 4
+
+    del g
+    _wait_until(lambda: buf_weak() is None)
+
+
+def test_memcpy_raw_ptrs_with_owners(init_cuda):
+    """Raw src/dst addresses: the graph retains both owners until destroyed."""
+    from cuda.core._utils.cuda_utils import driver, handle_return
+
+    _skip_if_no_mempool()
+    dev = Device()
+    mr = DeviceMemoryResource(dev)
+    src = mr.allocate(4, stream=dev.default_stream)
+    dst = mr.allocate(4, stream=dev.default_stream)
+    src.fill(0xCD, stream=dev.default_stream)
+    dev.default_stream.sync()
+    src_weak = weakref.ref(src)
+    dst_weak = weakref.ref(dst)
+    src_dptr = int(src.handle)
+    dst_dptr = int(dst.handle)
+
+    g = GraphDefinition()
+    g.memcpy(dst_dptr, src_dptr, 4, dst_owner=dst, src_owner=src)
+
+    del src, dst
+    gc.collect()
+    assert src_weak() is not None and dst_weak() is not None  # both owners retained
+
+    stream = dev.create_stream()
+    g.instantiate().launch(stream)
+    stream.sync()
+
+    out = (ctypes.c_uint8 * 4)(0)
+    handle_return(driver.cuMemcpyDtoH(out, dst_dptr, 4))
+    assert list(out) == [0xCD] * 4
+
+    del g
+    _wait_until(lambda: src_weak() is None and dst_weak() is None)
+
+
+def test_memcpy_mixed_buffer_and_raw_owner(init_cuda):
+    """Buffer dst is auto-retained; raw src uses src_owner. Both survive until destroyed."""
+    from cuda.core._utils.cuda_utils import driver, handle_return
+
+    _skip_if_no_mempool()
+    dev = Device()
+    mr = DeviceMemoryResource(dev)
+    src = mr.allocate(4, stream=dev.default_stream)
+    dst = mr.allocate(4, stream=dev.default_stream)
+    src.fill(0xCD, stream=dev.default_stream)
+    dev.default_stream.sync()
+    src_weak = weakref.ref(src)
+    dst_weak = weakref.ref(dst)
+    src_dptr = int(src.handle)
+    dst_dptr = int(dst.handle)
+
+    g = GraphDefinition()
+    g.memcpy(dst, src_dptr, 4, src_owner=src)
+
+    del src, dst
+    gc.collect()
+    assert src_weak() is not None and dst_weak() is not None  # explicit + auto owner
+
+    stream = dev.create_stream()
+    g.instantiate().launch(stream)
+    stream.sync()
+
+    out = (ctypes.c_uint8 * 4)(0)
+    handle_return(driver.cuMemcpyDtoH(out, dst_dptr, 4))
+    assert list(out) == [0xCD] * 4
+
+    del g
+    _wait_until(lambda: src_weak() is None and dst_weak() is None)
+
+
+def test_memcpy_buffer_and_dst_owner_rejected(init_cuda):
+    """dst_owner cannot be combined with a Buffer dst operand."""
+    _skip_if_no_mempool()
+    dev = Device()
+    mr = DeviceMemoryResource(dev)
+    buf = mr.allocate(4, stream=dev.default_stream)
+    dev.default_stream.sync()
+
+    g = GraphDefinition()
+    with pytest.raises(ValueError, match="dst_owner cannot be used when dst is a Buffer"):
+        g.memcpy(buf, buf, 4, dst_owner=object())
+
+
+def test_memcpy_buffer_and_src_owner_rejected(init_cuda):
+    """src_owner cannot be combined with a Buffer src operand."""
+    _skip_if_no_mempool()
+    dev = Device()
+    mr = DeviceMemoryResource(dev)
+    buf = mr.allocate(4, stream=dev.default_stream)
+    dev.default_stream.sync()
+
+    g = GraphDefinition()
+    with pytest.raises(ValueError, match="src_owner cannot be used when src is a Buffer"):
+        g.memcpy(buf, buf, 4, src_owner=object())