| #ifndef THP_STREAM_INC |
| #define THP_STREAM_INC |
| |
| #include <c10/core/Stream.h> |
| #include <c10/macros/Export.h> |
| #include <torch/csrc/python_headers.h> |
| |
| struct THPStream { |
| PyObject_HEAD int64_t stream_id; |
| int64_t device_type; |
| int64_t device_index; |
| }; |
| extern TORCH_API PyTypeObject* THPStreamClass; |
| |
| void THPStream_init(PyObject* module); |
| |
| inline bool THPStream_Check(PyObject* obj) { |
| return THPStreamClass && PyObject_IsInstance(obj, (PyObject*)THPStreamClass); |
| } |
| |
| PyObject* THPStream_Wrap(const c10::Stream& stream); |
| |
| #endif // THP_STREAM_INC |