Skip to content

Add versioning support to DLPack APIs #602

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.

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
merged 8 commits into from
Feb 9, 2024
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
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
96 changes: 73 additions & 23 deletions src/array_api_stubs/_draft/array_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,9 @@ def __complex__(self: array, /) -> complex:
"""

def __dlpack__(
self: array, /, *, stream: Optional[Union[int, Any]] = None
self: array, /, *,
stream: Optional[Union[int, Any]] = None,
max_version: Optional[tuple[int, int]] = None,
) -> PyCapsule:
"""
Exports the array for consumption by :func:`~array_api.from_dlpack` as a DLPack capsule.
Expand All @@ -298,45 +300,41 @@ def __dlpack__(
self: array
array instance.
stream: Optional[Union[int, Any]]
for CUDA and ROCm, a Python integer representing a pointer to a stream, on devices that support streams. ``stream`` is provided by the consumer to the producer to instruct the producer to ensure that operations can safely be performed on the array (e.g., by inserting a dependency between streams via "wait for event"). The pointer must be a positive integer or ``-1``. If ``stream`` is ``-1``, the value may be used by the consumer to signal "producer must not perform any synchronization". The ownership of the stream stays with the consumer. On CPU and other device types without streams, only ``None`` is accepted.
for CUDA and ROCm, a Python integer representing a pointer to a stream, on devices that support streams. ``stream`` is provided by the consumer to the producer to instruct the producer to ensure that operations can safely be performed on the array (e.g., by inserting a dependency between streams via "wait for event"). The pointer must be an integer larger than or equal to ``-1`` (see below for allowed values on each platform). If ``stream`` is ``-1``, the value may be used by the consumer to signal "producer must not perform any synchronization". The ownership of the stream stays with the consumer. On CPU and other device types without streams, only ``None`` is accepted.

For other device types which do have a stream, queue or similar synchronization mechanism, the most appropriate type to use for ``stream`` is not yet determined. E.g., for SYCL one may want to use an object containing an in-order ``cl::sycl::queue``. This is allowed when libraries agree on such a convention, and may be standardized in a future version of this API standard.
For other device types which do have a stream, queue, or similar synchronization/ordering mechanism, the most appropriate type to use for ``stream`` is not yet determined. E.g., for SYCL one may want to use an object containing an in-order ``cl::sycl::queue``. This is allowed when libraries agree on such a convention, and may be standardized in a future version of this API standard.

.. note::
Support for a ``stream`` value other than ``None`` is optional and implementation-dependent.

.. note::
Support for a ``stream`` value other than ``None`` is optional and implementation-dependent.


Device-specific notes:


.. admonition:: CUDA
:class: note
Device-specific values of ``stream`` for CUDA:

- ``None``: producer must assume the legacy default stream (default).
- ``1``: the legacy default stream.
- ``2``: the per-thread default stream.
- ``> 2``: stream number represented as a Python integer.
- ``0`` is disallowed due to its ambiguity: ``0`` could mean either ``None``, ``1``, or ``2``.


.. admonition:: ROCm
:class: note
Device-specific values of ``stream`` for ROCm:

- ``None``: producer must assume the legacy default stream (default).
- ``0``: the default stream.
- ``> 2``: stream number represented as a Python integer.
- Using ``1`` and ``2`` is not supported.

.. admonition:: Tip
:class: important

.. admonition:: Tip
:class: important

It is recommended that implementers explicitly handle streams. If
they use the legacy default stream, specifying ``1`` (CUDA) or ``0``
(ROCm) is preferred. ``None`` is a safe default for developers who do
not want to think about stream handling at all, potentially at the
cost of more synchronization than necessary.
It is recommended that implementers explicitly handle streams. If
they use the legacy default stream, specifying ``1`` (CUDA) or ``0``
(ROCm) is preferred. ``None`` is a safe default for developers who do
not want to think about stream handling at all, potentially at the
cost of more synchronizations than necessary.
max_version: Optional[tuple[int, int]]
The maximum DLPack version that the *consumer* (i.e., the caller of
``__dlpack__``) supports, in the form of a 2-tuple ``(major, minor)``.
This method may return a capsule of version ``max_version`` (recommended
if it does support that), or of a different version.

Returns
-------
Expand All @@ -353,9 +351,61 @@ def __dlpack__(

Notes
-----
The DLPack version scheme is SemVer, where the major DLPack versions
represent ABI breaks, and minor versions represent ABI-compatible additions
(e.g., new enum values for new data types or device types).

The ``max_version`` keyword was introduced in v2023.12, and goes
together with the ``DLManagedTensorVersioned`` struct added in DLPack
1.0. This keyword may not be used by consumers until a later time after
introduction, because producers may implement the support at a different
point in time.

It is recommended for the producer to use this logic in the implementation
of ``__dlpack__``:

.. code:: python

if max_version is None:
# Keep and use the DLPack 0.X implementation
# Note: from March 2025 onwards (but ideally as late as
# possible), it's okay to raise BufferError here
else:
# We get to produce `DLManagedTensorVersioned` now. Note that
# our_own_dlpack_version is the max version that the *producer*
# supports and fills in the `DLManagedTensorVersioned::version`
# field
if max_version >= our_own_dlpack_version:
# Consumer understands us, just return a Capsule with our max version
elif max_version[0] == our_own_dlpack_version[0]:
# major versions match, we should still be fine here -
# return our own max version
else:
# if we're at a higher major version internally, did we
# keep an implementation of the older major version around?
# For example, if the producer is on DLPack 1.x and the consumer
# is 0.y, can the producer still export a capsule containing
# DLManagedTensor and not DLManagedTensorVersioned?
# If so, use that. Else, the producer should raise a BufferError
# here to tell users that the consumer's max_version is too
# old to allow the data exchange to happen.

And this logic for the consumer in ``from_dlpack``:

.. code:: python

try:
x.__dlpack__(max_version=(1, 0))
# if it succeeds, store info from the capsule named "dltensor_versioned",
# and need to set the name to "used_dltensor_versioned" when we're done
except TypeError:
x.__dlpack__()

.. versionchanged:: 2022.12
Added BufferError.

.. versionchanged:: 2023.12
Added the ``max_version`` keyword.
"""

def __dlpack_device__(self: array, /) -> Tuple[Enum, int]:
Expand Down
5 changes: 5 additions & 0 deletions src/array_api_stubs/_draft/creation_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,11 @@ def from_dlpack(x: object, /) -> array:
If the ``__dlpack__`` and ``__dlpack_device__`` methods are not present
on the input array. This may happen for libraries that are never able
to export their data with DLPack.

Notes
-----
See :meth:`array.__dlpack__` for implementation suggestions for `from_dlpack` in
order to handle DLPack versioning correctly.
"""


Expand Down