mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
e35418b3be
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13342 This PR introduces a few new concepts: - DeviceGuardImplInterface, and implementations for CPU and CUDA, which provide a generic interface for interfacing with device and stream state, without requiring a direct dependency on the code in question. - InlineDeviceGuard, a general template for generating both specialized and dynamically dispatched device guard implementations. Dynamic dispatch is done by specializing it on a VirtualGuardImpl. - Provide a device-independent DeviceGuard class, which can be used even from CPU code. It uses the aforementioned dynamic dispatch. - CUDA-specialized CUDAGuard class, which doesn't have a dynamic dispatch but can only be used from CUDA. - StreamGuard, which is the same as above, but for streams rather than devices. - Optional variants of all the aforementioned guards, which are a no-op if no device/stream is specified - CUDAMultiStreamGuard, specifically for the case when we want to set a device on every guard. There are some subtle semantic changes, which have been thoroughly documented in the class definition. BC-breaking changes: - Move constructor/assignment have been removed from all device guard implementations. - In some cases where you previously wrote 'set_device' (or 'set_stream'), you now must write 'reset_device', because if you switch devices/device types, the stream/device on the previous device is unset. This is different from previous behavior. - CUDAGuard no longer handles streams, or multiple streams. Use CUDAStreamGuard or CUDAMultiStreamGuard as appropriate for your use case. Reviewed By: dzhulgakov Differential Revision: D12849620 fbshipit-source-id: f61956256f0b12be754b3234fcc73c2abc1be04e
106 lines
3.2 KiB
C++
106 lines
3.2 KiB
C++
#pragma once
|
|
|
|
#include "torch/csrc/python_headers.h"
|
|
|
|
#include "torch/csrc/Exceptions.h"
|
|
#include "torch/csrc/autograd/function.h"
|
|
#include "torch/csrc/autograd/variable.h"
|
|
#include "torch/csrc/autograd/saved_variable.h"
|
|
#include "torch/csrc/utils/object_ptr.h"
|
|
|
|
#include "c10/util/Optional.h"
|
|
#include "c10/DeviceGuard.h"
|
|
|
|
#include <vector>
|
|
#include <utility>
|
|
#include <memory>
|
|
|
|
namespace torch { namespace jit { struct Graph; }}
|
|
namespace torch { namespace autograd {
|
|
|
|
struct VariableInfo {
|
|
explicit VariableInfo(const Variable& var);
|
|
|
|
Variable zeros(at::OptionalDeviceGuard& device_guard) const;
|
|
|
|
at::Type* type;
|
|
at::Device device = at::kCPU;
|
|
std::vector<int64_t> size;
|
|
bool requires_grad;
|
|
};
|
|
|
|
// A Function which is implemented by a Python object (i.e., a THPFunction).
|
|
// Calls to 'apply' are forwarded to the Python method implementation.
|
|
struct PyFunction : public Function {
|
|
PyFunction(PyObject* obj) : obj(obj) {}
|
|
|
|
virtual variable_list apply(variable_list&& inputs) override;
|
|
variable_list legacy_apply(const variable_list& inputs);
|
|
|
|
virtual void release_variables() override;
|
|
virtual std::string name() const override;
|
|
virtual std::shared_ptr<Function> get_shared_ptr() override;
|
|
virtual bool is_traceable() override;
|
|
|
|
// THPFunction this Function is wrapping.
|
|
PyObject* obj;
|
|
};
|
|
|
|
/**
|
|
* Cast an object into a tuple, if it is not a tuple already. Returns true
|
|
* if the original object was not a tuple.
|
|
*/
|
|
inline bool ensure_tuple(THPObjectPtr& obj) {
|
|
if (PyTuple_Check(obj.get()))
|
|
return false;
|
|
|
|
PyObject *tuple = PyTuple_New(1);
|
|
if (!tuple) throw python_error();
|
|
PyTuple_SET_ITEM(tuple, 0, obj.release());
|
|
obj = tuple;
|
|
return true;
|
|
}
|
|
|
|
}} // namespace torch::autograd
|
|
|
|
struct THPFunction {
|
|
PyObject_HEAD
|
|
|
|
PyObject *needs_input_grad;
|
|
|
|
// Python tuple of tensors whose variables we should save. Set
|
|
// by Python with 'save_for_backward'. If nullptr, no tensors were
|
|
// saved.
|
|
PyObject *to_save;
|
|
// Python tuple of tensors which are not differentiable. Set by
|
|
// Python with 'mark_non_differentiable'. If nullptr, no tensors were
|
|
// non-differentiable.
|
|
PyObject *non_differentiable;
|
|
// Python tuple of tensors which had inplace updates in the forward()
|
|
// pass. Set by Python with 'mark_dirty'. If nullptr, no tensors were
|
|
// modified inplace.
|
|
PyObject *dirty_tensors;
|
|
|
|
std::vector<torch::autograd::VariableInfo> output_info;
|
|
std::vector<torch::autograd::VariableInfo> input_info;
|
|
std::vector<torch::autograd::SavedVariable> saved_variables;
|
|
// For each input, true if the input is a THPVariable
|
|
std::vector<bool> is_variable_input;
|
|
char has_freed_buffers;
|
|
|
|
// The C++ wrapper for this Python function.
|
|
// See a comment in THPFunction_asFunction for details about this field.
|
|
torch::autograd::PyFunction cdata;
|
|
};
|
|
|
|
bool THPFunction_initModule(PyObject *module);
|
|
extern PyTypeObject THPFunctionType;
|
|
extern PyObject *THPFunctionClass;
|
|
|
|
// XXX: this function requires the GIL (it can have side effects).
|
|
std::shared_ptr<torch::autograd::PyFunction> THPFunction_asFunction(THPFunction* self);
|
|
|
|
inline bool THPFunction_Check(PyObject* obj) {
|
|
return PyObject_IsInstance(obj, (PyObject*)&THPFunctionType);
|
|
}
|