ENH: Implement the DLPack Array API protocols for ndarray.

[hameerabbasi]

Fixes #19013

cc @rgommers @mattip @seberg

TODO

• from_dlpack function.
• dtype support in __dlpack__

TODO items added/edited by seberg:

• Move all dlpack specific code to a single C file, rather than methods.c
• I would feel much better if there is a PR to dlpack.h that clarifies alignment requirements (and use of byte_offset). After such a PR exists, we can review the current rules when to reject export. (if the rejection is not guaranteed to be fixed by arr.copy() we may have a problem.)
• As discussed in the meeting from_dlpack shall not be exposed as new, public API. There will be no public API to convert to NumPy from __dlpack__. It would be ok to have experimental/private, underscored API.
• (preferably) Make sure there is a PR to either dlpack.h or array-api that states that imported arrays shall be considered writable and exporters like NumPy should avoid exporting readonly arrays (truly readonly only?). Ideally, with a footnote that this means that users of immutable array libraries may be up for a big surprise.
  (mattip says:) Not sure how to resolve this. The data api issue is still open
• I am not quite sure if the ->deleter == NULL case is settled (I think it is, but maybe there needs to be some clarification in the header as well).
• Specify cleanup to know for a fact whether this implementation or the CuPy implementation is correct.

[rgommers]

Thanks Hameer, looks like a good start.

Of the points @seberg brought up on the issue, adding a version attribute seems the most important - because checking for version will be needed before any backwards-incompatible change (such as adding an extra field) can be done. I asked about it on dmlc/dlpack#34, and the suggestion was to add it as an attribute on the PyCapsule. Could you look into doing that?

[hameerabbasi]

Could you look into doing that?

Sure -- I commented there instead of here, so more of the relevant parties can see it.

[BvB93]

Hi Hameer, could you also add these two new methods to the ndarray stubs in numpy/__init__.pyi?
Their signatures are fairly simple, so that's a plus:

# NOTE: Every single import below is already present in `__init__.pyi`
# and doesn't have to be repeated
from typing import Any, Tuple
from typing_extensions import Literal as L

from numpy import number
from numpy.typing import NDArray

# `builtins.PyCapsule` unfortunately lacks annotations as of the moment;
# use `Any` as a stopgap measure
_PyCapsule = Any

def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ...
def __dlpack_device__(self) -> Tuple[L[1], L[0]]: ...

[seberg]

Had started looking at it, so a few small comments. It looks good, I have to think a bit more about the API, but that has nothing to do with the implementation.

I think if we put the header where you put it, it will be effectively public. I guess that is fine thoug? Otherwise I think placing it into common is good.

[BvB93]

Interestingly, if we could get builtins.PyCapsule in typeshed annotated as a parameterizable type,
then it would in principle be possible for static type checkers to read, e.g., the underlying shape and dtype.

The only thing that users or libraries would have to do here is declare the necessary annotations.

Perhaps something to consider for the future?

Examples

from typing import TypeVar, Any, Generic, Tuple as Tuple
import numpy as np

# Improvised `PyCapsule` annotation
_T = TypeVar("_T")
class PyCapsule(Generic[_T]): ...

# Construct a more compact `PyCapsule` alias; `Tuple` used herein to introduce 2 parameters 
# (there may be more appropriate types that can fulfill this functionality)
_Shape = TypeVar("_Shape", bound=Any)  # TODO: Wait for PEP 646's TypeVarTuple
_DType = TypeVar("_DType", bound=np.dtype[Any])
DLPackCapsule = PyCapsule[Tuple[_Shape, _DType]]

# A practical example
def from_dlpack(__x: DLPackCapsule[_Shape, _DType]) -> np.ndarray[_Shape, _DType]: ...

[hameerabbasi]

I'll consider this as out of scope of this PR for now, but will leave the conversation unresolved for visibility.

[hameerabbasi]

This name should be some standardized DLPack name? Or is this intentional for now to keep it effectively private right now?

Well, we have to be able to delete it to free the memory, but maybe we need some state as well, to ensure that we don't delete it twice if the consumer already deletes it without deleting the capsule (might be good for non-reference counted Python also?).

I thought I read that there was some capsule renaming involved (to signal that its "invalid" after deletion). I guess setting the data to NULL or so would be just as well. (If this is standardized, do that. If not setting data to NULL to guard against deleting twice? Once from python once from the consumer.)

There's a standardized name, see data-apis/array-api#186. I have modified the PR accordingly.

We need to also check for the DType being native. I assume unaligned is OK?

Unaligned is OK, see #19013 (comment). How does one check for native dtypes?

