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
|
# Owner(s): ["module: inductor"]
import torch
from torch import _dynamo as dynamo, _inductor as inductor
from torch._inductor.test_case import run_tests, TestCase
from torch._inductor.utils import gen_gm_and_inputs
from torch.fx import symbolic_trace
from torch.fx.experimental.proxy_tensor import make_fx
from torch.testing._internal.inductor_utils import HAS_CPU
class MyModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.a = torch.nn.Linear(10, 10)
self.b = torch.nn.Linear(10, 10)
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.relu(self.a(x))
x = torch.sigmoid(self.b(x))
return x
class MyModule2(MyModule):
def forward(self, x): # takes a dict of list
a, b = x["key"]
return {"result": super().forward(a) + b}
class MyModule3(MyModule):
def forward(self, x):
return (super().forward(x),)
class TestStandaloneInductor(TestCase):
"""
These test check that you can call TorchInductor directly without
going through TorchDynamo.
"""
def test_inductor_via_fx(self):
mod = MyModule3().eval()
inp = torch.randn(10)
correct = mod(inp)
mod_opt = inductor.compile(symbolic_trace(mod), [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_fx_tensor_return(self):
mod = MyModule().eval()
inp = torch.randn(10)
correct = mod(inp)
mod_opt = inductor.compile(symbolic_trace(mod), [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_fx_dict_input(self):
mod = MyModule2().eval()
inp = {"key": [torch.randn(10), torch.randn(10)]}
correct = mod(inp)
mod_opt = inductor.compile(symbolic_trace(mod), [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_make_fx(self):
mod = MyModule().eval()
inp = torch.randn(10)
correct = mod(inp)
mod_opt = inductor.compile(make_fx(mod)(inp), [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_bare_module(self):
mod = MyModule3().eval()
inp = torch.randn(10)
correct = mod(inp)
# no FX graph at all (mod must return list/tuple in this case)
mod_opt = inductor.compile(mod, [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_export1(self):
mod = MyModule3().eval()
inp = torch.randn(10)
correct = mod(inp)
gm, guards = dynamo.export(mod, inp, aten_graph=True)
mod_opt = inductor.compile(gm, [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_export2(self):
mod = MyModule2().eval()
inp = {"key": [torch.randn(10), torch.randn(10)]}
correct = mod(inp)
gm, guards = dynamo.export(mod, inp)
mod_opt = inductor.compile(gm, [inp])
actual = mod_opt(inp)
self.assertEqual(actual, correct)
def test_inductor_via_op_with_multiple_outputs(self):
x1 = torch.randn((2, 512, 128))
x2 = [128]
x3 = torch.randn(128)
x4 = torch.randn((128,))
x5 = 1e-6
mod, inp = gen_gm_and_inputs(
torch.ops.aten.native_layer_norm.default, (x1, x2, x3, x4, x5), {}
)
mod_opt = inductor.compile(mod, inp)
self.assertEqual(mod(*inp), mod_opt(*inp))
if __name__ == "__main__":
if HAS_CPU:
run_tests()
|
