Improve #449: Improve StridedMemoryView creation time

[mdboom]

Two changes:

1. Refactor the versioned/non-versioned paths to reduce the number of branches.
2. Create shape and strides tuples using Python/C API

The second change has a significant impact:

• Before: 1.23 us +- 0.06 us
• After: 1.02 us +- 0.06 us

Measured using this pyperf benchmark, based on the one in #449

import pyperf
import time


import cupy as cp
from cuda.core.experimental.utils import StridedMemoryView


inner_loops = 100_000


def bench_strided_memory_view(loops, x):
    range_it = range(loops * inner_loops)
    t0 = time.perf_counter()

    for _ in range_it:
        s = StridedMemoryView(x, -1)

    return time.perf_counter() - t0


runner = pyperf.Runner()
x = cp.empty((23, 4))
runner.bench_time_func('StridedMemoryView', bench_strided_memory_view, x, inner_loops=inner_loops)

Description

Improves #449

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Copilot]

Pull Request Overview

This PR improves the performance of StridedMemoryView creation by optimizing the tuple creation process and refactoring conditional logic. The changes achieve a ~17% performance improvement by using Python/C API calls instead of Python comprehensions for creating shape and strides tuples.

• Refactored versioned/non-versioned DLPack path handling to reduce branching complexity
• Replaced Python tuple comprehensions with direct Python/C API calls for creating shape and strides tuples
• Optimized stride conversion logic in the CAI (CUDA Array Interface) path

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[kkraus14]

The changes here LGTM! To trigger CI you need to comment "/ok to test", see here for more info: https://docs.gha-runners.nvidia.com/platform/apps/copy-pr-bot/faqs/

Most people are using the StridedMemoryView class as a way to generically handle input arrays that support dlpack and/or __cuda_array_interface__ right before kernel launches which is typically interfacing to some underlying C++ code. Currently when we initialize the class we eagerly populate python attributes for things like the ptr, shape, strides, dtype, etc. There's a non-trivial cost to this where I wonder if we should instead have C types and then lazily initialize the Python types as requested.

[mdboom]

/ok to test

[leofang]

Currently when we initialize the class we eagerly populate python attributes for things like the ptr, shape, strides, dtype, etc. There's a non-trivial cost to this where I wonder if we should instead have C types and then lazily initialize the Python types as requested.

Yes, I have been thinking these attributes should be populated lazily, which is possible because we still hold a reference to the dlpack capsule. @mdboom would you like to address this in a follow-up PR?

[leofang]

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]

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]

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.

[leofang]

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]

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]

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]

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.

[mdboom]

Currently when we initialize the class we eagerly populate python attributes for things like the ptr, shape, strides, dtype, etc. There's a non-trivial cost to this where I wonder if we should instead have C types and then lazily initialize the Python types as requested.

Yes, I have been thinking these attributes should be populated lazily, which is possible because we still hold a reference to the dlpack capsule. @mdboom would you like to address this in a follow-up PR?

Yeah, that makes sense. I'll do that in a separate PR and keep this just to the tuple-construction improvements.

[mdboom]

/ok to test

[mdboom]

/ok to test

[github-actions]

Doc Preview CI
Preview removed because the pull request was closed or merged.