Improve DLPack support for non CPU devices

.github/workflows/tests.yml

@@ -11,11 +11,16 @@ jobs:
           - python-version: '3.10'
             numpy-version: 'latest'
           - python-version: '3.10'
-            numpy-version: 'dev'            
+            numpy-version: 'dev'
           - python-version: '3.13'
             numpy-version: '1.26'
           - python-version: '3.14'
             numpy-version: '1.26'
+        include:
+          # exercise the DLPack interop tests, which need another array library
+          - python-version: '3.13'
+            numpy-version: 'latest'
+            torch: true
       fail-fast: false
     steps:
       - uses: actions/checkout@v7.0.1
@@ -34,6 +39,11 @@ jobs:
           fi
           python -m pip install pytest hypothesis
           python -c'import numpy as np; print(f"{np.__version__ = }")'
+      - name: Install PyTorch
+        if: matrix.torch
+        run: |
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          python -c'import torch; print(f"{torch.__version__ = }")'
       - name: Run Tests
         run: |
           pytest

array-api-tests-xfails.txt

@@ -46,9 +46,5 @@ array_api_tests/test_special_cases.py::test_unary[acosh(real(x_i) is +0 and imag
 
 array_api_tests/test_special_cases.py::test_unary[sqrt(real(x_i) is +infinity and isfinite(imag(x_i)) and imag(x_i) > 0) -> +0 + infinity j]
 
-# remove this when array-api-strict 2.6 is released: this is only for array-api-tests,
-# which runs self-tests against a released array-api-strict
-array_api_tests/test_dlpack.py::test_from_dlpack
-
 # NumPy 1.26 : NotImplementedError: The copy argument to __dlpack__ is not yet implemented
 array_api_tests/test_dlpack.py::test_dunder_dlpack

array_api_strict/_array_object.py

@@ -40,7 +40,13 @@
     _real_to_complex_map,
     _result_type,
 )
-from ._devices import CPU_DEVICE, Device, device_supports_dtype
+from ._devices import (
+    CPU_DEVICE,
+    Device,
+    _DLPACK_DEVICE_FOR,
+    _normalize_dl_device,
+    device_supports_dtype,
+)
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
 from ._typing import PyCapsule
 
@@ -602,6 +608,15 @@ def __dlpack__(
             if copy is not _undef:
                 raise ValueError("The copy argument to __dlpack__ requires at least version 2023.12 of the array API")
 
+        if self._device != CPU_DEVICE:
+            # Consistent with __buffer__ and __array__: data cannot leave a
+            # non-CPU device implicitly. Note that this also rules out consumers
+            # asking for the CPU device via dl_device: the array has to be moved
+            # to the CPU device first, with to_device or asarray.
+            raise BufferError(
+                f"Can't export array on the '{self._device}' device via DLPack."
+            )
+
         if np.lib.NumpyVersion(np.__version__) < '2.1.0':
             if max_version not in [_undef, None]:
                 raise NotImplementedError("The max_version argument to __dlpack__ is not yet implemented")
@@ -611,12 +626,20 @@ def __dlpack__(
                 raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented")
 
             return self._array.__dlpack__(stream=stream)
-        else:
-            kwargs = {'stream': stream}
-            if max_version is not _undef:
-                kwargs['max_version'] = max_version
-            if dl_device is not _undef:
-                kwargs['dl_device'] = dl_device
+
+        kwargs: dict[str, Any] = {'stream': stream}
+        if max_version is not _undef:
+            kwargs['max_version'] = max_version
+        if dl_device is not _undef:
+            if dl_device is not None:
+                requested = _normalize_dl_device(*dl_device)
+                if requested != _normalize_dl_device(*_DLPACK_DEVICE_FOR[self._device]):
+                    raise BufferError("unsupported device requested")
+                # The request is for the device the array is already on, which a
+                # plain export satisfies. NumPy is not told about it, as it only
+                # knows about its own spelling of the CPU device.
+                dl_device = None
+            kwargs['dl_device'] = dl_device
         if copy is not _undef:
             kwargs['copy'] = copy
         return self._array.__dlpack__(**kwargs)
@@ -625,8 +648,7 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]:
         """
         Performs the operation __dlpack_device__.
         """
-        # Note: device support is required for this
-        return self._array.__dlpack_device__()
+        return _DLPACK_DEVICE_FOR[self._device]
 
     def __eq__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
         """

array_api_strict/_creation_functions.py

@@ -7,7 +7,7 @@
 
 from ._dtypes import DType, _all_dtypes, _np_dtype, bool as xp_bool
 from ._devices import (
-    Device, device_supports_dtype, get_default_dtypes,
+    Device, device_supports_dtype, get_default_dtypes, _device_from_dlpack_device,
     check_device as _check_device
 )
 from ._flags import get_array_api_strict_flags
@@ -250,10 +250,28 @@ def from_dlpack(
         _check_device(device)
     else:
         device = None
-
-    if copy in [_undef, None]:
-        # numpy 1.26 does not have the copy= arg
-        return Array._new(np.from_dlpack(x), device=device)
+        if isinstance(x, Array):
+            # All the devices of this library share a DLPack device, so the
+            # logical device is read off the array itself.
+            device = x.device
+        elif hasattr(x, "__dlpack_device__"):
+            dl_type, dl_id = x.__dlpack_device__()
+            device = _device_from_dlpack_device(dl_type, dl_id)
+
+    if isinstance(x, Array):
+        # Arrays of this library are unwrapped instead of going through DLPack:
+        # the buffer is in host memory whatever the logical device is, and the
+        # DLPack export refuses arrays which are not on the CPU device.
+        x = x._array
+
+    if copy is _undef:
+        copy = None
+
+    if np.lib.NumpyVersion(np.__version__) < '2.1.0':
+        # numpy 1.26 does not have the copy= arg, and its from_dlpack never
+        # copies: copy=False needs nothing extra, copy=True is done here.
+        out = np.from_dlpack(x)
+        return Array._new(np.copy(out) if copy else out, device=device)
 
     return Array._new(np.from_dlpack(x, copy=copy), device=device)
 

array_api_strict/_devices.py

@@ -1,3 +1,4 @@
+from enum import IntEnum
 from typing import Final
 
 from ._dtypes import (
@@ -39,6 +40,64 @@ def __hash__(self) -> int:
 )
 
 
+class DLDeviceType(IntEnum):
+    """The DLPack device types, as defined by the DLPack ABI."""
+    kDLCPU = 1
+    kDLCUDA = 2
+    kDLCUDAHost = 3
+    kDLOpenCL = 4
+    kDLVulkan = 7
+    kDLMetal = 8
+    kDLVPI = 9
+    kDLROCM = 10
+    kDLROCMHost = 11
+    kDLExtDev = 12
+    kDLCUDAManaged = 13
+    kDLOneAPI = 14
+    kDLWebGPU = 15
+    kDLHexagon = 16
+    kDLMAIA = 17
+
+
+# All the devices of array_api_strict are fictitious and their data lives in host
+# memory, so they all report the CPU device, which is device number zero. Reporting
+# anything else makes consumers such as pytorch dispatch to the machinery of a
+# device they cannot actually reach, see
+# https://github.com/data-apis/array-api-strict/issues/219
+# The logical devices cannot be told apart through DLPack, which is fine: only the
+# CPU device can be exported at all, and from_dlpack recovers the logical device of
+# an array from this library without going through DLPack.
+_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = {
+    device: (DLDeviceType.kDLCPU, 0) for device in ALL_DEVICES
+}
+
+_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = {
+    (int(DLDeviceType.kDLCPU), 0): CPU_DEVICE,
+}
+
+
+def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]:
+    # `device_type` may be a member of another library's DLPack enum
+    return (int(device_type), device_id)
+
+
+def _device_from_dlpack_device(
+    device_type: IntEnum | int, device_id: int
+) -> Device:
+    key = _normalize_dl_device(device_type, device_id)
+    try:
+        return _DLPACK_DEVICE_TO_LOGICAL[key]
+    except KeyError:
+        try:
+            type_name = DLDeviceType(key[0]).name
+        except ValueError:
+            type_name = str(key[0])
+        raise BufferError(
+            f"No array_api_strict device matches the DLPack device "
+            f"({type_name}, {device_id})."
+        ) from None
+
+
 def check_device(device: Device | None) -> None:
     if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")

array_api_strict/tests/test_array_object.py

@@ -10,7 +10,9 @@
 
 from .. import ones, arange, reshape, asarray, result_type, all, equal, stack
 from .._array_object import Array
-from .._devices import CPU_DEVICE, Device
+from .._devices import (
+    ALL_DEVICES, CPU_DEVICE, Device, DLDeviceType, _DLPACK_DEVICE_FOR
+)
 from .._dtypes import (
     _all_dtypes,
     _boolean_dtypes,
@@ -759,6 +761,61 @@ def test_dlpack_2023_12(api_version):
         a.__dlpack__(copy=None)
 
 
+@pytest.mark.parametrize("device", ALL_DEVICES)
+def test_dlpack_device_numbers(device):
+    a = asarray([1, 2, 3], device=device)
+    # the data of every logical device lives in host memory, so they all report
+    # the CPU device, which is device number zero
+    assert a.__dlpack_device__() == (DLDeviceType.kDLCPU, 0)
+
+
+def test_dlpack_device_map_is_complete():
+    assert set(_DLPACK_DEVICE_FOR) == set(ALL_DEVICES)
+
+
+@pytest.mark.parametrize(
+    "device", [device for device in ALL_DEVICES if device != CPU_DEVICE]
+)
+def test_dlpack_export_from_non_cpu_device(device):
+    a = asarray([1, 2, 3], device=device)
+
+    with pytest.raises(BufferError):
+        a.__dlpack__()
+    with pytest.raises(BufferError):
+        np.from_dlpack(a)
+
+    if np.lib.NumpyVersion(np.__version__) < "2.1.0":
+        return
+
+    # asking for a device explicitly does not help: the array has to be moved
+    # to the CPU device first. Even asking for the CPU device, which is what
+    # __dlpack_device__ reports, is refused: these devices are meant to
+    # represent a GPU or other accelerator, so the consumer would end up with
+    # the data on a device the array is not logically on.
+    with pytest.raises(BufferError):
+        a.__dlpack__(dl_device=a.__dlpack_device__())
+    with pytest.raises(BufferError):
+        a.__dlpack__(dl_device=(DLDeviceType.kDLCPU, 0))
+    with pytest.raises(BufferError):
+        np.from_dlpack(a, device="cpu")
+
+    np.from_dlpack(a.to_device(CPU_DEVICE))
+
+
+def test_dlpack_export_from_cpu_device():
+    a = asarray([1, 2, 3])
+
+    a.__dlpack__()
+    np.from_dlpack(a)
+
+    if np.lib.NumpyVersion(np.__version__) < "2.1.0":
+        return
+
+    a.__dlpack__(dl_device=a.__dlpack_device__())
+    with pytest.raises(BufferError):
+        a.__dlpack__(dl_device=(DLDeviceType.kDLCUDA, 0))
+
+
 def test_pickle():
     """Check that arrays are pickleable (despite raising on `__new__`)"""
     a = ones(2)

array_api_strict/tests/test_creation_functions.py

@@ -24,7 +24,7 @@
 )
 from .._dtypes import float32, float64, complex64, int32, int64, bool as xp_bool
 from .._array_object import Array
-from .._devices import CPU_DEVICE, ALL_DEVICES, Device
+from .._devices import CPU_DEVICE, ALL_DEVICES, Device, DLDeviceType
 from .._info import __array_namespace_info__
 from .._flags import set_array_api_strict_flags
 
@@ -413,3 +413,29 @@ def test_from_dlpack_default_device():
     z = from_dlpack(np.asarray([1, 2, 3]))
     assert x.device == y.device == z.device == CPU_DEVICE
 
+
+@pytest.mark.parametrize("device", ALL_DEVICES)
+def test_from_dlpack_preserves_device(device):
+    x = asarray([1, 2, 3], device=device)
+    y = from_dlpack(x)
+    assert y.device == device
+
+
+def test_from_dlpack_unknown_device():
+    class ForeignArray:
+        """An array on a device which array_api_strict knows nothing about."""
+        def __init__(self):
+            self._array = np.asarray([1, 2, 3])
+
+        def __dlpack_device__(self):
+            return (DLDeviceType.kDLCUDA, 0)
+
+        def __dlpack__(self, **kwargs):
+            return self._array.__dlpack__(**kwargs)
+
+    with pytest.raises(BufferError):
+        from_dlpack(ForeignArray())
+
+    # an explicit device is a request to transfer, so nothing has to be inferred
+    assert from_dlpack(ForeignArray(), device=CPU_DEVICE).device == CPU_DEVICE
+

array_api_strict/tests/test_dlpack_interop.py

@@ -0,0 +1,44 @@
+"""Check how other libraries see arrays of array_api_strict via DLPack.
+
+Tests that libraries like NumPy and PyTorch can import arrays from
+array_api_strict via DLPack.
+
+Only run if torch is available.
+"""
+import numpy as np
+import pytest
+
+import array_api_strict as xp
+from .._devices import ALL_DEVICES, CPU_DEVICE
+
+torch = pytest.importorskip("torch")
+
+
+def test_export_from_cpu_device():
+    x = xp.asarray([1, 2, 3], device=CPU_DEVICE)
+
+    assert np.from_dlpack(x).tolist() == [1, 2, 3]
+    assert torch.from_dlpack(x).tolist() == [1, 2, 3]
+
+
+@pytest.mark.parametrize(
+    "device", [device for device in ALL_DEVICES if device != CPU_DEVICE]
+)
+def test_export_from_other_devices(device):
+    x = xp.asarray([1, 2, 3], device=device)
+
+    with pytest.raises(BufferError):
+        np.from_dlpack(x)
+    with pytest.raises(BufferError):
+        torch.from_dlpack(x)
+
+    assert torch.from_dlpack(x.to_device(CPU_DEVICE)).tolist() == [1, 2, 3]
+
+
+@pytest.mark.parametrize("device", ALL_DEVICES)
+def test_import_from_torch(device):
+    # int32 is the widest integer every device supports
+    x = xp.from_dlpack(torch.asarray([1, 2, 3], dtype=torch.int32), device=device)
+
+    assert x.device == device
+    assert xp.all(x == xp.asarray([1, 2, 3], dtype=xp.int32, device=device))