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
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
|
# Owner(s): ["module: fx"]
import copy
import unittest
from typing import Optional, Set, Type
import torch
import torch.fx
from torch.testing._internal.common_utils import IS_MACOS, TestCase
class TestDCE(TestCase):
def _custom_is_impure_node(self, node: torch.fx.Node) -> bool:
if node.is_impure():
return True
# a custom function that defines add operators as impure.
if node.target == torch.ops.aten.add:
return True
return False
def _has_nodes_without_users(self, m: torch.fx.GraphModule, custom: bool = False):
for node in m.graph.nodes:
if (not custom and node.is_impure()) or (
custom and self._custom_is_impure_node(node)
):
continue
if len(node.users) == 0:
return True
return False
def _get_num_placeholders(self, m: torch.fx.GraphModule) -> int:
count = 0
for node in m.graph.nodes:
if node.op == "placeholder":
count += 1
return count
def _run_dce_and_test(
self,
m: torch.nn.Module,
expect_dce_changes: bool,
modules_to_be_leafs: Optional[Set[Type]] = None,
custom: bool = False,
):
class TestTracer(torch.fx.Tracer):
def is_leaf_module(self, m, qualname):
if modules_to_be_leafs and type(m) in modules_to_be_leafs:
return True
return super().trace(m, qualname)
traced: torch.fx.GraphModule = torch.fx.GraphModule(m, TestTracer().trace(m))
print(str(traced.graph))
# Verify there are nodes without users (if expected).
has_nodes_without_users = self._has_nodes_without_users(traced, custom=custom)
if expect_dce_changes:
self.assertTrue(has_nodes_without_users)
else:
self.assertFalse(has_nodes_without_users)
# Get the original number of placeholders to verify it doesn't change
# during DCE.
orig_num_phs = self._get_num_placeholders(traced)
if custom:
changed = traced.graph.eliminate_dead_code(
is_impure_node=self._custom_is_impure_node
)
else:
changed = traced.graph.eliminate_dead_code()
self.assertTrue(changed if expect_dce_changes else not changed)
# Verify there are no nodes without users after DCE is run.
self.assertFalse(self._has_nodes_without_users(traced, custom=custom))
new_num_phs = self._get_num_placeholders(traced)
self.assertEqual(orig_num_phs, new_num_phs)
traced.recompile()
# Make sure we run and get the same results before/after DCE.
inputs = [torch.tensor([1.5])] * new_num_phs
inputs_copy = copy.deepcopy(inputs)
self.assertTrue(torch.equal(m(*inputs), traced(*inputs_copy)))
def test_simple(self):
"""
Tests that a single node in the graph is DCE'd correctly.
"""
class TestModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.attr_1 = torch.nn.Parameter(torch.tensor([-0.9]))
def forward(self, x):
a = x + 1
return x + self.attr_1
self._run_dce_and_test(TestModule(), expect_dce_changes=True)
def test_dead_chain(self):
"""
Tests that a chain of two nodes in the graph are DCE'd correctly.
"""
class TestModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.attr_1 = torch.nn.Parameter(torch.tensor([-0.9]))
def forward(self, x):
a = x + 1
b = a * 7
return x + self.attr_1
self._run_dce_and_test(TestModule(), expect_dce_changes=True)
def test_dead_getattr(self):
"""
Tests that a getatrr in the graph is DCE'd correctly.
"""
class TestModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.attr_1 = torch.nn.Parameter(torch.tensor([-0.9]))
def forward(self, x):
a = x + 1
b = a * self.attr_1
return x + 11
self._run_dce_and_test(TestModule(), expect_dce_changes=True)
def test_dead_placeholder(self):
"""
Tests that a placeholder in the graph is not DCE'd, as that would change
the function signature.
"""
class TestModule(torch.nn.Module):
def forward(self, x, y):
return x + 7
self._run_dce_and_test(TestModule(), expect_dce_changes=False)
def test_dead_placeholder_with_user(self):
"""
Tests that a placeholder in the graph is not DCE'd, as that would change
the function signature. Also verifies that a dead node that uses the
placeholder is DCE'd.
"""
class TestModule(torch.nn.Module):
def forward(self, x, y):
a = y + 2
return x + 7
self._run_dce_and_test(TestModule(), expect_dce_changes=True)
def test_keep_module_with_side_effects(self):
"""
Test that DCE doesn't remove a module if it's specified as having side effects.
"""
class ReLUImpure(torch.nn.ReLU):
_is_impure = True
class TestModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.relu = ReLUImpure()
def forward(self, a: torch.Tensor) -> torch.Tensor:
r = self.relu(a)
return a * 2
self._run_dce_and_test(
TestModule(), expect_dce_changes=False, modules_to_be_leafs={ReLUImpure}
)
def test_keep_torch_assert(self):
"""
Test that DCE doesn't remove torch._assert since it has side effects.
"""
class TestModule(torch.nn.Module):
def forward(self, a: torch.Tensor) -> torch.Tensor:
torch._assert(torch.equal(a, a), "a must equal a")
return a * 2
# Note: Don't need to specify torch._assert as having side effects
# because it's known to.
self._run_dce_and_test(TestModule(), expect_dce_changes=False)
def test_impure_nodes_args(self):
"""
Test that DCE doesn't remove call_function nodes with side effects.
"""
class TestModule(torch.nn.Module):
def forward(self, a: torch.Tensor) -> torch.Tensor:
torch._ops.ops.aten.add_.Tensor(a, 1)
return a * 2
# %add_ node should not be removed because it has side effects.
self._run_dce_and_test(TestModule(), expect_dce_changes=False)
def test_impure_kwargs(self):
"""
Test that DCE doesn't remove call_function nodes with side effects on kwargs.
"""
class TestModule(torch.nn.Module):
def forward(self, a: torch.Tensor) -> torch.Tensor:
b = a + 1
torch._ops.ops.aten.add.out(b, b, out=a, alpha=2)
return a
# %add_out node should not be removed because it has side effects.
self._run_dce_and_test(TestModule(), expect_dce_changes=False)
def test_impure_custom(self):
"""
Test that DCE doesn't remove nodes marked as impure by a custom function.
"""
class TestModule(torch.nn.Module):
def forward(self, a: torch.Tensor) -> torch.Tensor:
b = a + 1
c = torch._ops.ops.aten.add(b, b)
return a
# %add_out node should not be removed because it has side effects.
self._run_dce_and_test(TestModule(), expect_dce_changes=False, custom=True)
@unittest.skipIf(IS_MACOS, "Not working on macos")
def test_keep_collectives(self):
"""
Test that DCE doesn't remote collective ops even the results are not used.
"""
from torch.testing._internal.distributed.fake_pg import FakeStore
class TestModule(torch.nn.Module):
def forward(
self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
) -> torch.Tensor:
d = torch.ops.aten.mul.Tensor(a, b)
e = torch.ops.aten.mul.Tensor(a, c)
future = torch.ops._c10d_functional.all_reduce.default(e, "sum", "0")
synced_e = torch.ops._c10d_functional.wait_tensor.default(
future
) # synced_e is not used
return d
torch.distributed.init_process_group(
backend="fake",
world_size=2,
rank=0,
store=FakeStore(),
)
# collective nodes should not be removed because they have side effects.
self._run_dce_and_test(TestModule(), expect_dce_changes=False, custom=False)
torch.distributed.destroy_process_group()
@unittest.skipIf(IS_MACOS, "Not working on macos")
def test_keep_collectives_no_overload(self):
"""
Test that DCE doesn't remote collective ops (no overload version) even the results are not used.
"""
from torch.testing._internal.distributed.fake_pg import FakeStore
class TestModule(torch.nn.Module):
def forward(
self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
) -> torch.Tensor:
d = torch.ops.aten.mul(a, b)
e = torch.ops.aten.mul(a, c)
future = torch.ops._c10d_functional.all_reduce(e, "sum", "0")
synced_e = torch.ops._c10d_functional.wait_tensor(
future
) # synced_e is not used
return d
torch.distributed.init_process_group(
backend="fake",
world_size=2,
rank=0,
store=FakeStore(),
)
# collective nodes should not be removed because they have side effects.
self._run_dce_and_test(TestModule(), expect_dce_changes=False, custom=False)
torch.distributed.destroy_process_group()
|
