Skip to content

Improve #449: Improve StridedMemoryView creation time #838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mdboom merged 3 commits into from
Aug 15, 2025
Merged

Improve #449: Improve StridedMemoryView creation time #838

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
59d1b6b
Improve #449: Improve StridedMemoryView creation time
mdboom Aug 14, 2025
8a05be3
Add carray_int64_t_to_tuple function
mdboom Aug 15, 2025
df71f24
Move comment
mdboom Aug 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 20 additions & 16 deletions cuda_core/cuda/core/experimental/_memoryview.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,48 +184,52 @@ cdef StridedMemoryView view_as_dlpack(obj, stream_ptr, view=None):
stream=int(stream_ptr) if stream_ptr else None)

cdef void* data = NULL
cdef DLTensor* dl_tensor
cdef DLManagedTensorVersioned* dlm_tensor_ver
cdef DLManagedTensor* dlm_tensor
cdef const char *used_name
if cpython.PyCapsule_IsValid(
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME):
data = cpython.PyCapsule_GetPointer(
capsule, DLPACK_VERSIONED_TENSOR_UNUSED_NAME)
versioned = True
dlm_tensor_ver = <DLManagedTensorVersioned*>data
dl_tensor = &dlm_tensor_ver.dl_tensor
is_readonly = bool((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0)
used_name = DLPACK_VERSIONED_TENSOR_USED_NAME
elif cpython.PyCapsule_IsValid(
capsule, DLPACK_TENSOR_UNUSED_NAME):
data = cpython.PyCapsule_GetPointer(
capsule, DLPACK_TENSOR_UNUSED_NAME)
versioned = False
else:
assert False

cdef DLManagedTensor* dlm_tensor
cdef DLManagedTensorVersioned* dlm_tensor_ver
cdef DLTensor* dl_tensor
if versioned:
dlm_tensor_ver = <DLManagedTensorVersioned*>data
dl_tensor = &dlm_tensor_ver.dl_tensor
is_readonly = bool((dlm_tensor_ver.flags & DLPACK_FLAG_BITMASK_READ_ONLY) != 0)
else:
dlm_tensor = <DLManagedTensor*>data
dl_tensor = &dlm_tensor.dl_tensor
is_readonly = False
used_name = DLPACK_TENSOR_USED_NAME
else:
assert False

cdef StridedMemoryView buf = StridedMemoryView() if view is None else view
buf.ptr = <intptr_t>(dl_tensor.data)
buf.shape = tuple(int(dl_tensor.shape[i]) for i in range(dl_tensor.ndim))

# Construct shape and strides tuples using the Python/C API for speed

@leofang leofang Aug 15, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Q: Does this mean Cython is not generating efficient code for us? Could it be possible that I was using the generator form to construct a tuple and Cython is bad at this syntax?

@mdboom mdboom Aug 15, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, this is pretty inefficient. The generator causes Cython to make a separate C function which is then called through the Python iterator machinery, which is a lot less efficient than a C for loop (mainly because it makes a bunch of C function calls through C function pointers). Secondly, since the length of the tuple isn't known in advance, it causes it to be realloc'ed at least 3 times:

  • initial size of 0
  • overallocate to 10
  • truncate to the exact value of 2

It might be possible to write a cdef function to convert a C array pointer + known length into a tuple in this faster way, which would hide this use of the CPython API here. I'll look at that and measure it.

@mdboom mdboom Aug 15, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Abstracting this out to a cdef inline function seems to have no measurable overhead and definitely makes this more readable, so I've updated this PR to do that.

buf.shape = cpython.PyTuple_New(dl_tensor.ndim)
for i in range(dl_tensor.ndim):
cpython.PyTuple_SET_ITEM(buf.shape, i, cpython.PyLong_FromLong(dl_tensor.shape[i]))

@leofang leofang Aug 15, 2025 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we can use PyLong_FromLong because both shape and strides are int64_t
https://github.com/dmlc/dlpack/blob/7f393bbb86a0ddd71fde3e700fc2affa5cdce72d/include/dlpack/dlpack.h#L257-L263
we probably need to replace it with PyLong_FromLongLong

@leofang leofang Aug 15, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another question for my understanding: Both PyTuple_SET_ITEM and PyTuple_Set_Item steal the reference of the object o (the Python int that we convert from DLPack shape/stride members). Doesn't it mean that we should increment its refcount before calling setitem?

@mdboom mdboom Aug 15, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point about using LongLong. I will update.

As for reference counting, we don't need to increment here because we are "giving" the reference to the int to the tuple, and we don't need to hold on to the int ourselves to do anything else with it.

@leofang leofang Aug 15, 2025 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, so either constructor would return the obj with refcount 1 (for some reason I was thinking it's 0), and the ownership of this ref is transferred to the tuple.

if dl_tensor.strides:
buf.strides = tuple(
int(dl_tensor.strides[i]) for i in range(dl_tensor.ndim))
buf.strides = cpython.PyTuple_New(dl_tensor.ndim)
for i in range(dl_tensor.ndim):
cpython.PyTuple_SET_ITEM(buf.strides, i, cpython.PyLong_FromLong(dl_tensor.strides[i]))
else:
# C-order
buf.strides = None

buf.dtype = dtype_dlpack_to_numpy(&dl_tensor.dtype)
buf.device_id = device_id
buf.is_device_accessible = is_device_accessible
buf.readonly = is_readonly
buf.exporting_obj = obj

cdef const char* used_name = (
DLPACK_VERSIONED_TENSOR_USED_NAME if versioned else DLPACK_TENSOR_USED_NAME)
cpython.PyCapsule_SetName(capsule, used_name)

return buf
Expand Down
Loading