Skip to content

Commit 017e9cf

Browse files
tqchenrgommers
andauthored
[DLPack] Update to include DLPack C Exchange API (#984)
This PR updates to include DLPack C exchange API for fast exchange at the C extension level without going through Python. Co-authored-by: Ralf Gommers <ralf.gommers@gmail.com>
1 parent d88f7fe commit 017e9cf

File tree

3 files changed

+35
-0
lines changed

3 files changed

+35
-0
lines changed

spec/draft/API_specification/array_object.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,7 @@ Attributes
273273
array.shape
274274
array.size
275275
array.T
276+
array.__dlpack_c_exchange_api__
276277

277278
-------------------------------------------------
278279

spec/draft/design_topics/data_interchange.rst

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,19 @@ page gives a high-level specification for data exchange in Python using DLPack.
8585
below. They are not required to return an array object from ``from_dlpack``
8686
which conforms to this standard.
8787

88+
89+
DLPack C Exchange API
90+
---------------------
91+
92+
DLPack 1.3 introduces a C-level exchange API that can be used to speed up
93+
data exchange between arrays at the C-extension level. This API is available via
94+
the ``type(array_instance).__dlpack_c_exchange_api__`` attribute on the array type object.
95+
This is a static global object shared across all the array instances of the same type.
96+
For more details, see the `Python specification of DLPack <https://dmlc.github.io/dlpack/latest/python_spec.html>`__
97+
We recommend consumer libraries to start first using the Python-level ``__dlpack__`` first which will covers
98+
most cases, then start to use the C-level ``__dlpack_c_exchange_api__`` for performance critical cases.
99+
100+
88101
Non-supported use cases
89102
-----------------------
90103

src/array_api_stubs/_draft/array_object.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,27 @@ def T(self: array) -> array:
119119
Limiting the transpose to two-dimensional arrays (matrices) deviates from the NumPy et al practice of reversing all axes for arrays having more than two-dimensions. This is intentional, as reversing all axes was found to be problematic (e.g., conflicting with the mathematical definition of a transpose which is limited to matrices; not operating on batches of matrices; et cetera). In order to reverse all axes, one is recommended to use the functional ``permute_dims`` interface found in this specification.
120120
"""
121121

122+
@property
123+
def __dlpack_c_exchange_api__(self: array) -> PyCapsule:
124+
"""
125+
Object containing the DLPack C-API exchange API struct.
126+
127+
An optional static array type attribute stored in ``type(array_instance).__dlpack_c_exchange_api__``
128+
that can be used to retrieve the DLPack C-API exchange API struct in DLPack 1.3 or later to speed up
129+
exchange of array data at the C extension level without going through Python-level exchange.
130+
See :ref:`data-interchange` section for more details.
131+
132+
Returns
133+
-------
134+
out: PyCapsule
135+
The PyCapsule object containing the DLPack C-API exchange API struct.
136+
137+
138+
.. note::
139+
This is a static global object shared across all the array instances of the same type.
140+
It can be queried through the ``type(array_instance).__dlpack_c_exchange_api__`` attribute.
141+
"""
142+
122143
def __abs__(self: array, /) -> array:
123144
"""
124145
Calculates the absolute value for each element of an array instance.

0 commit comments

Comments
 (0)