def test_subgraph_rewriter_single_pattern_match(self):
class M(torch.nn.Module):
def forward(self, x):
val = torch.neg(x) + torch.relu(x)
return torch.add(val, val)
def pattern(x):
return torch.neg(x) + torch.relu(x)
def replacement(x):
return torch.relu(x)
def comparison(x):
val = torch.relu(x)
return torch.add(val, val)
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.rand(1, 3)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_output = comparison_fn(x)
test_output = traced.forward(x)
self.assertEqual(ref_output, test_output)
def test_match_node_with_primitive_type_values(self):
def f(x):
x = x + 3
return x
def pattern(x, num):
x = x + num
return x
def replacement(x, num):
x = x - num
return x
def comparison(x):
x = x - 3
return x
traced = symbolic_trace(f)
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_preserves_logic(self):
class M(torch.nn.Module):
def forward(self, x):
val = torch.neg(x) + torch.relu(x)
return torch.add(val, val)
def pattern(x):
return torch.neg(x) + torch.relu(x)
def comparison(x):
val = torch.neg(x) + torch.relu(x)
return torch.add(val, val)
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.rand(1, 3)
# Replace `pattern` with the same pattern (shouldn't change
# the underlying logic)
subgraph_rewriter.replace_pattern(traced, pattern, pattern)
traced.graph.lint()
ref_output = comparison_fn(x)
test_output = traced.forward(x)
self.assertEqual(ref_output, test_output)
def test_subgraph_writer_replace_consecutive_submodules(self):
def f(x):
x = torch.sigmoid(x)
x = torch.sigmoid(x)
return torch.sigmoid(x)
def pattern(x):
return torch.sigmoid(x)
def replacement(x):
return torch.exp(x)
def comparison(x):
x = torch.exp(x)
x = torch.exp(x)
return torch.exp(x)
traced = symbolic_trace(f)
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_pattern_output_pattern_node_can_have_users_that_are_not_matched(
self):
class M(torch.nn.Module):
def forward(self, x):
y = torch.relu(x)
return torch.neg(y) - y
def pattern(x):
return torch.relu(x)
def replacement(x):
return torch.sigmoid(x)
def comparison(x):
y = torch.sigmoid(x)
return torch.neg(y) - y
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_pattern_is_entire_graph(self):
class M(torch.nn.Module):
def forward(self, x):
a = torch.neg(x)
return torch.add(a, a)
def pattern(x):
a = torch.neg(x)
return torch.add(a, a)
def replacement(x):
a = torch.sigmoid(x)
return torch.cat([a, a])
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(replacement)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_traced_as_callable(self):
class M(torch.nn.Module):
def forward(self, x):
val = torch.neg(x) + torch.relu(x)
return torch.add(val, val)
class Pattern(torch.nn.Module):
def forward(self, x):
return torch.neg(x) + torch.relu(x)
class Replacement(torch.nn.Module):
def forward(self, x):
return torch.sigmoid(x)
def comparison(x):
val = torch.sigmoid(x)
return torch.add(val, val)
traced = symbolic_trace(M())
traced_pattern = symbolic_trace(Pattern())
traced_replacement = symbolic_trace(Replacement())
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, traced_pattern,
traced_replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_correct_output_replacement(self):
class M(torch.nn.Module):
def forward(self, x, y):
val = torch.neg(y) + torch.relu(x)
return torch.add(val, val)
def pattern(x):
return torch.relu(x)
def replacement(x):
return torch.neg(x)
def comparison(x, y):
val = torch.neg(y) + torch.neg(x)
return torch.add(val, val)
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.randn(4, 4)
y = torch.randn(4, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x, y)
test_outs = traced.forward(x, y)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_multiple_pattern_match(self):
class M(torch.nn.Module):
def forward(self, x, w1, w2):
m1 = torch.cat([w1, w2]).sum()
m2 = torch.cat([w1, w2]).sum()
return x + torch.max(m1) + torch.max(m2)
def pattern(w1, w2):
return torch.cat([w1, w2]).sum()
def replacement(w1, w2):
return torch.stack([w1, w2])
def comparison(x, w1, w2):
m1 = torch.stack([w1, w2])
m2 = torch.stack([w1, w2])
return x + torch.max(m1) + torch.max(m2)
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.rand(1, 3)
w1 = torch.rand(1, 3)
w2 = torch.rand(1, 3)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x, w1, w2)
test_outs = traced.forward(x, w1, w2)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_replaces_referenced_submodules(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
self.submod = torch.nn.ReLU()
def forward(self, x):
x = x + 1
return self.submod(self.sigmoid(x))
class Pattern(torch.nn.Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
self.submod = torch.nn.ReLU()
def forward(self, x):
return self.submod(self.sigmoid(x))
class Replacement(torch.nn.Module):
def __init__(self):
super().__init__()
self.id = torch.nn.Identity()
self.submod = torch.nn.ReLU()
def forward(self, x):
return self.submod(self.id(x))
class Comparison(torch.nn.Module):
def __init__(self):
super().__init__()
self.id = torch.nn.Identity()
self.submod = torch.nn.ReLU()
def forward(self, x):
x = x + 1
return self.submod(self.id(x))
traced = symbolic_trace(M())
comparison = Comparison()
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, Pattern(), Replacement())
traced.graph.lint()
ref_outs = comparison(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
traced.get_submodule("id")
with self.assertRaisesRegex(AttributeError, "has no attribute"):
traced.get_submodule("sigmoid")
submod = traced.get_submodule("submod")
self.assertEqual(type(submod), torch.nn.ReLU)
def test_subgraph_rewriter_internal_pattern_nodes_cannot_have_users_that_are_not_matched(
self):
class M(torch.nn.Module):
def forward(self, x, w1, w2, b1, b2):
m0 = torch.cat([w1, w2])
m1 = torch.cat([w1, w2])
m2 = torch.cat([x, b2])
t0 = torch.addmm(b1, m1, m2.t())
t1 = torch.sum(w1, 1)
t2 = torch.addmm(b1, m1, m2.t())
return torch.sum(t1), torch.sum(t2)
def pattern(x, w1, w2, b1, b2):
m1 = torch.cat([w1, w2])
m2 = torch.cat([x, b2])
return torch.addmm(b1, m1, m2.t())
def replacement(x, w1, w2, b1, b2):
return torch.cat([x, w1, w2])
traced = symbolic_trace(M())
# Result should be [] since no matches can be found
res = subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
self.assertEqual(res, [])
def test_match_attrs(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.attr = torch.tensor([3])
def forward(self, x):
return x + self.attr
class Pattern(torch.nn.Module):
def __init__(self):
super().__init__()
self.attr = torch.tensor([])
def forward(self, x):
return x + self.attr
class Replacement(torch.nn.Module):
def __init__(self):
super().__init__()
self.attr = torch.tensor([])
def forward(self, x):
return x - self.attr
class Comparison(torch.nn.Module):
def __init__(self):
super().__init__()
self.attr = torch.tensor([3])
def forward(self, x):
return x - self.attr
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(Comparison())
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, Pattern(), Replacement())
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_graph_argument_order(self):
class M(torch.nn.Module):
def forward(self, x, y):
return torch.mm(x, y)
def pattern(x, y):
return torch.mm(x, y)
def comparison(x, y):
return torch.mm(x, y)
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
y = torch.randn(4, 5)
subgraph_rewriter.replace_pattern(traced, pattern, pattern)
traced.graph.lint()
ref_outs = comparison_fn(x, y)
test_outs = traced.forward(x, y)
self.assertEqual(ref_outs, test_outs)
def test_subgraph_rewriter_placeholder_matching(self):
"""
This tests that a placeholder Node can be matched to a Node with
a different number of input Nodes. In the example below, the
original traced Module looks like this:
opcode target args kwargs
------------- ---------------------------------------------------------- ------------------------ --------
placeholder x () {}
call_function <built-in function add> (x, 3) {}
call_method dequantize (add,) {}
call_function <built-in method sigmoid of type object at 0x7f7c1f440fe0> (dequantize,) {}
call_method to (sigmoid, torch.float16) {}
output output (to,) {}
while the pattern we want to match looks like this:
opcode target args kwargs
------------- ---------------------------------------------------------- ------------------------ --------
placeholder x () {}
call_method dequantize (x,) {}
call_function <built-in method sigmoid of type object at 0x7f7c1f440fe0> (dequantize,) {}
call_method to (sigmoid, torch.float16) {}
output output (to,) {}
Here, we want to be able to match the original graph's
`call_function.add` Node with the pattern graph's
`plaeholder.x` Node.
Credit to Jerry Zhang (GitHub: jerryzh168) for this test case
"""
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.dtype = torch.float16
def forward(self, x):
x += 3
x = x.dequantize()
x = torch.sigmoid(x)
dtype = self.dtype
x = x.to(dtype)
return x
def pattern(x):
x = x.dequantize()
x = torch.sigmoid(x)
x = x.to(torch.float16)
return x
def replacement(x):
return x
def comparison(x):
return x + 3
traced = symbolic_trace(M())
comparison_fn = symbolic_trace(comparison)
x = torch.randn(3, 4)
subgraph_rewriter.replace_pattern(traced, pattern, replacement)
traced.graph.lint()
ref_outs = comparison_fn(x)
test_outs = traced.forward(x)
self.assertEqual(ref_outs, test_outs)