Skip to content

Commit fd15442

Browse files
committed
fix(cuda.core): keep DLPack owner alive for derived StridedMemoryView
Signed-off-by: Rui Luo <[email protected]>
1 parent 3429cba commit fd15442

2 files changed

Lines changed: 43 additions & 0 deletions

File tree

‎cuda_core/cuda/core/_memoryview.pyx‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -570,12 +570,18 @@ cdef class StridedMemoryView:
570570
"""
571571
Returns Buffer instance with the underlying data.
572572
If the SMV was created from a Buffer, it will return the same Buffer instance.
573+
If the SMV was imported via DLPack, it will return a new instance owned by this view.
573574
Otherwise, it will create a new instance with owner set to the exporting object.
574575
"""
575576
cdef object buffer
576577
if self._buffer is None:
577578
if isinstance(self.exporting_obj, Buffer):
578579
buffer = self.exporting_obj
580+
elif self.dl_tensor != NULL:
581+
# This view holds the consumed capsule and calls its deleter in
582+
# __dealloc__; the exporting object need not keep the memory alive.
583+
# Not cached: a cached Buffer owned by self would form a cycle.
584+
return Buffer.from_handle(self.ptr, 0, owner=self)
579585
else:
580586
buffer = Buffer.from_handle(self.ptr, 0, owner=self.exporting_obj)
581587
if self._buffer is None:

‎cuda_core/tests/test_utils_dlpack.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
"""
99

1010
import ctypes
11+
import gc
12+
import weakref
1113
from pathlib import Path
1214

1315
try:
@@ -105,6 +107,41 @@ def test_dlpack_view_of_buffer_reuses_exporting_buffer(init_cuda):
105107
buffer.close()
106108

107109

110+
class _EphemeralDLPackExporter:
111+
"""CPU exporter whose capsule is the only owner of a freshly created array."""
112+
113+
def __init__(self):
114+
self.array_ref = None
115+
116+
def __dlpack_device__(self):
117+
return (int(DLDeviceType.kDLCPU), 0)
118+
119+
def __dlpack__(self, **kwargs):
120+
array = np.arange(4, dtype=np.int32)
121+
self.array_ref = weakref.ref(array)
122+
return array.__dlpack__(**kwargs)
123+
124+
125+
@pytest.mark.agent_authored(model="claude-opus-5-5")
126+
def test_dlpack_derived_view_keeps_capsule_owned_storage_alive(init_cuda):
127+
"""A derived view keeps capsule-owned storage alive until the last view is released."""
128+
exporter = _EphemeralDLPackExporter()
129+
original = StridedMemoryView.from_dlpack(exporter, stream_ptr=-1)
130+
derived = original.view(dtype=original.dtype)
131+
ptr = original.ptr
132+
del original
133+
gc.collect()
134+
135+
# Check liveness before reading, so a regression fails without touching freed memory.
136+
assert exporter.array_ref() is not None
137+
assert derived.ptr == ptr
138+
assert np.from_dlpack(derived).tolist() == [0, 1, 2, 3]
139+
140+
del derived
141+
# No gc.collect(): the capsule must be released by reference counting, not a cycle.
142+
assert exporter.array_ref() is None
143+
144+
108145
def test_dlpack_export_unversioned_capsule_and_deleter():
109146
"""``__dlpack__()`` with no ``max_version`` yields an *unversioned* unused
110147
DLPack capsule; dropping it unconsumed runs ``_smv_pycapsule_deleter`` on

0 commit comments

Comments
 (0)