mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
cbb1ed2966
As described in OpenReg [next-steps](https://github.com/pytorch/pytorch/blob/main/test/cpp_extensions/open_registration_extension/README.md#next-steps), here we replace the current `open_registration_extension.cpp` test in PyTorch CI with openreg. The current `open_registration_extension.cpp` contains two parts: 1. Implentations to support `PrivateUse1` backend. 2. Helper functions used for UTs in `test_cpp_extensions_open_device_registration.py` and `test_transformers.py`. For the first part, we'll replace it with openreg. For the second part, we'll migrate them to ut files step by step. @albanD Pull Request resolved: https://github.com/pytorch/pytorch/pull/141815 Approved by: https://github.com/albanD
28 lines
684 B
C
28 lines
684 B
C
#ifndef THP_STREAM_INC
|
|
#define THP_STREAM_INC
|
|
|
|
#include <c10/core/Stream.h>
|
|
#include <c10/macros/Export.h>
|
|
#include <torch/csrc/Export.h>
|
|
#include <torch/csrc/python_headers.h>
|
|
|
|
struct THPStream {
|
|
PyObject_HEAD
|
|
int64_t stream_id;
|
|
int64_t device_type;
|
|
int64_t device_index;
|
|
// Used to switch stream context management, initialized lazily.
|
|
PyObject* context;
|
|
};
|
|
extern TORCH_API PyTypeObject* THPStreamClass;
|
|
|
|
void THPStream_init(PyObject* module);
|
|
|
|
inline bool THPStream_Check(PyObject* obj) {
|
|
return THPStreamClass && PyObject_IsInstance(obj, (PyObject*)THPStreamClass);
|
|
}
|
|
|
|
TORCH_PYTHON_API PyObject* THPStream_Wrap(const c10::Stream& stream);
|
|
|
|
#endif // THP_STREAM_INC
|