An improved version of subgraph matcher (pytorch#82090)

This new version of subgraph matcher further supports
- optionally match with pattern's placeholder and output nodes
- patterns with multiple outputs
- filtering out non-containing matches
- filtering out overlapping matches

TODOs:
- [x] Update replace_pattern() to use this matcher
- [x] Fix cases with identical anchor
- [x] Introduce wildcard matching, such Any, OneOf
- [ ] Improve node comparer to match args and kwargs values
Pull Request resolved: pytorch#82090
Approved by: https://github.com/ezyang

Author: SherlockNoMad

test/test_fx_passes.py

@@ -1,5 +1,6 @@
 # Owner(s): ["module: fx.passes"]
 
+from dataclasses import dataclass
 import operator
 import logging
 
@@ -9,6 +10,7 @@
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupport
 from torch.fx.passes.utils.fuser_utils import fuse_by_partitions
+from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
 from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
 from torch.testing._internal.jit_utils import JitTestCase
@@ -163,6 +165,8 @@ class MockOperatorSupport(OperatorSupport):
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         return node.op == "call_function" and node.target in {operator.add}
 
+
+@instantiate_parametrized_tests
 class TestFXGraphPasses(JitTestCase):
 
     @parametrize("fn, expected_partition", [
@@ -270,7 +274,372 @@ def test_fuser_util_xfail(self, partition):
         with self.assertRaises(Exception):
             fuse_by_partitions(gm, partitions)
 
-instantiate_parametrized_tests(TestFXGraphPasses)
+@dataclass
+class TestCase:
+    match_output: bool
+    match_placeholder: bool
+    num_matches: int
+    remove_overlapping_matches: bool = True
+
+class SingleNodePattern:
+    @staticmethod
+    def forward(x):
+        val = torch.neg(x)
+        return torch.add(val, val)
+
+    @staticmethod
+    def pattern(a):
+        return torch.neg(a)
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 0),
+        TestCase(False, True, 1),
+        TestCase(True, True, 0)
+    ]
+class SimplePattern:
+    @staticmethod
+    def forward(x, w1, w2):
+        m1 = torch.cat([w1, w2]).sum()
+        m2 = torch.cat([w2, w1]).sum()
+        m3 = torch.cat([m1, m2]).sum()
+        return x + torch.max(m1) + torch.max(m2) + m3
+
+    @staticmethod
+    def pattern(a, b):
+        return torch.cat([a, b]).sum()
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 3),
+        TestCase(True, False, 0),
+        TestCase(False, True, 2),
+        TestCase(True, True, 0)
+    ]
+
+class SimpleFullGraphMatching:
+    @staticmethod
+    def forward(x):
+        a = torch.neg(x)
+        return torch.add(a, a)
+
+    @staticmethod
+    def pattern(x):
+        a = torch.neg(x)
+        return torch.add(a, a)
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 1),
+        TestCase(False, True, 1),
+        TestCase(True, True, 1)
+    ]
+
+class DiamondShapePatternTestCase:
+    @staticmethod
+    def forward(x):
+        a = torch.neg(x)
+
+        a = a.relu()
+        left = a.sigmoid()
+        right = a.relu()
+        out = left + right
+
+        return out
+
+    @staticmethod
+    def pattern(a):
+        a = a.relu()
+        left = a.sigmoid()
+        right = a.relu()
+        out = left + right
+        return out
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 1),
+        TestCase(False, True, 0),
+        TestCase(True, True, 0)
+    ]
+
+class NonFullyContainedMatches:
+    @staticmethod
+    def forward(x, w1, w2, b1, b2):
+        # fully contained matched subgraph
+        m1 = torch.cat([w1, w2])
+        m2 = torch.cat([x, b2])
+        t0 = torch.addmm(b1, m1, m2.t())
+        t0_sum = torch.sum(t0)   # use of t0 is not leaking
+
+        # leaking matched subgraph, m3 is leaked
+        m3 = torch.cat([w1, w2])
+        m4 = torch.cat([x, b2])
+        t1 = torch.addmm(b1, m3, m4.t())
+        m3_sum = torch.sum(m3)
+
+        return t0_sum, m3_sum
+
+    @staticmethod
+    def pattern(x, w1, w2, b1, b2):
+        m1 = torch.cat([w1, w2])
+        m2 = torch.cat([x, b2])
+        return torch.addmm(b1, m1, m2.t())
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+
+        TestCase(True, False, 0),
+
+        TestCase(False, True, 1),     # leaked used of placeholder is not leaking
+    ]
+
+class ChainRepeatedPattern:
+    @staticmethod
+    def forward(x):
+        x = torch.sigmoid(x)
+        x = torch.sigmoid(x)
+        x = torch.sigmoid(x)
+        return torch.sigmoid(x)
+
+    @staticmethod
+    def pattern(x):
+        return torch.sigmoid(torch.sigmoid(x))
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 3, remove_overlapping_matches=False),
+        TestCase(False, False, 2, remove_overlapping_matches=True),
+        TestCase(True, False, 1),
+        TestCase(False, True, 1),
+        TestCase(True, True, 0)
+    ]
+
+class QuantizationModel:
+    @staticmethod
+    def forward(x):
+        x += 3
+        x = x.dequantize()
+        x = torch.sigmoid(x)
+        x = x.to(torch.float16)
+        return x
+
+    @staticmethod
+    def pattern(x):
+        x = x.dequantize()
+        x = torch.sigmoid(x)
+        x = x.to(torch.float16)
+        return x
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 1),
+        TestCase(False, True, 0),
+        TestCase(True, True, 0)
+    ]
+
+class MultipleOutputsWithDependency:
+    @staticmethod
+    def forward(x):
+        y = x.relu()
+        z = y.sigmoid()
+        return z, y
+
+    @staticmethod
+    def pattern(a):
+        b = a.relu()
+        c = b.sigmoid()
+        return b, c     # outputs have data dependency
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 0),
+        TestCase(False, True, 1),
+        TestCase(True, True, 0)
+    ]
+
+class MultipleOutputsWithoutDependency:
+    @staticmethod
+    def forward(x):
+        x = x + 1
+
+        # target subgraph to match
+        x = x.relu()
+        z = x.sum()
+        y = x.sigmoid()
+
+        out = y.sigmoid() + z.sum()
+        return out
+
+    @staticmethod
+    def pattern(a):
+        a = a.relu()
+        b = a.sigmoid()
+        c = a.sum()
+        return b, c
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 0),
+        TestCase(False, True, 0),
+        TestCase(True, True, 0)
+    ]
+
+class MultipleOutputsMultipleOverlappingMatches:
+    @staticmethod
+    def forward(x):
+        x = x + 1
+
+        # target subgraph to match
+        x = x.relu()
+        z = x.sum()
+        z1 = x.sum()
+        y = x.sigmoid()
+        y1 = x.sigmoid()
+
+        return z + z1 + y + y1
+
+    @staticmethod
+    def pattern(a):
+        a = a.relu()
+        b = a.sigmoid()
+        c = a.sum()
+        return a, b, c
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 4, remove_overlapping_matches=False),
+        TestCase(False, False, 1, remove_overlapping_matches=True),
+    ]
+
+class MultipleOutputsMultipleNonOverlappingMatches:
+    @staticmethod
+    def forward(x):
+        x = x + 1
+
+        # target subgraph to match
+        x = x.relu()
+        z = x.sum()
+        y = x.sigmoid()
+
+        x = x.relu()
+        z1 = x.sum()
+        y1 = x.sigmoid()
+
+        return z + z1 + y + y1
+
+    @staticmethod
+    def pattern(a):
+        a = a.relu()
+        b = a.sigmoid()
+        c = a.sum()
+        return b, c
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+    ]
+
+class MultipleOutputsIdenticalAnchor:
+    @staticmethod
+    def forward(x):
+        x = x + 1
+
+        # target subgraph to match
+        x = x.relu()
+        y = x.sigmoid()
+        y1 = x.sigmoid()
+
+        return y, y1
+
+    @staticmethod
+    def pattern(a):
+        a = a.relu()
+        b = a.sigmoid()
+        b1 = a.sigmoid()
+        return b, b1
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        # (False, False, 2),  # FIXME: currently still matches to 2, should fix to 1
+        TestCase(True, False, 1),
+        TestCase(False, True, 0),
+    ]
+
+
+class MultipleOutputsHorizontalPattern:
+    @staticmethod
+    def forward(x):
+        x = x + 1
+
+        # target subgraph to match
+        y1 = x.relu()
+        y2 = x.sigmoid()
+
+        return y1, y2
+
+    @staticmethod
+    def pattern(a):
+        b1 = a.relu()
+        b2 = a.sigmoid()
+
+        return b1, b2
+
+    test_cases = [
+        # match_output, match_placeholder, num_matches
+        TestCase(False, False, 1),
+        TestCase(True, False, 1),
+        TestCase(False, True, 0),
+        TestCase(True, True, 0)
+    ]
+
+
+@instantiate_parametrized_tests
+class TestFXMatcherUtils(JitTestCase):
+
+    @parametrize("test_model", [
+        SingleNodePattern,
+        SimplePattern,
+        SimpleFullGraphMatching,
+        DiamondShapePatternTestCase,
+        NonFullyContainedMatches,
+        ChainRepeatedPattern,
+        QuantizationModel,
+        MultipleOutputsWithDependency,
+        MultipleOutputsWithoutDependency,
+        MultipleOutputsMultipleOverlappingMatches,
+        MultipleOutputsMultipleNonOverlappingMatches,
+        MultipleOutputsIdenticalAnchor,
+        MultipleOutputsHorizontalPattern
+    ])
+    def test_subgraph_matcher(self, test_model):
+        traced = symbolic_trace(test_model.forward)
+        pattern_traced = symbolic_trace(test_model.pattern)
+
+        for test_case in test_model.test_cases:
+
+            matcher = SubgraphMatcher(pattern_traced.graph,
+                                      match_output=test_case.match_output,
+                                      match_placeholder=test_case.match_placeholder,
+                                      remove_overlapping_matches=test_case.remove_overlapping_matches)
+            matches = matcher.match(traced.graph)
+
+            assert len(matches) == test_case.num_matches
+
+            for match in matches:
+                for node in pattern_traced.graph.nodes:
+                    if not test_case.match_placeholder and node.op == "placeholder":
+                        continue
+                    if not test_case.match_output and node.op == "output":
+                        continue
+                    assert node in match.nodes_map
+
 
 if __name__ == "__main__":
     run_tests()