mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-05 00:14:54 +08:00
35f6a69191
Signed-off-by: Edward Z. Yang <ezyangfb.com> From @ezyang's original PR: There are a number of situations where we have non-backend kernels (e.g., CompositeImplicitAutograd, batching rules) which we would like to port to Python, but we have no way to integrate these ports with the overall system while using preexisting C++ registrations otherwise. This PR changes that by introducing a Python dispatcher (which can have its own kernels directly in Python), which can be interpose over ordinary C++ dispatch. The ingredients: We introduce a new PythonDispatcher dispatch key, that has the same tenor as FuncTorchDynamicLayerFrontMode: it works by getting triggered before every other dispatch key in the dispatch key, and shunting to a Python implementation The Python dispatcher is a per-interpreter global object that is enabled/disabled via the guard EnablePythonDispatcher/DisablePythonDispatcher. We don't make it compositional as I have no idea what a compositional version of this feature would look like. Because it is global, we don't need to memory manage it and so I use a simpler SafePyHandle (newly added) to control access to this pointer from non-Python C++. Like __torch_dispatch__, we use PyInterpreter to get to the Python interpreter to handle the dispatch. I need to reimplement dispatch table computation logic in Python. To do this, I expose a lot more helper functions for doing computations on alias dispatch keys and similar. I also improve the pybind11 handling for DispatchKey so that you can either accept the pybind11 bound enum or a string; this simplifies our binding code. See https://github.com/pybind/pybind11/issues/483#issuecomment-1237418106 for how this works; the technique is generally useful. I need to be able to call backend fallbacks. I do this by permitting you to call at a dispatch key which doesn't have a kernel for the operator; if the kernel doesn't exist, we check the backend fallback table instead. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/84826 Approved by: https://github.com/ezyang
90 lines
2.8 KiB
C++
90 lines
2.8 KiB
C++
#include <c10/core/SymIntArrayRef.h>
|
|
#include <c10/core/TensorImpl.h>
|
|
#include <c10/core/impl/PyInterpreter.h>
|
|
|
|
namespace c10 {
|
|
namespace impl {
|
|
|
|
struct NoopPyInterpreterVTable final : public PyInterpreterVTable {
|
|
std::string name() const override {
|
|
return "<unloaded interpreter>";
|
|
}
|
|
|
|
void decref(PyObject* pyobj, bool is_tensor) const override {} // do nothing
|
|
|
|
#define PANIC(m) \
|
|
TORCH_INTERNAL_ASSERT( \
|
|
0, \
|
|
"attempted to call " #m \
|
|
" on a Tensor with nontrivial PyObject after corresponding interpreter died")
|
|
|
|
c10::intrusive_ptr<TensorImpl> detach(const TensorImpl* self) const override {
|
|
PANIC(detach);
|
|
}
|
|
|
|
void dispatch(const c10::OperatorHandle& op, torch::jit::Stack* stack)
|
|
const override {
|
|
PANIC(dispatch);
|
|
}
|
|
|
|
void python_dispatcher(
|
|
const c10::OperatorHandle& op,
|
|
c10::DispatchKeySet,
|
|
torch::jit::Stack* stack) const override {
|
|
PANIC(python_dispatcher);
|
|
}
|
|
|
|
bool is_contiguous(const TensorImpl* self) const override {
|
|
PANIC(is_contiguous);
|
|
}
|
|
c10::Device device(const TensorImpl* self) const override {
|
|
PANIC(device);
|
|
}
|
|
int64_t dim(const TensorImpl* self) const override {
|
|
PANIC(dim);
|
|
}
|
|
c10::IntArrayRef strides(const TensorImpl* self) const override {
|
|
PANIC(strides);
|
|
}
|
|
c10::IntArrayRef sizes(const TensorImpl* self) const override {
|
|
PANIC(sizes);
|
|
}
|
|
c10::SymIntArrayRef sym_sizes(const TensorImpl* self) const override {
|
|
PANIC(sym_sizes);
|
|
}
|
|
c10::Layout layout(const TensorImpl* self) const override {
|
|
PANIC(layout);
|
|
}
|
|
c10::SymInt sym_numel(const TensorImpl* self) const override {
|
|
PANIC(sym_numel);
|
|
}
|
|
c10::SymIntArrayRef sym_strides(const TensorImpl* self) const override {
|
|
PANIC(sym_strides);
|
|
}
|
|
c10::SymInt sym_storage_offset(const TensorImpl* self) const override {
|
|
PANIC(sym_storage_offset);
|
|
}
|
|
|
|
// Just swallow the event, don't do anything
|
|
void trace_gpu_event_creation(uintptr_t event) const override {}
|
|
void trace_gpu_event_deletion(uintptr_t event) const override {}
|
|
void trace_gpu_event_record(uintptr_t event, uintptr_t stream)
|
|
const override {}
|
|
void trace_gpu_event_wait(uintptr_t event, uintptr_t stream) const override {}
|
|
void trace_gpu_memory_allocation(uintptr_t ptr) const override {}
|
|
void trace_gpu_memory_deallocation(uintptr_t ptr) const override {}
|
|
void trace_gpu_stream_creation(uintptr_t stream) const override {}
|
|
void trace_gpu_device_synchronization() const override {}
|
|
void trace_gpu_stream_synchronization(uintptr_t stream) const override {}
|
|
void trace_gpu_event_synchronization(uintptr_t event) const override {}
|
|
};
|
|
|
|
void PyInterpreter::disarm() noexcept {
|
|
// Intentionally leaked
|
|
static PyInterpreterVTable* noop_vtable = new NoopPyInterpreterVTable();
|
|
vtable_ = noop_vtable;
|
|
}
|
|
|
|
} // namespace impl
|
|
} // namespace c10
|