cuda.core: retain device allocations in graph memcpy/memset slots

Store DevicePtrHandle in slot table instead of Buffer wrappers so
reset/close cannot release memory while a graph still references it.
Add test-only weak_handle() for deterministic allocation lifetime
checks and extend graph lifetime tests accordingly.

Author: Andy-Jost

cuda_core/cuda/core/_resource_handles.pyi

@@ -24,4 +24,5 @@ FileDescriptorHandle = shared_ptr
 OpaqueArrayHandle = shared_ptr
 MipmappedArrayHandle = shared_ptr
 TexObjectHandle = shared_ptr
-SurfObjectHandle = shared_ptr
+SurfObjectHandle = shared_ptr
+OpaqueHandle = shared_ptr

cuda_core/cuda/core/_utils/_weak_handles.pyi

@@ -0,0 +1,55 @@
+# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_utils/_weak_handles.pyx
+
+"""Test-only weak handles for resource-handle lifetime checks.
+
+This module is **not** part of the public ``cuda.core`` API. It is built into
+the package (like other private ``_utils`` modules) purely so the test suite can
+observe, deterministically, when the strong references that keep a CUDA resource
+alive have all been released -- without relying on driver- or hardware-specific
+side effects (for example, whether freed device memory happens to remain
+readable).
+
+Every resource handle is owned by a C++ ``std::shared_ptr``. A **weak handle**
+is a non-owning ``std::weak_ptr`` observer of that control block: truthy while
+some strong owner remains, falsy once the last one is gone. Use :func:`weak_handle`
+to obtain a weak handle from a supported front-end object.
+
+To support another type, add a ``cdef _weak_from_<type>`` that reads its ``cdef``
+handle field (see ``*.pxd``), assigns to :ctype:`OpaqueHandle`, and extend the
+``isinstance`` chain in :func:`weak_handle`. Types whose slots hold arbitrary
+Python owners via ``make_opaque_py`` are not covered here -- use
+:class:`weakref.ref` on a weak-referenceable owner object in tests instead.
+"""
+from __future__ import annotations
+
+
+class WeakHandle:
+    """Non-owning weak handle for a resource's shared control block.
+
+    Truthy while some strong owner of the underlying resource handle remains,
+    falsy once the last strong reference is released. Obtain instances via
+    :func:`weak_handle` rather than constructing directly.
+    """
+
+    def __bool__(self):
+        ...
+
+    def expired(self):
+        """Return ``True`` once every strong owner of the handle is gone."""
+
+    def use_count(self):
+        """Number of strong owners currently sharing the handle."""
+
+def weak_handle(obj):
+    """Return a :class:`WeakHandle` observing the resource behind ``obj``.
+
+    Currently supports :class:`~cuda.core.Buffer` (device allocation handle).
+    See the module docstring for how to add more types.
+
+    Raises
+    ------
+    ValueError
+        If ``obj`` is a :class:`~cuda.core.Buffer` with no active allocation.
+    TypeError
+        If ``obj`` is not a supported type.
+    """

cuda_core/cuda/core/_utils/_weak_handles.pyx

