mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5531fafffe
This PR is on the way to getting compiled autograd's initial capture to stop specializing on Tensor metadata. This PR changes compiled autograd's initial capture to proxy an opaque (w.r.t. Dynamo) function into the graph for all built-in codegen'ed autograd nodes and validate_outputs. We changed each codegen'ed apply_with_saved (e.g. MulBackward0::apply_with_saved) to call into Python to proxy a function (compiled_autograd.ops.MulBackward0) into the graph. Then, we use the node's InputMetadata to "guess" at the properties of the output Tensors to create some new FakeTensors. Some details: - MulBackward0::apply_with_saved lives in libtorch_cpu, but needs to be call to Python via libtorch_python. There is an indirection (PyCompilerInterface) to do this. - MulBackward0::apply_with_saved passes a C++ function to Python. To make our lives easier, every codegen'ed apply_with_saved passes a C++ function with the same signature `(variable_list, ivalue_list) -> variable_list`. - We define how to pack arbitrary C++ types into IValue via a helper IValuePacker struct and codegen functional variants of each builtin C++ autograd node (e.g. MulBackward0_apply_functional_ivalue). MulBackward0 before this PR: https://gist.github.com/zou3519/a80381d5fa38e970e413fcd91b0530de MulBackward0 after this PR: https://gist.github.com/zou3519/0c2eee8b3d8d96232b51ef430b53c5b0 Test Plan: - existing tests Pull Request resolved: https://github.com/pytorch/pytorch/pull/143296 Approved by: https://github.com/jansel
28 lines
772 B
C++
28 lines
772 B
C++
#include <torch/csrc/autograd/engine.h>
|
|
#include <torch/csrc/dynamo/compiled_autograd.h>
|
|
|
|
namespace torch::dynamo::autograd {
|
|
|
|
std::unique_ptr<PyCompilerInterface> kPyCompilerInterface;
|
|
|
|
const std::unique_ptr<PyCompilerInterface>& getPyCompilerInterface() {
|
|
TORCH_INTERNAL_ASSERT(kPyCompilerInterface != nullptr);
|
|
return kPyCompilerInterface;
|
|
}
|
|
|
|
void setPyCompilerInterface(std::unique_ptr<PyCompilerInterface>&& impl) {
|
|
TORCH_INTERNAL_ASSERT(impl != nullptr);
|
|
kPyCompilerInterface = std::move(impl);
|
|
}
|
|
|
|
void resetPyCompilerInterface() {
|
|
kPyCompilerInterface.reset();
|
|
}
|
|
|
|
std::vector<std::optional<InputMetadata>> get_input_metadata(
|
|
const edge_list& edges) {
|
|
return torch::autograd::collect_input_metadata(edges);
|
|
}
|
|
|
|
} // namespace torch::dynamo::autograd
|