Skip to content

Commit e0d3c4e

Browse files
committed
fix(core): respect DLPack byte offsets
1 parent 757731a commit e0d3c4e

2 files changed

Lines changed: 42 additions & 1 deletion

File tree

‎cuda_core/cuda/core/_memoryview.pyx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1102,7 +1102,7 @@ cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None):
11021102
cdef StridedMemoryView buf = StridedMemoryView() if view is None else view
11031103
buf.dl_tensor = dl_tensor
11041104
buf.metadata = capsule
1105-
buf.ptr = <intptr_t>(dl_tensor.data)
1105+
buf.ptr = <intptr_t>(dl_tensor.data) + <intptr_t>(dl_tensor.byte_offset)
11061106
buf.device_id = device_id
11071107
buf.is_device_accessible = is_device_accessible
11081108
buf.readonly = is_readonly

‎cuda_core/tests/test_utils_dlpack.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,47 @@ class _DLManagedTensorVersioned(ctypes.Structure):
219219
]
220220

221221

222+
_PyCapsule_New = ctypes.pythonapi.PyCapsule_New
223+
_PyCapsule_New.argtypes = (ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p)
224+
_PyCapsule_New.restype = ctypes.py_object
225+
226+
227+
def test_from_dlpack_respects_byte_offset():
228+
"""Importing a DLPack capsule must include ``DLTensor.byte_offset`` in ptr."""
229+
src = np.arange(16, dtype=np.int32)
230+
offset = 8
231+
shape = (ctypes.c_int64 * 1)(len(src) - offset)
232+
strides = (ctypes.c_int64 * 1)(1)
233+
managed = _DLManagedTensorVersioned()
234+
managed.version = _DLPackVersion(1, 0)
235+
managed.manager_ctx = None
236+
managed.deleter = None
237+
managed.flags = 0
238+
managed.dl_tensor = _DLTensor(
239+
data=src.ctypes.data,
240+
device=_DLDevice(int(DLDeviceType.kDLCPU), 0),
241+
ndim=1,
242+
dtype=_DLDataType(0, 32, 1),
243+
shape=shape,
244+
strides=strides,
245+
byte_offset=offset * src.itemsize,
246+
)
247+
capsule = _PyCapsule_New(
248+
ctypes.addressof(managed), b"dltensor_versioned", None
249+
)
250+
251+
class _OffsetProducer:
252+
def __dlpack_device__(self):
253+
return (int(DLDeviceType.kDLCPU), 0)
254+
255+
def __dlpack__(self, **kwargs):
256+
return capsule
257+
258+
view = StridedMemoryView.from_dlpack(_OffsetProducer(), stream_ptr=-1)
259+
assert view.ptr == src.ctypes.data + offset * src.itemsize
260+
assert np.array_equal(np.from_dlpack(view), src[offset:])
261+
262+
222263
@pytest.mark.agent_authored(model="cursor-grok-4.5")
223264
@pytest.mark.parametrize(
224265
"max_version, capsule_name, managed_cls",

0 commit comments

Comments
 (0)