mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
d6d726f781
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55462 handles and symbolicate exception callstack thrown from backend. Objective of this diff is to achieve improve error reporting when exceptions are raised from lowered backend. We would effectively like to get the same model level stack trace that you would get without having lowered some module to backend. For example: ``` class AA(nn.Module): def forward(self, x, y): return x + y class A(nn.Module): def __init__(...): self.AA0 = AA() def forward(self, x, y): return self.AA0.forward(x, y) + 3 class B(nn.Module): def forward(self, x): return x + 2 class C(nn.Module): def __init__(...): self.A0 = A() self.B0 = B() def forward(self, x, y): return self.A0.forward(x, y) + self.B0.forward(x) ``` If the we then do C().forward(torch.rand((2,3)), torch.rand(14,2))) we will likely see error stack like: ``` C++ exception with description "The following operation failed in the TorchScript interpreter. Traceback of TorchScript (most recent call last): File "<string>", line 3, in forward def forward(self, x, y): return self.A0.forward(x, y) + self.B0.forward(x) ~~~~~~~~~~~~~~~ <--- HERE File "<string>", line 3, in forward def forward(self, x, y): return self.AA0.forward(x, y) + 3 ~~~~~~~~~~~~~~~~ <--- HERE File "<string>", line 3, in forward def forward(self, x, y): return x + y ~~~~~ <--- HERE ``` We would like to see the same error stack if we lowered C.A0 to some backend. With this diff we get something like: ``` Module hierarchy:top(C).A0(backend_with_compiler_demoLoweredModule).AA0(AA) Traceback of TorchScript (most recent call last): File "<string>", line 3, in FunctionName_UNKNOWN def forward(self, x, y): return self.A0.forward(x, y) + self.B0.forward(x) ~~~~~~~~~~~~~~~ <--- HERE File "<string>", line 5, in FunctionName_UNKNOWN typed_inputs: List[Any] = [x, y, ] if self.__backend.is_available() : _0, = self.__backend.execute(self.__handles["forward"], typed_inputs) ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE assert isinstance(_0, Tensor) return _0 File "<string>", line 3, in FunctionName_UNKNOWN def forward(self, x, y): return self.AA0.forward(x, y) + 3 ~~~~~~~~~~~~~~~~ <--- HERE File "<string>", line 3, in FunctionName_UNKNOWN def forward(self, x, y): return x + y ~~~~~ <--- HERE ``` This is achieved in 3 parts: Part 1: A. BackendDebugInfoRecorder: During backend lowering, in `to_backend`, before calling the preprocess function corresponding to the backend. This will facilitate recording of debug info (such as source range + inlined callstack) for the lowered module. B. Instantiate WithBackendDebugInfoRecorder with BackendDebugInfoRecorder. This initializes thread local pointer to BackendDebugInfoRecorder. C. generate_debug_handles: In preprocess function, the backend will call generate_debug_handles for each method being lowered separately. generate_debug_handles takes `Graph` of the method being lowered and returns a map of Node*-to-debug_handles. Backend is responsible for storing debug handles appropriately so as to raise exception (and later profiling) using debug handles when the exception being raised corresponds to particular Node that was lowered. Inside generate_debug_handles, we will query the current BackendDebugHandleInfoRecorder, that is issuing debug handles. This debug handle manager will issue debug handles as well as record debug_handles-to-<source range, inlined callstack> map. D. Back in `to_backend`, once the preprocess function is has finished lowering the module, we will call `stopRecord` on BackendDebugInfoRecorder. This will return the debug info map. This debug info is then stored inside the lowered module. Part 2: Serialization: During serialization for bytecode (lite interpreter), we will do two things: 1. Extract all the source ranges that are contained inside debug_handles-to-<source range, inlined callstack> map for lowered module. This will be source range corresponding to debug handles, including what is there is inlined callstack. Since we replaced original module with lowered module, we wont be serializing code for the original module and thus no source range. That is why the source range will have to be stored separately. We will lump all the source ranges for all the lowered modules in one single debug_pkl file. 2. Then we will serialize debug_handles-to-<source range, inlined callstack> map. Now during deserialization we will be able to reconstruct debug_handles-to-<source range, inlined callstack> map. Given all debug_handles are unique we would not need any module information. Test Plan: Tests are added in test_backend.cpp Tests are added in test_backend.cpp Imported from OSS Differential Revision: D27621330 D27621330 Reviewed By: raziel Pulled By: kimishpatel fbshipit-source-id: 0650ec68cda0df0a945864658cab226a97ba1890
78 lines
2.8 KiB
C++
78 lines
2.8 KiB
C++
#include <torch/csrc/jit/backends/backend.h>
|
|
#include <torch/csrc/jit/backends/backend_preprocess.h>
|
|
#include <torch/csrc/jit/backends/generate_debug_handles.h>
|
|
#include <torch/csrc/jit/passes/dead_code_elimination.h>
|
|
#include <torch/csrc/jit/passes/inliner.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace {
|
|
// For this backend, the actual compilation happens in preprocess function AOT.
|
|
// Put here for demonstration of backend
|
|
// as a whole piece. It's used when compilation is required. A dummy function
|
|
// can be passed when there's no usage of compilation in runtime backend lib.
|
|
c10::IValue preprocess(
|
|
const Module& mod,
|
|
const c10::Dict<IValue, IValue>& method_compile_spec) {
|
|
// The output of this process would produce a dictionary
|
|
// Key: method name.
|
|
// Val: compiled blob (represented by a string).
|
|
c10::Dict<IValue, IValue> compiled(StringType::get(), StringType::get());
|
|
|
|
for (const auto& method : mod.get_methods()) {
|
|
auto graph = method.function().graph()->copy();
|
|
// Must inline the graph for debug info map.
|
|
Inline(*graph);
|
|
// This is here because to test module hierarchy we will have
|
|
// getattr nodes which after inlining dont serve any purpose.
|
|
// Without removing them we will run into compilation errors.
|
|
// So eliminate deadcode just remove those getattr nodes.
|
|
EliminateDeadCode(graph);
|
|
// NOLINTNEXTLINE(performance-unnecessary-copy-initialization)
|
|
auto key = method.name();
|
|
auto node_debug_handles = generate_debug_handles(graph);
|
|
std::stringstream ss;
|
|
for (const auto& node : graph->nodes()) {
|
|
switch (node->kind()) {
|
|
case prim::Constant:
|
|
ss << node->kind().toDisplayString() << "#"
|
|
<< toIValue(node->output()).value();
|
|
ss << "<debug_handle>" << node_debug_handles[node];
|
|
break;
|
|
// NOLINTNEXTLINE(bugprone-branch-clone)
|
|
case aten::add:
|
|
ss << node->kind().toQualString();
|
|
ss << "<debug_handle>" << node_debug_handles[node];
|
|
break;
|
|
case aten::sub:
|
|
ss << node->kind().toQualString();
|
|
ss << "<debug_handle>" << node_debug_handles[node];
|
|
break;
|
|
default:
|
|
TORCH_CHECK(
|
|
false,
|
|
"The node of ",
|
|
node->kind().toQualString(),
|
|
" is not supported in this compiler. Source code: ",
|
|
node->sourceRange().str());
|
|
break;
|
|
}
|
|
ss << ",";
|
|
}
|
|
std::string blob = ss.str();
|
|
if (!blob.empty()) {
|
|
blob.pop_back();
|
|
}
|
|
compiled.insert(method.name(), blob);
|
|
}
|
|
return compiled;
|
|
}
|
|
|
|
constexpr auto backend_name = "backend_with_compiler_demo";
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
|
|
static auto pre_reg = backend_preprocess_register(backend_name, preprocess);
|
|
} // namespace
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|