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
|
// This file registers special JIT operators used to implement the PyTorch CUDA
// API in TorchScript.
#include <torch/csrc/api/include/torch/utils.h>
#include <torch/csrc/jit/cuda/cuda.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/runtime/custom_operator.h>
#include <torch/csrc/jit/runtime/operator.h>
namespace torch::jit {
namespace {
c10::AliasAnalysisKind aliasAnalysisFromSchema() {
return c10::AliasAnalysisKind::FROM_SCHEMA;
}
void _device_synchronize(int64_t device_index) {
// This is a helper API which synchronizes the device identified
// by the device index. The device index of the device is passed as an
// argument to this API.
auto current_device_index = c10::cuda::current_device();
// If the current_device and the device to synchronize are not
// the same, set the device to the device_index of the device
// to synchronize.
if (current_device_index != device_index) {
c10::cuda::set_device(device_index);
}
c10::cuda::device_synchronize();
// Reset the device to current_device before synchronizing.
if (current_device_index != device_index) {
c10::cuda::set_device(current_device_index);
}
}
RegisterOperators const reg({
Operator(
"cuda::current_stream.device(Device? device) -> __torch__.torch.classes.cuda.Stream",
[](Stack& stack) {
auto device = pop(stack).toOptional<c10::Device>();
c10::DeviceIndex device_index = device.has_value()
? device->index()
: c10::cuda::current_device();
auto s = c10::cuda::getCurrentCUDAStream(device_index);
auto st = make_custom_class<torch::jit::CUDAStream>(s);
push(stack, IValue(st));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::current_stream.int(int? val) -> __torch__.torch.classes.cuda.Stream",
[](Stack& stack) {
auto idx = pop(stack).toOptional<c10::DeviceIndex>();
c10::DeviceIndex device_index =
idx.has_value() ? idx.value() : c10::cuda::current_device();
auto s = c10::cuda::getCurrentCUDAStream(device_index);
auto st = make_custom_class<torch::jit::CUDAStream>(s);
push(stack, IValue(st));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::default_stream.device(Device? device) -> __torch__.torch.classes.cuda.Stream",
[](Stack& stack) {
auto device = pop(stack).toOptional<c10::Device>();
c10::DeviceIndex device_index = device.has_value()
? device->index()
: c10::cuda::current_device();
auto s = c10::cuda::getDefaultCUDAStream(device_index);
auto st = make_custom_class<torch::jit::CUDAStream>(s);
push(stack, IValue(st));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::default_stream.int(int? val) -> __torch__.torch.classes.cuda.Stream",
[](Stack& stack) {
auto idx = pop(stack).toOptional<c10::DeviceIndex>();
c10::DeviceIndex device_index =
idx.has_value() ? idx.value() : c10::cuda::current_device();
auto s = c10::cuda::getDefaultCUDAStream(device_index);
auto st = make_custom_class<torch::jit::CUDAStream>(s);
push(stack, IValue(st));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::_current_device() -> int",
[](Stack& stack) {
auto v = c10::cuda::current_device();
push(stack, static_cast<int>(v));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::_exchange_device(int64_t index) -> int",
[](Stack& stack) {
int64_t idx = -1;
pop(stack, idx);
if (idx < 0) {
push(stack, -1);
return;
}
auto prev_idx = c10::cuda::current_device();
c10::cuda::set_device(static_cast<c10::DeviceIndex>(idx));
push(stack, static_cast<int>(prev_idx));
},
// cuda::set_device has side effects.
c10::AliasAnalysisKind::CONSERVATIVE),
Operator(
"cuda::_maybe_exchange_device(int64_t index) -> int",
[](Stack& stack) {
int64_t idx = -1;
pop(stack, idx);
if (idx < 0) {
push(stack, -1);
return;
}
int prev_idx = c10::cuda::MaybeExchangeDevice(static_cast<int>(idx));
push(stack, prev_idx);
},
c10::AliasAnalysisKind::CONSERVATIVE),
Operator(
"cuda::_set_device(int64_t val) -> ()",
[](Stack& stack) {
int64_t idx = -1;
pop(stack, idx);
c10::cuda::set_device(static_cast<c10::DeviceIndex>(idx));
},
aliasAnalysisFromSchema()),
Operator(
"cuda::device_index(Device device) -> int",
[](Stack& stack) {
auto device = pop(stack);
auto idx = device.toDevice().index();
push(stack, idx);
},
aliasAnalysisFromSchema()),
Operator(
"cuda::device_count() -> int",
[](Stack& stack) { push(stack, at::cuda::device_count()); },
aliasAnalysisFromSchema()),
Operator(
"cuda::set_stream(__torch__.torch.classes.cuda.Stream stream) -> ()",
[](Stack& stack) {
auto v = pop(stack);
auto s = v.toCustomClass<torch::jit::CUDAStream>();
auto stream_device_idx = s->device_index();
auto cur_device_idx = c10::cuda::current_device();
// If the stream is not on the current device, change the
// device to the device of the stream.
if (cur_device_idx != stream_device_idx) {
c10::cuda::set_device(stream_device_idx);
}
// To set the current CUDA stream using
// c10::cuda::setCurrentCUDAStream, the jit::CUDAStream object needs
// to be converted to c10::cuda::CUDAStream. Since the latter cannot
// be returned from a class registered via TorchBind, this can only be
// achieved by packing the c10::cuda::CUDAStream instance contained
// inside the jit::CUDAStream object to a struct representation, and
// unpacking it inside this operator. The unpacked stream is then used
// to set the current CUDA stream.
auto unpacked = c10::cuda::CUDAStream::unpack3(
s->id(), stream_device_idx, c10::DeviceType::CUDA);
c10::cuda::setCurrentCUDAStream(unpacked);
},
aliasAnalysisFromSchema()),
Operator(
"cuda::synchronize() -> ()",
[](Stack& stack) { c10::cuda::device_synchronize(); },
aliasAnalysisFromSchema()),
Operator(
"cuda::synchronize.device(Device? device) -> ()",
[](Stack& stack) {
auto device = pop(stack).toOptional<c10::Device>();
c10::DeviceIndex device_index = device.has_value()
? device->index()
: c10::cuda::current_device();
_device_synchronize(device_index);
},
aliasAnalysisFromSchema()),
Operator(
"cuda::synchronize.int(int? val) -> ()",
[](Stack& stack) {
auto idx = pop(stack).toOptional<c10::DeviceIndex>();
c10::DeviceIndex device_index =
idx.has_value() ? idx.value() : c10::cuda::current_device();
_device_synchronize(device_index);
},
aliasAnalysisFromSchema()),
});
} // namespace
} // namespace torch::jit
|
