[Rewriter] Implement value/node level checkers for pattern matching infrastructure (#2459)

This PR extends the pattern matching infrastructure to support
value/node level checkers as requested in the issue. The implementation
allows for more sophisticated pattern matching by enabling custom
validation logic at both the node and value levels.

## Key Changes

### 1. Extended Pattern IR Classes
- **ValuePattern**: Added optional `_check` callable attribute via
`check` keyword argument
- **NodePattern**: Added optional `_check` callable attribute via
`check` keyword argument
- Both checkers take `(MatchContext, ir.Node/ir.Value)` and return
`bool` or `MatchResult`

### 2. Enhanced Pattern Building
- **_to_value_pattern**: Now accepts callable inputs, automatically
creating `ValuePattern` with checker
- **OpPatternBuilder.__call__**: Added `_check` parameter for node-level
validation

### 3. Extended MatchResult
- Added `node_bindings` property (similar to existing `value_bindings`)
- Provides access to pattern node → matched node mappings

### 4. Enhanced Pattern Matching
- **Pattern.match**: Now executes value/node level checks before
condition function
- Iterates through `node_bindings` and `value_bindings` to run
associated checkers
- Stops on first check failure with appropriate error handling

## Usage Examples

### Node-Level Checker
```python
def validated_add_checker(context, node):
    """Only accept Add nodes with no custom attributes."""
    return node.op_type == "Add" and len(node.attributes) == 0

def pattern(op, x, y):
    return op.Add(x, y, _check=validated_add_checker)
```

### Value-Level Checker
```python
def shape_checker(context, value):
    """Validate value has expected shape properties."""
    return hasattr(value, 'type') and hasattr(value.type, 'shape')

def pattern(op, x, y):
    validated_x = shape_checker  # Converted to ValuePattern with checker
    return op.Add(validated_x, y)
```

### Combined Checkers
```python
def pattern(op, x, y):
    validated_x = value_checker  # Value-level check
    return op.Add(validated_x, y, _check=node_checker)  # Node-level check
```

## Testing

Added comprehensive test suite (`ValueNodeCheckersTest`) covering:
- ✅ ValuePattern and NodePattern with checkers
- ✅ _to_value_pattern with callable inputs  
- ✅ OpPatternBuilder with _check parameter
- ✅ Pattern.match with successful node/value checkers
- ✅ Pattern.match with failing checkers (proper error handling)
- ✅ Backward compatibility with existing patterns

All existing tests continue to pass, ensuring no breaking changes.

Fixes #2458.

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

💬 Share your feedback on Copilot coding agent for the chance to win a
$200 gift card! Click
[here](https://survey.alchemer.com/s3/8343779/Copilot-Coding-agent) to
start the survey.

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: gramalingam <[email protected]>

Author: Copilot

docs/tutorial/rewriter/node_value_checkers.md

@@ -0,0 +1,187 @@
+(heading-target-checkers)=
+# Node and Value Level Checkers
+
+The pattern matching infrastructure supports custom validation logic at both the node and value levels through checker functions. These checkers allow for more sophisticated pattern matching by enabling additional constraints beyond basic operator and structure matching.
+
+## Value-Level Checkers
+
+Value-level checkers validate properties of specific values in the pattern. They are particularly useful for checking constants, shapes, or other value-specific properties.
+
+### Basic Usage
+
+A value checker is a function that takes a `MatchContext` and an `ir.Value`, and returns either a boolean or a `MatchResult`:
+
+```python
+def is_positive_constant(context, value: ir.Value):
+    """Check if a value is a positive constant."""
+    if value.const_value is not None:
+        # Get the numpy array from const_value
+        numpy_array = value.const_value.numpy()
+
+        # Check if it represents a single value and is positive
+        if numpy_array.size != 1:
+            return False
+
+        return float(numpy_array.item()) > 0
+
+    return False
+```
+
+You can use this checker directly in your pattern by passing the callable as an input:
+
+```python
+def add_pattern(op, x, y):
+    # Use callable as input to create ValuePattern with checker
+    return op.Add(is_positive_constant, y)
+```
+
+This pattern will only match `Add` operations where the first input is a positive constant value.
+
+### Example Usage
+
+```python
+from onnxscript.rewriter import pattern
+from onnxscript import ir, optimizer
+import onnx
+
+# Create a model with different Add operations
+model_proto = onnx.parser.parse_model("""
+    <ir_version: 7, opset_import: [ "" : 17]>
+    agraph (float[N] x, float[N] y) => (float[N] z1, float[N] z2, float[N] z3)
+    {
+        pos_const = Constant <value_float = 2.5> ()
+        neg_const = Constant <value_float = -1.5> ()
+        z1 = Add(x, y)           # non-constant first parameter
+        z2 = Add(pos_const, y)   # positive constant first parameter
+        z3 = Add(neg_const, y)   # negative constant first parameter
+    }
+""")
+model = ir.serde.deserialize_model(model_proto)
+
+# Apply constant propagation to set const_value fields
+optimizer.basic_constant_propagation(model.graph.all_nodes())
+
+# Create the pattern with value checker
+rule_pattern = pattern.Pattern(add_pattern)
+
+# Test matching against different Add nodes
+add_nodes = [node for node in model.graph if node.op_type == "Add"]
+
+# Non-constant first parameter - will not match
+match_result = rule_pattern.match(model, model.graph, add_nodes[0])
+print(f"Non-constant: {bool(match_result)}")  # False
+
+# Positive constant first parameter - will match
+match_result = rule_pattern.match(model, model.graph, add_nodes[1])
+print(f"Positive constant: {bool(match_result)}")  # True
+
+# Negative constant first parameter - will not match
+match_result = rule_pattern.match(model, model.graph, add_nodes[2])
+print(f"Negative constant: {bool(match_result)}")  # False
+```
+
+## Node-Level Checkers
+
+Node-level checkers validate properties of the operation nodes themselves, such as attributes, operation types, or other node-specific properties.
+
+### Basic Usage
+
+A node checker is a function that takes a `MatchContext` and an `ir.Node`, and returns either a boolean or a `MatchResult`:
+
+```python
+def shape_node_checker(context, node):
+    """Check if a Shape operation has start attribute equal to 0."""
+    return node.attributes.get_int("start", 0) == 0
+```
+
+You can use this checker by passing it to the `_check` parameter of an operation:
+
+```python
+def shape_pattern(op, x):
+    return op.Shape(x, _check=shape_node_checker)
+```
+
+This pattern will only match `Shape` operations where the `start` attribute is 0 (or not present, as the default is 0).
+
+### Example Usage
+
+```python
+from onnxscript.rewriter import pattern
+from onnxscript import ir
+import onnx
+
+# Create a model with different Shape operations
+model_proto = onnx.parser.parse_model("""
+    <ir_version: 7, opset_import: [ "" : 17]>
+    agraph (float[N, M] x) => (int64[2] z1, int64[2] z2, int64[1] z3)
+    {
+        z1 = Shape(x)
+        z2 = Shape <start: int = 0>(x)
+        z3 = Shape <start: int = 1>(x)
+    }
+""")
+model = ir.serde.deserialize_model(model_proto)
+
+# Create the pattern with node checker
+rule_pattern = pattern.Pattern(shape_pattern)
+
+# Test matching against different Shape nodes
+nodes = list(model.graph)
+shape_nodes = [node for node in nodes if node.op_type == "Shape"]
+
+# Shape without start attribute (default 0) - will match
+match_result = rule_pattern.match(model, model.graph, shape_nodes[0])
+print(f"No start attr: {bool(match_result)}")  # True
+
+# Shape with start=0 - will match
+match_result = rule_pattern.match(model, model.graph, shape_nodes[1])
+print(f"Start=0: {bool(match_result)}")  # True
+
+# Shape with start=1 - will not match
+match_result = rule_pattern.match(model, model.graph, shape_nodes[2])
+print(f"Start=1: {bool(match_result)}")  # False
+```
+
+## Combining Checkers
+
+You can combine both node-level and value-level checkers in the same pattern for more sophisticated matching:
+
+```python
+def complex_pattern(op, x, y):
+    # Value-level checker for first input
+    validated_x = is_positive_constant
+    # Node-level checker for the operation
+    return op.Add(validated_x, y, _check=lambda ctx, node: len(node.attributes) == 0)
+```
+
+This pattern will only match `Add` operations where:
+1. The first input is a positive constant (value-level check)
+2. The node has no custom attributes (node-level check)
+
+## Execution Timing and Limitations
+
+### When Checkers Are Called
+
+Node-level and value-level checkers are called **only at the end of the complete structural match**. This means:
+
+1. **Structural matching happens first**: The pattern matching engine first validates that the graph structure matches the pattern (correct operators, connections, etc.)
+2. **Checkers run after structural validation**: Only after the structural match succeeds do the node and value checkers execute
+3. **Order of execution**: Value-level checkers run first, followed by node-level checkers, and finally the pattern's condition function
+
+### Limitations with Pattern Disjunctions
+
+One important limitation of this design is that these checks don't compose well with pattern disjunctions (multiple alternative patterns). When searching among multiple value patterns:
+
+- **Only structural checking is performed initially**: If structural matching succeeds for the first alternative, other alternatives are not considered
+- **Checker failures don't trigger backtracking**: If a checker fails, the entire pattern match fails rather than trying the next alternative pattern
+
+This means you should be careful when designing patterns with multiple alternatives that rely on checkers, as the checker logic may prevent exploration of valid alternative matches.
+
+## 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
+
+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.

docs/tutorial/rewriter/rewrite_patterns.md

@@ -24,3 +24,6 @@ There are three main components needed when rewriting patterns in the graph:
 
 ```{include} commute.md
 ```
+
+```{include} node_value_checkers.md
+```

onnxscript/rewriter/_basics.py

@@ -188,6 +188,13 @@ def value_bindings(self) -> dict[_pattern_ir.ValuePattern, ir.Value]:
             raise ValueError("Value bindings can be accessed only at the top-level match.")
         return self._current_match.value_bindings
 
+    @property
+    def node_bindings(self) -> dict[_pattern_ir.NodePattern, ir.Node]:
+        """Returns the bindings for the node variables."""
+        if len(self._partial_matches) > 1:
+            raise ValueError("Node bindings can be accessed only at the top-level match.")
+        return self._current_match.node_bindings
+
     @property
     def outputs(self) -> MutableSequence[ir.Value]:
         """Returns the list of output values that matched the pattern."""

onnxscript/rewriter/_pattern_ir.py

@@ -224,6 +224,7 @@ def __call__(
         _outputs: int | list[str | None] = 1,
         _allow_other_attributes: bool | None = None,
         _allow_other_inputs: bool | None = None,
+        _check: Callable | None = None,
         **kwargs,
     ):
         if _version is not None:
@@ -255,6 +256,7 @@ def __call__(
             _outputs,
             allow_other_attributes=_allow_other_attributes,
             allow_other_inputs=_allow_other_inputs,
+            check=_check,
         )
         self.pattern_builder.add_node(node_pattern)
         output_values = node_pattern.outputs
@@ -266,7 +268,7 @@ def __call__(
 
 
 def _to_value_pattern(
-    x: ValuePattern | int | float | None,
+    x: ValuePattern | int | float | Callable | None,
 ) -> ValuePattern | None:
     """Promotes an input-value used to construct a NodePattern to a ValuePattern.
 
@@ -282,6 +284,8 @@ def _to_value_pattern(
     explicitly write this as:
     ::
         z = op.Add(x, op.Constant(0))
+
+    If a callable is provided, it will be converted to a ValuePattern with the callable as the check attribute.
     """
     if x is None or isinstance(x, ValuePattern):
         return x
@@ -291,6 +295,8 @@ def _to_value_pattern(
         if all(isinstance(i, (int, float)) for i in x):
             return Constant(x)
         raise ValueError("Only lists of int/float can be used as a ValuePattern")
+    if callable(x):
+        return ValuePattern(None, check=x)
 
     raise TypeError(f"Cannot convert {type(x)} to ValuePattern")
 
@@ -314,19 +320,24 @@ class ValuePattern:
     operations, so that we can write patterns like `x + 1` and `1 + x`.
     """
 
-    def __init__(self, name: str | None) -> None:
+    def __init__(self, name: str | None, *, check: Callable | None = None) -> None:
         self._name = name
+        self._check = check
         # Note: uses will be computed only when the full graph-pattern is constructed.
         self._uses: list[tuple[NodePattern, int]] = []
 
     def clone(self, node_map: dict[NodePattern, NodePattern]) -> ValuePattern:
         del node_map
-        return ValuePattern(self._name)
+        return ValuePattern(self._name, check=self._check)
 
     @property
     def name(self) -> str | None:
         return self._name
 
+    @property
+    def check_method(self) -> Callable | None:
+        return self._check
+
     def producer(self) -> NodePattern | None:
         return None
 
@@ -397,6 +408,7 @@ def __init__(
         *,
         allow_other_attributes: bool | None,
         allow_other_inputs: bool | None,
+        check: Callable | None = None,
     ):
         if allow_other_attributes is None:
             # Default behavior: allow other unmatched attributes in the node.
@@ -410,6 +422,7 @@ def __init__(
         self.attributes = attributes
         self.allow_other_attributes = allow_other_attributes
         self.allow_other_inputs = allow_other_inputs
+        self._check = check
         # In the common case, domain and op are constants, which can be used to optimize matching.
         if isinstance(op, str) and isinstance(domain, StringConstantPattern):
             # TODO(rama): support overloaded operators.
@@ -445,6 +458,10 @@ def op_identifier(self) -> ir.OperatorIdentifier | None:
     def op_type(self) -> str:
         return str(self.op)
 
+    @property
+    def check_method(self) -> Callable | None:
+        return self._check
+
     def matches(self, node: ir.Node, match: _basics.MatchResult) -> _basics.MatchResult:
         """Matches the pattern represented by self against a node.
 
@@ -498,6 +515,7 @@ def clone(self, node_map: dict[NodePattern, NodePattern], swap: bool) -> NodePat
             outputs,
             allow_other_attributes=self.allow_other_attributes,
             allow_other_inputs=self.allow_other_inputs,
+            check=self._check,
         )
         node_map[self] = copied
         return copied

onnxscript/rewriter/_pattern_ir_test.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import unittest
+
+from onnxscript.rewriter import _pattern_ir
+
+
+class PatternIRTest(unittest.TestCase):
+    """Test _pattern_ir module functionality."""
+
+    def test_value_pattern_with_check(self):
+        """Test ValuePattern with check attribute."""
+
+        def value_checker(context, value):
+            return True
+
+        # Test creating ValuePattern with check
+        value_pattern = _pattern_ir.ValuePattern("test_value", check=value_checker)
+        self.assertIs(value_pattern._check, value_checker)
+        self.assertEqual(value_pattern.name, "test_value")
+
+    def test_node_pattern_with_check(self):
+        """Test NodePattern with check attribute."""
+
+        def node_checker(context, node):
+            return True
+
+        # Test creating NodePattern with check
+        domain_pattern = _pattern_ir.StringConstantPattern("")
+        inputs = []
+        attributes = {}
+        outputs = ["output"]
+
+        node_pattern = _pattern_ir.NodePattern(
+            domain_pattern,
+            "Add",
+            inputs,
+            attributes,
+            outputs,
+            allow_other_attributes=True,
+            allow_other_inputs=True,
+            check=node_checker,
+        )
+        self.assertIs(node_pattern._check, node_checker)
+
+    def test_to_value_pattern_with_callable(self):
+        """Test _to_value_pattern function with callable input."""
+
+        def my_checker(context, value):
+            return True
+
+        result = _pattern_ir._to_value_pattern(my_checker)
+        self.assertIsInstance(result, _pattern_ir.ValuePattern)
+        self.assertIs(result._check, my_checker)
+        self.assertIsNone(result.name)
+
+    def test_op_pattern_builder_with_check(self):
+        """Test OpPatternBuilder with _check parameter."""
+
+        def node_checker(context, node):
+            return True
+
+        # Create OpPatternBuilder and call with _check parameter
+        opset_builder = _pattern_ir.OpsetPatternBuilder("")
+        result = opset_builder.Add(None, None, _check=node_checker)
+
+        # The result should be a NodeOutputPattern, and its producer should have the check
+        self.assertTrue(hasattr(result, "producer"))
+        producer = result.producer()
+        self.assertIsNotNone(producer)
+        self.assertTrue(hasattr(producer, "_check"))
+        self.assertIs(producer._check, node_checker)
+
+
+if __name__ == "__main__":
+    unittest.main()