Add versioning support to DLPack APIs

xref dmlc/dlpack#116

Author: rgommers

src/array_api_stubs/_draft/array_object.py

@@ -278,7 +278,9 @@ def __complex__(self: array, /) -> complex:
         """
 
     def __dlpack__(
-        self: array, /, *, stream: Optional[Union[int, Any]] = None
+        self: array, /, *,
+        max_version: Optional[tuple[int, int]] = None,
+        stream: Optional[Union[int, Any]] = None
     ) -> PyCapsule:
         """
         Exports the array for consumption by :func:`~array_api.from_dlpack` as a DLPack capsule.
@@ -287,46 +289,42 @@ def __dlpack__(
         ----------
         self: array
             array instance.
+        max_version: Optional[tuple[int, int]]
+            The maximum DLPack version that the consumer (i.e., the caller of
+            ``__dlpack__``) supports, in the form ``(major, minor)``.
+            This method may return that maximum version (recommended if it does
+            support that), or a different version.
         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 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.
 
+            .. 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 synchronization than necessary.
 
         Returns
         -------
@@ -343,9 +341,52 @@ def __dlpack__(
 
         Notes
         -----
+        Major DLPack versions represent ABI breaks, 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 for some time after
+        introduction. It is recommended 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: in >= 2 years from now (but ideally as late as
+                # possible), it's okay to raise BufferError here
+            else:
+                # We get to produce `DLManagedTensorVersioned` now
+                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?
+                    # If so, use that. Else, just return our max
+                    # version and let the consumer deal with it.
+
+        And this logic for the producer (i.e., in ``from_dlpack``):
+
+        .. code:: python
+
+            try:
+                x.__dlpack__(max_version=(1, 0))
+                # if it succeeds, store info about capsule name being "dltensor_versioned",
+                # and needing to set the capsule 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]:

src/array_api_stubs/_draft/creation_functions.py

@@ -212,6 +212,12 @@ def from_dlpack(x: object, /) -> array:
            :class: note
 
            The returned array may be either a copy or a view. See :ref:`data-interchange` for details.
+
+    Notes
+    -----
+    See :meth:`array.__dlpack__` for implementation suggestions for `from_dlpack` in
+    order to handle DLPack versioning correctly.
+
     """