diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 03cdd91..ab2dac2 100644 --- a/.github/workflows/tests.yml +++ b/.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 diff --git a/array-api-tests-xfails.txt b/array-api-tests-xfails.txt index b7042ff..21a4342 100644 --- a/array-api-tests-xfails.txt +++ b/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 diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py index 25e6595..e8b4e27 100644 --- a/array_api_strict/_array_object.py +++ b/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] """ diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py index 685044d..ca94e4d 100644 --- a/array_api_strict/_creation_functions.py +++ b/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) diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py index 1883199..79d8bdc 100644 --- a/array_api_strict/_devices.py +++ b/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}") diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py index 445ee80..8a048c4 100644 --- a/array_api_strict/tests/test_array_object.py +++ b/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) diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py index a28be64..4cbce07 100644 --- a/array_api_strict/tests/test_creation_functions.py +++ b/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 + diff --git a/array_api_strict/tests/test_dlpack_interop.py b/array_api_strict/tests/test_dlpack_interop.py new file mode 100644 index 0000000..4ac15c3 --- /dev/null +++ b/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))