mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
5103ecc5d8
Check all `.cpp` files except `jit` files for readability thoroughly. Pull Request resolved: https://github.com/pytorch/pytorch/pull/164561 Approved by: https://github.com/Skylion007
469 lines
16 KiB
C++
469 lines
16 KiB
C++
#include <pybind11/pybind11.h>
|
|
#include <torch/csrc/Device.h>
|
|
#include <torch/csrc/Event.h>
|
|
#include <torch/csrc/Stream.h>
|
|
#include <torch/csrc/THP.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
#include <torch/csrc/utils/pycfunction_helpers.h>
|
|
#include <torch/csrc/utils/python_arg_parser.h>
|
|
|
|
#include <c10/core/DeviceGuard.h>
|
|
#include <c10/core/Stream.h>
|
|
#include <c10/core/impl/DeviceGuardImplInterface.h>
|
|
#include <c10/util/Exception.h>
|
|
#include <c10/util/hash.h>
|
|
#include <structmember.h>
|
|
#include <cstdint>
|
|
|
|
PyTypeObject* THPStreamClass = nullptr;
|
|
|
|
static PyObject* THPStream_pynew(
|
|
PyTypeObject* type,
|
|
PyObject* args,
|
|
PyObject* kwargs) {
|
|
HANDLE_TH_ERRORS
|
|
|
|
int64_t stream_id = -1;
|
|
c10::DeviceType device_type{};
|
|
c10::DeviceIndex device_index{};
|
|
int64_t priority = 0;
|
|
|
|
static torch::PythonArgParser parser({
|
|
"Stream(Device device=None, *, int64_t priority=0)",
|
|
"Stream(int64_t stream_id, int64_t device_index, int64_t device_type, *, int64_t priority=0)",
|
|
});
|
|
|
|
torch::ParsedArgs<4> parsed_args;
|
|
auto r = parser.parse(args, kwargs, parsed_args);
|
|
|
|
std::unique_ptr<c10::DeviceGuard> device_guard_ptr;
|
|
|
|
if (r.idx == 0) {
|
|
auto default_accelerator = at::getAccelerator(false);
|
|
auto device = r.deviceOptional(0);
|
|
if (device.has_value()) {
|
|
device_type = device->type();
|
|
device_index = device->index();
|
|
// Initialize device guard if device is not None.
|
|
device_guard_ptr = std::make_unique<c10::DeviceGuard>(device.value());
|
|
} else {
|
|
// If device is None, we will use the current accelerator and index.
|
|
// If the current accelerator is not set, we will use the CPU as device
|
|
// type.
|
|
device_type = default_accelerator.value_or(c10::DeviceType::CPU);
|
|
c10::impl::VirtualGuardImpl impl{device_type};
|
|
const auto current_device = impl.getDevice();
|
|
device_index = current_device.index();
|
|
}
|
|
priority = r.toInt64WithDefault(1, 0);
|
|
} else if (r.idx == 1) {
|
|
stream_id = r.toInt64WithDefault(0, -1);
|
|
device_index = static_cast<c10::DeviceIndex>(r.toInt64WithDefault(1, 0));
|
|
device_type = static_cast<c10::DeviceType>(
|
|
r.toInt64WithDefault(2, static_cast<int64_t>(c10::DeviceType::CPU)));
|
|
priority = r.toInt64WithDefault(3, 0);
|
|
} else {
|
|
TORCH_CHECK(
|
|
false,
|
|
"parse stream arg fails please check the usage: ",
|
|
parser.get_signatures());
|
|
}
|
|
|
|
THPObjectPtr ptr(type->tp_alloc(type, 0));
|
|
if (!ptr) {
|
|
return nullptr;
|
|
}
|
|
|
|
THPStream* self = (THPStream*)ptr.get();
|
|
|
|
// If torch.Stream is not created from existing Stream, then create a new one.
|
|
// It requires other device backends override getNewStream method. How the new
|
|
// stream is created is backend specific. Backend should be able to correctly
|
|
// manage the lifetime of streams.
|
|
std::optional<c10::Stream> stream_opt;
|
|
if (r.idx == 0) {
|
|
c10::impl::VirtualGuardImpl impl{device_type};
|
|
stream_opt = impl.getNewStream(
|
|
c10::Device(device_type, device_index), static_cast<int>(priority));
|
|
} else {
|
|
stream_opt = c10::Stream::unpack3(stream_id, device_index, device_type);
|
|
}
|
|
|
|
TORCH_CHECK(stream_opt.has_value(), "Failed to create stream");
|
|
self->stream_id = static_cast<int64_t>(stream_opt->id());
|
|
// NOLINTNEXTLINE(bugprone-signed-char-misuse)
|
|
self->device_index = static_cast<int64_t>(stream_opt->device_index());
|
|
self->device_type = static_cast<int64_t>(stream_opt->device_type());
|
|
self->context = nullptr;
|
|
|
|
return (PyObject*)ptr.release();
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
PyObject* THPStream_Wrap(const c10::Stream& stream) {
|
|
HANDLE_TH_ERRORS
|
|
auto type = THPStreamClass;
|
|
THPObjectPtr ptr(type->tp_alloc(type, 0));
|
|
if (!ptr) {
|
|
throw python_error();
|
|
}
|
|
|
|
THPStream* self = (THPStream*)ptr.get();
|
|
self->stream_id = stream.id();
|
|
// NOLINTNEXTLINE(bugprone-signed-char-misuse)
|
|
self->device_index = static_cast<int64_t>(stream.device_index());
|
|
self->device_type = static_cast<int64_t>(stream.device_type());
|
|
self->context = nullptr;
|
|
return ptr.release();
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static void THPStream_dealloc(THPStream* self) {
|
|
Py_TYPE(self)->tp_free((PyObject*)self);
|
|
}
|
|
|
|
static PyObject* THPStream_get_device(THPStream* self, void* unused) {
|
|
HANDLE_TH_ERRORS
|
|
return THPDevice_New(c10::Device(
|
|
static_cast<c10::DeviceType>(self->device_type),
|
|
static_cast<c10::DeviceIndex>(self->device_index)));
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_query(PyObject* _self, PyObject* noargs) {
|
|
HANDLE_TH_ERRORS
|
|
auto self = (THPStream*)_self;
|
|
|
|
return PyBool_FromLong(c10::Stream::unpack3(
|
|
self->stream_id,
|
|
static_cast<c10::DeviceIndex>(self->device_index),
|
|
static_cast<c10::DeviceType>(self->device_type))
|
|
.query());
|
|
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_synchronize(PyObject* _self, PyObject* noargs) {
|
|
HANDLE_TH_ERRORS {
|
|
pybind11::gil_scoped_release no_gil;
|
|
auto self = (THPStream*)_self;
|
|
|
|
c10::Stream::unpack3(
|
|
self->stream_id,
|
|
static_cast<c10::DeviceIndex>(self->device_index),
|
|
static_cast<c10::DeviceType>(self->device_type))
|
|
.synchronize();
|
|
}
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_wait_event(PyObject* _self, PyObject* _event) {
|
|
HANDLE_TH_ERRORS {
|
|
auto self = (THPStream*)_self;
|
|
auto event = (THPEvent*)_event;
|
|
c10::Stream::unpack3(
|
|
self->stream_id,
|
|
static_cast<c10::DeviceIndex>(self->device_index),
|
|
static_cast<c10::DeviceType>(self->device_type))
|
|
.wait(event->event);
|
|
}
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_wait_stream(PyObject* _self, PyObject* _other) {
|
|
HANDLE_TH_ERRORS {
|
|
auto self = (THPStream*)_self;
|
|
auto other_stream = (THPStream*)_other;
|
|
c10::Event new_event(
|
|
static_cast<c10::DeviceType>(other_stream->device_type),
|
|
c10::EventFlag::PYTORCH_DEFAULT);
|
|
new_event.record(c10::Stream::unpack3(
|
|
other_stream->stream_id,
|
|
static_cast<c10::DeviceIndex>(other_stream->device_index),
|
|
static_cast<c10::DeviceType>(other_stream->device_type)));
|
|
c10::Stream::unpack3(
|
|
self->stream_id,
|
|
static_cast<c10::DeviceIndex>(self->device_index),
|
|
static_cast<c10::DeviceType>(self->device_type))
|
|
.wait(new_event);
|
|
}
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_record_event(
|
|
PyObject* _self,
|
|
PyObject* args,
|
|
PyObject* kwargs) {
|
|
HANDLE_TH_ERRORS
|
|
auto self = (THPStream*)_self;
|
|
PyObject* _new_event = nullptr;
|
|
PyObject* _event = Py_None;
|
|
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
|
|
constexpr const char* accepted_args[] = {"event", nullptr};
|
|
if (!PyArg_ParseTupleAndKeywords(
|
|
args,
|
|
kwargs,
|
|
"|O",
|
|
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
|
|
const_cast<char**>(accepted_args),
|
|
&_event)) {
|
|
TORCH_CHECK(false, "parse record_event arg fails");
|
|
}
|
|
if (_event != Py_None) {
|
|
// Increase the refcount of the event to avoid it being destroyed.
|
|
Py_INCREF(_event);
|
|
_new_event = _event;
|
|
} else {
|
|
_new_event = THPEvent_new(
|
|
static_cast<c10::DeviceType>(self->device_type),
|
|
c10::EventFlag::PYTORCH_DEFAULT);
|
|
}
|
|
auto new_event = (THPEvent*)_new_event;
|
|
TORCH_CHECK(new_event, "event must not be null");
|
|
new_event->event.record(c10::Stream::unpack3(
|
|
self->stream_id,
|
|
static_cast<c10::DeviceIndex>(self->device_index),
|
|
static_cast<c10::DeviceType>(self->device_type)));
|
|
return (PyObject*)new_event;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_repr(THPStream* self) {
|
|
HANDLE_TH_ERRORS
|
|
return THPUtils_packString(
|
|
"torch.Stream device_type=" +
|
|
c10::DeviceTypeName(
|
|
static_cast<c10::DeviceType>(self->device_type), true) +
|
|
", device_index=" + std::to_string(self->device_index) +
|
|
", stream_id=" + std::to_string(self->stream_id));
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static Py_hash_t THPStream_hash(THPStream* self) {
|
|
return static_cast<long>(at::hash_combine(
|
|
self->device_type,
|
|
(at::hash_combine(self->stream_id, self->device_index))));
|
|
}
|
|
|
|
static PyObject* THPStream_eq(THPStream* self, THPStream* other) {
|
|
HANDLE_TH_ERRORS
|
|
return PyBool_FromLong(
|
|
(self->stream_id == other->stream_id) &&
|
|
(self->device_index == other->device_index) &&
|
|
(self->device_type == other->device_type));
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_enter(PyObject* _self, PyObject* unused) {
|
|
HANDLE_TH_ERRORS
|
|
auto self = (THPStream*)_self;
|
|
c10::DeviceType stream_device_type =
|
|
static_cast<c10::DeviceType>(self->device_type);
|
|
// No operation is performed if the stream does not belong to an accelerator.
|
|
if (C10_UNLIKELY(!at::accelerator::isAccelerator(stream_device_type))) {
|
|
Py_INCREF(_self);
|
|
return _self;
|
|
}
|
|
c10::DeviceIndex cur_device_idx = at::accelerator::getDeviceIndex();
|
|
c10::DeviceIndex stream_device_idx =
|
|
static_cast<c10::DeviceIndex>(self->device_index);
|
|
// If the stream is not on the current device, switch the current device to
|
|
// the device of the stream.
|
|
if (stream_device_idx != cur_device_idx) {
|
|
at::accelerator::setDeviceIndex(stream_device_idx);
|
|
}
|
|
c10::Stream cur_stream = at::accelerator::getCurrentStream(stream_device_idx);
|
|
at::accelerator::setCurrentStream(c10::Stream::unpack3(
|
|
self->stream_id, stream_device_idx, stream_device_type));
|
|
// Save the current device index and previous stream to the context.
|
|
auto ctx_device_index =
|
|
THPObjectPtr(THPUtils_packDeviceIndex(cur_device_idx));
|
|
auto ctx_stream = THPObjectPtr(THPStream_Wrap(cur_stream));
|
|
TORCH_CHECK(!(self->context), "Stream's context should not be initialized.");
|
|
auto dict = THPObjectPtr(PyDict_New());
|
|
if (!dict) {
|
|
throw python_error();
|
|
}
|
|
self->context = dict.release();
|
|
if (PyDict_SetItemString(
|
|
self->context, "_ctx_device_index", ctx_device_index.get()) < 0) {
|
|
throw python_error();
|
|
}
|
|
if (PyDict_SetItemString(self->context, "_ctx_stream", ctx_stream.get()) <
|
|
0) {
|
|
throw python_error();
|
|
}
|
|
Py_INCREF(_self);
|
|
return _self;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_exit(PyObject* _self, PyObject* unused) {
|
|
HANDLE_TH_ERRORS
|
|
auto self = (THPStream*)_self;
|
|
// No operation is performed if the stream does not belong to an accelerator.
|
|
if (C10_UNLIKELY(!at::accelerator::isAccelerator(
|
|
static_cast<c10::DeviceType>(self->device_type)))) {
|
|
Py_RETURN_NONE;
|
|
}
|
|
PyObject* py_stream = nullptr;
|
|
if (PyDict_GetItemStringRef(self->context, "_ctx_stream", &py_stream) < 0) {
|
|
throw python_error();
|
|
}
|
|
auto ctx_stream = THPObjectPtr(py_stream);
|
|
PyObject* py_device_index = nullptr;
|
|
if (PyDict_GetItemStringRef(
|
|
self->context, "_ctx_device_index", &py_device_index) < 0) {
|
|
throw python_error();
|
|
}
|
|
auto ctx_device_index = THPObjectPtr(py_device_index);
|
|
TORCH_INTERNAL_ASSERT(
|
|
ctx_stream.get(), "ctx_stream should be present on the context dict.");
|
|
auto prev_stream = (THPStream*)(ctx_stream.get());
|
|
TORCH_INTERNAL_ASSERT(
|
|
ctx_device_index.get(),
|
|
"ctx_device_index should be present on the context dict.");
|
|
auto prev_device_index = THPUtils_unpackDeviceIndex(ctx_device_index.get());
|
|
at::accelerator::setCurrentStream(c10::Stream::unpack3(
|
|
prev_stream->stream_id,
|
|
static_cast<c10::DeviceIndex>(prev_stream->device_index),
|
|
static_cast<c10::DeviceType>(prev_stream->device_type)));
|
|
// Reset the current device to the previous device if they differ.
|
|
if (static_cast<c10::DeviceIndex>(self->device_index) != prev_device_index) {
|
|
at::accelerator::setDeviceIndex(prev_device_index);
|
|
}
|
|
Py_CLEAR(self->context);
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_ne(THPStream* self, THPStream* other) {
|
|
HANDLE_TH_ERRORS
|
|
return PyBool_FromLong(
|
|
(self->stream_id != other->stream_id) ||
|
|
(self->device_index != other->device_index) ||
|
|
(self->device_type != other->device_type));
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject* THPStream_richcompare(
|
|
PyObject* self,
|
|
PyObject* other,
|
|
int op) {
|
|
PyObject* result = nullptr;
|
|
if (other == Py_None) {
|
|
result = Py_False;
|
|
} else {
|
|
switch (op) {
|
|
case Py_EQ:
|
|
result = THPStream_eq((THPStream*)self, (THPStream*)other);
|
|
break;
|
|
case Py_NE:
|
|
result = THPStream_ne((THPStream*)self, (THPStream*)other);
|
|
break;
|
|
default:
|
|
result = Py_False;
|
|
break;
|
|
}
|
|
}
|
|
Py_XINCREF(result);
|
|
return result;
|
|
}
|
|
|
|
static const std::initializer_list<PyMemberDef> THPStream_members = {
|
|
{"stream_id",
|
|
T_LONGLONG,
|
|
offsetof(THPStream, stream_id),
|
|
READONLY,
|
|
nullptr},
|
|
{"device_index",
|
|
T_LONGLONG,
|
|
offsetof(THPStream, device_index),
|
|
READONLY,
|
|
nullptr},
|
|
{"device_type",
|
|
T_LONGLONG,
|
|
offsetof(THPStream, device_type),
|
|
READONLY,
|
|
nullptr},
|
|
{nullptr}};
|
|
|
|
static const std::initializer_list<PyGetSetDef> THPStream_properties = {
|
|
{"device", (getter)THPStream_get_device, nullptr, nullptr, nullptr},
|
|
{nullptr}};
|
|
|
|
static const std::initializer_list<PyMethodDef> THPStream_methods = {
|
|
{"query", THPStream_query, METH_NOARGS, nullptr},
|
|
{"synchronize", THPStream_synchronize, METH_NOARGS, nullptr},
|
|
{"wait_event", THPStream_wait_event, METH_O, nullptr},
|
|
{"wait_stream", THPStream_wait_stream, METH_O, nullptr},
|
|
{"record_event",
|
|
castPyCFunctionWithKeywords(THPStream_record_event),
|
|
METH_VARARGS | METH_KEYWORDS,
|
|
nullptr},
|
|
{"__eq__", (PyCFunction)THPStream_eq, METH_O, nullptr},
|
|
{"__enter__", THPStream_enter, METH_NOARGS, nullptr},
|
|
{"__exit__", THPStream_exit, METH_VARARGS, nullptr},
|
|
{nullptr}};
|
|
|
|
static PyTypeObject THPStreamType = {
|
|
PyVarObject_HEAD_INIT(nullptr, 0)
|
|
"torch.Stream", /* tp_name */
|
|
sizeof(THPStream), /* tp_basicsize */
|
|
0, /* tp_itemsize */
|
|
(destructor)THPStream_dealloc, /* tp_dealloc */
|
|
0, /* tp_vectorcall_offset */
|
|
nullptr, /* tp_getattr */
|
|
nullptr, /* tp_setattr */
|
|
nullptr, /* tp_reserved */
|
|
(reprfunc)THPStream_repr, /* tp_repr */
|
|
nullptr, /* tp_as_number */
|
|
nullptr, /* tp_as_sequence */
|
|
nullptr, /* tp_as_mapping */
|
|
(hashfunc)THPStream_hash, /* tp_hash */
|
|
nullptr, /* tp_call */
|
|
nullptr, /* tp_str */
|
|
nullptr, /* tp_getattro */
|
|
nullptr, /* tp_setattro */
|
|
nullptr, /* tp_as_buffer */
|
|
// NOLINTNEXTLINE(misc-redundant-expression)
|
|
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
|
|
nullptr, /* tp_doc */
|
|
nullptr, /* tp_traverse */
|
|
nullptr, /* tp_clear */
|
|
THPStream_richcompare, /* tp_richcompare */
|
|
0, /* tp_weaklistoffset */
|
|
nullptr, /* tp_iter */
|
|
nullptr, /* tp_iternext */
|
|
// NOLINTNEXTLINE(*const-cast)
|
|
const_cast<PyMethodDef*>(std::data(THPStream_methods)), /* tp_methods */
|
|
// NOLINTNEXTLINE(*const-cast)
|
|
const_cast<PyMemberDef*>(std::data(THPStream_members)), /* tp_members */
|
|
// NOLINTNEXTLINE(*const-cast)
|
|
const_cast<PyGetSetDef*>(std::data(THPStream_properties)), /* tp_getset */
|
|
nullptr, /* tp_base */
|
|
nullptr, /* tp_dict */
|
|
nullptr, /* tp_descr_get */
|
|
nullptr, /* tp_descr_set */
|
|
0, /* tp_dictoffset */
|
|
nullptr, /* tp_init */
|
|
nullptr, /* tp_alloc */
|
|
THPStream_pynew, /* tp_new */
|
|
};
|
|
|
|
void THPStream_init(PyObject* module) {
|
|
THPStreamClass = &THPStreamType;
|
|
Py_SET_TYPE(&THPStreamType, &PyType_Type);
|
|
if (PyType_Ready(&THPStreamType) < 0) {
|
|
throw python_error();
|
|
}
|
|
Py_INCREF(&THPStreamType);
|
|
if (PyModule_AddObject(module, "Stream", (PyObject*)&THPStreamType) < 0) {
|
|
throw python_error();
|
|
}
|
|
}
|