1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
#include <torch/csrc/jit/backends/backend.h>
#include <torch/csrc/jit/backends/backend_preprocess.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,
const BackendDebugHandleGenerator& generate_debug_handles) {
// 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 = toGraphFunction(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";
static auto pre_reg = backend_preprocess_register(backend_name, preprocess);
} // namespace
} // namespace jit
} // namespace torch
|
