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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
#include <torch/csrc/jit/passes/quantization/finalize.h>
#include <torch/csrc/jit/jit_log.h>
#include <torch/csrc/jit/passes/freeze_module.h>
#include <torch/csrc/jit/passes/prepack_folding.h>
#include <torch/csrc/jit/passes/quantization/quantization_patterns.h>
namespace torch {
namespace jit {
namespace {
void insertPrepackUnpackForLinear(std::shared_ptr<Graph>& graph) {
std::vector<QuantFusionInfo> patterns_and_replacements =
linear_prepack_unpack_patterns();
for (const auto& entry : patterns_and_replacements) {
SubgraphRewriter rewriter;
rewriter.RegisterRewritePattern(entry.pattern, entry.replacement);
rewriter.runOnGraph(graph, entry.filters);
}
}
void insertPrepackUnpackForConv(std::shared_ptr<Graph>& graph) {
std::vector<QuantFusionInfo> patterns_and_replacements =
conv_prepack_unpack_patterns();
for (const auto& entry : patterns_and_replacements) {
SubgraphRewriter rewriter;
rewriter.RegisterRewritePattern(entry.pattern, entry.replacement);
rewriter.runOnGraph(graph, entry.filters);
}
}
} // namespace
void QuantFusion(std::shared_ptr<Graph>& graph, QuantType quant_type) {
std::vector<QuantFusionInfo> patterns;
if (quant_type == QuantType::DYNAMIC) {
patterns = dynamic_quant_fusion_pattern_and_replacements();
} else {
patterns = quant_fusion_pattern_and_replacements();
}
for (const auto& info : patterns) {
SubgraphRewriter rewriter;
rewriter.RegisterRewritePattern(info.pattern, info.replacement);
rewriter.runOnGraph(graph, info.filters);
}
}
void InsertPrepackUnpack(std::shared_ptr<Graph>& graph) {
insertPrepackUnpackForLinear(graph);
insertPrepackUnpackForConv(graph);
}
void InsertPrepackUnpack(Module& module) {
for (auto& method : module.get_methods()) {
auto graph = method.graph();
InsertPrepackUnpack(graph);
}
for (Module m : module.children()) {
InsertPrepackUnpack(m);
}
}
void FoldQuantizedPrepackingOps(Module& module) {
auto filter_fn = [](const Node* n) -> bool {
return (
n->kind() == Symbol::fromQualString("quantized::linear_prepack") ||
n->kind() == Symbol::fromQualString("quantized::conv1d_prepack") ||
n->kind() == Symbol::fromQualString("quantized::conv2d_prepack") ||
n->kind() == Symbol::fromQualString("quantized::conv3d_prepack") ||
n->kind() ==
Symbol::fromQualString("quantized::conv_transpose1d_prepack") ||
n->kind() ==
Symbol::fromQualString("quantized::conv_transpose2d_prepack"));
};
PrePackingOpsFolder(module, filter_fn, "quantized");
}
Module Finalize(
Module& module,
QuantType quant_type,
const std::vector<std::string>& preserved_attrs) {
auto graph = module.get_method("forward").graph();
InsertPrepackUnpack(graph);
GRAPH_DUMP("Before QuantFusion:", graph);
QuantFusion(graph, quant_type);
auto frozen = freeze_module(module, preserved_attrs);
FoldQuantizedPrepackingOps(frozen);
return frozen;
}
} // namespace jit
} // namespace torch
|
