Skip to content

Add API for device support #96

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jan 12, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions spec/API_specification/array_object.md
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,17 @@ Data type of the array elements.

- array data type.

(attribute-device)=
### device

Hardware device the array data resides on.

#### Returns

- **out**: _<device>_

- a `device` object (see {ref}`device-support`).

(attribute-ndim)=
### ndim

Expand Down
66 changes: 55 additions & 11 deletions spec/API_specification/creation_functions.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ A conforming implementation of the array API standard must provide and support t
<!-- NOTE: please keep the functions in alphabetical order -->

(function-arange)=
### arange(start, /, *, stop=None, step=1, dtype=None)
### arange(start, /, *, stop=None, step=1, dtype=None, device=None)

Returns evenly spaced values within the half-open interval `[start, stop)` as a one-dimensional array.

Expand All @@ -39,14 +39,18 @@ This function cannot guarantee that the interval does not include the `stop` val

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- a one-dimensional array containing evenly spaced values. The length of the output array must be `ceil((stop-start)/step)`.

(function-empty)=
### empty(shape, /, *, dtype=None)
### empty(shape, /, *, dtype=None, device=None)

Returns an uninitialized array having a specified `shape`.

Expand All @@ -60,14 +64,18 @@ Returns an uninitialized array having a specified `shape`.

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array containing uninitialized data.

(function-empty_like)=
### empty_like(x, /, *, dtype=None)
### empty_like(x, /, *, dtype=None, device=None)
Copy link
Contributor

@leofang leofang Dec 19, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This reminds me of an interest inquiry we had a while ago: Should the _like*() functions honor the device where the input array x is on? (cupy/cupy#3457)

Now I look back, it seems also plausible if the output array is on the same device as x, but my rejection still holds: it is incompatible with any sane device management approaches (context manager, explicit function call (such as use_device(N)), etc). I suppose having the newly added argument device will encounter the same challenge. The most radical example is x on device 1, the argument device is set to 2, but the default/current device is 0.

Thoughts?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

but my rejection still holds: it is incompatible with any sane device management approaches

I agree with the argument you made on that issue.

I suppose having the newly added argument device will encounter the same challenge. The most radical example is x on device 1, the argument device is set to 2, but the default/current device is 0.

I don't see the conflict here. If a library has multiple ways of controlling device placement, the most explicit method should have the highest priority. So:

  1. If device= keyword is specified, that always takes precedence
  2. If device=None, then use the setting from a context manager, if set.
  3. If no context manager was used, then use the global default device/strategy

Your example seems very similar to the first example of https://pytorch.org/docs/stable/notes/cuda.html#cuda-semantics, which I think explains desired behaviour here.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The empty_like description seems clear enough: the _like is about shape and dtype only.

Copy link
Contributor

@leofang leofang Dec 28, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @rgommers I fully agree with everything you said above, but I noticed in the "Semantics" section that you're adding in this PR, the wording is a bit different and (arguably) less clear. Should we incorporate your above reply there? And it'd be nice to add a variant of the "_like is about shape and dtype only" emphasis to all the docstrings where device is an optional argument.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good ideas, done.


Returns an uninitialized array with the same `shape` as an input array `x`.

Expand All @@ -81,14 +89,18 @@ Returns an uninitialized array with the same `shape` as an input array `x`.

- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array having the same shape as `x` and containing uninitialized data.

(function-eye)=
### eye(N, /, *, M=None, k=0, dtype=None)
### eye(N, /, *, M=None, k=0, dtype=None, device=None)

Returns a two-dimensional array with ones on the `k`th diagonal and zeros elsewhere.

Expand All @@ -110,14 +122,18 @@ Returns a two-dimensional array with ones on the `k`th diagonal and zeros elsewh

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array where all elements are equal to zero, except for the `k`th diagonal, whose values are equal to one.

(function-full)=
### full(shape, fill_value, /, *, dtype=None)
### full(shape, fill_value, /, *, dtype=None, device=None)

Returns a new array having a specified `shape` and filled with `fill_value`.

Expand All @@ -135,14 +151,18 @@ Returns a new array having a specified `shape` and filled with `fill_value`.

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array where every element is equal to `fill_value`.

(function-full_like)=
### full_like(x, fill_value, /, *, dtype=None)
### full_like(x, fill_value, /, *, dtype=None, device=None)

Returns a new array filled with `fill_value` and having the same `shape` as an input array `x`.

Expand All @@ -160,14 +180,18 @@ Returns a new array filled with `fill_value` and having the same `shape` as an i

- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array having the same shape as `x` and where every element is equal to `fill_value`.

(function-linspace)=
### linspace(start, stop, num, /, *, dtype=None, endpoint=True)
### linspace(start, stop, num, /, *, dtype=None, device=None, endpoint=True)

Returns evenly spaced numbers over a specified interval.

Expand All @@ -194,6 +218,10 @@ Returns evenly spaced numbers over a specified interval.

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

- **endpoint**: _Optional\[ bool ]_

- boolean indicating whether to include `stop` in the interval. Default: `True`.
Expand All @@ -205,7 +233,7 @@ Returns evenly spaced numbers over a specified interval.
- a one-dimensional array containing evenly spaced values.

(function-ones)=
### ones(shape, /, *, dtype=None)
### ones(shape, /, *, dtype=None, device=None)

Returns a new array having a specified `shape` and filled with ones.

Expand All @@ -219,14 +247,18 @@ Returns a new array having a specified `shape` and filled with ones.

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array containing ones.

(function-ones_like)=
### ones_like(x, /, *, dtype=None)
### ones_like(x, /, *, dtype=None, device=None)

Returns a new array filled with ones and having the same `shape` as an input array `x`.

Expand All @@ -240,14 +272,18 @@ Returns a new array filled with ones and having the same `shape` as an input arr

- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array having the same shape as `x` and filled with ones.

(function-zeros)=
### zeros(shape, /, *, dtype=None)
### zeros(shape, /, *, dtype=None, device=None)

Returns a new array having a specified `shape` and filled with zeros.

Expand All @@ -261,14 +297,18 @@ Returns a new array having a specified `shape` and filled with zeros.

- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_

- an array containing zeros.

(function-zeros_like)=
### zeros_like(x, /, *, dtype=None)
### zeros_like(x, /, *, dtype=None, device=None)

Returns a new array filled with zeros and having the same `shape` as an input array `x`.

Expand All @@ -282,6 +322,10 @@ Returns a new array filled with zeros and having the same `shape` as an input ar

- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.

- **device**: _Optional\[ &lt;device&gt; ]_

- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.

#### Returns

- **out**: _&lt;array&gt;_
Expand Down
100 changes: 99 additions & 1 deletion spec/design_topics/device_support.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,5 +2,103 @@

# Device support

TODO. See https://github.com/data-apis/array-api/issues/39
For libraries that support execution on more than a single hardware device - e.g. CPU and GPU, or multiple GPUs - it is important to be able to control on which device newly created arrays get placed and where execution happens. Attempting to be fully implicit doesn't always scale well to situations with multiple GPUs.

Existing libraries employ one or more of these three methods to exert such control:
1. A global default device, which may be fixed or user-switchable.
2. A context manager to control device assignment within its scope.
3. Local control via explicit keywords and a method to transfer arrays to another device.

This standard chooses to add support for method 3 (local control), because it's the most explicit and granular, with its only downside being verbosity. A context manager may be added in the future - see {ref}`device-out-of-scope` for details.


## Intended usage

The intended usage for the device support in the current version of the
standard is _device handling in library code_. The assumed pattern is that
users create arrays (for which they can use all the relevant device syntax
that the library they use provides), and that they then pass those arrays
into library code which may have to do the following:

- Create new arrays on the same device as an array that's passed in.
- Determine whether two input arrays are present on the same device or not.
- Move an array from one device to another.
- Create output arrays on the same device as the input arrays.
- Pass on a specified device to other library code.

```{note}
Given that there is not much that's currently common in terms of
device-related syntax between different array libraries, the syntax included
in the standard is kept as minimal as possible while enabling the
above-listed use cases.
```

## Syntax for device assignment

The array API will offer the following syntax for device assignment and
cross-device data transfer:

1. A `.device` property on the array object, which returns a `Device` object
representing the device the data in the array is stored on, and supports
comparing devices for equality with `==` and `!=` within the same library
(e.g., by implementing `__eq__`); comparing device objects from different
libraries is out of scope).
2. A `device=None` keyword for array creation functions, which takes an
instance of a `Device` object.
3. A `.to_device(device)` method on the array object, with `device` again being
a `Device` object, to move an array to a different device.

