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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
|
# mypy: allow-untyped-defs
import functools
import torch
from torch._inductor.compile_fx import fake_tensor_prop
from torch._inductor.utils import GPU_TYPES
from ..._dynamo.utils import counters
from .. import config
from ..pattern_matcher import (
_return_true,
CallFunction,
fwd_only,
Ignored,
init_once_fakemode,
KeywordArg,
Match,
PatternMatcherPass,
register_graph_pattern,
register_replacement,
stable_topological_sort,
)
aten = torch.ops.aten
# First pass_patterns[0] are applied, then [1], then [2]
pass_patterns = [
PatternMatcherPass(),
PatternMatcherPass(),
PatternMatcherPass(),
]
binary_folding_pass = PatternMatcherPass()
def freezing_passes(gm: torch.fx.GraphModule, aot_example_inputs):
"""
Passes that are applied to the graph to freeze pass.
"""
from ..freezing import constant_fold
lazy_init()
# We need a few rounds of binary folding to get rid of all the
# unnecessary nodes, but may need a good method to chose the rounds number.
# works like: conv+binary+binary.
binary_folding = counters["inductor"]["binary_folding"]
fake_tensor_prop(gm, aot_example_inputs, True)
torch._inductor.fx_passes.binary_folding.mark_mixed_dtype_allowed_computation_ops(
gm
)
for _ in range(4):
constant_fold(gm)
# Make sure meta['val'] is properly set for all nodes
fake_tensor_prop(gm, aot_example_inputs, True)
binary_folding_pass.apply(gm.graph) # type: ignore[arg-type]
# If we don't have binary folding, we don't need to run the pass again.
# TODO: remove the need to run fake_tensor_prop on the whole model.
if counters["inductor"]["binary_folding"] == binary_folding:
break
binary_folding = counters["inductor"]["binary_folding"]
torch._inductor.fx_passes.binary_folding.recover_original_precision_folded_computation_ops(
gm
)
constant_fold(gm)
fake_tensor_prop(gm, aot_example_inputs, True)
for pattern in pass_patterns:
pattern.apply(gm.graph) # type: ignore[arg-type]
# The CPU weight packing always assume the conv's weight is channels last,
# So make sure the layout_optimization is on when doing it.
if (
torch._C._has_mkldnn
and config.cpp.weight_prepack
and config.layout_optimization
):
from .mkldnn_fusion import _eliminate_duplicate_packed_nodes
_eliminate_duplicate_packed_nodes(gm)
stable_topological_sort(gm.graph)
gm.recompile()
gm.graph.lint()
@init_once_fakemode
def lazy_init():
if torch._C._has_mkldnn and config.cpp.weight_prepack:
from .mkldnn_fusion import _mkldnn_weight_pack_init
_mkldnn_weight_pack_init()
from .binary_folding import binary_folding_init
addmm_patterns_init()
binary_folding_init()
def register_freezing_graph_pattern(pattern, extra_check=_return_true, pass_number=0):
return register_graph_pattern(
pattern,
extra_check=extra_check,
pass_dict=pass_patterns[pass_number],
)
def register_binary_folding_pattern(pattern, extra_check=_return_true):
return register_graph_pattern(
pattern,
extra_check=extra_check,
pass_dict=binary_folding_pass,
)
@functools.lru_cache(None)
def addmm_patterns_init():
device = next(
(gpu for gpu in GPU_TYPES if getattr(torch, gpu).is_available()), "cpu"
)
val = functools.partial(torch.empty, (10, 10), device=device, requires_grad=False)
def check_concat_weights(match):
is_cpu = match.kwargs["inp"].meta["val"].is_cpu
if is_cpu and not config.cpp.enable_concat_linear:
return False
weight_inputs = ["w1", "w2"]
if "w3" in match.kwargs:
weight_inputs.append("w3")
equal_shape_inputs = [weight_inputs]
if "b1" in match.kwargs:
bias_inputs = ["b1", "b2"]
if "b3" in match.kwargs:
bias_inputs.append("b3")
equal_shape_inputs.append(bias_inputs)
for equal_shape_group in equal_shape_inputs:
inps = [match.kwargs[name] for name in equal_shape_group]
if not all(
inp.op == "get_attr"
and inp.meta["val"].shape == inps[0].meta["val"].shape
for inp in inps
):
return False
return True
def matmul_fuse_pattern(inp, w1, w2, w3):
return (inp @ w1, inp @ w2, inp @ w3)
def matmul_replacement(inp, w1, w2, w3):
cat_t = torch.cat((w1, w2, w3), dim=1)
mm = inp @ cat_t
return mm.chunk(3, dim=1)
register_replacement(
matmul_fuse_pattern,
matmul_replacement,
[val(), val(), val(), val()],
fwd_only,
pass_patterns[0],
extra_check=check_concat_weights,
exclusive_arg_names=("w1", "w2", "w3"),
)
def matmul_fuse_pattern_two(inp, w1, w2):
return (inp @ w1, inp @ w2)
def matmul_replacement_two(inp, w1, w2):
cat_t = torch.cat((w1, w2), dim=1)
mm = inp @ cat_t
return mm.chunk(2, dim=1)
register_replacement(
matmul_fuse_pattern_two,
matmul_replacement_two,
[val(), val(), val()],
fwd_only,
pass_patterns[0],
extra_check=check_concat_weights,
exclusive_arg_names=("w1", "w2"),
)
def addmm_fuse_pattern_second(inp, w1, w2, w3, b1, b2, b3):
return (
aten.addmm(b1, inp, w1),
aten.addmm(b2, inp, w2),
aten.addmm(b3, inp, w3),
)
def addmm_fuse_replacement_second(inp, w1, w2, w3, b1, b2, b3):
cat_w = torch.cat((w1, w2, w3), dim=1)
cat_b = torch.cat((b1, b2, b3))
return aten.addmm(cat_b, inp, cat_w).chunk(3, dim=1)
register_replacement(
addmm_fuse_pattern_second,
addmm_fuse_replacement_second,
[val() for _ in range(7)],
fwd_only,
pass_patterns[0],
extra_check=check_concat_weights,
exclusive_arg_names=("w1", "w2", "w3", "b1", "b2", "b3"),
)
def same_dtype(match):
return match.output_node().args[0].meta["val"].dtype == match.kwargs["dtype"]
@register_graph_pattern(
CallFunction(
torch.ops.prims.convert_element_type.default,
Ignored(),
KeywordArg("dtype"),
),
pass_dict=pass_patterns[0],
extra_check=same_dtype,
)
def unnecessary_dtype_convert(match: Match, **kwargs):
"""Remove unnecessary dtype conversion op, probably left as a result of Conv-Bn folding"""
graph = match.graph
node = match.output_node()
node.replace_all_uses_with(node.args[0]) # type: ignore[arg-type]
graph.erase_node(node)
|
