-
Notifications
You must be signed in to change notification settings - Fork 21.3k
/
CUDAStream.cpp
377 lines (327 loc) · 12.3 KB
/
CUDAStream.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAFunctions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/util/Exception.h>
#include <array>
#include <atomic>
#include <cstdint>
#include <mutex>
#include <vector>
#include <iostream>
namespace c10 {
namespace cuda {
namespace {
// Internal implementation that leaks the stream. It's not intended to be used
// outside of this file.
struct LeakyStreamInternals {
LeakyStreamInternals() = default;
C10_DISABLE_COPY_AND_ASSIGN(LeakyStreamInternals);
~LeakyStreamInternals() {
// NB: this code is invoked only in the destruction of global variables
// (since we never shrink the corresponding vectors). At this point the CUDA
// runtime might be already destroyed and invoking cudaStreamDestroy leads
// to a crash. It's likely an issue in CUDA, but to be safe - let's just
// "forget" the destruction.
// if (stream) cudaStreamDestroy(stream);
}
DeviceIndex device_index = -1;
int32_t stream_id = -1;
cudaStream_t stream = nullptr;
};
// Global stream state and constants
static DeviceIndex num_gpus = -1;
static constexpr int kStreamsPerPoolBits = 5;
static constexpr int kStreamsPerPool = 1 << kStreamsPerPoolBits;
static constexpr unsigned int kDefaultFlags = cudaStreamNonBlocking;
// Note: stream priority is not supported by HIP
// Note: lower numbers are higher priorities, zero is default priority
#ifndef __HIP_PLATFORM_HCC__
static int kHighPriority = -1;
static int kLowPriority = 0;
#endif // __HIP_PLATFORM_HCC__
// Default streams
static std::once_flag init_flag;
static LeakyStreamInternals default_streams[C10_COMPILE_TIME_MAX_GPUS];
// Non-default streams
// Note: the number of CUDA devices is determined at run time,
// and the low and high priority pools are lazily initialized
// when the first stream is requested for a device.
// The device flags track the initialization of each device, while
// the low and high priority counters track, for each device, the next stream
// in the pool to be returned when a stream is requested (round-robin fashion
// , see the note in CUDAStream.h).
//
// unique_ptr<T[]> is used instead of vector<T> because T might be non-moveable
// and non-copyable.
static std::once_flag device_flags[C10_COMPILE_TIME_MAX_GPUS];
static std::atomic<uint32_t> low_priority_counters[C10_COMPILE_TIME_MAX_GPUS];
static std::atomic<uint32_t> high_priority_counters[C10_COMPILE_TIME_MAX_GPUS];
static std::array<LeakyStreamInternals, kStreamsPerPool>
low_priority_streams[C10_COMPILE_TIME_MAX_GPUS];
static std::array<LeakyStreamInternals, kStreamsPerPool>
high_priority_streams[C10_COMPILE_TIME_MAX_GPUS];
// Note [StreamId assignment]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~
// How do we assign stream IDs?
//
// -- 25 bits -- -- 2 bits -- -- 5 bits -----
// zeros StreamIdType stream id index
//
// Where StreamIdType:
// 00 = default stream
// 01 = low priority stream
// 10 = high priority stream
//
// This is not really for efficiency; it's just easier to write the code
// to extract the index if we do this with bitmasks :)
//
// We are obligated to treat the stream ID 0 as the default stream, per the
// invariant specified in c10::Stream. However, all other numbers are entirely
// an internal implementation detail, we reserve the right to renumber streams
// however we like.
//
// Note that it is really important that the MSB is zero; StreamId is a
// *signed* integer, and unsigned to signed conversion outside of the
// bounds of signed integer representation is undefined behavior. You
// could work around this with something like
// https://stackoverflow.com/questions/13150449/efficient-unsigned-to-signed-cast-avoiding-implementation-defined-behavior
// but it seems a bit overkill for this.
enum class StreamIdType : uint8_t {
DEFAULT = 0x0,
LOW = 0x1,
HIGH = 0x2,
};
std::ostream& operator<<(std::ostream& stream, StreamIdType s) {
switch (s) {
case StreamIdType::DEFAULT:
stream << "DEFAULT";
break;
case StreamIdType::LOW:
stream << "LOW";
break;
case StreamIdType::HIGH:
stream << "HIGH";
break;
default:
stream << static_cast<uint8_t>(s);
break;
}
return stream;
}
// StreamId is 32-bit, so we can just rely on regular promotion rules.
// We rely on streamIdIndex and streamIdType being non-negative;
// see Note [Hazard when concatenating signed integers]
static inline StreamIdType streamIdType(StreamId s) {
return static_cast<StreamIdType>(s >> kStreamsPerPoolBits);
}
static inline size_t streamIdIndex(StreamId s) {
return static_cast<size_t>(s & ((1 << kStreamsPerPoolBits) - 1));
}
StreamId makeStreamId(StreamIdType st, size_t si) {
return (static_cast<StreamId>(st) << kStreamsPerPoolBits) |
static_cast<StreamId>(si);
}
template <typename T, typename A>
static bool pointer_within(const T* ptr, const A& arr) {
return std::greater_equal<const T*>()(ptr, arr.data()) &&
std::less<const T*>()(ptr, arr.data() + arr.size());
}
static StreamId CUDAStream_getStreamId(const LeakyStreamInternals* ptr) {
// Hypothetically, we could store the stream ID in the stream. But that
// introduces a degree of freedom which could lead to bugs (where we
// misnumber streams in the pool, or overwrite the number). Better
// to just compute it based on the metric that actually matters,
// which is how we map IDs back into the vectors.
DeviceIndex device_index = ptr->device_index;
// Check if it's the default stream
if (ptr == &default_streams[device_index]) {
return makeStreamId(StreamIdType::DEFAULT, 0);
}
// Check if it's a low priority stream
// NB: Because ptr may not necessarily lie within the array, we must use
// std::less and similar templates to avoid UB that arises when
// doing an operator< comparison.
if (pointer_within<LeakyStreamInternals>(
ptr, low_priority_streams[device_index])) {
return makeStreamId(
StreamIdType::LOW, ptr - low_priority_streams[device_index].data());
}
// Check if it's a high priority stream
if (pointer_within<LeakyStreamInternals>(
ptr, high_priority_streams[device_index])) {
return makeStreamId(
StreamIdType::HIGH, ptr - high_priority_streams[device_index].data());
}
AT_ASSERTM(
0,
"Could not compute stream ID for ",
ptr,
" on device ",
device_index,
" (something has gone horribly wrong!)");
}
// Thread-local current streams
static thread_local LeakyStreamInternals** current_streams = nullptr;
// Populates global values and creates a default stream for each device.
// Note: the default stream on each device is signified by a nullptr,
// and so is not created as usual.
// In particular, we don't need to switch devices when creating the
// streams.
// Warning: this function must only be called once!
static void initGlobalStreamState() {
num_gpus = device_count();
// Check if the number of GPUs matches the expected compile-time max number
// of GPUs.
AT_ASSERTM(
num_gpus <= C10_COMPILE_TIME_MAX_GPUS,
"Number of CUDA devices on the machine is larger than the compiled "
"max number of gpus expected (",
C10_COMPILE_TIME_MAX_GPUS,
"). Increase that and recompile.");
// Initializes default streams
for (auto i = decltype(num_gpus){0}; i < num_gpus; ++i) {
default_streams[i].device_index = i;
low_priority_counters[i] = 0;
high_priority_counters[i] = 0;
}
}
// Creates the low and high priority stream pools for the specified device
// Warning: only call once per device!
static void initDeviceStreamState(DeviceIndex device_index) {
// Switches to the requested device so streams are properly associated
// with it.
CUDAGuard device_guard{device_index};
for (auto i = decltype(kStreamsPerPool){0}; i < kStreamsPerPool; ++i) {
auto& lowpri_stream = low_priority_streams[device_index][i];
auto& hipri_stream = high_priority_streams[device_index][i];
lowpri_stream.device_index = device_index;
hipri_stream.device_index = device_index;
#ifndef __HIP_PLATFORM_HCC__
C10_CUDA_CHECK(cudaStreamCreateWithPriority(
&lowpri_stream.stream, kDefaultFlags, kLowPriority));
C10_CUDA_CHECK(cudaStreamCreateWithPriority(
&hipri_stream.stream, kDefaultFlags, kHighPriority));
#else
C10_CUDA_CHECK(
cudaStreamCreateWithFlags(&lowpri_stream.stream, kDefaultFlags));
C10_CUDA_CHECK(
cudaStreamCreateWithFlags(&hipri_stream.stream, kDefaultFlags));
#endif // __HIP_PLATFORM_HCC__
}
}
// Init front-end to ensure initialization only occurs once
static void initCUDAStreamsOnce() {
// Inits default streams (once, globally)
std::call_once(init_flag, initGlobalStreamState);
if (current_streams) {
return;
}
// Inits current streams (thread local) to default streams
current_streams =
(LeakyStreamInternals**)malloc(num_gpus * sizeof(LeakyStreamInternals*));
for (auto i = decltype(num_gpus){0}; i < num_gpus; ++i) {
current_streams[i] = &default_streams[i];
}
}
// Helper to verify the GPU index is valid
static inline void check_gpu(DeviceIndex device_index) {
AT_ASSERT(device_index >= 0 && device_index < num_gpus);
}
// Helper to determine the index of the stream to return
// Note: Streams are returned round-robin (see note in CUDAStream.h)
static uint32_t get_idx(std::atomic<uint32_t>& counter) {
auto raw_idx = counter++;
return raw_idx % kStreamsPerPool;
}
// See Note [StreamId assignment]
LeakyStreamInternals* CUDAStream_internals(CUDAStream s) {
c10::DeviceIndex device_index = s.device_index();
StreamIdType st = streamIdType(s.unwrap().id());
size_t si = streamIdIndex(s.unwrap().id());
switch (st) {
case StreamIdType::DEFAULT:
AT_ASSERTM(
si == 0,
"Unrecognized stream ",
s.unwrap(),
" (I think this should be the default stream, but I got a non-zero index ",
si,
").",
" Did you manufacture the StreamId yourself? Don't do that; use the",
" official API like c10::cuda::getStreamFromPool() to get a new stream.");
return &default_streams[device_index];
case StreamIdType::LOW:
return &low_priority_streams[device_index][si];
case StreamIdType::HIGH:
return &high_priority_streams[device_index][si];
default:
AT_ASSERTM(
0,
"Unrecognized stream ",
s.unwrap(),
" (I didn't recognize the stream type, ",
st,
")");
}
}
CUDAStream CUDAStream_fromInternals(const LeakyStreamInternals* ptr) {
return CUDAStream(
CUDAStream::UNCHECKED,
Stream(
Stream::UNSAFE,
c10::Device(DeviceType::CUDA, ptr->device_index),
CUDAStream_getStreamId(ptr)));
}
} // anonymous namespace
cudaStream_t CUDAStream::stream() const {
auto ptr = CUDAStream_internals(*this);
AT_ASSERT(ptr);
return ptr->stream;
}
// Returns a stream from the requested pool
// Note: when called the first time on a device, this will create the
// stream pools for that device.
CUDAStream getStreamFromPool(
const bool isHighPriority,
DeviceIndex device_index) {
initCUDAStreamsOnce();
if (device_index == -1)
device_index = current_device();
check_gpu(device_index);
// Initializes the stream pools (once)
std::call_once(
device_flags[device_index], initDeviceStreamState, device_index);
if (isHighPriority) {
const auto idx = get_idx(high_priority_counters[device_index]);
return CUDAStream_fromInternals(&high_priority_streams[device_index][idx]);
}
const auto idx = get_idx(low_priority_counters[device_index]);
return CUDAStream_fromInternals(&low_priority_streams[device_index][idx]);
}
CUDAStream getDefaultCUDAStream(DeviceIndex device_index) {
initCUDAStreamsOnce();
if (device_index == -1) {
device_index = current_device();
}
check_gpu(device_index);
return CUDAStream_fromInternals(&default_streams[device_index]);
}
CUDAStream getCurrentCUDAStream(DeviceIndex device_index) {
initCUDAStreamsOnce();
if (device_index == -1) {
device_index = current_device();
}
check_gpu(device_index);
return CUDAStream_fromInternals(current_streams[device_index]);
}
void setCurrentCUDAStream(CUDAStream stream) {
initCUDAStreamsOnce();
auto ptr = CUDAStream_internals(stream);
AT_ASSERT(ptr);
current_streams[ptr->device_index] = ptr;
}
std::ostream& operator<<(std::ostream& stream, const CUDAStream& s) {
return stream << s.unwrap();
}
} // namespace cuda
} // namespace c10