1*da0073e9SAndroid Build Coastguard Worker #include <pybind11/pybind11.h>
2*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/Device.h>
3*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/Event.h>
4*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/Stream.h>
5*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/THP.h>
6*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/utils/pybind.h>
7*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/utils/pycfunction_helpers.h>
8*da0073e9SAndroid Build Coastguard Worker #include <torch/csrc/utils/python_arg_parser.h>
9*da0073e9SAndroid Build Coastguard Worker
10*da0073e9SAndroid Build Coastguard Worker #include <c10/core/DeviceGuard.h>
11*da0073e9SAndroid Build Coastguard Worker #include <c10/core/Stream.h>
12*da0073e9SAndroid Build Coastguard Worker #include <c10/core/impl/DeviceGuardImplInterface.h>
13*da0073e9SAndroid Build Coastguard Worker #include <c10/util/Exception.h>
14*da0073e9SAndroid Build Coastguard Worker #include <c10/util/hash.h>
15*da0073e9SAndroid Build Coastguard Worker #include <structmember.h>
16*da0073e9SAndroid Build Coastguard Worker #include <cstdint>
17*da0073e9SAndroid Build Coastguard Worker
18*da0073e9SAndroid Build Coastguard Worker PyTypeObject* THPStreamClass = nullptr;
19*da0073e9SAndroid Build Coastguard Worker
THPStream_pynew(PyTypeObject * type,PyObject * args,PyObject * kwargs)20*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_pynew(
21*da0073e9SAndroid Build Coastguard Worker PyTypeObject* type,
22*da0073e9SAndroid Build Coastguard Worker PyObject* args,
23*da0073e9SAndroid Build Coastguard Worker PyObject* kwargs) {
24*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
25*da0073e9SAndroid Build Coastguard Worker
26*da0073e9SAndroid Build Coastguard Worker int64_t stream_id = -1;
27*da0073e9SAndroid Build Coastguard Worker int64_t device_type = 0;
28*da0073e9SAndroid Build Coastguard Worker int64_t device_index = 0;
29*da0073e9SAndroid Build Coastguard Worker int64_t priority = 0;
30*da0073e9SAndroid Build Coastguard Worker
31*da0073e9SAndroid Build Coastguard Worker static torch::PythonArgParser parser({
32*da0073e9SAndroid Build Coastguard Worker "Stream(Device device=None, *, int64_t priority=0)",
33*da0073e9SAndroid Build Coastguard Worker "Stream(int64_t stream_id, int64_t device_index, int64_t device_type, *, int64_t priority=0)",
34*da0073e9SAndroid Build Coastguard Worker });
35*da0073e9SAndroid Build Coastguard Worker
36*da0073e9SAndroid Build Coastguard Worker torch::ParsedArgs<4> parsed_args;
37*da0073e9SAndroid Build Coastguard Worker auto r = parser.parse(args, kwargs, parsed_args);
38*da0073e9SAndroid Build Coastguard Worker
39*da0073e9SAndroid Build Coastguard Worker std::unique_ptr<c10::DeviceGuard> device_guard_ptr;
40*da0073e9SAndroid Build Coastguard Worker
41*da0073e9SAndroid Build Coastguard Worker if (r.idx == 0) {
42*da0073e9SAndroid Build Coastguard Worker auto default_accelerator = at::getAccelerator(false);
43*da0073e9SAndroid Build Coastguard Worker auto device = r.deviceOptional(0);
44*da0073e9SAndroid Build Coastguard Worker if (device.has_value()) {
45*da0073e9SAndroid Build Coastguard Worker device_type = static_cast<int64_t>(device->type());
46*da0073e9SAndroid Build Coastguard Worker device_index = static_cast<int64_t>(device->index());
47*da0073e9SAndroid Build Coastguard Worker // Initialize device guard if device is not None.
48*da0073e9SAndroid Build Coastguard Worker device_guard_ptr = std::make_unique<c10::DeviceGuard>(device.value());
49*da0073e9SAndroid Build Coastguard Worker } else {
50*da0073e9SAndroid Build Coastguard Worker // If device is None, we will use the current accelerator and index.
51*da0073e9SAndroid Build Coastguard Worker // If the current accelerator is not set, we will use the CPU as device
52*da0073e9SAndroid Build Coastguard Worker // type.
53*da0073e9SAndroid Build Coastguard Worker device_type = static_cast<int64_t>(
54*da0073e9SAndroid Build Coastguard Worker default_accelerator.value_or(c10::DeviceType::CPU));
55*da0073e9SAndroid Build Coastguard Worker c10::impl::VirtualGuardImpl impl{
56*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(device_type)};
57*da0073e9SAndroid Build Coastguard Worker const auto current_device = impl.getDevice();
58*da0073e9SAndroid Build Coastguard Worker device_index = current_device.index();
59*da0073e9SAndroid Build Coastguard Worker }
60*da0073e9SAndroid Build Coastguard Worker priority = r.toInt64WithDefault(1, 0);
61*da0073e9SAndroid Build Coastguard Worker } else if (r.idx == 1) {
62*da0073e9SAndroid Build Coastguard Worker stream_id = r.toInt64WithDefault(0, -1);
63*da0073e9SAndroid Build Coastguard Worker device_index = r.toInt64WithDefault(1, 0);
64*da0073e9SAndroid Build Coastguard Worker device_type =
65*da0073e9SAndroid Build Coastguard Worker r.toInt64WithDefault(2, static_cast<int64_t>(c10::DeviceType::CPU));
66*da0073e9SAndroid Build Coastguard Worker priority = r.toInt64WithDefault(3, 0);
67*da0073e9SAndroid Build Coastguard Worker } else {
68*da0073e9SAndroid Build Coastguard Worker TORCH_CHECK(
69*da0073e9SAndroid Build Coastguard Worker false,
70*da0073e9SAndroid Build Coastguard Worker "parse stream arg fails please check the usage: ",
71*da0073e9SAndroid Build Coastguard Worker parser.get_signatures());
72*da0073e9SAndroid Build Coastguard Worker }
73*da0073e9SAndroid Build Coastguard Worker
74*da0073e9SAndroid Build Coastguard Worker THPObjectPtr ptr(type->tp_alloc(type, 0));
75*da0073e9SAndroid Build Coastguard Worker if (!ptr) {
76*da0073e9SAndroid Build Coastguard Worker return nullptr;
77*da0073e9SAndroid Build Coastguard Worker }
78*da0073e9SAndroid Build Coastguard Worker
79*da0073e9SAndroid Build Coastguard Worker THPStream* self = (THPStream*)ptr.get();
80*da0073e9SAndroid Build Coastguard Worker
81*da0073e9SAndroid Build Coastguard Worker // If torch.Stream is not created from existing Stream, then create a new one.
82*da0073e9SAndroid Build Coastguard Worker // It requires other device backends override getNewStream method. How the new
83*da0073e9SAndroid Build Coastguard Worker // stream is created is backend specific. Backend should be able to correctly
84*da0073e9SAndroid Build Coastguard Worker // manage the lifetime of streams.
85*da0073e9SAndroid Build Coastguard Worker std::optional<c10::Stream> stream_opt;
86*da0073e9SAndroid Build Coastguard Worker if (r.idx == 0) {
87*da0073e9SAndroid Build Coastguard Worker c10::impl::VirtualGuardImpl impl{static_cast<c10::DeviceType>(device_type)};
88*da0073e9SAndroid Build Coastguard Worker stream_opt = impl.getNewStream(
89*da0073e9SAndroid Build Coastguard Worker c10::Device(static_cast<c10::DeviceType>(device_type), device_index),
90*da0073e9SAndroid Build Coastguard Worker static_cast<int>(priority));
91*da0073e9SAndroid Build Coastguard Worker } else {
92*da0073e9SAndroid Build Coastguard Worker stream_opt = c10::Stream::unpack3(
93*da0073e9SAndroid Build Coastguard Worker stream_id,
94*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceIndex>(device_index),
95*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(device_type));
96*da0073e9SAndroid Build Coastguard Worker }
97*da0073e9SAndroid Build Coastguard Worker
98*da0073e9SAndroid Build Coastguard Worker TORCH_CHECK(stream_opt.has_value(), "Failed to create stream");
99*da0073e9SAndroid Build Coastguard Worker self->stream_id = static_cast<int64_t>(stream_opt->id());
100*da0073e9SAndroid Build Coastguard Worker self->device_index = static_cast<int64_t>(stream_opt->device_index());
101*da0073e9SAndroid Build Coastguard Worker self->device_type = static_cast<int64_t>(stream_opt->device_type());
102*da0073e9SAndroid Build Coastguard Worker
103*da0073e9SAndroid Build Coastguard Worker return (PyObject*)ptr.release();
104*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
105*da0073e9SAndroid Build Coastguard Worker }
106*da0073e9SAndroid Build Coastguard Worker
THPStream_Wrap(const c10::Stream & stream)107*da0073e9SAndroid Build Coastguard Worker PyObject* THPStream_Wrap(const c10::Stream& stream) {
108*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
109*da0073e9SAndroid Build Coastguard Worker auto type = (PyTypeObject*)THPStreamClass;
110*da0073e9SAndroid Build Coastguard Worker THPObjectPtr ptr(type->tp_alloc(type, 0));
111*da0073e9SAndroid Build Coastguard Worker if (!ptr) {
112*da0073e9SAndroid Build Coastguard Worker throw python_error();
113*da0073e9SAndroid Build Coastguard Worker }
114*da0073e9SAndroid Build Coastguard Worker
115*da0073e9SAndroid Build Coastguard Worker THPStream* self = (THPStream*)ptr.get();
116*da0073e9SAndroid Build Coastguard Worker self->stream_id = stream.id();
117*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(bugprone-signed-char-misuse)
118*da0073e9SAndroid Build Coastguard Worker self->device_index = static_cast<int64_t>(stream.device_index());
119*da0073e9SAndroid Build Coastguard Worker self->device_type = static_cast<int64_t>(stream.device_type());
120*da0073e9SAndroid Build Coastguard Worker return ptr.release();
121*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
122*da0073e9SAndroid Build Coastguard Worker }
123*da0073e9SAndroid Build Coastguard Worker
THPStream_dealloc(THPStream * self)124*da0073e9SAndroid Build Coastguard Worker static void THPStream_dealloc(THPStream* self) {
125*da0073e9SAndroid Build Coastguard Worker Py_TYPE(self)->tp_free((PyObject*)self);
126*da0073e9SAndroid Build Coastguard Worker }
127*da0073e9SAndroid Build Coastguard Worker
THPStream_get_device(THPStream * self,void * unused)128*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_get_device(THPStream* self, void* unused) {
129*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
130*da0073e9SAndroid Build Coastguard Worker return THPDevice_New(c10::Device(
131*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type),
132*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceIndex>(self->device_index)));
133*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
134*da0073e9SAndroid Build Coastguard Worker }
135*da0073e9SAndroid Build Coastguard Worker
THPStream_query(PyObject * _self,PyObject * noargs)136*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_query(PyObject* _self, PyObject* noargs) {
137*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
138*da0073e9SAndroid Build Coastguard Worker auto self = (THPStream*)_self;
139*da0073e9SAndroid Build Coastguard Worker
140*da0073e9SAndroid Build Coastguard Worker return PyBool_FromLong(c10::Stream::unpack3(
141*da0073e9SAndroid Build Coastguard Worker self->stream_id,
142*da0073e9SAndroid Build Coastguard Worker self->device_index,
143*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type))
144*da0073e9SAndroid Build Coastguard Worker .query());
145*da0073e9SAndroid Build Coastguard Worker
146*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
147*da0073e9SAndroid Build Coastguard Worker }
148*da0073e9SAndroid Build Coastguard Worker
THPStream_synchronize(PyObject * _self,PyObject * noargs)149*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_synchronize(PyObject* _self, PyObject* noargs) {
150*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS {
151*da0073e9SAndroid Build Coastguard Worker pybind11::gil_scoped_release no_gil;
152*da0073e9SAndroid Build Coastguard Worker auto self = (THPStream*)_self;
153*da0073e9SAndroid Build Coastguard Worker
154*da0073e9SAndroid Build Coastguard Worker c10::Stream::unpack3(
155*da0073e9SAndroid Build Coastguard Worker self->stream_id,
156*da0073e9SAndroid Build Coastguard Worker self->device_index,
157*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type))
158*da0073e9SAndroid Build Coastguard Worker .synchronize();
159*da0073e9SAndroid Build Coastguard Worker }
160*da0073e9SAndroid Build Coastguard Worker Py_RETURN_NONE;
161*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
162*da0073e9SAndroid Build Coastguard Worker }
163*da0073e9SAndroid Build Coastguard Worker
THPStream_wait_event(PyObject * _self,PyObject * _event)164*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_wait_event(PyObject* _self, PyObject* _event) {
165*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS {
166*da0073e9SAndroid Build Coastguard Worker auto self = (THPStream*)_self;
167*da0073e9SAndroid Build Coastguard Worker auto event = (THPEvent*)_event;
168*da0073e9SAndroid Build Coastguard Worker c10::Stream::unpack3(
169*da0073e9SAndroid Build Coastguard Worker self->stream_id,
170*da0073e9SAndroid Build Coastguard Worker self->device_index,
171*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type))
172*da0073e9SAndroid Build Coastguard Worker .wait(event->event);
173*da0073e9SAndroid Build Coastguard Worker }
174*da0073e9SAndroid Build Coastguard Worker Py_RETURN_NONE;
175*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
176*da0073e9SAndroid Build Coastguard Worker }
177*da0073e9SAndroid Build Coastguard Worker
THPStream_wait_stream(PyObject * _self,PyObject * _other)178*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_wait_stream(PyObject* _self, PyObject* _other) {
179*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS {
180*da0073e9SAndroid Build Coastguard Worker auto self = (THPStream*)_self;
181*da0073e9SAndroid Build Coastguard Worker auto other_stream = (THPStream*)_other;
182*da0073e9SAndroid Build Coastguard Worker c10::Event new_event(
183*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(other_stream->device_type),
184*da0073e9SAndroid Build Coastguard Worker c10::EventFlag::PYTORCH_DEFAULT);
185*da0073e9SAndroid Build Coastguard Worker new_event.record(c10::Stream::unpack3(
186*da0073e9SAndroid Build Coastguard Worker other_stream->stream_id,
187*da0073e9SAndroid Build Coastguard Worker other_stream->device_index,
188*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(other_stream->device_type)));
189*da0073e9SAndroid Build Coastguard Worker c10::Stream::unpack3(
190*da0073e9SAndroid Build Coastguard Worker self->stream_id,
191*da0073e9SAndroid Build Coastguard Worker self->device_index,
192*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type))
193*da0073e9SAndroid Build Coastguard Worker .wait(new_event);
194*da0073e9SAndroid Build Coastguard Worker }
195*da0073e9SAndroid Build Coastguard Worker Py_RETURN_NONE;
196*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
197*da0073e9SAndroid Build Coastguard Worker }
198*da0073e9SAndroid Build Coastguard Worker
THPStream_record_event(PyObject * _self,PyObject * args,PyObject * kwargs)199*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_record_event(
200*da0073e9SAndroid Build Coastguard Worker PyObject* _self,
201*da0073e9SAndroid Build Coastguard Worker PyObject* args,
202*da0073e9SAndroid Build Coastguard Worker PyObject* kwargs) {
203*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
204*da0073e9SAndroid Build Coastguard Worker auto self = (THPStream*)_self;
205*da0073e9SAndroid Build Coastguard Worker PyObject* _new_event;
206*da0073e9SAndroid Build Coastguard Worker PyObject* _event = Py_None;
207*da0073e9SAndroid Build Coastguard Worker
208*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
209*da0073e9SAndroid Build Coastguard Worker constexpr const char* accepted_args[] = {"event", nullptr};
210*da0073e9SAndroid Build Coastguard Worker if (!PyArg_ParseTupleAndKeywords(
211*da0073e9SAndroid Build Coastguard Worker args,
212*da0073e9SAndroid Build Coastguard Worker kwargs,
213*da0073e9SAndroid Build Coastguard Worker "|O",
214*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
215*da0073e9SAndroid Build Coastguard Worker const_cast<char**>(accepted_args),
216*da0073e9SAndroid Build Coastguard Worker &_event)) {
217*da0073e9SAndroid Build Coastguard Worker TORCH_CHECK(false, "parse record_event arg fails");
218*da0073e9SAndroid Build Coastguard Worker }
219*da0073e9SAndroid Build Coastguard Worker if (_event != Py_None) {
220*da0073e9SAndroid Build Coastguard Worker // Increase the refcount of the event to avoid it being destroyed.
221*da0073e9SAndroid Build Coastguard Worker Py_INCREF(_event);
222*da0073e9SAndroid Build Coastguard Worker _new_event = _event;
223*da0073e9SAndroid Build Coastguard Worker } else {
224*da0073e9SAndroid Build Coastguard Worker _new_event = THPEvent_new(
225*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type),
226*da0073e9SAndroid Build Coastguard Worker c10::EventFlag::PYTORCH_DEFAULT);
227*da0073e9SAndroid Build Coastguard Worker }
228*da0073e9SAndroid Build Coastguard Worker auto new_event = (THPEvent*)_new_event;
229*da0073e9SAndroid Build Coastguard Worker TORCH_CHECK(new_event, "event must not be null");
230*da0073e9SAndroid Build Coastguard Worker new_event->event.record(c10::Stream::unpack3(
231*da0073e9SAndroid Build Coastguard Worker self->stream_id,
232*da0073e9SAndroid Build Coastguard Worker self->device_index,
233*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type)));
234*da0073e9SAndroid Build Coastguard Worker return (PyObject*)new_event;
235*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
236*da0073e9SAndroid Build Coastguard Worker }
237*da0073e9SAndroid Build Coastguard Worker
THPStream_repr(THPStream * self)238*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_repr(THPStream* self) {
239*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
240*da0073e9SAndroid Build Coastguard Worker return THPUtils_packString(
241*da0073e9SAndroid Build Coastguard Worker "torch.Stream device_type=" +
242*da0073e9SAndroid Build Coastguard Worker c10::DeviceTypeName(
243*da0073e9SAndroid Build Coastguard Worker static_cast<c10::DeviceType>(self->device_type), true) +
244*da0073e9SAndroid Build Coastguard Worker ", device_index=" + std::to_string(self->device_index) +
245*da0073e9SAndroid Build Coastguard Worker ", stream_id=" + std::to_string(self->stream_id));
246*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
247*da0073e9SAndroid Build Coastguard Worker }
248*da0073e9SAndroid Build Coastguard Worker
THPStream_hash(THPStream * self)249*da0073e9SAndroid Build Coastguard Worker static Py_hash_t THPStream_hash(THPStream* self) {
250*da0073e9SAndroid Build Coastguard Worker return static_cast<long>(at::hash_combine(
251*da0073e9SAndroid Build Coastguard Worker self->device_type,
252*da0073e9SAndroid Build Coastguard Worker (at::hash_combine(self->stream_id, self->device_index))));
253*da0073e9SAndroid Build Coastguard Worker }
254*da0073e9SAndroid Build Coastguard Worker
THPStream_eq(THPStream * self,THPStream * other)255*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_eq(THPStream* self, THPStream* other) {
256*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
257*da0073e9SAndroid Build Coastguard Worker return PyBool_FromLong(
258*da0073e9SAndroid Build Coastguard Worker (self->stream_id == other->stream_id) &&
259*da0073e9SAndroid Build Coastguard Worker (self->device_index == other->device_index) &&
260*da0073e9SAndroid Build Coastguard Worker (self->device_type == other->device_type));
261*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
262*da0073e9SAndroid Build Coastguard Worker }
263*da0073e9SAndroid Build Coastguard Worker
THPStream_ne(THPStream * self,THPStream * other)264*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_ne(THPStream* self, THPStream* other) {
265*da0073e9SAndroid Build Coastguard Worker HANDLE_TH_ERRORS
266*da0073e9SAndroid Build Coastguard Worker return PyBool_FromLong(
267*da0073e9SAndroid Build Coastguard Worker (self->stream_id != other->stream_id) ||
268*da0073e9SAndroid Build Coastguard Worker (self->device_index != other->device_index) ||
269*da0073e9SAndroid Build Coastguard Worker (self->device_type != other->device_type));
270*da0073e9SAndroid Build Coastguard Worker END_HANDLE_TH_ERRORS
271*da0073e9SAndroid Build Coastguard Worker }
272*da0073e9SAndroid Build Coastguard Worker
THPStream_richcompare(PyObject * self,PyObject * other,int op)273*da0073e9SAndroid Build Coastguard Worker static PyObject* THPStream_richcompare(
274*da0073e9SAndroid Build Coastguard Worker PyObject* self,
275*da0073e9SAndroid Build Coastguard Worker PyObject* other,
276*da0073e9SAndroid Build Coastguard Worker int op) {
277*da0073e9SAndroid Build Coastguard Worker PyObject* result = NULL;
278*da0073e9SAndroid Build Coastguard Worker if (other == Py_None) {
279*da0073e9SAndroid Build Coastguard Worker result = Py_False;
280*da0073e9SAndroid Build Coastguard Worker } else {
281*da0073e9SAndroid Build Coastguard Worker switch (op) {
282*da0073e9SAndroid Build Coastguard Worker case Py_EQ:
283*da0073e9SAndroid Build Coastguard Worker result = THPStream_eq((THPStream*)self, (THPStream*)other);
284*da0073e9SAndroid Build Coastguard Worker break;
285*da0073e9SAndroid Build Coastguard Worker case Py_NE:
286*da0073e9SAndroid Build Coastguard Worker result = THPStream_ne((THPStream*)self, (THPStream*)other);
287*da0073e9SAndroid Build Coastguard Worker break;
288*da0073e9SAndroid Build Coastguard Worker default:
289*da0073e9SAndroid Build Coastguard Worker result = Py_False;
290*da0073e9SAndroid Build Coastguard Worker break;
291*da0073e9SAndroid Build Coastguard Worker }
292*da0073e9SAndroid Build Coastguard Worker }
293*da0073e9SAndroid Build Coastguard Worker Py_XINCREF(result);
294*da0073e9SAndroid Build Coastguard Worker return result;
295*da0073e9SAndroid Build Coastguard Worker }
296*da0073e9SAndroid Build Coastguard Worker
297*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
298*da0073e9SAndroid Build Coastguard Worker static struct PyMemberDef THPStream_members[] = {
299*da0073e9SAndroid Build Coastguard Worker {"stream_id",
300*da0073e9SAndroid Build Coastguard Worker T_LONGLONG,
301*da0073e9SAndroid Build Coastguard Worker offsetof(THPStream, stream_id),
302*da0073e9SAndroid Build Coastguard Worker READONLY,
303*da0073e9SAndroid Build Coastguard Worker nullptr},
304*da0073e9SAndroid Build Coastguard Worker {"device_index",
305*da0073e9SAndroid Build Coastguard Worker T_LONGLONG,
306*da0073e9SAndroid Build Coastguard Worker offsetof(THPStream, device_index),
307*da0073e9SAndroid Build Coastguard Worker READONLY,
308*da0073e9SAndroid Build Coastguard Worker nullptr},
309*da0073e9SAndroid Build Coastguard Worker {"device_type",
310*da0073e9SAndroid Build Coastguard Worker T_LONGLONG,
311*da0073e9SAndroid Build Coastguard Worker offsetof(THPStream, device_type),
312*da0073e9SAndroid Build Coastguard Worker READONLY,
313*da0073e9SAndroid Build Coastguard Worker nullptr},
314*da0073e9SAndroid Build Coastguard Worker {nullptr}};
315*da0073e9SAndroid Build Coastguard Worker
316*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
317*da0073e9SAndroid Build Coastguard Worker static struct PyGetSetDef THPStream_properties[] = {
318*da0073e9SAndroid Build Coastguard Worker {"device", (getter)THPStream_get_device, nullptr, nullptr, nullptr},
319*da0073e9SAndroid Build Coastguard Worker {nullptr}};
320*da0073e9SAndroid Build Coastguard Worker
321*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
322*da0073e9SAndroid Build Coastguard Worker static PyMethodDef THPStream_methods[] = {
323*da0073e9SAndroid Build Coastguard Worker {"query", THPStream_query, METH_NOARGS, nullptr},
324*da0073e9SAndroid Build Coastguard Worker {"synchronize", THPStream_synchronize, METH_NOARGS, nullptr},
325*da0073e9SAndroid Build Coastguard Worker {"wait_event", THPStream_wait_event, METH_O, nullptr},
326*da0073e9SAndroid Build Coastguard Worker {"wait_stream", THPStream_wait_stream, METH_O, nullptr},
327*da0073e9SAndroid Build Coastguard Worker {"record_event",
328*da0073e9SAndroid Build Coastguard Worker castPyCFunctionWithKeywords(THPStream_record_event),
329*da0073e9SAndroid Build Coastguard Worker METH_VARARGS | METH_KEYWORDS,
330*da0073e9SAndroid Build Coastguard Worker nullptr},
331*da0073e9SAndroid Build Coastguard Worker {"__eq__", (PyCFunction)THPStream_eq, METH_O, nullptr},
332*da0073e9SAndroid Build Coastguard Worker {nullptr}};
333*da0073e9SAndroid Build Coastguard Worker
334*da0073e9SAndroid Build Coastguard Worker PyTypeObject THPStreamType = {
335*da0073e9SAndroid Build Coastguard Worker PyVarObject_HEAD_INIT(nullptr, 0) "torch.Stream", /* tp_name */
336*da0073e9SAndroid Build Coastguard Worker sizeof(THPStream), /* tp_basicsize */
337*da0073e9SAndroid Build Coastguard Worker 0, /* tp_itemsize */
338*da0073e9SAndroid Build Coastguard Worker (destructor)THPStream_dealloc, /* tp_dealloc */
339*da0073e9SAndroid Build Coastguard Worker 0, /* tp_vectorcall_offset */
340*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_getattr */
341*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_setattr */
342*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_reserved */
343*da0073e9SAndroid Build Coastguard Worker (reprfunc)THPStream_repr, /* tp_repr */
344*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_as_number */
345*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_as_sequence */
346*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_as_mapping */
347*da0073e9SAndroid Build Coastguard Worker (hashfunc)THPStream_hash, /* tp_hash */
348*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_call */
349*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_str */
350*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_getattro */
351*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_setattro */
352*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_as_buffer */
353*da0073e9SAndroid Build Coastguard Worker // NOLINTNEXTLINE(misc-redundant-expression)
354*da0073e9SAndroid Build Coastguard Worker Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
355*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_doc */
356*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_traverse */
357*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_clear */
358*da0073e9SAndroid Build Coastguard Worker THPStream_richcompare, /* tp_richcompare */
359*da0073e9SAndroid Build Coastguard Worker 0, /* tp_weaklistoffset */
360*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_iter */
361*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_iternext */
362*da0073e9SAndroid Build Coastguard Worker THPStream_methods, /* tp_methods */
363*da0073e9SAndroid Build Coastguard Worker THPStream_members, /* tp_members */
364*da0073e9SAndroid Build Coastguard Worker THPStream_properties, /* tp_getset */
365*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_base */
366*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_dict */
367*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_descr_get */
368*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_descr_set */
369*da0073e9SAndroid Build Coastguard Worker 0, /* tp_dictoffset */
370*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_init */
371*da0073e9SAndroid Build Coastguard Worker nullptr, /* tp_alloc */
372*da0073e9SAndroid Build Coastguard Worker THPStream_pynew, /* tp_new */
373*da0073e9SAndroid Build Coastguard Worker };
374*da0073e9SAndroid Build Coastguard Worker
THPStream_init(PyObject * module)375*da0073e9SAndroid Build Coastguard Worker void THPStream_init(PyObject* module) {
376*da0073e9SAndroid Build Coastguard Worker THPStreamClass = &THPStreamType;
377*da0073e9SAndroid Build Coastguard Worker Py_SET_TYPE(&THPStreamType, &PyType_Type);
378*da0073e9SAndroid Build Coastguard Worker if (PyType_Ready(&THPStreamType) < 0) {
379*da0073e9SAndroid Build Coastguard Worker throw python_error();
380*da0073e9SAndroid Build Coastguard Worker }
381*da0073e9SAndroid Build Coastguard Worker Py_INCREF(&THPStreamType);
382*da0073e9SAndroid Build Coastguard Worker if (PyModule_AddObject(module, "Stream", (PyObject*)&THPStreamType) < 0) {
383*da0073e9SAndroid Build Coastguard Worker throw python_error();
384*da0073e9SAndroid Build Coastguard Worker }
385*da0073e9SAndroid Build Coastguard Worker }
386