19
19
20
20
21
21
class Device :
22
+ """Represent a GPU and act as an entry point for cuda.core features.
22
23
24
+ This is a singleton object that helps ensure interoperability
25
+ across multiple libraries imported in the process to both see
26
+ and use the same GPU device.
27
+
28
+ While acting as the entry point, many other CUDA resources can be
29
+ allocated such as streams and buffers. Any :obj:`Context` dependent
30
+ resource created through this device, will continue to refer to
31
+ this device's context.
32
+
33
+ Newly returend :obj:`Device` object are is a thread-local singleton
34
+ for a specified device.
35
+
36
+ Note
37
+ ----
38
+ Will not initialize the GPU.
39
+
40
+ Parameters
41
+ ----------
42
+ device_id : int, optional
43
+ Device ordinal to return a :obj:`Device` object for.
44
+ Default value of `None` return the currently used device.
45
+
46
+ """
23
47
__slots__ = ("_id" , "_mr" , "_has_inited" )
24
48
25
49
def __new__ (cls , device_id = None ):
@@ -54,15 +78,29 @@ def _check_context_initialized(self, *args, **kwargs):
54
78
55
79
@property
56
80
def device_id (self ) -> int :
81
+ """Return device ordinal."""
57
82
return self ._id
58
83
59
84
@property
60
85
def pci_bus_id (self ) -> str :
86
+ """Return a PCI Bus Id string for this device."""
61
87
bus_id = handle_return (cudart .cudaDeviceGetPCIBusId (13 , self ._id ))
62
88
return bus_id [:12 ].decode ()
63
89
64
90
@property
65
91
def uuid (self ) -> str :
92
+ """Return a UUID for the device.
93
+
94
+ Returns 16-octets identifying the device. If the device is in
95
+ MIG mode, returns its MIG UUID which uniquely identifies the
96
+ subscribed MIG compute instance.
97
+
98
+ Note
99
+ ----
100
+ MIG UUID is only returned when device is in MIG mode and the
101
+ driver is older than CUDA 11.4.
102
+
103
+ """
66
104
driver_ver = handle_return (cuda .cuDriverGetVersion ())
67
105
if driver_ver >= 11040 :
68
106
uuid = handle_return (cuda .cuDeviceGetUuid_v2 (self ._id ))
@@ -74,19 +112,21 @@ def uuid(self) -> str:
74
112
75
113
@property
76
114
def name (self ) -> str :
77
- # assuming a GPU name is less than 128 characters...
78
- name = handle_return (cuda .cuDeviceGetName (128 , self ._id ))
115
+ """Return the device name."""
116
+ # Use 256 characters to be consistent with CUDA Runtime
117
+ name = handle_return (cuda .cuDeviceGetName (256 , self ._id ))
79
118
name = name .split (b'\0 ' )[0 ]
80
119
return name .decode ()
81
120
82
121
@property
83
122
def properties (self ) -> dict :
123
+ """Return information about the compute-device."""
84
124
# TODO: pythonize the key names
85
125
return handle_return (cudart .cudaGetDeviceProperties (self ._id ))
86
126
87
127
@property
88
128
def compute_capability (self ) -> ComputeCapability :
89
- """Returns a named tuple with 2 fields: major and minor. """
129
+ """Return a named tuple with 2 fields: major and minor."""
90
130
major = handle_return (cudart .cudaDeviceGetAttribute (
91
131
cudart .cudaDeviceAttr .cudaDevAttrComputeCapabilityMajor , self ._id ))
92
132
minor = handle_return (cudart .cudaDeviceGetAttribute (
@@ -96,12 +136,20 @@ def compute_capability(self) -> ComputeCapability:
96
136
@property
97
137
@precondition (_check_context_initialized )
98
138
def context (self ) -> Context :
139
+ """Return the current :obj:`Context` associated with this device.
140
+
141
+ Note
142
+ ----
143
+ Device must be initialized.
144
+
145
+ """
99
146
ctx = handle_return (cuda .cuCtxGetCurrent ())
100
147
assert int (ctx ) != 0
101
148
return Context ._from_ctx (ctx , self ._id )
102
149
103
150
@property
104
151
def memory_resource (self ) -> MemoryResource :
152
+ """Return :obj:`MemoryResource` associated with this device."""
105
153
return self ._mr
106
154
107
155
@memory_resource .setter
@@ -112,27 +160,53 @@ def memory_resource(self, mr):
112
160
113
161
@property
114
162
def default_stream (self ) -> Stream :
163
+ """Return default CUDA :obj:`Stream` associated with this device.
164
+
165
+ The type of default stream returned depends on if the environment
166
+ variable CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM is set.
167
+
168
+ If set, returns a per-thread default stream. Otherwise returns
169
+ the legacy stream.
170
+
171
+ """
115
172
return default_stream ()
116
173
117
174
def __int__ (self ):
175
+ """Return device_id."""
118
176
return self ._id
119
177
120
178
def __repr__ (self ):
121
179
return f"<Device { self ._id } ({ self .name } )>"
122
180
123
181
def set_current (self , ctx : Context = None ) -> Union [Context , None ]:
124
- """
125
- Entry point of this object. Users always start a code by
182
+ """Set device to be used for GPU executions.
183
+
184
+ Initializes CUDA and sets the calling thread to a valid CUDA
185
+ context. By default the primary context is used, but optional `ctx`
186
+ parameter can be used to explicitly supply a :obj:`Context` object.
187
+
188
+ Providing a `ctx` causes the previous set context to be popped and returned.
189
+
190
+ Parameters
191
+ ----------
192
+ ctx : :obj:`Context`, optional
193
+ Optional context to push onto this device's current thread stack.
194
+
195
+ Returns
196
+ -------
197
+ Union[:obj:`Context`, None], optional
198
+ Popped context.
199
+
200
+ Examples
201
+ --------
202
+ Acts as an entry point of this object. Users always start a code by
126
203
calling this method, e.g.
127
-
204
+
128
205
>>> from cuda.core.experimental import Device
129
206
>>> dev0 = Device(0)
130
207
>>> dev0.set_current()
131
208
>>> # ... do work on device 0 ...
132
-
133
- The optional ctx argument is for advanced users to bind a
134
- CUDA context with the device. In this case, the previously
135
- set context is popped and returned to the user.
209
+
136
210
"""
137
211
if ctx is not None :
138
212
if not isinstance (ctx , Context ):
@@ -163,25 +237,94 @@ def set_current(self, ctx: Context=None) -> Union[Context, None]:
163
237
self ._has_inited = True
164
238
165
239
def create_context (self , options : ContextOptions = None ) -> Context :
166
- # Create a Context object (but do NOT set it current yet!).
167
- # ContextOptions is a dataclass for setting e.g. affinity or CIG
168
- # options.
240
+ """Create a new :obj:`Context` object.
241
+
242
+ Note
243
+ ----
244
+ The newly context will not be set as current.
245
+
246
+ Parameters
247
+ ----------
248
+ options : :obj:`ContextOptions`, optional
249
+ Customizable dataclass for context creation options.
250
+
251
+ Returns
252
+ -------
253
+ :obj:`Context`
254
+ Newly created context object.
255
+
256
+ """
169
257
raise NotImplementedError ("TODO" )
170
258
171
259
@precondition (_check_context_initialized )
172
260
def create_stream (self , obj = None , options : StreamOptions = None ) -> Stream :
173
- # Create a Stream object by either holding a newly created
174
- # CUDA stream or wrapping an existing foreign object supporting
175
- # the __cuda_stream__ protocol. In the latter case, a reference
176
- # to obj is held internally so that its lifetime is managed.
261
+ """Create a Stream object.
262
+
263
+ New stream objects can be created in two different ways:
264
+
265
+ 1) Create a new CUDA stream with customizable `options`.
266
+ 2) Wrap an existing foreign `obj` supporting the __cuda_stream__ protocol.
267
+
268
+ Option (2) internally holds a reference to the foreign object
269
+ such that the lifetime is managed.
270
+
271
+ Note
272
+ ----
273
+ Device must be initialized.
274
+
275
+ Parameters
276
+ ----------
277
+ obj : Any, optional
278
+ Any object supporting the __cuda_stream__ protocol.
279
+ options : :obj:`StreamOptions`, optional
280
+ Customizable dataclass for stream creation options.
281
+
282
+ Returns
283
+ -------
284
+ :obj:`Stream`
285
+ Newly created stream object.
286
+
287
+ """
177
288
return Stream ._init (obj = obj , options = options )
178
289
179
290
@precondition (_check_context_initialized )
180
291
def allocate (self , size , stream = None ) -> Buffer :
292
+ """Allocate device memory from a specified stream.
293
+
294
+ Allocates device memory of `size` bytes on the specified `stream`
295
+ using the memory resource currently associated with this Device.
296
+
297
+ Parameter `stream` is optional, using a default stream by default.
298
+
299
+ Note
300
+ ----
301
+ Device must be initialized.
302
+
303
+ Parameters
304
+ ----------
305
+ size : int
306
+ Number of bytes to allocate.
307
+ stream : :obj:`Stream`, optional
308
+ The stream establishing the stream ordering semantic.
309
+ Default value of `None` uses default stream.
310
+
311
+ Returns
312
+ -------
313
+ :obj:`Buffer`
314
+ Newly created buffer object.
315
+
316
+ """
181
317
if stream is None :
182
318
stream = default_stream ()
183
319
return self ._mr .allocate (size , stream )
184
320
185
321
@precondition (_check_context_initialized )
186
322
def sync (self ):
323
+ """Synchronize the device.
324
+
325
+ Note
326
+ ----
327
+ Device must be initialized.
328
+
329
+ """
187
330
handle_return (cudart .cudaDeviceSynchronize ())
0 commit comments