Expose match functionality of rewrite-rule by extracting base classes

[Copilot]

This PR extracts the pattern matching functionality from rewrite rules into standalone base classes, allowing users to use pattern matching without needing the replacement functionality.

Changes

New Base Classes

PatternImpl: Core pattern matching functionality

• Encapsulates _target_pattern, _matcher, and _condition_function
• Provides match() method that returns MatchResult or None
• Can be used standalone for pattern matching without rewriting

PatternBase: Base class for class-based pattern definition

• Provides abstract pattern() method for defining patterns
• Provides optional check() method for condition functions
• Includes create_pattern_impl() method to generate PatternImpl instances

Updated Classes

RewriteRule: Now inherits from PatternImpl

• Maintains all existing functionality
• Gains access to standalone pattern matching capabilities
• Uses inherited match() method in try_rewrite()

RewriteRuleClassBase: Now inherits from PatternBase

• Maintains all existing functionality
• Gains access to pattern-only capabilities
• Still provides rule() class method to create RewriteRule instances

Usage Examples

Standalone Pattern Matching

from onnxscript.rewriter import pattern

# Define a pattern
def identity_pattern(op, x):
    return op.Identity(x)

# Create a pattern matcher (no replacement needed)
pattern_matcher = pattern.PatternImpl(identity_pattern, name="IdentityMatcher")

# Use it to check if a node matches the pattern
match_result = pattern_matcher.match(model, graph, node)
if match_result:
    print(f"Pattern matched! Found {len(match_result.nodes)} nodes")

Class-Based Pattern Definition

class MyPattern(pattern.PatternBase):
    def pattern(self, op, x):
        return op.Identity(x)
    
    def check(self, context, x):
        # Custom condition logic
        return pattern.MatchResult()

# Create a pattern implementation
my_pattern = MyPattern()
pattern_impl = my_pattern.create_pattern_impl()

Existing Functionality Preserved

# RewriteRule still works exactly as before
rule = pattern.RewriteRule(target_pattern, replacement_pattern)

# But now it can also be used for just pattern matching
match_result = rule.match(model, graph, node)  # New capability
count = rule.apply_to_model(model)  # Existing functionality

Backward Compatibility

All existing functionality is preserved. The changes are purely additive - existing code using RewriteRule and RewriteRuleClassBase will continue to work without modification.

Testing

• All existing tests pass (34/34 tests successful)
• Added comprehensive test suite for new base classes
• Created example demonstrating standalone pattern matching usage
• Verified inheritance relationships work correctly

Fixes #2446.

---

💡 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 in the docs.

[codecov]

❌ 7 Tests Failed:

Tests completed	Failed	Passed	Skipped
14153	7	14146	1879

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0837_test_reduce_log_sum_asc_axes

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_reduce_log_sum_asc_axes'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_reduce_log_sum_asc_axes' (e=No module named 'tests.onnx_backend_test_code.test_reduce_log_sum_asc_axes') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_log_sum_asc_axes.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_log_sum_asc_axes.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_log_sum_asc_axes(data: FLOAT[3,4,5], axes: INT64[2]) -> (FLOAT[5]):
E       reduced = opset18.ReduceLogSum(data, axes, keepdims=0)
E       return reduced

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0941_test_reshape_negative_extended_dims

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_reshape_negative_extended_dims'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_reshape_negative_extended_dims' (e=No module named 'tests.onnx_backend_test_code.test_reshape_negative_extended_dims') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reshape_negative_extended_dims.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reshape_negative_extended_dims.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset21
E   
E   @script()
E   def bck_test_reshape_negative_extended_dims(data: FLOAT[2,3,4], shape: INT64[4]) -> (FLOAT[1,2,3,4]):
E       reshaped = opset21.Reshape(data, shape)
E       return reshaped

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0641_test_max_uint64

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_max_uint64'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_max_uint64' (e=No module named 'tests.onnx_backend_test_code.test_max_uint64') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_max_uint64.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_max_uint64.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import UINT64
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_max_uint64(data_0: UINT64[3], data_1: UINT64[3]) -> (UINT64[3]):
E       result = opset13.Max(data_0, data_1)
E       return result

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

Do we want to expose the PatternImpl class? Or would there be a better name? Exposing a class with Impl seems a little strange.

Agreed! I've renamed CompiledPattern to Pattern which provides a much cleaner public API. "Pattern" is more intuitive and doesn't expose implementation details. Fixed in commit fadd8c6.

[justinchuby]

LGTM - I think a follow up is needed: The way "Pattern" and "PatternBase" are related can be confusing. It isn't clear if they are related at all even though they have similar names.

[gramalingam]

LGTM - I think a follow up is needed: The way "Pattern" and "PatternBase" are related can be confusing. It isn't clear if they are related at all even though they have similar names.

Yes. This can be addressed later, but I think it stems from a similar confusion between RewriteRuleClassBase and RewriteRule (the latter being the original, and the former becoming preferred for more complex rules).