onnxscript 0.1.0__py3-none-any.whl

onnxscript/rewriter/pattern.py

@@ -0,0 +1,1714 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import abc
+import contextlib
+import dataclasses
+import enum
+import inspect
+import itertools
+import math
+from collections import defaultdict
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    MutableSequence,
+    Protocol,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import onnxscript.optimizer
+from onnxscript import ir
+from onnxscript.ir import _convenience, _tape
+
+T = TypeVar("T")
+
+
+class Pattern(Protocol[T]):  # type: ignore[misc]
+    """This is essentially a Predicate[T], that is, a Callable[[T], bool] bound to the name "matches"."""
+
+    def matches(self, item: T) -> bool: ...
+
+
+class StringPattern(abc.ABC, Pattern[str]):
+    """Abstract base class for string patterns."""
+
+    @abc.abstractmethod
+    def matches(self, item: str) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def __str__(self) -> str:
+        pass
+
+
+class StringConstantPattern(StringPattern):
+    """Matches strings with given value."""
+
+    def __init__(self, value: str):
+        self._value = value
+
+    def matches(self, item: str) -> bool:
+        return item == self._value
+
+    def __str__(self) -> str:
+        return self._value
+
+    def value(self) -> str:
+        return self._value
+
+
+class PrefixPattern(StringPattern):
+    """Matches strings with a given prefix."""
+
+    def __init__(self, value: str) -> None:
+        self._value = value
+
+    def matches(self, value: str) -> bool:
+        return value.startswith(self._value)
+
+    def __str__(self) -> str:
+        return f"{self._value}*"
+
+
+class AttrPattern(Pattern[Union[ir.Attr, ir.RefAttr]]):
+    """Base class for an attribute pattern. Matches any attribute value by default."""
+
+    def __init__(self, name: str | None):
+        self._name = name
+
+    @property
+    def name(self) -> str | None:
+        return self._name
+
+    def matches(self, attr: ir.Attr | ir.RefAttr) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return self._name if self._name is not None else "anonymous:" + str(id(self))
+
+
+# TODO: Support tensors. Align with usage elsewhere.
+SupportedAttrTypes = Union[
+    int,
+    float,
+    str,
+    Sequence[int],
+    Sequence[float],
+    Sequence[str],
+]
+
+
+class AttrConstantPattern(AttrPattern):
+    """Matches attributes with given value.
+
+    Uses standard equality for matching. For list-valued attributes, the order of elements matters.
+    If order is immaterial, we need to define a separate pattern for that.
+    """
+
+    def __init__(self, value: SupportedAttrTypes):
+        super().__init__(None)
+        self._value = value
+
+    def matches(self, attr: ir.Attr | ir.RefAttr) -> bool:
+        return isinstance(attr, ir.Attr) and attr.value == self._value
+
+    def __str__(self) -> str:
+        return str(self._value)
+
+
+def _to_attr_pattern(value: AttrPattern | ValuePattern | SupportedAttrTypes) -> AttrPattern:
+    """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,
+        # 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.)
+        return AttrPattern(value.name)
+    if isinstance(value, (int, float, str)):
+        return AttrConstantPattern(value)
+    if isinstance(value, Sequence):
+        if all(isinstance(i, (int, float)) for i in value):
+            return AttrConstantPattern(value)
+        if all(isinstance(i, str) for i in value):
+            return AttrConstantPattern(value)
+        raise ValueError("Only lists of int/float/str can be used as an AttrPattern")
+    raise TypeError(f"Cannot convert {type(value)} to AttrPattern")
+
+
+class OpsetPatternBuilder:
+    """Represents an opset pattern and a pattern builder.
+
+    (i) It is used to create a NodePattern (via OpPatternBuilder).
+    Example usage:
+    ::
+
+        z = op.Matmul(x, y)
+
+    Here, `op` is an instance of OpsetPatternBuilder and `op.Matmul` is an instance
+    of OpPatternBuilder, and  `op.Matmul(x, y)` is an instance of NodePattern.
+
+    (ii) It contains a domain pattern matched against the actual opset domain used in the
+    input model.
+    """
+
+    def __init__(self, domain: StringPattern | str, record: bool = False) -> None:
+        if isinstance(domain, str):
+            domain = StringConstantPattern(domain)
+        self._domain_pattern = domain
+        if record:
+            self._nodes: list[NodePattern] | None = []
+        else:
+            self._nodes = None
+
+    def domain_pattern(self) -> StringPattern:
+        return self._domain_pattern
+
+    def __getattr__(self, op_name: str) -> OpPatternBuilder:
+        return OpPatternBuilder(self, op_name)
+
+    def submodule(self, name: str) -> OpPatternBuilder:
+        """This method is used to match against submodule ops with prefix."""
+        return OpPatternBuilder(self, PrefixPattern(name))
+
+    def __str__(self) -> str:
+        return str(self._domain_pattern)
+
+    def add_node(self, node: NodePattern) -> None:
+        if self._nodes is not None:
+            self._nodes.append(node)
+
+    def nodes(self) -> Sequence[NodePattern]:
+        if self._nodes is None:
+            raise ValueError("Nodes were not recorded.")
+        return self._nodes
+
+
+onnxop = OpsetPatternBuilder("")
+
+torch_module_op = OpsetPatternBuilder(PrefixPattern("pkg.torch"))
+
+
+class OpPatternBuilder:
+    """A utility class to build a NodePattern.
+
+    It is used primarily to create a NodePattern.
+    Example usage:
+    ::
+
+        z = op.Matmul(x, y)
+
+    Here, `op` is an instance of OpsetPatternBuilder and `op.Matmul` is an instance
+    of OpPatternBuilder, and  `op.Matmul(x, y)` is an instance of NodePattern.
+
+    """
+
+    def __init__(
+        self,
+        pattern_builder: OpsetPatternBuilder,
+        op_name: str | Pattern[str],
+    ) -> None:
+        self.pattern_builder = pattern_builder
+        self.op_name = op_name
+
+    def __call__(
+        self,
+        *args,
+        _domain: str | None = None,
+        _version: int | None = None,
+        _outputs: int | list[str | None] = 1,
+        _allow_other_attributes: bool | None = None,
+        _allow_other_inputs: bool | None = None,
+        **kwargs,
+    ):
+        if _version is not None:
+            raise ValueError(
+                "The pattern builder does not support '_version' keyword argument. "
+                "Version restrictions should be handled by rewrite rules."
+            )
+        if _domain is None:
+            opset_pattern = self.pattern_builder.domain_pattern()
+        elif isinstance(_domain, str):
+            opset_pattern = StringConstantPattern(_domain)
+        else:
+            # TODO(rama): allow OpsetPatternBuilder as _domain.
+            raise TypeError("_domain must be a string.")
+
+        if isinstance(_outputs, int):
+            _outputs = [None for _ in range(_outputs)]
+        elif not isinstance(_outputs, Sequence) or not all(
+            isinstance(x, (str, type(None))) for x in _outputs
+        ):
+            raise ValueError("_outputs must be an int or a list[str|None].")
+        inputs = [_to_value_pattern(x) for x in args]
+        attributes = {name: _to_attr_pattern(value) for (name, value) in kwargs.items()}
+        node_pattern = NodePattern(
+            opset_pattern,
+            self.op_name,
+            inputs,
+            attributes,
+            _outputs,
+            allow_other_attributes=_allow_other_attributes,
+            allow_other_inputs=_allow_other_inputs,
+        )
+        self.pattern_builder.add_node(node_pattern)
+        output_values = node_pattern.outputs
+        # Unpack outputs if there is only one output, the common case.
+        if len(output_values) == 1:
+            return output_values[0]
+        else:
+            return output_values
+
+
+def _to_value_pattern(
+    x: ValuePattern | int | float | None,
+) -> ValuePattern | None:
+    """Promotes an input-value used to construct a NodePattern to a ValuePattern.
+
+    Example usage:
+    ::
+        x = op.MatMul(a, b)
+        z = op.Add(x, 0)
+
+    In this example, `a, `b`, and `x` are ValuePatterns used to construct a NodePattern.
+    `0` is a constant (int) value, and is automatically promoted to a ValuePattern.
+
+    Note that this is a shorthand for creating a Constant pattern. The user can more
+    explicitly write this as:
+    ::
+        z = op.Add(x, op.Constant(0))
+    """
+    if x is None or isinstance(x, ValuePattern):
+        return x
+    if isinstance(x, (int, float)):
+        return Constant(x)
+    if isinstance(x, Sequence):
+        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")
+
+    raise TypeError(f"Cannot convert {type(x)} to ValuePattern")
+
+
+class MatchResult:
+    """Represents the result of a match operation.
+
+    A match can either succeed or fail.
+    If it succeeds, it returns a list of nodes that matched the pattern
+    and a set of bindings for the variables in the pattern.
+
+    Example:
+    ::
+        def pattern(x, shape1, shape2):
+            t1 = op.Reshape(x, shape1)
+            t2 = op.Reshape(t1, shape2)
+            return t2
+    The above pattern matches a sequence of two Reshape ops.
+    The matched_nodes will contain the two Reshape ops, and the bindings will
+    contain the values that are bound to the variables `x`, `shape1`, and `shape2`.
+    """
+
+    def __init__(self) -> None:
+        self._success: bool = True
+        # For a successful match, _matched_nodes is a list of values that matched the pattern.
+        # These include the internal nodes of the pattern that were matched, but not
+        # the leaves (sub-trees) that match against the variables in the pattern.
+        # These represent the values that will be replaced by the replacement pattern.
+        self._matched_nodes: MutableSequence[ir.Node] = []
+        # For a successful match, bindings is a dictionary of mapping pattern-variable-names
+        # to values.
+        self.bindings: dict[str, Any] = {}
+        self.outputs: list[ir.Value] = []
+        # For a failed match, _reason is a string that describes the reason for the failure.
+        self._reason: str = ""
+        # Track the node that caused the failure.
+        # TODO: May be useful to extend this to be a collection of Nodes and Values.
+        self._failure_node: ir.Node | None = None
+
+    def __bool__(self):
+        return self._success
+
+    def fail(self, reason: str = "", node: ir.Node | None = None) -> MatchResult:
+        self._success = False
+        self._reason = reason
+        self._failure_node = node
+        return self
+
+    @property
+    def reason(self) -> str:
+        return self._reason
+
+    @property
+    def nodes(self) -> MutableSequence[ir.Node]:
+        return self._matched_nodes
+
+    def bind(self, var: str, value: Any) -> bool:
+        """Binds a pattern variable name to a value from the matched IR.
+
+        Returns True if the binding is successful, False otherwise (when the binding is inconsistent).
+        """
+        if var in self.bindings:
+            # TODO(rama): Use appropriate equality-check here.
+            if self.bindings[var] == value:
+                return True
+            self._success = False
+            return False
+        self.bindings[var] = value
+        return True
+
+    def extend(self, other: MatchResult | bool):
+        if not self._success:
+            return
+        if not other:
+            self._success = False
+            return
+        if isinstance(other, bool):
+            return
+        for var, val in other.bindings.items():
+            if var in self.bindings:
+                # TODO: handle attribute var bindings
+                if self.bindings[var] != val:
+                    self._success = False
+                    return
+            else:
+                self.bindings[var] = val
+        assert self._matched_nodes is not None, "_matched_nodes should not be None."
+        self._matched_nodes.extend(other._matched_nodes)  # type: ignore[attr-defined]
+
+
+_pattern_builder: OpsetPatternBuilder = onnxop
+
+
+@contextlib.contextmanager
+def pattern_builder(builder: OpsetPatternBuilder):
+    global _pattern_builder
+    prev_builder = _pattern_builder
+    _pattern_builder = builder
+    yield
+    _pattern_builder = prev_builder
+
+
+class ValuePattern:
+    """Base class for all patterns that match against IR values.
+
+    This is used primarily to provide operator overloadings for arithmetic
+    operations, so that we can write patterns like `x + 1` and `1 + x`.
+    """
+
+    def __init__(self, name: str | None) -> None:
+        self._name = name
+        # 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)
+
+    @property
+    def name(self) -> str | None:
+        return self._name
+
+    def producer(self) -> NodePattern | None:
+        return None
+
+    def uses(self) -> Sequence[tuple[NodePattern, int]]:
+        return self._uses
+
+    def append_use(self, node: NodePattern, index: int):
+        self._uses.append((node, index))
+
+    def __repr__(self) -> str:
+        return f"ValuePattern({self._name!r})"
+
+    def __add__(self, other):
+        return _pattern_builder.Add(self, other)
+
+    def __radd__(self, other):
+        return _pattern_builder.Add(other, self)
+
+    def __sub__(self, other):
+        return _pattern_builder.Sub(self, other)
+
+    def __rsub__(self, other):
+        return _pattern_builder.Sub(other, self)
+
+    def __mul__(self, other):
+        return _pattern_builder.Mul(self, other)
+
+    def __rmul__(self, other):
+        return _pattern_builder.Mul(other, self)
+
+    def __truediv__(self, other):
+        return _pattern_builder.Div(self, other)
+
+    def __rtruediv__(self, other):
+        return _pattern_builder.Div(other, self)
+
+    def __pow__(self, other):
+        return _pattern_builder.Pow(self, other)
+
+    def __str__(self) -> str:
+        return self._name if self._name is not None else "anonymous:" + str(id(self))
+
+
+class NodePattern:
+    """Represents a pattern that matches against a Node.
+
+    This differs from a NodeOutputPattern in that it matches against a node (which
+    may produce 1 or more outputs), whereas a NodeOutputPattern matches against
+    a specific output of a node.
+
+    Args:
+        domain: pattern to match against the domain of the node.
+        op: pattern or string constant to match against the op_type of the node.
+        inputs: sequence of ValuePatterns (or constants) to match against the inputs of the node.
+        attributes: dictionary of attribute patterns to match against the attributes of the node.
+        outputs: specifies pattern-variable-name for outputs (or None)
+        allow_other_attributes: specifies whether other attributes (not mentioned in `attributes`)
+          are allowed in the node.
+    """
+
+    def __init__(
+        self,
+        domain: StringPattern,
+        op: str | Pattern[str],
+        inputs: Sequence[int | float | ValuePattern | None],
+        attributes: dict[str, AttrPattern],
+        outputs: Sequence[str | None],
+        *,
+        allow_other_attributes: bool | None,
+        allow_other_inputs: bool | None,
+    ):
+        if allow_other_attributes is None:
+            # Default behavior: allow other unmatched attributes in the node.
+            allow_other_attributes = True
+        if allow_other_inputs is None:
+            # TODO(rama): Should we default to True? For now, we preserve the current behavior.
+            allow_other_inputs = False
+        self.domain = domain
+        self.op = StringConstantPattern(op) if isinstance(op, str) else op
+        self.inputs = [_to_value_pattern(x) for x in inputs]
+        self.attributes = attributes
+        self.allow_other_attributes = allow_other_attributes
+        self.allow_other_inputs = allow_other_inputs
+        # 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.
+            overload = ""
+            self._op_identifier: tuple[str, str, str] | None = (
+                domain.value(),
+                op,
+                overload,
+            )
+        else:
+            self._op_identifier = None
+        self.outputs = [NodeOutputPattern(self, i, name) for i, name in enumerate(outputs)]
+
+        # Update uses for inputs.
+        for index, value in enumerate(self.inputs):
+            if value is not None:
+                value.append_use(self, index)
+
+    def __str__(self) -> str:
+        inputs = ", ".join(str(v) for v in self.inputs)
+        outputs = ", ".join(str(v) for v in self.outputs)
+        attributes = ", ".join(f"{k}={v}" for k, v in self.attributes.items())
+        op = str(self.op)
+        domain = str(self.domain)
+        qualified_op = f"{domain}.{op}" if domain else op
+        inputs_and_attributes = f"{inputs}, {attributes}" if attributes else inputs
+        return f"{outputs} = {qualified_op} ({inputs_and_attributes})"
+
+    def op_identifier(self) -> Tuple[str, str, str] | None:
+        return self._op_identifier
+
+    @property
+    def op_type(self) -> str:
+        return str(self.op)
+
+    def matches(self, node: ir.Node, match: MatchResult) -> MatchResult:
+        """Matches the pattern represented by self against a node.
+
+        This is purely a local node-level match, and does not consider the subgraph rooted at the node.
+        We check the domain, op_type, and attributes of the node, but not the inputs.
+        """
+        # TODO(rama): Ensure we handle "" and "onnx.ai" correctly.
+        if not self.op.matches(node.op_type):
+            return match.fail(
+                f"OpType mismatch: expected {self.op}, got {node.op_type}.", node
+            )
+        if not self.domain.matches(node.domain):
+            return match.fail(
+                f"Domain mismatch: expected {self.domain}, got {node.domain}.", node
+            )
+
+        for name, attr_pattern in self.attributes.items():
+            attr_value = node.attributes.get(name)
+            if attr_value is None:
+                return match.fail(f"Attribute {name} not found in node.", node)
+            if not attr_pattern.matches(attr_value):
+                return match.fail(
+                    f"Attribute {name} mismatch: expected {attr_pattern}, got {attr_value}.",
+                    node,
+                )
+            if attr_pattern.name is not None:
+                if not match.bind(attr_pattern.name, attr_value):
+                    return match
+
+        if not self.allow_other_attributes:
+            for name in node.attributes:
+                # TODO: Support matching default nodes for attributes.
+                if name not in self.attributes:
+                    return match.fail(f"Attribute {name} not expected in node.", node)
+
+        return match
+
+    def clone(self, node_map: dict[NodePattern, NodePattern], swap: bool) -> NodePattern:
+        inputs = [(v.clone(node_map) if v is not None else None) for v in self.inputs]
+        if swap:
+            assert len(inputs) == 2, (
+                "Internal error: commutative swap applies only to binary ops."
+            )
+            inputs = [inputs[1], inputs[0]]
+        outputs = [value.name for value in self.outputs]
+        copied = NodePattern(
+            self.domain,
+            self.op,
+            inputs,
+            self.attributes,
+            outputs,
+            allow_other_attributes=self.allow_other_attributes,
+            allow_other_inputs=self.allow_other_inputs,
+        )
+        node_map[self] = copied
+        return copied
+
+
+class NodeOutputPattern(ValuePattern):
+    """Represents a pattern that matches against a specific output of a Node.
+
+    This is the primary pattern used to match against computed values, that
+    is values computed using a specific op.
+    """
+
+    def __init__(
+        self, producer: NodePattern, output_index: int, name: str | None = None
+    ) -> None:
+        super().__init__(name)
+        self._producer = producer
+        self._output_index = output_index
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> NodeOutputPattern:
+        return node_map[self._producer].outputs[self._output_index]
+        # return NodeOutputPattern(node_map[self._producer], self._output_index, self._name)
+
+    @property
+    def output_index(self) -> int:
+        return self._output_index
+
+    def producer(self) -> NodePattern:
+        return self._producer
+
+
+Var = ValuePattern
+
+
+def _is_pattern_variable(x: Any) -> bool:
+    # The derived classes of ValuePattern represent constant patterns and node-output patterns.
+    return type(x) is ValuePattern
+
+
+class Constant(ValuePattern):
+    """Represents a pattern that matches against a scalar constant value."""
+
+    def __init__(
+        self,
+        value: int | float | Sequence[int] | Sequence[float],
+        rel_tol: float = 1e-5,
+        abs_tol: float = 1e-8,
+    ) -> None:
+        super().__init__(None)
+        self._value = list(value) if isinstance(value, Sequence) else value
+        self._rel_tol = rel_tol
+        self._abs_tol = abs_tol
+
+    def clone(self, node_map: dict[NodePattern, NodePattern]) -> Constant:
+        del node_map
+        return Constant(self._value, self._rel_tol, self._abs_tol)
+
+    @property
+    def value(self) -> int | float | list[int] | list[float]:
+        return self._value
+
+    def matches(self, value: ir.Value, match: MatchResult) -> MatchResult:
+        constant_value = value.const_value
+        if constant_value is None:
+            return match.fail(f"Value is not a constant, expecting {self.value}.")
+
+        constant_value_numpy = constant_value.numpy()
+        if isinstance(self._value, list):
+            if constant_value_numpy.shape != (len(self._value),):
+                return match.fail(f"Value has mismatching shape, expecting ({self.value},).")
+            if not all(
+                math.isclose(
+                    constant_value_numpy.item(i),
+                    self._value[i],
+                    rel_tol=self._rel_tol,
+                    abs_tol=self._abs_tol,
+                )
+                for i in range(len(self._value))
+            ):
+                return match.fail(
+                    f"Value mismatch: expected {self._value}, got {constant_value_numpy}."
+                )
+            return match
+
+        # Scalar constant case:
+        # TODO (rama): allow users to specify shape requirement, if desired.
+        if constant_value_numpy.size != 1:
+            return match.fail(f"Value is not a scalar, expecting {self.value}.")
+
+        if not math.isclose(
+            constant_value_numpy.item(),
+            self._value,
+            rel_tol=self._rel_tol,
+            abs_tol=self._abs_tol,
+        ):
+            match.fail(
+                f"Value mismatch: expected {self._value}, got {constant_value_numpy.item()}."
+            )
+
+        # Note: If the value is produced by a Constant node, we could include
+        # the Constant node in the return_value list. However, we don't do that.
+        # Instead, we will rely on DCE to remove the constant node if it is not
+        # used elsewhere.
+        return match
+
+    def __str__(self) -> str:
+        return str(self._value)
+
+
+def _nodes_in_pattern(outputs: Sequence[ValuePattern]) -> list[NodePattern]:
+    """Returns all nodes used in a pattern, given the outputs of the pattern."""
+    node_patterns: list[NodePattern] = []
+
+    def visit(value_patterns: Sequence[ValuePattern | None]) -> None:
+        for value_pattern in value_patterns:
+            if isinstance(value_pattern, NodeOutputPattern):
+                node_pattern = value_pattern.producer()
+                if node_pattern not in node_patterns:
+                    node_patterns.append(node_pattern)
+                    visit(node_pattern.inputs)
+
+    visit(outputs)
+    node_patterns.reverse()
+    return node_patterns
+
+
+def _add_backward_slice(node: NodePattern, backward_slice: set[NodePattern]) -> None:
+    """Adds all nodes in the backward slice of given node to the set `backward_slice`.
+
+    The backward slice of a node is the set of all nodes that are reachable from the node
+    in a backward traversal from the given node.
+    """
+    if node in backward_slice:
+        return
+    backward_slice.add(node)
+    for value_pattern in node.inputs:
+        if isinstance(value_pattern, NodeOutputPattern):
+            _add_backward_slice(value_pattern.producer(), backward_slice)
+
+
+class GraphPattern:
+    """Represents a pattern that can be matched against a subgraph."""
+
+    def __init__(
+        self,
+        inputs: Sequence[ValuePattern],
+        outputs: Sequence[ValuePattern],
+        nodes: Sequence[NodePattern],
+    ) -> None:
+        self._inputs = inputs
+        self._outputs = outputs
+        if len(outputs) == 0:
+            raise ValueError("GraphPattern must have at least one output")
+        self._nodes = nodes  # _nodes_in_pattern(outputs)
+
+        # Determine the output nodes of the pattern. These are a minimal set of nodes
+        # whose backward-slices cover the entire pattern.
+        output_nodes: set[NodePattern] = set()
+        covered: set[NodePattern] = set()
+        for value_pattern in outputs:
+            if not isinstance(value_pattern, ValuePattern):
+                raise TypeError(
+                    f"Invalid type {type(value_pattern)} for graph pattern output."
+                )
+            if isinstance(value_pattern, Constant):
+                raise NotImplementedError(
+                    "Constant values are not allowed as graph pattern outputs."
+                )
+            if isinstance(value_pattern, NodeOutputPattern):
+                candidate = value_pattern.producer()
+                if candidate not in covered:
+                    output_nodes.add(candidate)
+                    _add_backward_slice(candidate, covered)
+
+        self.output_nodes: list[NodePattern] = list(output_nodes)
+
+    @property
+    def output_node(self) -> NodePattern:
+        if len(self.output_nodes) != 1:
+            raise ValueError("GraphPattern does not have unique output node.")
+        return self.output_nodes[0]
+
+    def node(self, index: int) -> NodePattern:
+        return self._nodes[index]
+
+    def num_nodes(self) -> int:
+        return len(self._nodes)
+
+    def __len__(self) -> int:
+        return self.num_nodes()
+
+    @property
+    def inputs(self) -> Sequence[ValuePattern]:
+        return self._inputs
+
+    @property
+    def outputs(self) -> Sequence[ValuePattern]:
+        return self._outputs
+
+    def __iter__(self) -> Iterator[NodePattern]:
+        return iter(self._nodes)
+
+    def __reversed__(self) -> Iterator[NodePattern]:
+        return reversed(self._nodes)
+
+    @property
+    def has_single_output_node(self) -> bool:
+        return len(self.output_nodes) == 1
+
+    @property
+    def num_outputs(self) -> int:
+        return len(self._outputs)
+
+    def commute(self) -> Sequence[GraphPattern]:
+        def commute_node(node: NodePattern) -> Iterable[bool]:
+            if node.op_identifier() == ("", "Add", "") or node.op_identifier() == (
+                "",
+                "Mul",
+                "",
+            ):
+                # Try with and without swapping inputs.
+                return [False, True]
+            # No swapping of inputs
+            return [False]
+
+        iteration_space = [commute_node(node) for node in self._nodes]
+
+        def copy_graph(swap_list: Iterable[bool]) -> GraphPattern:
+            if not any(swap_list):
+                # No need to swap inputs of any node
+                return self
+            # Create a copy of the graph, with swapped inputs for the nodes that need it.
+            node_map: dict[NodePattern, NodePattern] = {}
+            new_inputs = [v.clone(node_map) for v in self._inputs]
+            new_nodes = [
+                node.clone(node_map, swap) for node, swap in zip(self._nodes, swap_list)
+            ]
+            new_outputs = [v.clone(node_map) for v in self._outputs]
+            return GraphPattern(new_inputs, new_outputs, new_nodes)
+
+        return [copy_graph(swap_list) for swap_list in itertools.product(*iteration_space)]
+
+    def __str__(self) -> str:
+        inputs = ", ".join(str(v) for v in self._inputs)
+        outputs = ", ".join(str(v) for v in self._outputs)
+        nodes = "\n   ".join(str(n) for n in self._nodes)
+        return f"pattern ({inputs}) {{\n   {nodes}\n   return {outputs}\n}}"
+
+
+def _to_graph_pattern(pattern_constructor: Callable) -> GraphPattern:
+    """Convert a pattern-construction function to a GraphPattern.
+
+    A pattern-construction function will return values as below:
+    ::
+        def pattern(op, x: Var, shape1: Var, shape2: Var):
+            ...
+            return outputs
+
+    We create a pattern graph by creating pattern-variables for each parameter of the function,
+    and calling the function. The returned values are normalized to a list of ValuePatterns,
+    which represent the outputs of the pattern graph.
+
+    Args:
+        pattern_constructor: Callable
+
+    Returns:
+        GraphPattern: A representation of the pattern that can be matched against a subgraph.
+    """
+    _pattern_vars = inspect.signature(pattern_constructor).parameters
+    pattern_inputs = [Var(v) for v in _pattern_vars][1:]  # Skip the first parameter
+    builder = OpsetPatternBuilder("", record=True)
+    with pattern_builder(builder):
+        pattern_outputs = pattern_constructor(builder, *pattern_inputs)
+    # TODO(rama): classify inputs as value/attribute vars
+    # Returned value could be a single ValuePattern or a list of ValuePatterns.
+    # Normalize representation to a list of ValuePatterns.
+    if isinstance(pattern_outputs, ValuePattern):
+        pattern_outputs = [pattern_outputs]
+    return GraphPattern(pattern_inputs, pattern_outputs, builder.nodes())
+
+
+def _valid_to_replace(
+    matched_nodes: Sequence[ir.Node], output_values: Sequence[ir.Value]
+) -> bool:
+    """Check that values computed by the matched_nodes, except for output_values, are used only by the matched_nodes."""
+    # * Must check that all values matched by pattern are used only by pattern,
+    # except for the value that is replaced.
+    # * Must ensure that replacement subgraph does not use any of the deleted
+    # (intermediate) values. (Not necessary for now. Guaranteed.)
+    for n in matched_nodes:
+        for v in n.outputs:
+            if v in output_values:
+                continue
+            if v.is_graph_output():
+                # value is an output-value of the graph/function.
+                return False
+            for consumer, _ in v.uses():
+                if consumer not in matched_nodes:
+                    return False
+    return True
+
+
+RewriterContext = _tape.Builder
+
+
+@dataclasses.dataclass
+class ReplacementSubgraph:
+    """A subgraph that will replace the matched pattern."""
+
+    match: MatchResult
+    new_outputs: Sequence[ir.Value]
+    new_nodes: Sequence[ir.Node]
+    new_initializers: Sequence[ir.Value]
+    used_opsets: _tape.UsedOpsets
+
+
+def always_true(*args, **kwargs) -> bool:
+    """A condition function that always returns True.
+
+    This is used when no condition function is provided for a rewrite rule.
+    """
+    return True
+
+
+class ReplacementPatternFunction:
+    """The replacement pattern that will replace the targeted pattern.
+
+    Attributes:
+        function (Callable): The replacement function that will be used to replace the matched pattern.
+    """
+
+    def __init__(self, function) -> None:
+        self._function = function
+
+    def get_replacement(self, match: MatchResult) -> ReplacementSubgraph | None:
+        context = RewriterContext()
+        new_outputs = self._function(context, **match.bindings)
+        if new_outputs is None:
+            return None  # Failed to create replacement subgraph
+        if not isinstance(new_outputs, Sequence):
+            new_outputs = [new_outputs]
+        return ReplacementSubgraph(
+            match, new_outputs, context.nodes, context.initializers, context.used_opsets
+        )
+
+
+def _update_opset_imports(
+    graph_or_function: ir.Graph | ir.Function, delta: ReplacementSubgraph
+):
+    imports = graph_or_function.opset_imports
+    for domain, version in delta.used_opsets:
+        if domain not in imports:
+            # use 1 as default version if not explicitly specified
+            imports[domain] = version if version is not None else 1
+        elif version is not None and version != imports[domain]:
+            raise ValueError(
+                f"Multiple versions of opset {domain} used. "
+                f"Expected version {imports[domain]}, but got {version}."
+            )
+
+
+class PatternMatcher(abc.ABC):
+    def __init__(self, pattern: GraphPattern) -> None:
+        self.pattern = pattern
+
+    @abc.abstractmethod
+    def match(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        node: ir.Node,
+        *,
+        verbose: int = 0,
+        remove_nodes: bool = True,
+        tracer: MatchingTracer | None = None,
+    ) -> MatchResult:
+        """Match the pattern against the subgraph ending at the given node."""
+
+    def __str__(self) -> str:
+        return str(self.pattern)
+
+
+class SimplePatternMatcher(PatternMatcher):
+    def __init__(self, pattern: GraphPattern) -> None:
+        super().__init__(pattern)
+        self._current_node: ir.Node | None = None
+
+    def fail(self, reason: str, node: ir.Node | None = None) -> bool:
+        if self._verbose:
+            if self._matched:  # Print only if at least one node successfully matched.
+                count = len(self._matched)
+                print(f"Match failed after {count} nodes: {reason}")
+        self._match.fail(reason, node or self._current_node)
+        return False
+
+    def _match_constant(self, pattern_constant: Constant, value: ir.Value) -> bool:
+        """Match a Constant pattern against a value.
+
+        If the constant value is produced by a Constant node, we do not include
+        the constant node as part of the matched graph. Thus, it will not be deleted,
+        if subgraph replacement happens. But subsequent DCE will remove the constant
+        node if it is not used elsewhere.
+        """
+        constant_value = value.const_value
+        if constant_value is None:
+            return self.fail(
+                f"Value {value.name} is not a constant, expecting {pattern_constant.value}.",
+            )
+
+        try:
+            constant_value_numpy = constant_value.numpy()
+        except FileNotFoundError:
+            return self.fail(f"Constant value of {value.name} not available.")
+
+        pattern_constant_value = pattern_constant._value
+
+        if isinstance(pattern_constant_value, list):
+            expected_shape = (len(pattern_constant_value),)
+            if constant_value_numpy.shape != expected_shape:
+                return self.fail(f"Value has mismatching shape, expecting {expected_shape}.")
+            if not all(
+                math.isclose(
+                    constant_value_numpy.item(i),
+                    pattern_constant_value[i],
+                    rel_tol=pattern_constant._rel_tol,
+                    abs_tol=pattern_constant._abs_tol,
+                )
+                for i in range(len(pattern_constant_value))
+            ):
+                return self.fail(
+                    f"Value mismatch: expected {pattern_constant_value}, got {constant_value_numpy}."
+                )
+            return True
+
+        # TODO (rama): allow users to specify shape requirement, if desired.
+        if constant_value_numpy.size != 1:
+            return self.fail(
+                f"Value {value.name} is not a scalar, expecting {pattern_constant_value}.",
+            )
+
+        if not math.isclose(
+            constant_value_numpy.item(),
+            pattern_constant_value,
+            rel_tol=pattern_constant._rel_tol,
+            abs_tol=pattern_constant._abs_tol,
+        ):
+            return self.fail(
+                f"Constant value mismatch: expected {pattern_constant_value}, got {constant_value_numpy.item()}.",
+            )
+
+        return True
+
+    def _match_node(self, pattern_node: NodePattern, node: ir.Node) -> bool:
+        """Matches a pattern subgraph against subgraph rooted at node."""
+        self._current_node = node
+        # Graph-matching: we do not allow the same pattern node to be matched against
+        # different graph nodes.
+        if pattern_node in self._matched:
+            if self._matched[pattern_node] is not node:
+                return self.fail("Same pattern node is matched against different graph nodes.")
+            return True
+        match = self._match
+        if not pattern_node.matches(node, match):
+            return self.fail(match.reason)
+
+        if self._verbose:
+            print(f"Matched: {node.op_type}")
+
+        match.nodes.append(node)
+        self._matched[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):
+                return self.fail(
+                    f"Number of inputs ({len(node.inputs)}) is less than expected ({len(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)}"
+                )
+
+        for arg_value, arg_pattern in zip(node.inputs, pattern_node.inputs):
+            # arg_pattern could be a Var, if it's the original arg.
+            if arg_pattern is None:
+                if arg_value is None:
+                    continue
+                else:
+                    return self.fail("(Optional) input is expected to be None but is not.")
+            if not self._match_value(arg_pattern, arg_value):
+                return False
+
+        for i, output_value_pattern in enumerate(pattern_node.outputs):
+            if not self._bind_value(output_value_pattern, node.outputs[i]):
+                return False
+
+        return True
+
+    def _bind_value(self, pattern_value: ValuePattern, value: ir.Value | None) -> bool:
+        """Bind a ValuePattern var to ir Value."""
+        if pattern_value.name is not None:
+            match = self._match
+            if pattern_value.name in match.bindings:
+                # TODO(rama): Use appropriate equality-check here: future extension possibility.
+                if match.bindings[pattern_value.name] == value:
+                    return True
+                return self.fail(f"Variable {pattern_value.name} is bound to multiple values.")
+            match.bindings[pattern_value.name] = value
+        return True
+
+    def _match_value(self, pattern_value: ValuePattern, value: ir.Value | None) -> bool:
+        """Match an IR value against a ValuePattern instance."""
+        if not self._bind_value(pattern_value, value):
+            return False
+
+        if isinstance(pattern_value, NodeOutputPattern):
+            if value is None:
+                return self.fail("Mismatch: Computed node pattern does not match None.")
+            return self._match_node_output(pattern_value, value)
+        if isinstance(pattern_value, Constant):
+            if value is None:
+                return self.fail("Mismatch: Constant pattern does not match None.")
+            return self._match_constant(pattern_value, value)
+        return True
+
+    def _match_node_output(self, pattern_value: NodeOutputPattern, value: ir.Value) -> bool:
+        """Match an IR value against a NodeOutputPattern instance."""
+        node = value.producer()
+        if node is None:
+            return self.fail(
+                "Mismatch: Computed node pattern does not match uncomputed IR value."
+            )
+        if value.index() != pattern_value.output_index:
+            return self.fail(
+                f"Node output index mismatch: expected {pattern_value._output_index}, got {value.index()}."
+            )
+        return self._match_node(pattern_value.producer(), node)
+
+    def _init_match(self, verbose: int) -> None:
+        """Initialize the match state. Invoked before starting a new match."""
+        self._verbose = verbose
+        self._matched: dict[NodePattern, ir.Node] = {}
+        self._match: MatchResult = MatchResult()
+        self._current_node = None
+
+    def _get_output_values(self) -> list[ir.Value] | None:
+        """Get values bound to the output variables of the pattern."""
+        output_values: list[ir.Value] = []
+        unbound_values: list[str] = []
+        for j, value_pattern in enumerate(self.pattern.outputs):
+            if value_pattern.name is not None:
+                if value_pattern.name in self._match.bindings:
+                    output_values.append(self._match.bindings[value_pattern.name])
+                else:
+                    unbound_values.append(value_pattern.name)
+            elif isinstance(value_pattern, NodeOutputPattern):
+                i = value_pattern.output_index
+                node = value_pattern.producer()
+                if node in self._matched:
+                    output_values.append(self._matched[node].outputs[i])
+                else:
+                    unbound_values.append(f"output_{j}")
+            elif isinstance(value_pattern, Constant):
+                raise NotImplementedError("Constant values as return-values not supported.")
+        if unbound_values:
+            self._match.fail(f"Error: Output values not found: {unbound_values}")
+            return None
+        return output_values
+
+    def _match_single_output_node(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        node: ir.Node,
+        check_removable: bool,
+    ) -> MatchResult:
+        del model
+        del graph_or_function
+
+        pattern = self.pattern
+        match = self._match
+
+        if not pattern.has_single_output_node:
+            return match.fail(
+                "Internal Error: SimplePatternMatcher should not be used for patterns with multiple output nodes."
+            )
+
+        if not self._match_node(pattern.output_node, node):
+            return match
+
+        output_values = self._get_output_values()
+        if output_values is None:
+            # TODO(rama): Is this a valid (useful) case?
+            return match
+        if check_removable and not _valid_to_replace(match.nodes, output_values):
+            # TODO(rama): Match status should be updated to reflect failure reason.
+            return match.fail("Matched nodes have other uses preventing replacement.")
+
+        match.outputs.extend(output_values)
+        return match
+
+    def _multi_match(self, candidate: Iterable[ir.Node], check_removable: bool) -> MatchResult:
+        """Find a match for a pattern with multiple output nodes.
+
+        For a pattern with K output nodes, the input candidate should specify K nodes
+        in the graph that will be matched against the pattern output nodes.
+
+        Args:
+            candidate: An iterable of nodes that will be matched against the pattern output nodes.
+            check_removable: If True, check that the matched nodes can be removed (that is, that
+                they are not used elsewhere in the graph).
+        """
+        match = self._match
+        for pattern_node, node in zip(self.pattern.output_nodes, candidate):
+            if not self._match_node(pattern_node, node):
+                return match
+        output_values = self._get_output_values()
+        if output_values is None:
+            return match
+
+        if check_removable and not _valid_to_replace(match.nodes, output_values):
+            return match.fail("Matched nodes have other uses preventing replacement.")
+
+        match.outputs.extend(output_values)
+        return match
+
+    def match(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        node: ir.Node,
+        *,
+        verbose: int = 0,
+        remove_nodes: bool = True,
+        tracer: MatchingTracer | None = None,
+    ) -> MatchResult:
+        """Match the pattern against the subgraph ending at the given node.
+
+        For patterns with multiple output nodes, the given node is matched
+        against the first output node in the pattern. For the remaining
+        output nodes in the pattern, we use a brute-force algorithm that
+        enumerates all possible combinations of nodes from the graph (with
+        a filter based on op-type).
+
+        TODO: Consider omitting parameters model and graph_or_function. With
+        the new IR, the graph can be obtained from the node, and the model is
+        not used. But this is a shared abstract method of the Matcher interface,
+        so other matcher implementation also needs to be updated. More importantly,
+        matching in the presence of subgraphs (control-flow) can introduce some
+        complications which require careful consideration.
+        """
+        self._tracer = tracer
+        if self.pattern.has_single_output_node:
+            self._init_match(verbose)
+            return self._match_single_output_node(
+                model, graph_or_function, node, check_removable=remove_nodes
+            )
+        else:
+            # Note: This is a potentially expensive algorithm for matching patterns with
+            # multiple output nodes. For patterns with N output nodes, we try all possible
+            # combinations of N nodes from the graph, and check if they match the pattern.
+            # The first node is fixed to the node argument in this method call. We do
+            # some simple filtering by restricting the candidates for each remaining
+            # output nodes to graph nodes with the same op_type as the corresponding pattern
+            # node. For now, this is intended to be a simple, but robust, implementation
+            # that can be used for debugging and testing. The GenericPatternMatcher is a
+            # more sophisticated implementation, but incomplete.
+            pattern_output_nodes = self.pattern.output_nodes
+            op_to_nodes: dict[tuple[str, str, str], list[ir.Node]] = {}
+            for n in graph_or_function:
+                op_to_nodes.setdefault(n.op_identifier(), []).append(n)
+            all_nodes = iter(graph_or_function)
+
+            def get_nodes(pattern_node):
+                id = pattern_node.op_identifier()
+                if id is None:
+                    return all_nodes
+                return op_to_nodes.get(id, [])
+
+            candidates = [iter([node])] + [get_nodes(pn) for pn in pattern_output_nodes[1:]]
+            match = None
+            for combination in itertools.product(*candidates):
+                self._init_match(verbose)
+                match = self._multi_match(combination, check_removable=remove_nodes)
+                if match:
+                    return match
+            if match is None:
+                return MatchResult().fail("No match found.")
+            return match
+
+
+class RewriteRule:
+    def __init__(
+        self,
+        target_pattern: GraphPattern | Callable,
+        replacement_pattern: ReplacementPatternFunction | Callable,
+        condition_function: Callable | None = None,
+        matcher: PatternMatcher | Callable[[GraphPattern], PatternMatcher] | None = None,
+        verbose: int = 0,
+        name: str | None = None,
+        remove_nodes: bool = True,
+        graph_pre_visitor: Callable[[], None] | None = None,
+        graph_post_visitor: Callable[[], None] | None = None,
+    ) -> None:
+        """Create a rewrite rule.
+
+        Args:
+            target_pattern: The GraphPattern that will be matched against the IR.
+                If a callable is provided, it will be converted to a GraphPattern.
+            replacement_pattern: The ReplacementPatternFunction that will be used to
+                replace the matched pattern. If a callable is provided, it will be
+                converted to a ReplacementPatternFunction.
+            condition_function: The condition function that will be used to check if
+                the pattern match found should be rewritten.
+            matcher: The pattern matcher that will be used to match the pattern.
+                If not provided, a default matcher will be used.
+            verbose: The verbosity level of the rule.
+            name: An optional name for the pattern that will show up in verbose logging.
+            remove_nodes: If True, the matched nodes will be removed from the graph.
+            graph_pre_visitor: A function that will be called before applying the
+                rewriting to the top-level graph or a function.
+            graph_post_visitor: A function that will be called after the rewriting
+                is complete for a graph or function.
+        """
+
+        if not isinstance(target_pattern, GraphPattern):
+            target_pattern = _to_graph_pattern(target_pattern)
+        self._target_pattern = target_pattern
+
+        if not isinstance(replacement_pattern, ReplacementPatternFunction):
+            replacement_pattern = ReplacementPatternFunction(replacement_pattern)
+        self._replacement_pattern = replacement_pattern
+        self._condition_function = condition_function or always_true
+        if isinstance(matcher, PatternMatcher):
+            self._matcher = matcher
+        elif matcher is None:
+            if target_pattern.has_single_output_node:
+                self._matcher = SimplePatternMatcher(self._target_pattern)
+            else:
+                import onnxscript.rewriter.generic_pattern as generic_pattern
+
+                self._matcher = generic_pattern.GenericPatternMatcher(self._target_pattern)
+        else:
+            self._matcher = matcher(self._target_pattern)
+        self._verbose = verbose
+        self.name = name
+        self.remove_nodes = remove_nodes
+        self.graph_pre_visitor = graph_pre_visitor
+        self.graph_post_visitor = graph_post_visitor
+
+    def __str__(self) -> str:
+        return self.name if self.name else "Anonymous Rule"
+
+    def try_rewrite(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        node: ir.Node,
+        *,
+        verbose: int | None = None,
+        tracer: MatchingTracer | None = None,
+    ) -> ReplacementSubgraph | None:
+        """If the node matches the pattern, then replace the node with the replacement pattern."""
+        if verbose and verbose > 2:
+            print(f"[try_rewrite] {self}")
+        verbose = verbose if verbose is not None else self._verbose
+        match = self._matcher.match(
+            model, graph_or_function, node, verbose=verbose, remove_nodes=self.remove_nodes
+        )
+        if match:
+            context = None  # TODO(rama)
+            for var in self._target_pattern.inputs:
+                if var.name is not None:
+                    if var.name not in match.bindings:
+                        match.bindings[var.name] = None
+            if not self._condition_function(context, **match.bindings):
+                if tracer:
+                    tracer.log(
+                        self, graph_or_function, node, match, MatchStatus.CONDITION_FAILED
+                    )
+                return None
+            replacement_subgraph = self._replacement_pattern.get_replacement(match)
+            if replacement_subgraph is None:
+                if tracer:
+                    tracer.log(
+                        self, graph_or_function, node, match, MatchStatus.REPLACEMENT_FAILED
+                    )
+                return None
+            if len(replacement_subgraph.new_outputs) != self._target_pattern.num_outputs:
+                raise ValueError(
+                    f"Number of outputs from replacement function does not match the number of outputs from the target pattern. "
+                    f"Expected {self._target_pattern.num_outputs}, but got {len(replacement_subgraph.new_outputs)}."
+                )
+            # TODO(rama): Remove the opset imports from deleted nodes?
+            _update_opset_imports(graph_or_function, replacement_subgraph)
+            _update_opset_imports(model.graph, replacement_subgraph)
+            if tracer:
+                tracer.log(self, graph_or_function, node, match, MatchStatus.SUCCESS)
+            return replacement_subgraph
+        if tracer:
+            tracer.log(self, graph_or_function, node, match, MatchStatus.NO_MATCH)
+        return None
+
+    def apply_to_model(
+        self,
+        model: ir.Model,
+        *,
+        commute: bool = False,
+        verbose: int | None = None,
+        debug: bool = False,
+    ):
+        # A convenience method to apply the rule to a model. We use a RewriteRuleSet to
+        # handle commutative rules.
+        return RewriteRuleSet([self], commute=commute).apply_to_model(
+            model, verbose=verbose, debug=debug
+        )
+
+    def commute(self) -> Sequence[RewriteRule]:
+        def replace_pattern(new_pattern):
+            """Return a shallow copy of self with node_pattern replaced by new_pattern."""
+            # TODO(rama): Maybe we should use a better alternative to construct new matcher.
+            matcher_class = type(self._matcher)
+            return RewriteRule(
+                new_pattern,
+                self._replacement_pattern,
+                self._condition_function,
+                matcher_class(new_pattern),
+                self._verbose,
+                self.name,
+                self.remove_nodes,
+                self.graph_pre_visitor,
+                self.graph_post_visitor,
+            )
+
+        return [replace_pattern(p) for p in self._target_pattern.commute()]
+
+
+class RewriteRuleAsClass:
+    """Defines a class grouping method pattern, rewrite, check.
+    This class is then given to function :func:`make_rewrite_rule_from_class`
+    to define a new rule.
+    """
+
+    @classmethod
+    def pattern(cls, op, *_) -> Any:
+        raise NotImplementedError("Method 'pattern' must be overwritten.")
+
+    @classmethod
+    def rewrite(cls, op, *_) -> Any:
+        raise NotImplementedError("Method 'rewrite' must be overwritten.")
+
+    @classmethod
+    def check(cls, context, *_, **__) -> bool:
+        return True
+
+
+def make_rewrite_rule_from_class(
+    rule_class: type | RewriteRuleAsClass, generic: bool = False
+) -> RewriteRule:
+    """Creates a RewriteRule from a class defining the function
+    pattern, rewrite, check with class method. It makes it is easier
+    to read when a module contains multiple patterns.
+
+    Example::
+
+        class TransposeIdentity(RewriteRuleAsClass):
+            @classmethod
+            def pattern(cls, op, x, perm):
+                return op.Transpose(x, perm=perm)
+
+            @classmethod
+            def check(cls, context, x: ir.Value, perm: ir.Attr | ir.RefAttr) -> bool:
+                if isinstance(perm, ir.RefAttr):
+                    return False
+                if perm.type == ir.AttributeType.INTS:
+                    if perm.value == list(range(len(perm.value))):
+                        return True
+                return False
+
+            @classmethod
+            def rewrite(cls, op, x: ir.Value, perm: ir.Attr | None = None):
+                return op.Identity(x)
+
+        transpose_identity_rule = make_rewrite_rule_from_class(TransposeIdentity)
+    """
+    assert hasattr(rule_class, "pattern"), f"Method 'pattern' is missing from {rule_class!r}."
+    assert hasattr(rule_class, "rewrite"), f"Method 'rewrite' is missing from {rule_class!r}."
+    assert hasattr(rule_class, "check"), f"Method 'check' is missing from {rule_class!r}."
+    if generic:
+        import onnxscript.rewriter.generic_pattern as orpp
+
+        return RewriteRule(
+            rule_class.pattern,
+            rule_class.rewrite,
+            rule_class.check,
+            orpp.GenericPatternMatcher,
+            name=rule_class.__name__,  # type: ignore[union-attr]
+        )
+    return RewriteRule(
+        rule_class.pattern,
+        rule_class.rewrite,
+        rule_class.check,
+        name=rule_class.__name__,  # type: ignore[union-attr]
+    )
+
+
+# Variation of RewriteRuleAsClass that is based on instance methods instead of class methods.
+# Useful to implement a family of rules to support pattern variations.
+# TODO: cleanup the naming conventions for these inter-related classes.
+class RewriteRuleClassBase:
+    @classmethod
+    def rule(cls, *args, **kwargs):
+        instance = cls(*args, **kwargs)
+        setup = instance.setup if hasattr(instance, "setup") else None
+        cleanup = instance.cleanup if hasattr(instance, "cleanup") else None
+        return RewriteRule(
+            instance.pattern,
+            instance.rewrite,
+            instance.check,
+            name=instance.name,
+            remove_nodes=instance.remove_nodes,
+            graph_pre_visitor=setup,
+            graph_post_visitor=cleanup,
+        )
+
+    def __init__(self, name: str | None = None, remove_nodes: bool = True) -> None:
+        self.name = name or self.__class__.__name__
+        self.remove_nodes = remove_nodes
+
+    def pattern(self, op, *args, **kwargs):
+        raise NotImplementedError("Method 'pattern' must be implemented by derived class.")
+
+    def check(self, op, *args, **kwargs):
+        # Default check function that always returns True.
+        return True
+
+    def rewrite(self, op, *args, **kwargs):
+        raise NotImplementedError("Method 'rewrite' must be implemented by derived class.")
+
+
+class RewriteRuleSet:
+    def __init__(self, rules: Sequence[RewriteRule], *, commute: bool = False) -> None:
+        if commute:
+            rules = list(itertools.chain.from_iterable([rule.commute() for rule in rules]))
+        self.rules = rules
+
+    def _apply_to_graph_or_function(
+        self,
+        model: ir.Model,
+        graph_or_function: ir.Graph | ir.Function,
+        *,
+        verbose: int | None,
+        tracer: MatchingTracer | None = None,
+    ) -> int:
+        """
+        Apply the rewrite rules to the given graph or function.
+
+        Args:
+            model: The model to which the rewrite rules are applied.
+            graph_or_function: The graph or function to which the rewrite rules are applied.
+            verbose: The verbosity level. Defaults to None.
+            tracer: The tracer for debugging. Defaults to None.
+
+        Returns:
+            The number of rewrite rules applied.
+        """
+        count = 0
+
+        # NOTE: Rules should be prioritized in the order they are added to the RewriteRuleSet.
+        # And the graph is applied in order.
+        for rule in self.rules:
+            if rule.graph_pre_visitor:
+                rule.graph_pre_visitor()
+            for node in graph_or_function:
+                delta = rule.try_rewrite(
+                    model, graph_or_function, node, verbose=verbose, tracer=tracer
+                )
+                if delta is None or tracer is not None:
+                    continue
+                assert isinstance(delta, ReplacementSubgraph)
+                if delta.new_initializers:
+                    if isinstance(graph_or_function, ir.Function):
+                        # TODO(rama): Can't add initializers to functions. But currently this is not
+                        # an issue, as we apply inlining before applying rewrite rules.
+                        if verbose:
+                            print(
+                                f"Rewrites adding initializers not supported for functions: {rule}"
+                            )
+                        continue
+                    initializers = graph_or_function.initializers
+                    for initializer in delta.new_initializers:
+                        if initializer.name in initializers:
+                            if verbose:
+                                print(f"Initializer {initializer.name} already exists.")
+                            continue
+                    for initializer in delta.new_initializers:
+                        initializers[initializer.name] = initializer  # type: ignore[index]
+                # TODO: This does not yet handle the problem of determining the correct insertion point
+                # for inserted nodes in the case of patterns with multiple output-nodes. The following
+                # is sufficient for patterns with a single output-node "node", which can serve as the
+                # insertion-point.
+                onnxscript.optimizer.basic_constant_propagation(delta.new_nodes)
+                _convenience.replace_nodes_and_values(
+                    graph_or_function,
+                    node,
+                    delta.match.nodes if rule.remove_nodes else [],
+                    delta.new_nodes,
+                    delta.match.outputs,
+                    delta.new_outputs,
+                )
+                count += 1
+            if rule.graph_post_visitor:
+                rule.graph_post_visitor()
+
+        return count
+
+    def apply_to_model(
+        self, model: ir.Model, *, verbose: int | None = None, debug: bool = False
+    ) -> int:
+        """Apply the rewrite rules in the set to the model.
+
+        Args:
+            model: The model to which the rewrite rules are applied.
+            verbose: The verbosity level of messages. Defaults to None.
+            debug: Whether to enable debugging. Defaults to False. In the
+                debug mode, no changes are made to the model, only a report is produced at
+                the end about the best matches found.
+
+        Returns:
+            The number of applications of rewrite rules.
+        """
+        assert isinstance(model, ir.Model)
+        tracer = MatchingTracer() if debug else None
+        onnxscript.optimizer.basic_constant_propagation(model.graph)
+        count = self._apply_to_graph_or_function(
+            model, model.graph, verbose=verbose, tracer=tracer
+        )
+        for function in model.functions.values():
+            onnxscript.optimizer.basic_constant_propagation(function)
+            count += self._apply_to_graph_or_function(
+                model, function, verbose=verbose, tracer=tracer
+            )
+        if tracer:
+            tracer.report()
+        return count
+
+    def __iter__(self):
+        yield from self.rules
+
+
+class MatchStatus(enum.IntEnum):
+    """The status of a pattern-matching operation."""
+
+    NO_MATCH = 0  # No successful match found for entire pattern graph
+    CONDITION_FAILED = 1  # Subsequent validation check failed
+    REPLACEMENT_FAILED = 2  # Replacement subgraph could not be created
+    SUCCESS = 3  # A successful match was found
+
+
+@dataclasses.dataclass
+class MatchInfo:
+    """The status of a pattern-matching operation. An extension of MatchResult."""
+
+    match_result: MatchResult
+    root_node: ir.Node
+    container: ir.Graph | ir.Function
+    status: MatchStatus
+
+    def score(self) -> int:
+        """Return a score for the match."""
+        return len(self.match_result.nodes) + int(self.status.value) * 100
+
+
+class MatchingTracer:
+    """A debugging helper class to trace the matching of a pattern against a graph.
+
+    This is used to track the best matches found for each rule, and to report the
+    results at the end of the matching.
+    """
+
+    def __init__(self) -> None:
+        self._log: dict[RewriteRule, list[MatchInfo]] = defaultdict(list)
+
+    def log(
+        self,
+        rule: RewriteRule,
+        container: ir.Graph | ir.Function,
+        node: ir.Node,
+        match_result: MatchResult,
+        status: MatchStatus,
+    ) -> None:
+        this_match = MatchInfo(match_result, node, container, status)
+        this_score = this_match.score()
+        if this_score == 0:
+            return
+        best_matches = self._log[rule]
+        if best_matches:
+            if this_score < best_matches[0].score():
+                return
+            if this_score > best_matches[0].score():
+                best_matches.clear()
+        best_matches.append(this_match)
+
+    def report(self) -> None:
+        import onnxscript.rewriter._ir_utils as ir_utils
+
+        print("===")
+        for rule, matches in self._log.items():
+            if not matches:
+                continue
+            print(f"Rule: {rule}")
+            print(f"Best score: {matches[0].score()}")
+            for match in matches:
+                print(f"Status: {match.status}")
+                if match.status == MatchStatus.NO_MATCH:
+                    print("Graph matching failed: " + match.match_result.reason)
+                    node = match.match_result._failure_node
+                    if node:
+                        print("Failure at or around node:")
+                        node.display()
+                print("Matched nodes:")
+                ir_utils.display_nodes(match.match_result.nodes)
+                print("===")