array-api-compat icon indicating copy to clipboard operation
array-api-compat copied to clipboard

Remove "cpu" device

Open asmeurer opened this issue 5 months ago • 8 comments

Right now we have and support "cpu" as a device for all supported libraries in the to_device() helper. In particular, we allow to_device(cupy_array, "cpu") to convert a CuPy array to a NumPy array.

This was implemented as a way to portably move an array to the host device, as per https://github.com/data-apis/array-api/issues/626. However, it looks like that issue is going to be resolved in a different way, by adding a device keyword to from_dlpack https://github.com/data-apis/array-api/pull/741.

So I propose we remove this behavior, as it extends things in a way that isn't specified by the standard. Since I'm not sure who is using this, we may want to issue a deprecation warning first.

CC @tylerjereddy @leofang

asmeurer avatar Feb 08 '24 20:02 asmeurer

we aren't using it in SciPy at least - we now just use the test infra from the libraries we are testing with to remain on the same device.

lucascolley avatar Feb 08 '24 20:02 lucascolley

+1 for slow deprecation in favor of the more explicit cross-device / namespace API provided by from_dlpack once it is officially part of the spec and reasonably well implemented/adopted.

ogrisel avatar Feb 09 '24 13:02 ogrisel

I have a deprecation for using "cpu" in to_device for CuPy open at https://github.com/data-apis/array-api-compat/pull/87, but I don't think it's a good idea to merge it until the from_dlpack device support is actually implemented. Though if others disagree I can merge it sooner (it's just a deprecation warning).

asmeurer avatar Mar 08 '24 20:03 asmeurer

Agreed, it seems too early to merge that.

rgommers avatar Mar 10 '24 20:03 rgommers

BTW, is there an issue/PR to track implementation of the device= and copy= kwarg and more generally cross-device transfers in xp.from_dlpack?

dlpack itself has been released (in RC) with an updated spec that says:

  • https://github.com/dmlc/dlpack/blob/main/docs/source/python_spec.rst

Starting Python array API standard v2023, a copy can be explicitly requested (or disabled) through the new copy argument of from_dlpack(). When a copy is made, the producer must set the DLPACK_FLAG_BITMASK_IS_COPIED bit flag. It is also possible to request cross-device copies through the new device argument, though the v2023 standard only mandates the support of kDLCPU.

but then array-api-compat does not seem to have been updated yet:

>>> import array_api_compat.torch as xp_torch
>>> import array_api_compat.numpy as xp_np
>>> xp_np.from_dlpack(xp_torch.ones(3, device="mps"))
Traceback (most recent call last):
  Cell In[40], line 1
    xp_np.from_dlpack(xp_torch.ones(3, device="mps"))
  File ~/miniforge3/envs/dev/lib/python3.11/site-packages/torch/_tensor.py:1450 in __dlpack__
    return torch.to_dlpack(self)
RuntimeError: Cannot pack tensors on mps:0

Similarly with array-api-strict:

>>> import array_api_strict as xp_strict
>>> xp_strict.from_dlpack(xp_torch.ones(3, device="mps"))
Traceback (most recent call last):
  Cell In[42], line 1
    xp_strict.from_dlpack(xp_torch.ones(3, device="mps"))
  File ~/miniforge3/envs/dev/lib/python3.11/site-packages/array_api_strict/_creation_functions.py:185 in from_dlpack
    return Array._new(np.from_dlpack(x))
  File ~/miniforge3/envs/dev/lib/python3.11/site-packages/torch/_tensor.py:1450 in __dlpack__
    return torch.to_dlpack(self)
RuntimeError: Cannot pack tensors on mps:0

I think the "cpu" device should not be deprecated while it's not possible to use from_dlpack to convert stuff back to CPU/numpy.

ogrisel avatar Mar 20 '24 14:03 ogrisel

I'm definitely not going to deprecate "cpu" until this gets implemented upstream.

Is it possible to implement wrapping support for this in this library? My understanding was that we would need to just wait for upstream support. Keep in mind that array-api-compat should consist of pure Python wrappers only.

By the way, the "cpu" device that is being deprecated here only refers to using CuPy with the to_device() helper, where we manually support it as a device. The use of "cpu" with numpy or pytorch will continue to work because that's just a feature of the underlying library.

asmeurer avatar Mar 20 '24 18:03 asmeurer

Indeed there is nothing to do in array-api-compat besides waiting for upstream. For the specific case of being able to replace to_device(cupy_array, "cpu") by np.from_dlpack(cupy_array), I am not sure which would need to be updated: numpy.from_dlpack or cupy's __dlpack__ / __dlpack_device__ or both.

For array-api-strict on the other hand this will require adding the new kwargs (copy and device) to implement the new spec fully.

ogrisel avatar Mar 22 '24 17:03 ogrisel

There was some discussion at the meeting Thursday that there were some things that could be done here, but I don't know if I completely followed it. I will need to rewatch the recording. There was a reference to https://github.com/data-apis/array-api/pull/741#discussion_r1486470895.

asmeurer avatar Mar 25 '24 22:03 asmeurer