xref: /aosp_15_r20/external/pytorch/torch/csrc/Stream.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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