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
|
# Owner(s): ["oncall: jit"]
from torch.testing._internal.jit_utils import JitTestCase
from torch._C import parse_ir
import torch
if __name__ == '__main__':
raise RuntimeError("This test file is not meant to be run directly, use:\n\n"
"\tpython test/test_jit.py TESTNAME\n\n"
"instead.")
class TestAliasAnalysis(JitTestCase):
def test_becomes_wildcard_annotations(self):
graph_str = """
graph(%a.1 : Tensor, %b.1 : Tensor):
%11 : NoneType = prim::Constant()
%8 : int = prim::Constant[value=0]()
%7 : int = prim::Constant[value=1]()
%x.1 : Tensor = aten::add(%a.1, %b.1, %7)
%y.1 : Tensor[] = aten::split(%x.1, %7, %8)
return ()
"""
graph = parse_ir(graph_str)
alias_db = graph.alias_db()
split_node = graph.findNode("aten::split")
# split input enters wildcard set, list initalized as containing wildcard set
self.assertTrue(alias_db.may_contain_alias(next(split_node.inputs()), split_node.output()))
# because %x.1 enters wildcard set, it now aliases other members of wildcard set (graph inputs)
self.assertTrue(alias_db.may_contain_alias(next(split_node.inputs()), next(graph.inputs())))
def test_nested_list_construct_not_wildcard(self):
@torch.jit.script
def foo(x):
y = torch.rand([2, 2])
return [y]
graph = foo.graph
graph.alias_db()
alias_db = graph.alias_db()
ten_construct = graph.findNode("aten::rand").output()
output = next(graph.outputs())
self.assertTrue(alias_db.may_contain_alias(ten_construct, output))
self.assertFalse(alias_db.may_contain_alias(next(graph.inputs()), ten_construct))
def test_recursive_calls(self):
@torch.jit.script
def foo(x, y):
x.add_(1)
return x + y
@torch.jit.script
def caller():
a = torch.rand([2, 2])
b = torch.ones([2, 2])
out1 = foo(a, b)
c = torch.rand([1])
d = torch.ones([2])
out2 = foo(d, c)
return out1, out2
isFrozen = False
descend_function_calls = True
alias_db = caller.graph.alias_db(isFrozen, descend_function_calls)
func_calls = caller.graph.findAllNodes("prim::CallFunction")
self.assertEqual(len(func_calls), 2)
for node in func_calls:
inps = list(node.inputs())
self.assertTrue(alias_db.has_writers(inps[1]))
self.assertFalse(alias_db.has_writers(inps[2]))
class Mod(torch.nn.Module):
def forward(self):
a = torch.rand([2, 2])
b = torch.ones([2, 2])
out1 = self.foo2(a, b)
c = torch.rand([1])
d = torch.ones([2])
out2 = self.foo2(d, c)
return out1, out2
def foo2(self, x, y):
x.add_(1)
return x + y
mod = torch.jit.script(Mod())
alias_db = mod.graph.alias_db(isFrozen, descend_function_calls)
func_calls = mod.graph.findAllNodes("prim::CallMethod")
self.assertEqual(len(func_calls), 2)
for node in func_calls:
inps = list(node.inputs())
self.assertTrue(alias_db.has_writers(inps[1]))
self.assertFalse(alias_db.has_writers(inps[2]))
|