@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Test-only weak handles for resource-handle lifetime checks.
+
+This module is **not** part of the public ``cuda.core`` API. It is built into
+the package (like other private ``_utils`` modules) purely so the test suite can
+observe, deterministically, when the strong references that keep a CUDA resource
+alive have all been released -- without relying on driver- or hardware-specific
+side effects (for example, whether freed device memory happens to remain
+readable).
+
+Every resource handle is owned by a C++ ``std::shared_ptr``. A **weak handle**
+is a non-owning ``std::weak_ptr`` observer of that control block: truthy while
+some strong owner remains, falsy once the last one is gone. Use :func:`weak_handle`
+to obtain a weak handle from a supported front-end object.
+
+To support another type, add a ``cdef _weak_from_<type>`` that reads its ``cdef``
+handle field (see ``*.pxd``), assigns to :ctype:`OpaqueHandle`, and extend the
+``isinstance`` chain in :func:`weak_handle`. Types whose slots hold arbitrary
+Python owners via ``make_opaque_py`` are not covered here -- use
+:class:`weakref.ref` on a weak-referenceable owner object in tests instead.
+"""
+
+from cuda.core._memory._buffer cimport Buffer
+from cuda.core._resource_handles cimport OpaqueHandle
+
+
+# Cython cannot spell ``weak_ptr[const void]`` inline (the ``const void``
+# template argument fails to parse), so the weak type and its one constructor
+# are provided by a small inline C++ shim local to this test-only module. This
+# keeps the production resource_handles translation units untouched.
+cdef extern from *:
+    """
+    #include <memory>
+    namespace cuda_core_test {
+    using OpaqueWeakHandle = std::weak_ptr<const void>;
+    static inline OpaqueWeakHandle make_weak(const std::shared_ptr<const void>& h) {
+        return OpaqueWeakHandle(h);
+    }
+    }  // namespace cuda_core_test
+    """
+    cppclass OpaqueWeakHandle "cuda_core_test::OpaqueWeakHandle":
+        OpaqueWeakHandle()
+        bint expired()
+        long use_count()
+    OpaqueWeakHandle make_weak "cuda_core_test::make_weak" (const OpaqueHandle& h)
+
+
+cdef class WeakHandle:
+    """Non-owning weak handle for a resource's shared control block.
+
+    Truthy while some strong owner of the underlying resource handle remains,
+    falsy once the last strong reference is released. Obtain instances via
+    :func:`weak_handle` rather than constructing directly.
+    """
+
+    cdef OpaqueWeakHandle _w
+
+    def __bool__(self):
+        return not self._w.expired()
+
+    def expired(self):
+        """Return ``True`` once every strong owner of the handle is gone."""
+        return self._w.expired()
+
+    def use_count(self):
+        """Number of strong owners currently sharing the handle."""
+        return self._w.use_count()
+
+
+cdef WeakHandle _weak_from_opaque(OpaqueHandle h):
+    # Build the weak handle from a (temporary) strong handle. The strong copy
+    # lives only for the duration of this call, so it does not perturb the
+    # reference count the weak handle later reports.
+    cdef WeakHandle wh = WeakHandle.__new__(WeakHandle)
+    wh._w = make_weak(h)
+    return wh
+
+
+cdef WeakHandle _weak_from_buffer(Buffer buf):
+    cdef OpaqueHandle h = buf._h_ptr
+    if not h:
+        raise ValueError("Buffer has no active allocation")
+    return _weak_from_opaque(h)
+
+
+def weak_handle(obj):
+    """Return a :class:`WeakHandle` observing the resource behind ``obj``.
+
+    Currently supports :class:`~cuda.core.Buffer` (device allocation handle).
+    See the module docstring for how to add more types.
+
+    Raises
+    ------
+    ValueError
+        If ``obj`` is a :class:`~cuda.core.Buffer` with no active allocation.
+    TypeError
+        If ``obj`` is not a supported type.
+    """
+    if isinstance(obj, Buffer):
+        return _weak_from_buffer(obj)
+    raise TypeError(
+        f"weak_handle() does not support {type(obj).__name__!r}; "
+        "supported types: Buffer"
+    )

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -106,6 +106,18 @@ class GraphBuilder:
     to ambiguity. New graph builders should instead be created through a
     :obj:`~_device.Device`, or a :obj:`~_stream.stream` object.
 
+    .. note::
+
+        Operations recorded during capture reference your memory but do not
+        take ownership of it. As with ordinary stream work, you must keep the
+        operands alive for as long as the completed graph may execute -- for
+        example, the :obj:`~_memory.Buffer` objects passed to :func:`~launch`
+        or :meth:`~_memory.Buffer.copy_to`. Host callbacks added with
+        :meth:`callback` are the exception: the callable (and any copied
+        ``user_data``) are retained for the graph's lifetime. This differs from
+        building a graph explicitly with :class:`~graph.GraphDefinition`, which
+        retains the operands it is given.
+
     """
 
     def __init__(self):

cuda_core/cuda/core/graph/_graph_definition.pyi

@@ -6,6 +6,7 @@ from __future__ import annotations
 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._utils.cuda_utils import driver
 from cuda.core.graph._graph_builder import (Graph, GraphCompleteOptions,
@@ -85,7 +86,7 @@ class GraphDefinition:
         See :meth:`GraphNode.deallocate` for full documentation.
         """
 
-    def memset(self, dst: int, value, width: int, height: int=1, pitch: int=0) -> MemsetNode:
+    def memset(self, dst: Buffer | int, value, width: int, *, height: int=1, pitch: int=0, dst_owner=None) -> MemsetNode:
         """Add an entry-point memset node (no dependencies).
 
         See :meth:`GraphNode.memset` for full documentation.
@@ -120,7 +121,7 @@ class GraphDefinition:
             A new EmptyNode that depends on all input nodes.
         """
 
-    def memcpy(self, dst: int, src: int, size: int) -> MemcpyNode:
+    def memcpy(self, dst: Buffer | int, src: Buffer | int, size: int, *, dst_owner=None, src_owner=None) -> MemcpyNode:
         """Add an entry-point memcpy node (no dependencies).
 
         See :meth:`GraphNode.memcpy` for full documentation.

cuda_core/cuda/core/graph/_graph_node.pyi

@@ -9,6 +9,7 @@ from collections.abc import Iterable
 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._utils.cuda_utils import driver
 from cuda.core.graph._adjacency_set_proxy import AdjacencySetProxy
@@ -178,13 +179,16 @@ class GraphNode:
             A new FreeNode representing the free operation.
         """
 
-    def memset(self, dst: int, value, width: int, height: int=1, pitch: int=0) -> MemsetNode:
+    def memset(self, dst: Buffer | int, value, width: int, *, height: int=1, pitch: int=0, dst_owner=None) -> MemsetNode:
         """Add a memset node depending on this node.
 
         Parameters
         ----------
-        dst : int
-            Destination device pointer.
+        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.
@@ -194,14 +198,23 @@ 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. 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``.
         """
 
-    def memcpy(self, dst: int, src: int, size: int) -> MemcpyNode:
+    def memcpy(self, dst: Buffer | int, src: Buffer | int, size: int, *, 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
@@ -210,17 +223,35 @@ class GraphNode:
 
         Parameters
         ----------
-        dst : int
-            Destination pointer (device or pinned host).
-        src : int
-            Source pointer (device or pinned host).
+        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.
         """
 
     def embed(self, child: GraphDefinition) -> ChildGraphNode: