Extend basic matcher to handle multiple-output-nodes

onnxscript/rewriter/generic_pattern.py

@@ -298,7 +298,7 @@
             return self.none(starting_node, inspect.currentframe().f_lineno)
 
         for graph_input, pattern_input in zip(graph_node.inputs, pattern_node.inputs):
-            if len(list(graph_input.uses())) != len(list(pattern_input.uses())):
+            if len(graph_input.uses()) != len(pattern_input.uses()):
                 self._hint(
                     "BACKWARD: one input is used outside the pattern",
                     "-- pattern",
@@ -423,12 +423,12 @@
             return match_count
         if len(free) < len(pattern_node_users_not_matched):
             # Not enough successors to match the remaining patterns.
-            return self.none(node, inspect.currentframe().f_lineno)
+            return self.none(starting_node, inspect.currentframe().f_lineno)
         if len(pattern_node_users_not_matched) == len(free) == 1:
             # Only one option again.
             graph_node = free[0]
             if pattern_node_users_not_matched[0].op_identifier() != graph_node.op_identifier():
-                return self.none(node, inspect.currentframe().f_lineno)
+                return self.none(starting_node, inspect.currentframe().f_lineno)
 
             key = pattern_node_users_not_matched[0]
             if self.verbose >= 10:
@@ -461,11 +461,11 @@
                 "-- model-matched",
                 pattern_node_users_matched,
             )
-            return self.none(node, inspect.currentframe().f_lineno)
+            return self.none(starting_node, inspect.currentframe().f_lineno)
         for k, v in ec.items():
             if gc[k] < v:
                 # Not enough types to match.
-                return self.none(node, inspect.currentframe().f_lineno)
+                return self.none(starting_node, inspect.currentframe().f_lineno)
 
         # At this stage, we know matching the types is possible.
         # We first mark whatever is possible.

onnxscript/rewriter/generic_pattern_test.py

@@ -12,13 +12,21 @@
 import onnx.parser
 import onnx.reference
 import onnxruntime as ort
+import parameterized
 
 from onnxscript import ir
 from onnxscript.rewriter import generic_pattern, pattern
 
 FLOAT = onnx.TensorProto.FLOAT
 
 
+@parameterized.parameterized_class(
+    ("matcher_algo",),
+    [
+        (generic_pattern.GenericPatternMatcher,),
+        (pattern.SimplePatternMatcher,),
+    ],
+)
 class GenericPatternTest(unittest.TestCase):
     def _range(self, *shape, bias: float | None = None):
         n = np.prod(shape)
@@ -48,7 +56,7 @@ def validate_mapping(context, x, y, z, **_) -> bool:
             match_pattern,
             apply_pattern,
             validate_mapping,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
         )
 
         class AddAdd(onnx.reference.op_run.OpRun):
@@ -128,7 +136,7 @@ def validate_mapping(context, **_) -> bool:
             match_pattern,
             apply_pattern,
             validate_mapping,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
             verbose=10,
         )
 
@@ -256,11 +264,7 @@ def match_pattern(op, x):
         def apply_pattern(op, x, **_):
             return op.SinCos(x, domain="com.microsoft", outputs=2)
 
-        rule = pattern.RewriteRule(
-            match_pattern,
-            apply_pattern,
-            matcher=generic_pattern.GenericPatternMatcher,
-        )
+        rule = pattern.RewriteRule(match_pattern, apply_pattern, matcher=self.matcher_algo)
         model_proto = onnx.parser.parse_model(
             """
             <ir_version: 7, opset_import: [ "" : 17]>
@@ -281,8 +285,10 @@ def apply_pattern(op, x, **_):
         self.assertEqual(len(graph.node), 2)
         self.assertEqual(graph.node[0].op_type, "SinCos")
 
-    @unittest.skip("Input variable reuse not supported yet")
     def test_shared_root_value_extra_use(self):
+        if self.matcher_algo is generic_pattern.GenericPatternMatcher:
+            raise unittest.SkipTest("GenericPatternMatcher does not support extra uses yet.")
+
         def match_pattern(op, x):
             t1 = op.Sin(x)
             t2 = op.Cos(x)
@@ -294,7 +300,7 @@ def apply_pattern(op, x, **_):
         rule = pattern.RewriteRule(
             match_pattern,
             apply_pattern,
-            matcher=generic_pattern.GenericPatternMatcher,
+            matcher=self.matcher_algo,
         )
         model_proto = onnx.parser.parse_model(
             """
@@ -314,7 +320,7 @@ def apply_pattern(op, x, **_):
         rule.apply_to_model(ir_model)
         graph = ir_model.graph
         self.assertEqual(len(graph), 3)
-        self.assertEqual(graph.node[0].op_type, "SinCos")
+        self.assertEqual(graph.node(0).op_type, "SinCos")
 
     def test_rotary_embedding(self):
         # The test work on a model if it has the expected name.
@@ -367,7 +373,7 @@ def apply_pattern(op, x, pos_ids, axis, **_):
             match_pattern,
             apply_pattern,
             validate_mapping,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
             verbose=10,
         )
 
@@ -389,7 +395,8 @@ def apply_pattern(op, x, pos_ids, axis, **_):
         self.assertEqual(expected, [n.op_type for n in rewriten_model.graph.node])
         out = buffer.getvalue()
         # TODO(Rama): What is this assertion testing? Is it to check that `verbose` is working?
-        self.assertIn("[GenericPatternMatcher.match", out)
+        if self.matcher_algo is generic_pattern.GenericPatternMatcher:
+            self.assertIn("[GenericPatternMatcher.match", out)
 
     def test_rotary_embedding_onnxscript(self):
         # The test work on a model if it has the expected name.
@@ -432,7 +439,7 @@ def rotary_apply_pattern(op, x, pos_ids, axis, **_):
             rotary_match_pattern,
             rotary_apply_pattern,
             validate_rotary_mapping,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
             verbose=10,
         )
 
@@ -454,7 +461,8 @@ def rotary_apply_pattern(op, x, pos_ids, axis, **_):
         self.assertEqual(expected, [n.op_type for n in rewriten_model.graph.node])
         out = buffer.getvalue()
         # TODO(justinchuby): Remove this assert - capturing stdout is not robust
-        self.assertIn("[GenericPatternMatcher.match", out)
+        if self.matcher_algo is generic_pattern.GenericPatternMatcher:
+            self.assertIn("[GenericPatternMatcher.match", out)
 
     def test_rotary_emb_file_onnxscript(self):
         # The test work on a model if it has the expected name.
@@ -504,7 +512,7 @@ def rotary_apply_pattern(op, x, pos_ids, axis):
             rotary_match_pattern,
             rotary_apply_pattern,
             validate_rotary_mapping,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
             verbose=10,
         )
 
@@ -561,7 +569,7 @@ def transpose_transpose_apply_pattern(op, X, XT: ir.Value, Y, **_):
             transpose_transpose_pattern,
             transpose_transpose_apply_pattern,
             transpose_transpose_check,
-            generic_pattern.GenericPatternMatcher,
+            self.matcher_algo,
             verbose=0,
         )