Skip to content

Commit eb276c1

Browse files
authored
Add API for device support (#96)
* Complete the "device support" section Closes gh-39 * Add device keywords to the creation functions * Add the device object and a device array attribute * Update for review comments, narrow scope of syntax and semantics * Remove device object from API specification * Address review comments about control method priority, and `*_like` * A small tweak on wording for `.device` Assumes a bit less about implementation. A string like `'cpu'` should meet the requirements, and it doesn't have `__neq__`.
1 parent 17b50f5 commit eb276c1

File tree

3 files changed

+165
-12
lines changed

3 files changed

+165
-12
lines changed

spec/API_specification/array_object.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,17 @@ Data type of the array elements.
201201

202202
- array data type.
203203

204+
(attribute-device)=
205+
### device
206+
207+
Hardware device the array data resides on.
208+
209+
#### Returns
210+
211+
- **out**: _<device>_
212+
213+
- a `device` object (see {ref}`device-support`).
214+
204215
(attribute-ndim)=
205216
### ndim
206217

spec/API_specification/creation_functions.md

Lines changed: 55 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ A conforming implementation of the array API standard must provide and support t
1212
<!-- NOTE: please keep the functions in alphabetical order -->
1313

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

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

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

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

42+
- **device**: _Optional\[ &lt;device&gt; ]_
43+
44+
- device to place the created array on, if given. Default: `None`.
45+
4246
#### Returns
4347

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

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

4852
(function-empty)=
49-
### empty(shape, /, *, dtype=None)
53+
### empty(shape, /, *, dtype=None, device=None)
5054

5155
Returns an uninitialized array having a specified `shape`.
5256

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

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

67+
- **device**: _Optional\[ &lt;device&gt; ]_
68+
69+
- device to place the created array on, if given. Default: `None`.
70+
6371
#### Returns
6472

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

6775
- an array containing uninitialized data.
6876

6977
(function-empty_like)=
70-
### empty_like(x, /, *, dtype=None)
78+
### empty_like(x, /, *, dtype=None, device=None)
7179

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

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

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

92+
- **device**: _Optional\[ &lt;device&gt; ]_
93+
94+
- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.
95+
8496
#### Returns
8597

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

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

90102
(function-eye)=
91-
### eye(N, /, *, M=None, k=0, dtype=None)
103+
### eye(N, /, *, M=None, k=0, dtype=None, device=None)
92104

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

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

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

125+
- **device**: _Optional\[ &lt;device&gt; ]_
126+
127+
- device to place the created array on, if given. Default: `None`.
128+
113129
#### Returns
114130

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

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

119135
(function-full)=
120-
### full(shape, fill_value, /, *, dtype=None)
136+
### full(shape, fill_value, /, *, dtype=None, device=None)
121137

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

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

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

154+
- **device**: _Optional\[ &lt;device&gt; ]_
155+
156+
- device to place the created array on, if given. Default: `None`.
157+
138158
#### Returns
139159

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

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

144164
(function-full_like)=
145-
### full_like(x, fill_value, /, *, dtype=None)
165+
### full_like(x, fill_value, /, *, dtype=None, device=None)
146166

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

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

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

183+
- **device**: _Optional\[ &lt;device&gt; ]_
184+
185+
- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.
186+
163187
#### Returns
164188

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

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

169193
(function-linspace)=
170-
### linspace(start, stop, num, /, *, dtype=None, endpoint=True)
194+
### linspace(start, stop, num, /, *, dtype=None, device=None, endpoint=True)
171195

172196
Returns evenly spaced numbers over a specified interval.
173197

@@ -194,6 +218,10 @@ Returns evenly spaced numbers over a specified interval.
194218
195219
- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.
196220
221+
- **device**: _Optional\[ &lt;device&gt; ]_
222+
223+
- device to place the created array on, if given. Default: `None`.
224+
197225
- **endpoint**: _Optional\[ bool ]_
198226
199227
- boolean indicating whether to include `stop` in the interval. Default: `True`.
@@ -205,7 +233,7 @@ Returns evenly spaced numbers over a specified interval.
205233
- a one-dimensional array containing evenly spaced values.
206234
207235
(function-ones)=
208-
### ones(shape, /, *, dtype=None)
236+
### ones(shape, /, *, dtype=None, device=None)
209237
210238
Returns a new array having a specified `shape` and filled with ones.
211239
@@ -219,14 +247,18 @@ Returns a new array having a specified `shape` and filled with ones.
219247
220248
- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.
221249
250+
- **device**: _Optional\[ &lt;device&gt; ]_
251+
252+
- device to place the created array on, if given. Default: `None`.
253+
222254
#### Returns
223255
224256
- **out**: _&lt;array&gt;_
225257
226258
- an array containing ones.
227259
228260
(function-ones_like)=
229-
### ones_like(x, /, *, dtype=None)
261+
### ones_like(x, /, *, dtype=None, device=None)
230262
231263
Returns a new array filled with ones and having the same `shape` as an input array `x`.
232264
@@ -240,14 +272,18 @@ Returns a new array filled with ones and having the same `shape` as an input arr
240272
241273
- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.
242274
275+
- **device**: _Optional\[ &lt;device&gt; ]_
276+
277+
- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.
278+
243279
#### Returns
244280
245281
- **out**: _&lt;array&gt;_
246282
247283
- an array having the same shape as `x` and filled with ones.
248284
249285
(function-zeros)=
250-
### zeros(shape, /, *, dtype=None)
286+
### zeros(shape, /, *, dtype=None, device=None)
251287
252288
Returns a new array having a specified `shape` and filled with zeros.
253289
@@ -261,14 +297,18 @@ Returns a new array having a specified `shape` and filled with zeros.
261297
262298
- output array data type. If `dtype` is `None`, the output array data type must be the default floating-point data type. Default: `None`.
263299
300+
- **device**: _Optional\[ &lt;device&gt; ]_
301+
302+
- device to place the created array on, if given. Default: `None`.
303+
264304
#### Returns
265305
266306
- **out**: _&lt;array&gt;_
267307
268308
- an array containing zeros.
269309
270310
(function-zeros_like)=
271-
### zeros_like(x, /, *, dtype=None)
311+
### zeros_like(x, /, *, dtype=None, device=None)
272312
273313
Returns a new array filled with zeros and having the same `shape` as an input array `x`.
274314
@@ -282,6 +322,10 @@ Returns a new array filled with zeros and having the same `shape` as an input ar
282322
283323
- output array data type. If `dtype` is `None`, the output array data type must be inferred from `x`. Default: `None`.
284324
325+
- **device**: _Optional\[ &lt;device&gt; ]_
326+
327+
- device to place the created array on, if given. If `device` is `None`, the default device must be used, not `x.device`. Default: `None`.
328+
285329
#### Returns
286330
287331
- **out**: _&lt;array&gt;_

spec/design_topics/device_support.md

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,103 @@
22

33
# Device support
44

5-
TODO. See https://github.com/data-apis/array-api/issues/39
5+
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.
66

7+
Existing libraries employ one or more of these three methods to exert such control:
8+
1. A global default device, which may be fixed or user-switchable.
9+
2. A context manager to control device assignment within its scope.
10+
3. Local control via explicit keywords and a method to transfer arrays to another device.
11+
12+
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.
13+
14+
15+
## Intended usage
16+
17+
The intended usage for the device support in the current version of the
18+
standard is _device handling in library code_. The assumed pattern is that
19+
users create arrays (for which they can use all the relevant device syntax
20+
that the library they use provides), and that they then pass those arrays
21+
into library code which may have to do the following:
22+
23+
- Create new arrays on the same device as an array that's passed in.
24+
- Determine whether two input arrays are present on the same device or not.
25+
- Move an array from one device to another.
26+
- Create output arrays on the same device as the input arrays.
27+
- Pass on a specified device to other library code.
28+
29+
```{note}
30+
Given that there is not much that's currently common in terms of
31+
device-related syntax between different array libraries, the syntax included
32+
in the standard is kept as minimal as possible while enabling the
33+
above-listed use cases.
34+
```
35+
36+
## Syntax for device assignment
37+
38+
The array API will offer the following syntax for device assignment and
39+
cross-device data transfer:
40+
41+
1. A `.device` property on the array object, which returns a `Device` object
42+
representing the device the data in the array is stored on, and supports
43+
comparing devices for equality with `==` and `!=` within the same library
44+
(e.g., by implementing `__eq__`); comparing device objects from different
45+
libraries is out of scope).
46+
2. A `device=None` keyword for array creation functions, which takes an
47+
instance of a `Device` object.
48+
3. A `.to_device(device)` method on the array object, with `device` again being
49+
a `Device` object, to move an array to a different device.
50+
51+
```{note}
52+
The only way to obtain a `Device` object is from the `.device` property on
53+
the array object, hence there is no `Device` object in the array API itself
54+
that can be instantiated to point to a specific physical or logical device.
55+
```
56+
57+
58+
## Semantics
59+
60+
Handling devices is complex, and some frameworks have elaborate policies for
61+
handling device placement. Therefore this section only gives recommendations,
62+
rather than hard requirements:
63+
64+
- Respect explicit device assignment (i.e. if the input to the `device=` keyword
65+
is not `None`, guarantee that the array is created on the given device, and
66+
raise an exception otherwise).
67+
- Preserve device assignment as much as possible (e.g. output arrays from a
68+
function are expected to be on the same device as input arrays to the
69+
function).
70+
- Raise an exception if an operation involves arrays on different devices
71+
(i.e. avoid implicit data transfer between devices).
72+
- Use a default for `device=None` which is consistent between functions
73+
within the same library.
74+
- If a library has multiple ways of controlling device placement, the most
75+
explicit method should have the highest priority. For example:
76+
1. If `device=` keyword is specified, that always takes precedence
77+
2. If `device=None`, then use the setting from a context manager, if set.
78+
3. If no context manager was used, then use the global default device/strategy
79+
80+
81+
(device-out-of-scope)=
82+
83+
## Out of scope for device support
84+
85+
Individual libraries may offers APIs for one or more of the following topics,
86+
however those are out of scope for this standard:
87+
88+
- Identifying a specific physical or logical device across libraries
89+
- Setting a default device globally
90+
- Stream/queue control
91+
- Distributed allocation
92+
- Memory pinning
93+
- A context manager for device control
94+
95+
```{note}
96+
A context manager for controlling the default device is present in most existing array
97+
libraries (NumPy being the exception). There are concerns with using a
98+
context manager however. A context manager can be tricky to use at a high
99+
level, since it may affect library code below function calls (non-local
100+
effects). See, e.g., [this PyTorch issue](https://github.com/pytorch/pytorch/issues/27878)
101+
for a discussion on a good context manager API.
102+
103+
Adding a context manager may be considered in a future version of this API standard.
104+
```

0 commit comments

Comments
 (0)