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
|
# mypy: allow-untyped-defs
from contextlib import contextmanager
from torch.fx.graph_module import (
_format_import_block,
GraphModule,
reduce_graph_module,
reduce_package_graph_module,
)
from torch.package import PackageExporter, sys_importer
from ._compatibility import compatibility
_use_lazy_graph_module_flag = False
_force_skip_lazy_graph_module_flag = False
@compatibility(is_backward_compatible=False)
@contextmanager
def _force_skip_lazy_graph_module():
"""
Skip using lazy graph module disregarding the setting of _use_lazy_graph_module.
Use to skip _LazyGraphModule when testing inductor torchscript related backend.
torch.jit.script a _LazyGraphModule results in following error:
https://gist.github.com/shunting314/5143654c8084aed84ecd19b818258a69
"""
try:
global _force_skip_lazy_graph_module_flag
prior = _force_skip_lazy_graph_module_flag
_force_skip_lazy_graph_module_flag = True
yield
finally:
_force_skip_lazy_graph_module_flag = prior
@compatibility(is_backward_compatible=False)
@contextmanager
def _use_lazy_graph_module(should_use: bool):
try:
global _use_lazy_graph_module_flag
prior = _use_lazy_graph_module_flag
_use_lazy_graph_module_flag = (
should_use and not _force_skip_lazy_graph_module_flag
)
yield
finally:
_use_lazy_graph_module_flag = prior
@compatibility(is_backward_compatible=False)
def _get_graph_module_cls():
return _LazyGraphModule if _use_lazy_graph_module_flag else GraphModule
def _make_graph_module(*args, graph_module_cls=None, **kwargs):
if graph_module_cls is None:
graph_module_cls = _get_graph_module_cls()
return graph_module_cls(*args, **kwargs)
@compatibility(is_backward_compatible=False)
class _LazyGraphModule(GraphModule):
"""
The main difference between _LazyGraphModule and GraphModule is how recompile happens.
GraphModule will do a 'recompile' call to generate python code and the forward method when it's
constructed. Later on if the graph get updated, recompile method can be called again to refresh
the saved python code and forward method.
However in some cases especially in inductor, the recompilation can be a waste since we never
check the python code for the graph module or call its forward method. A few more concreate
examples regarding pattern matching fx passes in inductor:
1. some passes will update the graph to be compiled and then call recompile on the GraphModule.
2. some passes will trace small pattern function to search it in the graph being compiled and
replace the match with the traced graph of a replacement function. The pattern graph and
replacement graph are quite small but there are large amount of them. Doing GraphModule.recompile
for them in GraphModule.__init__ is also a waste of time.
However simply skip calling GraphModule.recompile in these scenarios is also dangeruous.
People may want to check the python code or call the GraphModule's forward method for debugging purposes.
The way _LazyGraphModule solves it is, we override the recompile method to just mark the
need for recompilation but does not do the actual recompilation. Later on if people really
access the compiled python code or call the GraphModule's forward method, we do the real
recompilation.
"""
@classmethod
def from_graphmodule(cls, gm: GraphModule):
if isinstance(gm, _LazyGraphModule):
return gm
else:
return _LazyGraphModule(gm, gm.graph)
@staticmethod
def force_recompile(gm):
"""
Sometimes we need force a recompile as a workaround
- we want to do the real recompilation before symbolic_trace to avoid error:
https://gist.github.com/shunting314/75549c2e82ae07ac1139c94a3583d259
"""
if isinstance(gm, _LazyGraphModule):
gm.real_recompile()
def real_recompile(self):
if self._needs_recompile():
self._real_recompile()
@classmethod
def _needs_recompile(cls):
return cls.forward is cls._lazy_forward
def _lazy_forward(self, *args, **kwargs):
# Call self.real_recompile() rather than self._real_recompile() here.
# The _lazy_forward method may be saved and call repeatedly.
# Calling self.real_recompile can make sure we skip recompilation if
# we have already done so.
self.real_recompile()
assert not self._needs_recompile()
# call `__call__` rather than 'forward' since recompilation may
# install a wrapper for `__call__` to provide a customized error
# message.
return self(*args, **kwargs)
forward = _lazy_forward
# TODO: we shold handle __reduce_deploy__ the same way as __reduce_package__,
# or __reduce__ by calling _real_recompile. But I don't find a good way
# to test __reduce_deploy__ out. Also it's very unlikely that LazyGraphModule
# will be used in torch::deploy. So it's skipped for now.
def __reduce_package__(self, exporter: PackageExporter):
"""
Follow GraphModule.__reduce__ but call 'self._real_recompile' rather
than 'self.recompile' since for a _LazyGraphModule, self.recompile just
mark the need of recompilation and does not return the PythonCode object.
"""
python_code = self._real_recompile()
dict_without_graph = self.__dict__.copy()
dict_without_graph["_graphmodule_cls_name"] = self.__class__.__name__
del dict_without_graph["_graph"]
generated_module_name = f"fx-generated._{exporter.get_unique_id()}"
import_block = _format_import_block(python_code.globals, exporter.importer)
module_code = import_block + self.code
exporter.save_source_string(generated_module_name, module_code)
return (
reduce_package_graph_module,
(dict_without_graph, generated_module_name),
)
def __reduce__(self):
"""
Follow GraphModule.__reduce__ but call 'self._real_recompile' rather
than 'self.recompile' since for a _LazyGraphModule, self.recompile just
mark the need of recompilation and does not return the PythonCode object.
"""
python_code = self._real_recompile()
dict_without_graph = self.__dict__.copy()
import_block = _format_import_block(python_code.globals, sys_importer)
del dict_without_graph["_graph"]
return (reduce_graph_module, (dict_without_graph, import_block))
def _real_recompile(self):
return super().recompile()
@classmethod
def recompile(cls):
cls.forward = cls._lazy_forward
@property
def code(self) -> str:
self.real_recompile()
return super().code
def __str__(self) -> str:
"""
str(GraphModule) will access the _code attribute. Make sure recompile
happens so _code attribute is available.
"""
self.real_recompile()
return super().__str__()
|
