Narrow multi-output match candidates to consumers of bound inputs

onnxscript/rewriter/_matcher.py

@@ -391,18 +391,56 @@ def match(
             # that can be used for debugging and testing. The GenericPatternMatcher is a
             # more sophisticated implementation, but incomplete.
             pattern_output_nodes = self.pattern.output_nodes
-            op_to_nodes: dict[tuple[str, str, str], list[ir.Node]] = {}
-            for n in graph_or_function:
-                op_to_nodes.setdefault(n.op_identifier(), []).append(n)
-            all_nodes = iter(graph_or_function)
+
+            # The first output node is matched first in every combination, so its result
+            # does not depend on the other candidates. Match it once: if it fails, no
+            # combination can succeed, and if it succeeds, its bindings narrow down the
+            # candidates for the other output nodes.
+            self._init_match(verbose)
+            if not self._match_node(pattern_output_nodes[0], node):
+                return self._match
+            first_match = self._match
+
+            def bound_value(pattern_value) -> ir.Value | None:
+                # Named pattern variables are bound by name, the others by the pattern value.
+                if pattern_value.name is not None:
+                    value = first_match.bindings.get(pattern_value.name)
+                else:
+                    value = first_match.value_bindings.get(pattern_value)
+                return value if isinstance(value, ir.Value) else None
+
+            op_to_nodes: dict[tuple[str, str, str], list[ir.Node]] | None = None
 
             def get_nodes(pattern_node):
+                nonlocal op_to_nodes
                 id = pattern_node.op_identifier()
+                # A graph node can only match if each of its inputs equals the value
+                # already bound to the corresponding pattern input, so it must be a
+                # consumer of that value.
+                for index, input_pattern in enumerate(pattern_node.inputs):
+                    if input_pattern is None:
+                        continue
+                    bound = bound_value(input_pattern)
+                    if bound is None:
+                        continue
+                    # A dict keeps the first-use order and drops repeated uses of a node.
+                    consumers = {
+                        user: None
+                        for user, user_index in bound.uses()
+                        if user_index == index
+                        and user.graph is self._graph
+                        and (id is None or user.op_identifier() == id)
+                    }
+                    return list(consumers)
                 if id is None:
-                    return all_nodes
+                    return list(graph_or_function)
+                if op_to_nodes is None:
+                    op_to_nodes = {}
+                    for n in graph_or_function:
+                        op_to_nodes.setdefault(n.op_identifier(), []).append(n)
                 return op_to_nodes.get(id, [])
 
-            candidates = [iter([node])] + [get_nodes(pn) for pn in pattern_output_nodes[1:]]
+            candidates = [[node]] + [get_nodes(pn) for pn in pattern_output_nodes[1:]]
             match = None
             for combination in itertools.product(*candidates):
                 self._init_match(verbose)

onnxscript/rewriter/pattern_test.py

@@ -4,6 +4,7 @@
 import io
 import logging
 import unittest
+import unittest.mock
 
 import numpy as np
 import onnx.checker
@@ -1055,5 +1056,75 @@ def div_replacement(op, x, y):
         )
 
 
+class MultiOutputPatternTest(unittest.TestCase):
+    @staticmethod
+    def _rule() -> pattern.RewriteRule:
+        def relu_neg_pattern(op, x):
+            return op.Relu(x), op.Neg(x)
+
+        def relu_neg_replacement(op, x):
+            return op.Relu(x), op.Mul(x, op.Constant(value_float=-1.0))
+
+        return pattern.RewriteRule(relu_neg_pattern, relu_neg_replacement)
+
+    def test_other_outputs_are_matched_against_consumers_of_bound_inputs(self):
+        model_proto = onnx.parser.parse_model(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N] a, float[N] b, float[N] c) => (float[N] y1, float[N] y2, float[N] y3, float[N] y4)
+            {
+                y1 = Relu(a)
+                y2 = Neg(b)
+                y3 = Neg(c)
+                y4 = Relu(c)
+            }
+            """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+
+        count = self._rule().apply_to_model(model)
+
+        self.assertEqual(count, 1)
+        # Only the Relu and Neg that share the input c are rewritten.
+        self.assertEqual(
+            sorted(n.op_type for n in model.graph), ["Constant", "Mul", "Neg", "Relu", "Relu"]
+        )
+        (mul,) = [n for n in model.graph if n.op_type == "Mul"]
+        self.assertEqual(mul.inputs[0].name, "c")
+
+    def test_candidates_are_limited_to_consumers_of_bound_inputs(self):
+        # Each Relu has its own input, so no Neg can complete a match with it. Only the
+        # Neg that reads the same input is tried, instead of every Neg in the graph.
+        n = 20
+        inputs = ", ".join(f"float[N] x{i}" for i in range(n))
+        outputs = ", ".join(f"float[N] r{i}, float[N] s{i}" for i in range(n))
+        body = "\n".join(f"r{i} = Relu(x{i})\ns{i} = Neg(x{(i + 1) % n})" for i in range(n))
+        model_proto = onnx.parser.parse_model(
+            f"""
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph ({inputs}) => ({outputs})
+            {{
+                {body}
+            }}
+            """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        matcher_class = onnxscript.rewriter._matcher.SimplePatternMatcher
+        original_multi_match = matcher_class._multi_match
+        calls = 0
+
+        def counting_multi_match(self, *args, **kwargs):
+            nonlocal calls
+            calls += 1
+            return original_multi_match(self, *args, **kwargs)
+
+        with unittest.mock.patch.object(matcher_class, "_multi_match", counting_multi_match):
+            count = self._rule().apply_to_model(model)
+
+        # Every Relu shares its input with exactly one Neg, so each pair is rewritten.
+        self.assertEqual(count, n)
+        self.assertLessEqual(calls, 2 * n)
+
+
 if __name__ == "__main__":
     unittest.main()