Unify failure-handling in rewrite-rule (#2866)

Unify failure-handling in rewrite-rule

---------

Signed-off-by: G Ramalingam <grama@microsoft.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: gramalingam

docs/tutorial/rewriter/node_value_checkers.md

@@ -179,9 +179,16 @@ This means you should be careful when designing patterns with multiple alternati
 
 ## Error Handling
 
-Checkers can return either:
-- `True`: Check passed, continue matching
-- `False`: Check failed, pattern does not match
-- `MatchResult`: More detailed result with potential failure reasons
+Both check functions (including condition functions and node/value-level checkers) and
+rewrite functions support the same conventions for indicating failure:
 
-If a checker raises an exception, it will be caught and treated as a match failure, allowing patterns to fail gracefully when encountering unexpected conditions.
+- **`MatchResult` with `.fail()`** *(recommended)*: Return `MatchResult().fail("reason", source)` to indicate failure with a descriptive reason and optional source node/value. This provides the most useful diagnostic information for debugging.
+- **Raise `MatchFailureError`** *(recommended)*: Import it as `from onnxscript.rewriter.rewriter import MatchFailureError` and raise `MatchFailureError("reason", source1, source2, ...)` to indicate failure associated with one or more `ir.Node` or `ir.Value` objects. Each source should be passed as a separate positional argument (do not pass a list as a single argument). This is especially convenient in utility functions called from a check or rewrite, since it avoids having to explicitly propagate failure status through the call chain.
+- **Return `None` or `False`**: These indicate failure without providing a reason. They are supported but not recommended, since a failure reason is valuable for debugging why a rule did not apply.
+
+Including a descriptive failure reason is strongly encouraged. The rewriter's tracing infrastructure
+uses these reasons to report why rules failed to match, which is essential for diagnosing
+issues when developing or debugging rewrite rules.
+
+For **check functions**, success is indicated by returning `True` or a truthy `MatchResult`.
+For **rewrite functions**, success is indicated by returning one or more `ir.Value` results.

onnxscript/rewriter/_rewrite_rule.py

@@ -217,9 +217,34 @@ def __init__(self, function) -> None:
 
     def get_replacement(self, match: _basics.MatchResult) -> ReplacementSubgraph | None:
         context = RewriterContext()
-        new_outputs = self._function(context, **match.bindings)
-        if new_outputs is None:
-            return None  # Failed to create replacement subgraph
+        try:
+            new_outputs = self._function(context, **match.bindings)
+        except _basics.MatchFailureError as e:
+            match.fail(e.reason, list(e.failure_sources))
+            return None
+        # Support the same failure conventions as check functions for uniformity:
+        # - None or False: indicates failure without a reason (not recommended)
+        # - Falsy MatchResult: failure with reason/source info (recommended)
+        # - MatchFailureError exception: failure with reason/source info (recommended)
+        if new_outputs is None or new_outputs is False:
+            return None
+        if isinstance(new_outputs, _basics.MatchResult):
+            if not new_outputs:
+                # A falsy MatchResult is the recommended way to signal failure with
+                # reason/source information from a replacement function.
+                match.fail(
+                    new_outputs.reason,
+                    new_outputs.failure_nodes_and_values,
+                )
+                return None
+            # A truthy MatchResult should never be returned from a replacement
+            # function. Treat this as a programmer error to avoid silent failures.
+            raise TypeError(
+                "Replacement function returned a truthy MatchResult. "
+                "Replacement functions should either return None/False for a "
+                "generic failure, return a *falsy* MatchResult to provide "
+                "failure details, or raise MatchFailureError."
+            )
         if not isinstance(new_outputs, Sequence):
             new_outputs = [new_outputs]
         return ReplacementSubgraph(

onnxscript/rewriter/pattern_base_test.py

@@ -6,6 +6,7 @@
 
 from onnxscript import ir
 from onnxscript.rewriter import pattern
+from onnxscript.rewriter._basics import MatchFailureError
 
 
 class PatternTest(unittest.TestCase):
@@ -249,5 +250,97 @@ def rewrite(self, op, x):
         self.assertEqual(rule.name, "SimpleIdentityRule")
 
 
+class RewriteFailureConventionsTest(unittest.TestCase):
+    """Test that rewrite functions support the same failure conventions as check functions.
+
+    The check side supports three failure conventions:
+    - Return a falsy MatchResult (via .fail())
+    - Raise MatchFailureError
+    - Return None or False
+
+    The rewrite side should support all three as well, for uniformity.
+    """
+
+    _IDENTITY_MODEL_TEXT = """
+        <ir_version: 7, opset_import: [ "" : 17]>
+        agraph (float[N] x) => (float[N] z)
+        {
+            z = Identity(x)
+        }
+    """
+
+    def _apply_rewrite_rule(self, rewrite_fn):
+        """Helper that builds a RewriteRule with the given rewrite function and applies it."""
+
+        def identity_pattern(op, x):
+            return op.Identity(x)
+
+        rule = pattern.RewriteRule(identity_pattern, rewrite_fn, name="TestRule")
+        model = ir.from_onnx_text(self._IDENTITY_MODEL_TEXT)
+        count = rule.apply_to_model(model)
+        return count
+
+    def test_rewrite_returning_none_is_treated_as_failure(self):
+        """Returning None from rewrite indicates failure (no replacement made)."""
+
+        def rewrite_returns_none(op, x):
+            return None
+
+        count = self._apply_rewrite_rule(rewrite_returns_none)
+        self.assertEqual(count, 0)
+
+    def test_rewrite_returning_false_is_treated_as_failure(self):
+        """Returning False from rewrite indicates failure (deprecated but supported)."""
+
+        def rewrite_returns_false(op, x):
+            return False
+
+        count = self._apply_rewrite_rule(rewrite_returns_false)
+        self.assertEqual(count, 0)
+
+    def test_rewrite_raising_match_failure_error_is_treated_as_failure(self):
+        """Raising MatchFailureError from rewrite indicates failure."""
+
+        def rewrite_raises_error(op, x):
+            raise MatchFailureError("Cannot rewrite this node")
+
+        count = self._apply_rewrite_rule(rewrite_raises_error)
+        self.assertEqual(count, 0)
+
+    def test_rewrite_returning_falsy_match_result_is_treated_as_failure(self):
+        """Returning a falsy MatchResult from rewrite indicates failure."""
+
+        def rewrite_returns_failed_match_result(op, x):
+            result = pattern.MatchResult()
+            return result.fail("Rewrite not applicable")
+
+        count = self._apply_rewrite_rule(rewrite_returns_failed_match_result)
+        self.assertEqual(count, 0)
+
+    def test_rewrite_returning_ir_value_succeeds(self):
+        """Returning an ir.Value from rewrite is success (the normal case)."""
+
+        # Use a non-self-referential pattern to avoid infinite rewrite loops
+        def add_zero_pattern(op, x):
+            return op.Add(x, 0.0)
+
+        def identity_replacement(op, x):
+            return op.Identity(x)
+
+        rule = pattern.RewriteRule(add_zero_pattern, identity_replacement, name="TestRule")
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N] x) => (float[N] z)
+            {
+                c0 = Constant<value_float = 0.0>()
+                z = Add(x, c0)
+            }
+            """
+        )
+        count = rule.apply_to_model(model)
+        self.assertGreaterEqual(count, 1)
+
+
 if __name__ == "__main__":
     unittest.main()