Maybe we should just explicitly list them? We cannot include longdouble here, since the C99 standard doesn't specify it (I am not even sure IBM double-double is strictly standard conform). And I assume Float in DLPack means IEEE compatible float format of defined size. float16, float32, and float64 are fine, for the others maybe DLPack should add an extended float field or so, if we want it?

I just don't see longdoublel to be much good, considering its already system dependend on the CPU alone...

Can we verify this?

[hameerabbasi]

Can we verify this?

Ignore, it does mean IEEE float. Does np.float80 or np.float128 not respect the IEEE convention? In that case, it makes sense to limit both complex and floats. Otherwise, the receiving library can deny the incoming datatype by reading the bits field.

[eric-wieser]

np.float80 etc are all just aliases for np.longdouble based on what C produces for sizeof(long double)

[seberg]

I guess if longdouble is 64bit, you can be sure its just double. That we could check for if we want. There are a couple of macros in numpy/core/src/common/npy_fpmath.h such as HAVE_LDOUBLE_IEEE_QUAD_BE. I don't know how well they work/are tested, but I guess technically you could check for that to catch the rare cases where we have proper quad precision.

Technically, most current hardware uses something like 80bit (IEEE?) stored as 96bit or 128bit. But DLPack can't describe it (its an old thing that our float128 should relaly be float128_80 or something else than float128...). Something, that will eventually be important if quad precision becomes a real thing.

[seberg]

How does one check for native dtypes?

You can use PyDataType_ISNOTSWAPPED.

[leofang]

clear hierarchal structure (ndarray <- DLPackMemory <- capsule)

I guess you can see it that way, but I see no reason for DLPackMemory to hold on to the original capsule! At least unless there is some reasons due to stream handling.

No, it has nothing to do with stream handling (which is done at the Python level during handshaking), it's about keeping the ref count of the underlying resource. Conceptually we could definitely use a capsule to back an ndarray, in practice though cupy.ndarray expects a few extra methods from its data member to return device, size, pointer, etc, hence wrapping the capsule with a very thin layer (DLPackMemory).

[seberg]

Not really the right place to continue discussing. But no: The capsule owns the DLManagedTensor which owns its own memory and is not reference counted (there is exactly one reference that cannot be shared). Renaming the capsule steals ownership from it (and you must rename it).
At that point, the capsule does not own anything (check the capsule destructor it does nothing after renaming). So you simply can delete the cdef object dltensor and self.dltensor = dltensor line. The only tricky thing is memory handling (i.e. renaming the capsule at the right time), and I think cupy should move that directly after the line getting the pointer from the capsule (at that point DLPackMemory has ownership). That looks like a bug, for example if you are confronted with an unknown dtype.

[mattip]

This code path can only be hit if we are re-wrapping a dlpack capsule with an offset. It is currently not tested, which should be done in a follow-on PR.

[seberg]

Yeah, don't worry about it. This is nonsense anyway, I deleted the whole thing (the logic was probably incomplete anyway) and replaced it with the following comment:

    /*
     * Note: the `dlpack.h` header suggests/standardizes that `data` must be
     * 256-byte aligned.  We ignore this intentionally, because `__dlpack__`
     * standardizes that `byte_offset` must be 0 (for now) to not break pytorch:
     * https://github.com/data-apis/array-api/issues/293#issuecomment-964111413
     *
     * We further assume that exporting fully unaligned data is OK even without
     * `byte_offset` since the standard does not reject it.
     * Presumably, pytorch will support importing `byte_offset != 0` and NumPy
     * can choose to use it starting about 2023.  At that point, it may be
     * that NumPy MUST use `byte_offset` to adhere to the standard (as
     * specified in the header)!
     */

[seberg]

I will put this in once the tests pass, but will then open a few issues about dlpack, with pretty harmless bugs or general followups. The release notes get reviewed before the release anyway.

Thanks everyone and especially Matti for doing the follow-ups.

[rgommers]

🎉 thanks @mattip, @hameerabbasi, @seberg, @leofang & all other reviewers!

[leofang]

Silly question: Can someone remind me why we renamed it to np._from_dlpack() (prefixed with an underscore)?

[seberg]

This was pushed in a ~1 week before branching with not quite settled things in dlpack itself (as I remember). Adding the underscore was the way to not having to discuss all of that in a short time frame.

[leofang]

I see, thanks for the context @seberg. Sounds like a decision made offline (which I am fine with). Is there any plan to make it a public API (by removing the underscore)?

[seberg]

Yes, we can discuss this obviously. It may have been discussed in a meeting, in which case it is publically noted in the numpy/archive repo. I feel the PR had the underscore initially, because the first step was to make it available for the _array_api namespace.

[mattip]

The only mention of dlpack in the docs is in the release notes, so maybe we could discuss making it public together with some documentation.

[leofang]

Thanks, guys. I created an issue to track this change: #20743.