Handle matching against None explicitly (#2460)

Provide a way to indicate that a pattern-variable can match successfully
against a None-valued input. Cleanup current handling which was
inconsistent in one place. Add test cases.

---------

Signed-off-by: Ganesan Ramalingam <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: gramalingam

onnxscript/rewriter/_matcher.py

@@ -149,18 +149,21 @@ def _match_node(self, pattern_node: _pattern_ir.NodePattern, node: ir.Node) -> b
         match.bind_node(pattern_node, node)
 
         # TODO: Revisit this to handle optional trailing inputs better.
-        if pattern_node.allow_other_inputs:
-            if len(node.inputs) < len(pattern_node.inputs):
+
+        if len(node.inputs) > len(pattern_node.inputs):
+            if not pattern_node.allow_other_inputs:
                 return self.fail(
-                    f"Number of inputs ({len(node.inputs)}) is less than expected ({len(pattern_node.inputs)})"
+                    f"Number of inputs ({len(node.inputs)}) is greater than expected ({len(pattern_node.inputs)})"
                 )
+            checked_inputs = zip(node.inputs, pattern_node.inputs)
         else:
-            if len(node.inputs) != len(pattern_node.inputs):
-                return self.fail(
-                    f"Input nums mismatch. {len(node.inputs)} vs {len(pattern_node.inputs)}"
-                )
+            # In ONNX, trailing Nones can be omitted in the inputs of a node. So, we extend actual
+            # node inputs with None values to match the pattern node inputs length when zipping.
+            checked_inputs = itertools.zip_longest(
+                node.inputs, pattern_node.inputs, fillvalue=None
+            )
 
-        for arg_value, arg_pattern in zip(node.inputs, pattern_node.inputs):
+        for arg_value, arg_pattern in checked_inputs:
             # arg_pattern could be a Var, if it's the original arg.
             if arg_pattern is None:
                 if arg_value is None:
@@ -216,6 +219,11 @@ def _match_value(
                 if pattern_value.tag_var is not None:
                     self._match.bind(pattern_value.tag_var, i)
             return result
+        # Default case: a plain pattern variable (ValuePattern)
+        if value is None and not pattern_value.can_match_none:
+            return self.fail(
+                f"Mismatch: pattern variable {pattern_value} does not match None."
+            )
         return True
 
     def _match_node_output(

onnxscript/rewriter/_pattern_ir.py

@@ -123,12 +123,16 @@ def _to_attr_pattern(value: AttrPattern | ValuePattern | SupportedAttrTypes) ->
     """Represents promotion of values allowed as keyword-arguments in a pattern-builder call to an AttrPattern."""
     if isinstance(value, AttrPattern):
         return value
-    if type(value) is ValuePattern:
-        # This is a hack. Currently, when we create pattern-variables, we create them as ValuePattern,
+    if isinstance(value, Var):
+        # This is a hack. Currently, when we create pattern-variables, we create them as Var,
         # and change them to AttrPattern if/when used in an attribute context. We could use type
         # annotations to distinguish between ValuePattern and AttrPattern, but forces users to
         # use these type annotations.
         # TODO: check for misuse at rule-creation time. (Currently will be caught by matcher at match-time.)
+        if value.can_match_none or value.check_method is not None:
+            raise ValueError(
+                "Pattern variables used in attributes must not have can_match_none or check_method set."
+            )
         return AttrPattern(value.name)
     if isinstance(value, (int, float, str)):
         return AttrConstantPattern(value)
@@ -320,9 +324,12 @@ class ValuePattern:
     operations, so that we can write patterns like `x + 1` and `1 + x`.
     """
 
-    def __init__(self, name: str | None, *, check: Callable | None = None) -> None:
+    def __init__(
+        self, name: str | None, *, check: Callable | None = None, can_match_none: bool = False
+    ) -> None:
         self._name = name
         self._check = check
+        self._can_match_none = can_match_none
         # Note: uses will be computed only when the full graph-pattern is constructed.
         self._uses: list[tuple[NodePattern, int]] = []
 
@@ -338,6 +345,11 @@ def name(self) -> str | None:
     def check_method(self) -> Callable | None:
         return self._check
 
+    @property
+    def can_match_none(self) -> bool:
+        """Indicates whether this variable can match a None input."""
+        return self._can_match_none
+
     def producer(self) -> NodePattern | None:
         return None
 
@@ -547,7 +559,17 @@ def producer(self) -> NodePattern:
         return self._producer
 
 
-Var = ValuePattern
+class Var(ValuePattern):
+    """Represents a pattern-variable."""
+
+    def __init__(
+        self, name: str | None, *, check: Callable | None = None, can_match_none: bool = False
+    ) -> None:
+        super().__init__(name, check=check, can_match_none=can_match_none)
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> Var:
+        """Clones the pattern-variable, preserving its name and check method."""
+        return Var(self.name, check=self.check_method, can_match_none=self.can_match_none)
 
 
 class AnyValue(ValuePattern):

onnxscript/rewriter/ort_fusions/attention.py

@@ -34,7 +34,6 @@ def pattern(
         qkv_bias,
         # mask_index,
         past,
-        attention_bias,
         num_heads,
         # scale,
         start1,
@@ -106,7 +105,7 @@ def pattern(
                 value_BSD,
                 qkv_bias,
                 None,  # key_padding_mask
-                attention_bias,
+                pattern.Var("attention_bias", can_match_none=True),
                 past_key,
                 past_value,
                 num_heads=num_heads,
@@ -127,7 +126,7 @@ def pattern(
                 value_BSD,
                 qkv_bias,
                 None,  # key_padding_mask
-                attention_bias,
+                pattern.Var("attention_bias", can_match_none=True),
                 None,  # past_key
                 None,  # past_value
                 num_heads=num_heads,

onnxscript/rewriter/ort_fusions/fuse_mha_bias.py

@@ -52,9 +52,9 @@ def pattern(
             value_BSD,
             None,  # bias
             None,  # key padding mask
-            mask,  # attention mask/bias
-            past_key,
-            past_value,
+            pattern.Var("mask", can_match_none=True),  # attention mask/bias
+            pattern.Var("past_key", can_match_none=True),
+            pattern.Var("past_value", can_match_none=True),
             num_heads=num_heads,
             # scale=scale,
             _domain="com.microsoft",

onnxscript/rewriter/pattern.py

@@ -10,6 +10,7 @@
     Constant,
     OpsetPatternBuilder,
     OrValue,
+    Var,
     pattern_builder,
     torch_module_op,
 )
@@ -41,4 +42,5 @@
     "PatternMatcher",
     "SimplePatternMatcher",
     "torch_module_op",
+    "Var",
 ]

onnxscript/rewriter/pattern_test.py

@@ -450,8 +450,9 @@ def test_model(x: FLOAT[1024]) -> FLOAT[1024]:
         self.assertEqual(model.graph.node(1).op_type, "Original")
 
     def test_match_optional_input(self):
-        def none_pattern(op, optional_input, x):
+        def none_pattern(op, x):
             # match against a call to Original where the first input may or may not be None
+            optional_input = pattern.Var("optional_input", can_match_none=True)
             return op.Original(optional_input, x)
 
         def replacement(op, optional_input, x):
@@ -478,6 +479,44 @@ def test_model(x: FLOAT[1024]) -> FLOAT[1024]:
         self.assertEqual(model.graph.node(0).op_type, "ReplacedNone")
         self.assertEqual(model.graph.node(1).op_type, "ReplacedNotNone")
 
+    def test_mismatched_number_of_inputs(self):
+        def var_length_pattern(op):
+            # match against a call to Original where the first input may or may not be None
+            input1 = pattern.Var("input1", can_match_none=False)
+            input2 = pattern.Var("input2", can_match_none=True)
+            return op.Original(input1, input2)
+
+        def replacement(op, input1, input2):
+            return op.Replaced(input1, input2)
+
+        rule = pattern.RewriteRule(var_length_pattern, replacement)
+
+        @script()
+        def test_model(x: FLOAT[1024], y: FLOAT[1024], z: FLOAT[1024]) -> FLOAT[1024]:
+            # Pattern should NOT match following 2 calls, since pattern requires first input to be non-None
+            t0 = op.Original()
+            t1 = op.Original(None, x)
+
+            # Pattern should match following 3 calls, since second input can be None
+            t2 = op.Original(x)
+            t3 = op.Original(x, None)
+            t4 = op.Original(x, y)
+
+            # Pattern should NOT match following call, since it has more than 2 inputs
+            t5 = op.Original(x, y, z)
+            return op.All(t0, t1, t2, t3, t4, t5)
+
+        model_proto = test_model.to_model_proto()
+        model = ir.serde.deserialize_model(model_proto)
+
+        count = rule.apply_to_model(model)
+        self.assertEqual(count, 3)
+        self.assertEqual(len(model.graph), 7)
+        self.assertEqual(
+            [n.op_type for n in model.graph],
+            ["Original", "Original", "Replaced", "Replaced", "Replaced", "Original", "All"],
+        )
+
     def test_graph_visitor(self):
         class ReplaceFoo(pattern.RewriteRuleClassBase):
             def __init__(self):