Implement MatchContext class for rewriter pattern matching (#2455)

This PR introduces the `PatternMatchContext` class to provide context
information during pattern matching in the ONNX rewriter system.

## Changes Made

### Core Implementation
- **Added `PatternMatchContext` class** in
`onnxscript/rewriter/_basics.py` with read-only properties:
  - `model`: The model being matched
  - `graph_or_function`: The graph or function being matched  
  - `main_root_node`: The main root node of the matching subgraph
  - `output_values`: The output values of the matching subgraph
  - `nodes`: All nodes of the matching subgraph

- **Updated pattern matching logic** in
`onnxscript/rewriter/_rewrite_rule.py` at line 134 to create and pass
`PatternMatchContext` instances to condition functions

- **Exported the new class** in the rewriter module's `__all__` list for
external use

### Usage Example
```python
def condition_with_context(context, x, y):
    # Access match context information
    model = context.model
    main_node = context.main_root_node
    matched_nodes = context.nodes
    outputs = context.output_values
    
    # Use context for advanced pattern validation
    if main_node.op_type == "Mul" and len(matched_nodes) > 1:
        return True
    return False

rule = pattern.RewriteRule(
    target_pattern, 
    replacement_pattern,
    condition_function=condition_with_context
)
```

### Testing
- **Comprehensive test suite** in
`onnxscript/rewriter/pattern_match_context_test.py` covering:
  - Property access and type validation
  - Read-only behavior enforcement
  - Backward compatibility with existing condition functions
  - Practical usage scenarios in real pattern matching

### Backward Compatibility
- All existing condition functions continue to work unchanged
- The `context` parameter is passed as the first argument, maintaining
the existing `**match.bindings` pattern
- No breaking changes to the existing API

## Validation
- All existing rewriter tests pass (39/39 tests in pattern-related
modules)
- New functionality validated with 4 comprehensive test cases
- Integration testing confirms proper context creation and usage

Fixes #2454.

<!-- START COPILOT CODING AGENT TIPS -->
---

💡 You can make Copilot smarter by setting up custom instructions,
customizing its development environment and configuring Model Context
Protocol (MCP) servers. Learn more [Copilot coding agent
tips](https://gh.io/copilot-coding-agent-tips) in the docs.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: gramalingam <10075881+gramalingam@users.noreply.github.com>

Author: Copilot

.lintrunner.toml

@@ -114,9 +114,8 @@ include_patterns = [
     '**/*.py',
 ]
 exclude_patterns = [
-    'examples/**',  # TODO: Merge with docs/examples
-    'docs/examples/**',
-    'docs/tutorial/examples/**',
+    'examples/**',
+    'docs/**',
     'onnxscript/converter_test.py',
     'tests/functions/**',
     'tests/models/**',

docs/tutorial/rewriter/conditional_rewrite.md

@@ -50,3 +50,54 @@ With all the necessary components in place, the pattern rewrite rule with the `m
 The final graph with the applied rewrite looks as follows:
 
 ![broadcast_rewrite](examples/img/broadcast_02.png){align=center}
+
+# Using MatchContext for Advanced Condition Checking
+
+The `context` parameter passed to condition functions is an instance of {py:class}`onnxscript.rewriter.MatchContext`, which provides access to additional information about the pattern match that can be useful for sophisticated condition checking.
+
+## MatchContext Properties
+
+The MatchContext provides the following read-only properties:
+
+- `model`: The entire ONNX model being matched
+- `graph_or_function`: The specific graph or function being matched
+- `root`: The root node of the matching subgraph
+- `output_values`: The output values of the matching subgraph
+- `nodes`: All nodes that are part of the matching subgraph
+
+## Example Usage
+
+Here's an example showing how to use the MatchContext to implement more sophisticated condition checking:
+
+```python
+def advanced_condition_check(context, x, y, **_):
+    """Example condition function using MatchContext."""
+
+    # Access the main node of the pattern match
+    main_node = context.root
+
+    # Check that the main_node does not have an attribute called "alpha"
+    if "alpha" in main_node.attributes:
+        return False
+
+    # Access the broader graph context and check that x occurs as a graph-input
+    model = context.model
+    if x not in model.graph.inputs:
+        return False
+
+    # You can inspect the matched nodes for advanced validation
+    for node in context.nodes:
+        if node.op_type == "Constant":
+            # Check properties of constant nodes in the match
+            pass
+
+    # Access output values for shape/type validation
+    outputs = context.output_values
+    if len(outputs) > 0 and outputs[0].shape is not None:
+        # Validate output shapes
+        pass
+
+    return True
+```
+
+This context information enables condition functions to make decisions based on the broader graph structure, the specific nodes involved in the match, and relationships between matched patterns and the rest of the model.

docs/tutorial/rewriter/examples/broadcast_matmul.py

@@ -79,8 +79,6 @@ def check_if_not_need_reshape(
     Returns:
         True if we need to replace the pattern, False otherwise.
     """
-    del context  # Reserved for future extensions
-
     input_a_shape = input_a.shape
     input_b_shape = input_b.shape
     shape_c_tensor = shape_c.const_value

onnxscript/rewriter/__init__.py

@@ -9,6 +9,7 @@
     "rewrite",
     "RewritePass",
     "MatchResult",
+    "MatchContext",
     "RewriteRule",
     "RewriteRuleClassBase",
     "RewriteRuleSet",
@@ -31,7 +32,7 @@
     pattern,
     redundant_scatter_nd,
 )
-from onnxscript.rewriter._basics import MatchingTracer, MatchResult, MatchStatus
+from onnxscript.rewriter._basics import MatchContext, MatchingTracer, MatchResult, MatchStatus
 from onnxscript.rewriter._rewrite_rule import (
     RewriterContext,
     RewriteRule,

onnxscript/rewriter/_basics.py

@@ -340,6 +340,82 @@ def print(self):
         print(separator)
 
 
+class MatchContext:
+    """A read-only context containing information about a pattern match.
+
+    This class captures information about the context describing a match to a given pattern,
+    providing access to the model, graph/function, root node, output values, and all
+    nodes of the matching subgraph.
+    """
+
+    def __init__(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        root: ir.Node,
+        match_result: MatchResult,
+    ) -> None:
+        """Initialize the pattern match context.
+
+        Args:
+            model: The model being matched.
+            graph_or_function: The graph or function being matched.
+            root: The root node of the matching subgraph.
+            match_result: The match result containing matched nodes and outputs.
+        """
+        self._model = model
+        self._graph_or_function = graph_or_function
+        self._root = root
+        self._match_result = match_result
+
+    @property
+    def model(self) -> ir.Model:
+        """The model being matched."""
+        return self._model
+
+    @property
+    def graph_or_function(self) -> ir.Graph | ir.Function:
+        """The graph or function being matched."""
+        return self._graph_or_function
+
+    @property
+    def root(self) -> ir.Node:
+        """The root node of the matching subgraph."""
+        return self._root
+
+    @property
+    def output_values(self) -> Sequence[ir.Value]:
+        """The output values of the matching subgraph."""
+        return self._match_result.outputs
+
+    @property
+    def nodes(self) -> Sequence[ir.Node]:
+        """All the nodes of the matching subgraph."""
+        return self._match_result.nodes
+
+    def display(self, *, in_graph_order: bool = True) -> None:
+        """Display the nodes in the pattern match context.
+
+        Args:
+            in_graph_order: If True, display nodes in the order they appear in the
+                graph/function. If False, display nodes in the order they appear
+                in the match result.
+        """
+        nodes = self.nodes
+        if not nodes:
+            return
+
+        if in_graph_order:
+            # Display nodes in same order as in graph/function
+            for node in self._graph_or_function:
+                if node in nodes:
+                    node.display()
+        else:
+            # Display nodes in match order
+            for node in nodes:
+                node.display()
+
+
 class MatchingTracer:
     """A debugging helper class to trace the matching of a pattern against a graph.
 

onnxscript/rewriter/_rewrite_rule.py

@@ -131,7 +131,7 @@ def match(
             remove_nodes=check_nodes_are_removable,
         )
         if match:
-            context = None  # TODO(rama)
+            context = _basics.MatchContext(model, graph_or_function, node, match)
             for var in self._target_pattern.inputs:
                 if var.name is not None:
                     if var.name not in match.bindings:

onnxscript/rewriter/match_context_test.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Test for MatchContext functionality."""
+
+import unittest
+
+import onnx.parser
+
+from onnxscript import ir
+from onnxscript.rewriter import pattern
+
+
+class MatchContextTest(unittest.TestCase):
+    def test_context_usage_in_condition_function(self):
+        """Test that MatchContext can be meaningfully used in condition functions."""
+
+        model_proto = onnx.parser.parse_model(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[N] x, float[N] y) => (float[N] z)
+            {
+                c1 = Constant<value_float = 1.0>()
+                t1 = Div(c1, x)
+                z = Mul(t1, y)
+            }
+        """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+
+        def condition_using_context(context, x, y):
+            # Use context to check properties of the match
+            self.assertIs(context.model, model)
+            self.assertIs(context.graph_or_function, model.graph)
+            self.assertIs(context.root, model.graph[2])
+
+            # Verify that we can inspect the matched nodes
+            self.assertEqual(len(context.nodes), 2)
+
+            return True  # Allow the rewrite
+
+        def reciprocal_mul_pattern(op, x, y):
+            return (1 / x) * y
+
+        def replacement(op, x, y):
+            return op.Div(y, x)
+
+        rule = pattern.RewriteRule(
+            reciprocal_mul_pattern, replacement, condition_function=condition_using_context
+        )
+
+        count = rule.apply_to_model(model)
+        self.assertEqual(count, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()