```{note}
The only way to obtain a `Device` object is from the `.device` property on
the array object, hence there is no `Device` object in the array API itself
that can be instantiated to point to a specific physical or logical device.
```


## Semantics

Handling devices is complex, and some frameworks have elaborate policies for
handling device placement. Therefore this section only gives recommendations,
rather than hard requirements:

- Respect explicit device assignment (i.e. if the input to the `device=` keyword
is not `None`, guarantee that the array is created on the given device, and
raise an exception otherwise).
- Preserve device assignment as much as possible (e.g. output arrays from a
function are expected to be on the same device as input arrays to the
function).
- Raise an exception if an operation involves arrays on different devices
(i.e. avoid implicit data transfer between devices).
- Use a default for `device=None` which is consistent between functions
within the same library.
- If a library has multiple ways of controlling device placement, the most
explicit method should have the highest priority. For example:
1. If `device=` keyword is specified, that always takes precedence
2. If `device=None`, then use the setting from a context manager, if set.
3. If no context manager was used, then use the global default device/strategy


(device-out-of-scope)=

## Out of scope for device support

Individual libraries may offers APIs for one or more of the following topics,
however those are out of scope for this standard:

- Identifying a specific physical or logical device across libraries
- Setting a default device globally
- Stream/queue control
- Distributed allocation
- Memory pinning
- A context manager for device control

```{note}
A context manager for controlling the default device is present in most existing array
libraries (NumPy being the exception). There are concerns with using a
context manager however. A context manager can be tricky to use at a high
level, since it may affect library code below function calls (non-local
effects). See, e.g., [this PyTorch issue](https://github.com/pytorch/pytorch/issues/27878)
for a discussion on a good context manager API.

Adding a context manager may be considered in a future version of this API standard.
```