diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md new file mode 100644 index 0000000000000..95f21704e9fb8 --- /dev/null +++ b/docs/Pattern_Rewriter.md @@ -0,0 +1,1540 @@ +# Writing C++ graph rewrite rules + +This document describes how to express fusion optimizations as rewrite rules. + +## Example: composing two Transposes + +The following example replaces `Transpose(Transpose(x, p), q)` with one Transpose. + +```cpp +#include "core/optimizer/pattern_rewriter.h" + +using namespace onnxruntime; +using namespace onnxruntime::rewrite; + +bool ComposeTransposePermutations(const Graph&, Match& match) { + // Validate that the permutations are composable and compute the "composed" permutation + // ... +} + +Rule MakeComposeTransposeRule() { + Pattern pattern(R"( + (x, p : ints, q : ints) => { + a = Transpose(x, perm = p) + return Transpose(a, perm = q) + })"); + Replacement replacement("(x, composed) => Transpose(x, perm = composed)"); + return Rule{"ComposeTranspose", pattern, replacement, + ComposeTransposePermutations, {{kOnnxDomain, 13, 25}}}; +} +``` + +A rule combines a pattern, semantic check-and-bind function, replacement, +and supported opset ranges. +* The pattern describes the subgraph to be fused +* The replacement describes the new subgraph that will replace the fused subgraph +* The check-and-bind function (called `ComposeTransposePermutations` in above +example) checks for extra conditions that must hold for the fusion to be +applied. It can also compute other information (eg., new constant values or +attribute values) needed to construct the replacement. In this example, +it **checks** known rank and valid explicit permutations and **binds** the +composed permutation for the replacement. For inner permutation `p` and +outer `q`, composition is `composed[i] = p[q[i]]`, not `q[p[i]]`. +* The opset ranges specify the ONNX opsets for which the fusion is valid. + +The complete definition of `ComposeTransposePermutations` is: +```cpp +bool ComposeTransposePermutations(const Graph&, Match& match) { + const NodeArg* x = match.valueOf("x"); + const auto* shape = x->Shape(); + if (!shape) return match.Reject("Transpose composition requires known rank."); + const int rank = shape->dim_size(); + const ONNX_NAMESPACE::AttributeProto& p = match.attributeOf("p"); + const ONNX_NAMESPACE::AttributeProto& q = match.attributeOf("q"); + auto valid = [rank](const ONNX_NAMESPACE::AttributeProto& perm) { + if (perm.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_INTS || + perm.ints_size() != rank) return false; + InlinedHashSet seen; + for (int64_t axis : perm.ints()) + if (axis < 0 || axis >= rank || !seen.insert(axis).second) return false; + return true; + }; + if (!valid(p) || !valid(q)) return match.Reject("Invalid transpose permutation."); + ONNX_NAMESPACE::AttributeProto composed; + composed.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); + for (int64_t axis : q.ints()) composed.add_ints(p.ints(static_cast(axis))); + match.update("composed", std::move(composed)); + return true; +} +``` + +The pattern requires explicit `perm` attributes; it does not match default +reverse-axis Transposes. `attributeOf` returns a const reference to the captured +`AttributeProto` rather than copying its potentially large payload; the reference +borrows the Match. The pattern guarantees that `p` and `q` exist before the +checker runs. This example deliberately uses mandatory captures. An annotation +such as `p : ints | None` would also match an omitted `perm`; its checker would +use `optionalAttributeOf("p")` and explicitly implement the default reverse +permutation. See [optional attribute captures](#optional-attribute-captures). +Known rank is enough: dimension sizes need not be +known. Identity composition remains a Transpose with the identity permutation, +not a zero-node replacement. The built-in boundary check will not apply the +fusion if the inner Transpose result is exposed as a graph output or is +consumed elsewhere. The rank-three +noncommuting permutations `[1, 0, 2]` then `[0, 2, 1]` compose to `[1, 2, 0]`. + +## Patterns + +`Pattern` is an owning, copyable parsed graph description. Repeated names bind +the **same `NodeArg`**, not merely equal tensor contents. It has a standalone +execution API; neither a `Rule`, replacement, nor `PatternRewriter` is needed: + +```cpp +Pattern pattern("(x) => { a = Identity(x); return Sigmoid(a) }"); +std::optional match; +ORT_RETURN_IF_ERROR(pattern.MatchAt(graph, root.Index(), match)); +if (match) { + const NodeArg* input = match->valueOf("x"); + // Inspect input; do not mutate graph. + (void)input; +} +``` + +`MatchAt` matches at exactly the specified root node. It returns a match, +or `std::nullopt` for an ordinary nonmatch (clearing any previous result). +Status errors and configuration +exceptions leave the caller's result unchanged. The result owns its selected +plan but borrows the unchanged graph, as described in [Match](#the-match-object). +The graph must already be resolved; an unresolved graph or invalid root returns +an error. Malformed/non-rooted pattern configuration throws. + +This standalone API uses the same structural, attribute, constant, type/shape, +layer and **deletion-boundary** matcher as the engine. It finds rewrite-safe +patterns, not arbitrary subgraph occurrences. It does not apply rule opset-range +or execution-provider policy, run a checker, construct a replacement, or traverse +other roots/subgraphs. After introducing pattern disjunction below, +[matching alternatives at one root](#matching-alternatives-at-one-root) describes +first-match ordering and the separate ordered visitor API. + +### Pattern syntax and shared replacement notation + +Definitions start with named inputs and `=>`. Braces around the body are optional; +use them to make multi-statement rules easier to read. Assignments create named +operator outputs; `return` either calls the final operator (the root) or names an +ordered list of previously defined operator outputs. Whitespace separates calls, and semicolons are +optional separators. A trailing semicolon before the closing brace is allowed. +Only whitespace and comments may follow the closing brace. In both patterns and +replacements, `#` starts a comment through the end of the line (or end of text). +Comments act as whitespace between tokens, including in headers and bodies; +they cannot split an identifier, number, or multi-character token. +Hashes inside single- or double-quoted attribute strings remain literal: + +```text +# A complete line comment +(x) => { + a = Identity(x) # An inline comment + return Identity(a) +} # A trailing comment +``` + +```cpp +Pattern pattern(R"( + (x) => { + add = Add(x, 3.0) + clip = Clip(add, 0.0, 6.0) + product = Mul(clip, x) + result = Div(product, 6.0) + return result; + } +)"); +Replacement replacement("(x) => HardSwish(x)"); +``` + +The example assigns one root output, named `result`, then returns it. Assignment +and return are separate statements: `return result = Div(...)` is rejected. +A final call without an assignment, such as `return Div(product, 6.0)`, also describes one output, +but does not give that output a lookup name. Replacements additionally accept +a final call without the `return` keyword, as above. + +Use `domain.Operator(...)` (ONNX text-format notation) or +`domain::Operator(...)` for qualified operators in patterns and replacements: + +```cpp +Pattern pattern("(x, y) => { mm = com.microsoft.FusedMatMul(x, y); return Identity(mm); }"); +Replacement replacement("(x, y) => com.microsoft.FusedMatMul(x, y)"); +``` + +An unqualified operator uses the standard ONNX domain. `ai.onnx::Identity` +and `ai.onnx.Identity` are equivalent to `Identity`. Dotted qualification +splits at the last dot, so `com.microsoft.FusedMatMul` has domain +`com.microsoft` and operator name `FusedMatMul`. + +Names must be unique and defined before use. Every declared pattern input must +participate, and every operator must reach the root. Inputs to different named +variables may alias; repeating one variable requires identical binding. +`match.valueOf("x")` and `match.valueOf("clip")` return bound `const NodeArg*` +values. Unknown names and unbound values raise explicit errors. Read-only +`Pattern::NamedValue`, `Nodes`, `Constants`, `Outputs`, `ValueCount`, `IsInput`, +`IsOptional`, and `AttributeNames` +are available for introspection. + +### Optional type and shape annotations + +Pattern header variables and assigned outputs may carry ` : type` annotations. +Unannotated variables retain their existing unconstrained behavior: + +```cpp +Pattern pattern(R"( + (x : float[M, N], y : float[N, K]) => { + product : float[M, K] = MatMul(x, y); + return Identity(product); + } +)"); +``` + +Dimension names such as `M`, `N`, and `K` are **declared implicitly** by their +first appearance in an annotation. Do not list them as separate header inputs. +They cannot collide with tensor, attribute, or element-type variable names. + +| Annotation | Constraint | +|---|---| +| `float` | Rank-zero float32 tensor (known scalar shape) | +| `float[]` or `float[...]` | Float32 tensor of any rank, including unknown rank | +| `float[M, N]` | Known rank 2, with named dimension bindings | +| `float[2, 3]` | Known rank 2 and exact static dimensions 2 and 3 | +| `float[M, ...]` | Known rank at least 1, binding its first dimension | +| `float[_, N]` | Known rank 2; first dimension unconstrained and unnamed | +| `tensor[M, N]` | Known rank 2, without an element-type constraint | +| `tensor[M, N]` | Known rank 2, with shared element-type variable `T` | +| `tensor[]` | Any rank, with shared element-type variable `T` | + +Tensor element names follow ONNX, for example `float`, `double`, `float16`, +`bfloat16`, `int64`, `int32`, `uint8`, `bool`, and `string`. Only tensor types +are supported; sequences, maps, and optional ONNX value types are not annotations. +For a tensor binding, `None` in a type union describes an absent operator slot, not an ONNX optional +value type: `(x : float[], low : float | None, high : float | None)` allows absent +scalar Clip bounds. An absent slot matches only its `None` alternative, without +binding any dimensions or element type. + +The matcher uses the model's existing metadata; it does not infer shapes or +solve equations. Rank-specific constraints reject unknown rank, and fixed +dimensions require known equal integers. A named dimension binds to an integer +for a static dimension or a nonempty string for a symbolic dimension. +A name occurring **once** in the expanded selected annotation combination is +a capture: an anonymous unknown source dimension is allowed and leaves the +capture unbound. A **repeated** name is an equality constraint: every occurrence +must have a known integer or nonempty symbolic name, and all must agree exactly. +Unknown-plus-known and unknown-plus-unknown both reject before the checker, +regardless of traversal order. Integer 2 differs from symbolic `"batch"`, and +`"batch"` differs from `"other"`. `_` accepts any dimension, including unknown, +without creating a capture or equality constraint. + +For example, `(x : float[M, K], y : float[M, N])` requires both `M` dimensions to be +known and equal; the singleton `K` and `N` dimensions may be anonymous unknowns. +`float[M, M]` imposes equality between two positions of one tensor as well. +Counts include only the selected alternative at each annotation occurrence, +not mutually exclusive alternatives or failed search branches. Calls preserve +hygienic dimension names and each conjunctive annotation occurrence. Merely +forwarding an interface formal through a call does not duplicate its annotation; +an explicit interface output annotation is an additional constraint, even if +the output aliases an input. Repeating a name there therefore also requires +known metadata. `Options::on_failure` reports a type/shape failure naming the +unknown repeated dimension or conflicting bindings. +Likewise, a missing element type leaves `T` unbound, while a concrete annotation +such as `float` requires a known matching element type. + +In a check-and-bind or replacement callback, `match.dimensionOf("M")` returns a +`Dimension` (`std::variant`), and +`match.elementTypeOf("T")` returns an optional ONNX element-type enum value +represented as `int32_t`. Unknown declared bindings are empty; undeclared names +raise an error. These bindings are local to each match, cannot be overwritten +with computed attributes, and survive clearing computed attributes. +`Pattern::TypeConstraints()` and `IsAttribute()` expose read-only introspection. + +Attribute captures can be declared and annotated in the header: + +```cpp +Pattern pattern("(x : float[M, N], axes : ints) => { t = Transpose(x, perm = axes); return Identity(t); }"); +Replacement replacement("(x, axes) => Transpose(x, perm = axes)"); +``` + +A header name used in keyword position denotes an attribute, not a tensor. +Attribute annotations are `int`, `float`, `string`, `ints`, `floats`, `strings`, +or `tensor`, matching the corresponding concrete `AttributeProto` type. +For example, `float` means a FLOAT attribute in keyword position but a scalar +float32 tensor in tensor position; use `int64` for an integer tensor and `int` +for an INT attribute. Attribute variables cannot also be tensor inputs or carry +shapes. An attribute annotation permits one concrete type, optionally combined +with `None`, as described below; arbitrary attribute-type unions remain unsupported. +Undeclared captures remain supported and untyped, and require a present attribute. + +Annotations are constraints on **patterns only**; replacements continue to use +unannotated inputs and derive types/shapes from preserved outputs and schema +inference. Annotations on returns or replacements are rejected rather than +silently ignored. Each annotation allows at most 64 dimensions, with at most +256 distinct implicit dimension/element-type variables per pattern. + +### Type alternatives + +Use `|` for a disjunction of tensor constraints: + +```cpp +Pattern pattern("(x : float[M] | float[M, N]) => return Identity(x)"); +Pattern optional("(x, low : float | None) => return Clip(x, low)"); +``` + +Type alternatives are tried left to right, with annotated variables processed +in definition order. Matching finds the first **globally consistent** combination: +if a later constraint conflicts with a tentative dimension or element-type +binding, earlier annotation alternatives are reconsidered. Failed alternatives +leave no bindings behind. Variables mentioned only in unchosen alternatives +remain unbound. Check-and-bind callbacks run once after a consistent combination is selected; +Callback rejection does not trigger another combination. + +The product of alternative counts across annotations is limited to 64 at parse time. +Empty branches, duplicate constraints, and `None` alone are rejected. +`None | float[M]` is also accepted. `None` is permitted in header tensor inputs +and optional attribute captures, not assigned outputs. Replacements remain untyped. +The former experimental `x?` syntax is removed; use `x : tensor[...] | None` for +an unconstrained optional tensor slot. + +### Optional input slots + +Use `None` in an operator input position to require an absent input. It matches +both an empty ONNX input name and an omitted trailing slot. Mark a header +variable with a tensor type union including `None` to allow either a present value or an absent slot: + +```cpp +Pattern pattern("(x, low : tensor[...] | None, high : tensor[...] | None) => { a = Identity(x); return Clip(a, low, high); }"); +Replacement replacement("(x, low, high) => Clip(x, low, high)"); +Pattern absent_min("(x, high) => { a = Identity(x); return Clip(a, None, high); }"); +``` + +An ordinary header variable requires a present input. Repeated optional variables +must bind identically, including absence: `(x, v : tensor[...] | None) => Clip(x, v, v)` does not match +one present and one absent slot. `match.valueOf("low")` returns `nullptr` when +the optional binding is absent; check-and-bind callbacks must check it before dereferencing. +Replacements preserve an absent binding as an empty input name. Explicit `None` +is also supported in replacements. Schema validation still rejects missing +required inputs. This is slot matching, not variadic matching: extra actual +input slots are not accepted. `None` is reserved and cannot be a variable name. + +### Constants and attributes + +Positional numeric literals in patterns match **rank-zero constant initializers** +of the exact value and element type: `3.0`, `3e0`, and `0.5` are float32; +`3`, `-1`, and `0` are int64. Thus `2` and `2.0` are distinct constraints. +A singleton vector is not a scalar. External tensor data and overridable +initializers do not satisfy these constraints. Constant operators require +an earlier constant-folding pass. + +A bracketed positional literal is a **rank-one tensor**, not an attribute: + +```cpp +Pattern pattern("(x) => { reshaped = Reshape(x, [2, 3]); return Identity(reshaped); }"); +Replacement replacement("(x) => Reshape(x, [2, 3])"); +``` + +`[1, 2]` is an int64 tensor of shape `[2]`; `[1.0, 2.0]` is float32 of the same +shape. `[1]` has shape `[1]` and does not match scalar `1` or a matrix `[[1]]`. +Matching requires exact rank, length, element type, order, and values. Nested +lists are not supported. All elements must have the same numeric type spelling: +mixing integer and floating-point literals is rejected instead of silently +converting them. Float elements must be finite. `[]` is an empty int64 tensor +of shape `[0]`. Replacement tensor lists synthesize immutable initializers. + +Keyword arguments constrain attributes, not tensor inputs: + +```cpp +Pattern pattern("(x) => { h = HardSigmoid(x, alpha = 0.25, beta = 0.5); return Mul(h, x); }"); +``` + +Attribute equality compares the protobuf attribute type as well as its value. +For example, `axis = -1` is an integer attribute, while `epsilon = 1e-5` is a float. +Single- or double-quoted string attributes are supported; strings may escape +quotes and backslashes. Other attributes are unconstrained. A listed attribute +must be present in the resolved node's attribute map. Staged replacement nodes +receive schema default attributes before verification and requeueing, so later +rules and callbacks see defaults without an intervening `Graph::Resolve()`. + +In keyword positions the same numeric list syntax instead denotes an attribute: + +```cpp +Pattern pattern("(x) => { t = Transpose(x, perm = [1, 0]); return Identity(t); }"); +Replacement replacement("(x) => Transpose(x, perm = [1, 0])"); +``` + +Integer lists become `INTS` attributes; float lists become `FLOATS`. +Matching compares the exact attribute type, length, order, and each value. +`[]` denotes an empty `INTS` attribute. Scalar and list attributes are distinct. +Pattern syntax supports scalar integer, float, and string attribute equality, +and numeric list equality, but not custom predicate syntax or string lists. +Inspect `match.nodes` in the C++ check-and-bind callback for ranges, shapes, or other semantic conditions. +Positional inputs must precede keyword attributes. + +Instead of a literal, a pattern keyword may capture an attribute by name: + +```cpp +Pattern pattern("(x) => { t = Transpose(x, perm = axes); return Identity(t); }"); +Replacement replacement("(x, axes) => Transpose(x, perm = axes)"); +``` + +The attribute must be present and concrete. `match.attributeOf("axes")` exposes +the captured `AttributeProto` to check-and-bind callbacks and text replacements without extra preparation. +Capture names are separate from tensor values and cannot collide with them; +declared header captures are classified as attributes rather than tensor inputs. +Repeating a capture requires identical protobuf attribute content, ignoring the +original attribute name and documentation string; type and list order matter. +Captures can hold any concrete attribute type, but replacement schemas decide +whether that type is valid. Captures are per match, survive clearing computed +attributes, and cannot be overwritten with `match.update`. Matching reads only the resolved node's attribute map, including any defaults +already inserted by `Graph::Resolve`. + +### Optional attribute captures + +Declare `name : type | None` (or `name : None | type`) to capture either a +concrete attribute or its absence. The concrete type can be `int`, `float`, +`string`, `ints`, `floats`, `strings`, or `tensor`; shape/element-type +annotations and `const` are not attribute qualifiers. `None` alone and unions +of multiple concrete attribute types are rejected. Classification follows use: +`float | None` in a keyword position describes an optional FLOAT attribute, +whereas in a tensor position it describes an optional scalar float32 tensor. + +```cpp +Pattern pattern("(x, perm : ints | None) => return Transpose(x, perm = perm)"); +std::optional match; +ORT_RETURN_IF_ERROR(pattern.MatchAt(graph, root_index, match)); +if (match) { + const ONNX_NAMESPACE::AttributeProto* perm = match->optionalAttributeOf("perm"); + if (perm) { + // An explicit permutation, including an empty list for rank zero. + } else { + // Transpose uses reverse-axis order; interpretation belongs to the checker. + } +} +``` + +`optionalAttributeOf` returns a pointer to a captured/prepared attribute, or +`nullptr` for a successfully captured absence. It throws for unknown names, +tensor bindings and captures unavailable in the selected plan. `attributeOf` +still returns a const reference and throws for absence. Both lookups are +available on `BindingView` with the same relative-name/scoping rules. +References/pointers borrow the Match. Copying a Match preserves presence state; +`ClearUpdates` does not clear captures or captured absence, and `update` cannot +overwrite an optional capture even when absent. Prepared attributes still +require concrete values: there is no prepared-None `update` API. + +Repeating a capture requires identical presence at every occurrence and, when +present, identical normalized attribute contents. Missing/present is a mismatch +in either order; an empty list is present, not absent. Each concrete candidate +has fresh presence state. Ordinary nested calls and alternatives remap optional +attribute bindings hygienically, just like mandatory captures. An unselected +branch does not fabricate evidence of an absent attribute. + +Presence means presence in the **resolved map**, not original model spelling. +For example, omitted LeakyRelu `alpha` becomes a present default after Resolve; +Transpose `perm` has no schema default and may remain absent. Optional capture +does not implement omitted-attribute semantics or prove replacement equivalence. + +A replacement keyword referring to captured absence is **omitted**, not written +as an UNDEFINED/empty protobuf attribute: + +```cpp +Rule preserve{"PreserveTranspose", + "(x, perm : ints | None) => return Transpose(x, perm = perm)", + "(x, perm) => Transpose(x, perm = perm)", {}, + {{kOnnxDomain, 13, 25}}}; +``` + +Staging then inserts target-schema defaults and validates required attributes. +Omitting a required attribute fails and rolls back staging; it does not retry +another candidate. The example illustrates propagation, not an optimization +(use a one-iteration limit if executing this self-reproducing rule). + +Scalar numeric literals in replacements synthesize rank-zero float32 or int64 +initializers; list literals synthesize rank-one tensors. Replacement keyword +literals set scalar or list attributes on the new node. + +### Branched patterns and multiple outputs + +Assignments can describe a DAG rather than only a chain: + +```cpp +Pattern pattern(R"( + (x, y) => { + sum = Add(x, y) + product = Mul(sum, x) + return Add(sum, product) + } +)"); +``` + +Here both uses of `sum` must have the same producer. The matcher does not search +for equivalent computations or reorder commutative inputs. + +Comma-separated assignment names declare the exact ordered output count: + +```cpp +Pattern pattern("(x) => { first, second = Split(x, axis = 0); return Concat(first, second, axis = 0); }"); +Replacement replacement("(x) => Identity(x)"); +``` + +For a multi-output root, assign its outputs first, then return their names: + +```cpp +Pattern pattern(R"( + (x) => { + copied = Identity(x); + left, right = Split(copied, axis = 0); + return left, right; + } +)"); +Replacement replacement("(x) => { left, right = Split(x, axis = 0); return left, right; }"); +``` + +This replacement preserves the two explicitly returned original `NodeArg`s in order, +including their names, types, shapes, graph-output roles, and consumers. +The replacement names are local labels, not new externally visible graph names. +`return Split(x)` describes an operator with one declared output; use assignment +names to describe its exact arity. Named returns may select and reorder outputs, +including outputs from different nodes: + +```cpp +Pattern pattern("(x) => { a = Identity(x); y = Add(a, a); z = Mul(y, y); return y, z; }"); +Replacement replacement("(x) => { y = Add(x, x); z = Mul(y, y); return y, z; }"); +``` + +The final operator remains the unique sink used to anchor backward matching. +Every defined operator must reach it, independently of the returned-value order. +Multiple independent sinks are rejected. Return position defines correspondence: +the first replacement return inherits the first pattern return's `NodeArg`, +and so on. The returned counts must match. Returning a leaf input, duplicate +names, or undefined values is rejected. + +Only **explicitly returned** values may have graph-output roles, external +consumers, or implicit subgraph captures. All other outputs, including unreturned +sink outputs, must remain within the deletion set or be unused. Internally +consumed returned values are correctly wired into the replacement graph. +If an external consumer of a returned intermediate feeds a pattern input, +the staged producer dependencies are checked for cycles before deletion. +A replacement that introduces a cycle is rejected with a status and rolled +back; a mapping that preserves an acyclic dependency graph is allowed. + +## Composable pattern libraries + +`PatternLibrary` parses an owning library of interfaces and named patterns. +Every concrete named pattern is callable; graph bodies and alternatives are +two definition forms in one namespace, not separate declaration kinds. +Definitions may refer forward to other definitions. +`library.GetPattern("Fuse")` returns a copyable `Pattern` usable by the standalone +matching APIs and every `Rule` constructor. The library and its source text need +not outlive the returned pattern. + +```cpp +PatternLibrary library(R"( + pattern activation(g : float[M, N], u : float[M, N], + alpha : float[...] | None) -> (product : float[M, N]) { + scaled = (Mul(g, alpha) | g) + s = Sigmoid(scaled) + a = Mul(g, s) + product = Mul(a, u) + return product + } + + pattern Fuse(g, u, alpha : float[...] | None) { + product = activation(g, u, alpha) + return Identity(product) + } +)"); +``` + +Ordinary patterns declare typed positional inputs and, optionally, positional +outputs using `-> (product : float[M, N])`. `return product` supplies that result; +multiple results use `return first, second`. Output annotations constrain the +returned values and share the definition's dimension variables with inputs; +they do not manufacture shape evidence. Output names must not shadow a +different local value. Without an output signature, names and annotations +are inferred from returned variables. An alias may inherit an input annotation +while another expression branch returns an unannotated operator result; these +are compatible unconstrained-result alternatives. Each selected plan keeps +only its actual annotation occurrences, so reconciliation does not create +additional repeated-dimension constraints. + +Ordinary inputs retain their tensor/attribute classification from use, including +typed attribute formals such as +`pattern P(x, axis : int) -> (y) { y = Softmax(x, axis = axis); return y }`. +Adding an output signature does not impose an interface's all-tensor input +contract. Declared outputs themselves must be tensors. Output labels cannot +replace a different existing tensor binding, including a direct alias or a +positional call result; reusing a label for the same handle remains valid. + +`product = activation(g, u, alpha)` binds a tensor directly. For multiple results +write `first, second = Pair(x)`, just as for a multi-output ONNX operator. +Explicit multi-target calls require exact positional arity and distinct, +previously unused target names. No `product.port` notation is needed. + +### Expressions, alternatives, and branch-local inputs + +The example factors optional scaling outside the shared Sigmoid/Mul suffix: +`scaled = (Mul(g, alpha) | g)`. An assignment `scaled = g` is an alias to the exact +same `NodeArg`, not an inserted Identity. Reusing both names therefore imposes +the same binding-identity requirement as reusing either name. + +Alternatives are tried in source order, expanding into complete concrete +plans. Put multiplication first to prefer it when both alternatives can match. +There is no automatic longest match. Checker rejection may try the next +complete plan, with fresh bindings. Every included operator must contribute +to a returned value in **each** plan. Thus +`alpha = Add(a, b); scaled = (Mul(g, alpha) | g); return Sigmoid(scaled)` is rejected +because the Add is disconnected in the identity plan. Connecting that Add +through another returned computation makes the definition valid. + +Call-expression alternatives may invoke definitions with different input +signatures. Each call checks its own actual arguments; only positional result +arity and declared result annotations must be compatible. No named interface +is required: + +```text +pattern plain(g : float[M, N], u : float[M, N]) -> (product : float[M, N]) { + s = Sigmoid(g); a = Mul(g, s); product = Mul(a, u); return product +} +pattern scaled(g : float[M, N], u : float[M, N], alpha : float[...]) -> (product : float[M, N]) { + v = Mul(g, alpha); s = Sigmoid(v); a = Mul(g, s); product = Mul(a, u); return product +} +pattern choice(g, u, alpha : float[...] | None) { + product = plain(g, u) | scaled(g, u, alpha) + return product +} +pattern discovered(g, u) { + product = plain(g, u) | scaled(g, u, _) + return product +} +``` + +`_` as a call actual binds a **real, present tensor** without a caller-level +name. It is not ONNX's omitted input slot (`None`). Each `_` is independent; +callee type/shape/constancy and repeated-formal binding requirements still +apply. The callee-local binding remains available to scoped C++ views. +Wildcards for attribute formals are rejected. + +A declared input may be unused in a selected branch only if its annotation +permits `None`. In `choice` the plain plan leaves `alpha` unbound. Another +reachable use may bind it, in which case its tensor annotation and all uses' +`NodeArg` identity must agree. Required inputs remain required; `None` cannot +disable an actual required operator slot. `match.optionalValueOf("alpha")` +returns `nullptr` for a declared unbound optional tensor and throws for an +unknown name or attribute. Existing `valueOf` retains its optional-null +behavior; required unbound lookups remain errors. Select an appropriate +pre-parsed replacement in C++ when an optional capture is absent: consuming +it in a required replacement input is an explicit schema-validation error, +not a successful rewrite or candidate retry. + +Result compatibility compares the ordered declared constraints structurally, +including dimension symbol spellings, but ignores result labels. An +unannotated result is unconstrained. Parenthesized expressions may nest; +calls and aliases preserve source-order selection identities. Expression +occurrences expose their selected call/operator/alias name via +`match.instance("product").alternative()`, with `alternativeView()` descending +into a selected call's metadata. Operator and value-alias leaves have no +callee-local scope or deeper `alternativeView()`; grouped branches retain +their nested selection metadata under the `` tag. +See [instance binding views](#instance-binding-views-and-handwritten-c-structs) +for scoped lookups and a complete checker/replacement example. + +Interfaces remain optional reusable tensor contracts, not callable graphs: +`interface Unary(x : float[N]) -> (y : float[N])` and +`pattern Copy implements Unary { y = Identity(x); return y }` remain supported. +The older explicit `(port = local_value)` returns, colon-signature extra +discovered inputs, and scoped DSL `p.port` accesses are compatibility syntax, +not necessary for new positional examples. An unnamed returned operator call +still exposes `output` (or `output0`, `output1`, etc.). + +### Calls versus standalone entries + +Reusable bodies may return identity aliases or independent output branches. +Every required declared input must be used, and every operator must reach at least +one returned value, but those returned values need not have a single root. +`GetPattern(name)` validates **every fully expanded plan** as a standalone +entry: it must have an operator root, all operators must feed that final +operator, and returned values cannot be leaf/identity aliases or duplicate +the same value through different output names. Callable bodies may still +export multiple aliases of the same value. An entry violation raises a +use-time diagnostic naming the requested entry. Inherited interface inputs +and mapped outputs must remain tensor bindings, checked during library +construction; extra body-only inputs may capture attributes. Mapped interface +outputs are always required to be present, even when unannotated and aliased +to an optional input. Interface tensor-kind requirements survive calls and +aliases; callers cannot reclassify these mapped values as attribute captures +or pass them to attribute formals. Conflicts fail during library construction. +Thus a library +may contain useful callable patterns that cannot independently serve as rules: + +```text +pattern identity(x) -> (y) { return x } +pattern branches(x, z) { + left = Identity(x) + right = Neg(z) + return left, right +} +pattern join(x, z) { + left, right = branches(x, z) + value = identity(left) + return Add(value, right) +} +``` + +`GetPattern("join")` succeeds; `GetPattern("identity")` and +`GetPattern("branches")` fail. These checks apply equally to body patterns, +aliases and alternative definitions. Legacy standalone `Pattern("(x) => ...")` +syntax remains supported; its operator-root and output-boundary checks are unchanged. + +The legacy colon signature extends an interface implementation with external +inputs discovered by matching. For example, +`pattern extra implements Unary : (alpha : float[...]) -> { y = Mul(x, alpha); return y }` +is called with just `extra(x)`, while `alpha` binds to the matched Mul's input. +It remains a tensor, not an automatically decoded float. +Additional attribute captures can also be declared in that signature using +the existing attribute-annotation syntax. There is no separate capture +language. Both `float[]` and `float[...]` continue to mean **any rank**, not +scalar. + +An **alternative pattern** supplies ordered named members. Members can be +body patterns, other alternative patterns, or single aliases; interfaces +cannot be members. They must have the same named interface (or all have no +interface) and identical ordered input/output port names and annotations. +This is a strict structural annotation check (including symbol spellings), +not semantic type equivalence; equivalent any-rank `[]`/`[...]` are normalized. +Extra discovered inputs of an `implements` body are not common call ports. +A named call `value = Forms(x)` and the legacy anonymous call +`value = (first | second)(x)` have the same matching behavior when `Forms = first | second`. +Unlike call-expression alternatives, these definition alternatives retain +the common input signature requirement. Each occurrence +has an independent selected immediate member, available through `alternative()`; +no user-defined tag name is required. Nesting and single aliases preserve +their selection identities: + +```text +pattern first(x) { value = Identity(x); return value } +pattern second(x) { value = Neg(x); return value } +pattern Inner = first | second +pattern Outer = Inner | first +pattern Alias = Outer +pattern nested(x) { + value = Alias(x) + return Identity(value) +} +``` + +For a selection of `second`, `value` reports definition `Alias`, member `Outer`; +successive `alternativeView()` calls report `Outer`/`Inner`, `Inner`/`second`, +then body `second` with no member. Legacy anonymous alternatives use the same +representation and retain their immediate named members too. Patterns may +call other patterns within the expansion and dependency-nesting bounds. +The separate `|` type-annotation syntax is unchanged. + +Calls bind positional output variables. The **first output name** is also +the C++ occurrence-scope handle for `match.instance("first")`; this explicit +introspection API is separate from `match.valueOf("first")`, which returns +the actual result tensor. Legacy one-name calls to multi-output patterns +remain instance-only: `p = Pair(x)` requires `p.port` to address a result and +`valueOf("p")`/tensor use of `p` fails, rather than selecting an arbitrary +output. Tensor handles, attribute captures, +dimensions and element-type variables are remapped into the call occurrence. +Two instances do not share a capture just because its local spelling is the +same. Formal inputs and returned aliases bind to the caller's actual +`NodeArg`s; no integer handle from another definition is reused. +Caller and callee annotations are conjunctive, while alternatives within +one annotation remain disjunctive. A call inserts its callee's ordered +constraint occurrences at the call site, preserving callee order even when +actual arguments are reordered or aliased to the same value. +An annotation classified as a tensor in a compiled callee stays a tensor +after composition. A caller cannot turn that mapped value into an attribute +through a later keyword capture or attribute-formal call; this unsupported +conversion fails during construction, including for ordinary identity aliases. +Typed attribute formals already classified as attributes in the callee remain +supported. + +Graph alternatives expand to at most **64 complete concrete plans per +definition**, in source order (the rightmost call varies fastest). +Each plan retains the existing limits of 64 nodes, 256 values and 64 +type-annotation combinations. Each plan also permits at most **256 expanded +call occurrences**, counted cumulatively through nested bodies, including +zero-node identity calls. This bounds scoped metadata even when composition +adds no nodes or values; overflow is rejected before copying callee metadata. +Libraries contain at most 256 definitions. A dependency path may contain at +most **64 concrete pattern definitions**, including its entry and any named +alternative/alias wrappers (interfaces do not count). This bound is independent +of declaration order and is enforced before deeper recursive compilation to +protect native stack usage. Forward-reference chains at 64 succeed; 65 or +more fail with an explicit nesting-limit diagnostic. +Expression nesting is also bounded to 64 (including call parentheses), sharing the native +recursion budget with active dependency compilation. No expression has more +than 64 branches, and expression alternatives share the 64-complete-plan limit. +Recursive calls, duplicate definitions/ports, missing members, incompatible +interfaces and invalid mappings fail during construction. An unresolved +unqualified call must name a registered standard ONNX operator; use an +explicit domain for other operators. This disambiguates misspelled pattern +calls from ONNX calls without changing standalone `Pattern` parsing. + +Each complete candidate passes the existing structural, type/constant, +deletion-boundary and layer checks. In engine use, provider and opset checks +also apply before C++ acceptance. If the checker rejects it, the next concrete +candidate is tried with fresh captures and computed attributes. Earlier rules +retain priority over every alternative of a later rule. Search is bounded +enumeration, not a recursive producer/permutation matcher. Replacement or callback +errors propagate; they never cause a retry. Match-only statistics count passing +concrete plans, so multiple alternatives of one logical rule can contribute matches. +See [the rewriter engine](#the-rewriter-engine-applying-collections-of-rules) +for rule application policy. + +Graph alternatives and type alternatives have different retry behavior. +For each concrete graph candidate, the bounded annotation search selects the +first globally consistent type/shape combination. Checker rejection does not +try another annotation combination for that candidate; it advances to the +next concrete graph candidate. Neither search permutes operands or explores +different producers. + +### Matching alternatives at one root + +Disjunction can produce multiple matches for the same pattern and root. +For example, when an Identity feeds a Sigmoid, both branches below can match: +the first binds `x` to the Identity input, while the second binds `x` to the +Sigmoid input. + +```cpp +PatternLibrary library(R"( + pattern E(x) { + value = Identity(x) | x + return Sigmoid(value) + } +)"); +Pattern pattern = library.GetPattern("E"); +size_t count = 0; +ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& match) { + ++count; + const NodeArg* input = match.valueOf("x"); + // Inspect input; do not mutate graph. + (void)input; + return true; // Continue to the next concrete candidate. +})); +``` + +`MatchAt` returns only the first successful concrete plan in source order. +It does not run a checker or search another candidate after a later checker +rejection. `VisitMatches` anchors at the same root in an **already resolved** +graph and visits successful expanded graph plans in source order, each with +fresh bindings. Both APIs use the standalone eligibility checks described in +[Patterns](#patterns), without rule opset-range or provider filtering. +`true` means **continue**, `false` means **stop**; these are traversal controls, +not check-and-bind acceptance. Ordinary nonmatches/exhaustion return OK without +invoking the visitor. Unresolved graphs, invalid root indices and empty visitors +return errors; malformed/non-rooted pattern configuration throws. Visitor +exceptions propagate, never silently retry. The callback may copy a candidate +but must not mutate the graph. + +To try the next concrete plan after checker rejection, adapt the checker result +explicitly (using a checker appropriate to the supplied pattern). The +[check-and-bind section](#the-check-and-bind-function) describes checkers in detail: + +```cpp +std::optional accepted; +ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& candidate) { + if (!check_and_bind(graph, candidate)) return true; // Reject: try next. + accepted = candidate; // Keeps prepared attributes plus selected-plan ownership. + return false; // Accept: stop enumeration. +})); +``` + +Do **not** directly return the checker boolean from a visitor: its traversal +meaning is opposite. The first consistent annotation combination is selected +within each concrete graph plan; rejecting that candidate does not search +other annotation combinations. + +### Optional immutable tensor constraints and constant readers + +The prefix qualifier `const`, as in `top_k : const int64[1]`, requires an +immutable constant initializer in addition to the tensor annotation. +It does not decode contents, bind a scalar, or require local payload storage. +Runtime inputs and overridable initializers do not satisfy it. +It can qualify individual tensor alternatives, including +`const float[1] | float[] | None`. It cannot qualify an attribute type. + +`ReadConstantTensor(graph, value, data, available)` has typed overloads for +`InlinedVector` and `InlinedVector`. A missing value, +nonconstant/overridable initializer, wrong element type, or external payload +returns OK with `available = false`, leaving `data` unchanged. A malformed +available payload returns an error, leaving both outputs unchanged. +Success decodes raw or typed protobuf storage and sets `available = true`. +C++ checks rank, supported values and all dependent relationships; decoding +does not implicitly enforce scalar shape. Propagate reader errors, for +example with `ORT_THROW_IF_ERROR` in a bool checker, rather than turning +corruption into an ordinary alternative rejection. + +There are no pattern-valued parameters, construction-time scalar parameters, +`where` clauses, tensor-content capture declarations, or C++ code generation. +Semantic acceptance, tensor decoding, arithmetic and kernel support checks +remain C++ responsibilities. + +## Replacements + +`Replacement` owns an immutable parsed graph-construction recipe. Copying shares +the recipe, not per-call tensors or prepared attributes. Its syntax uses the +same operators, literals, explicit optional slots, assignments and positional +returns illustrated above; composition/expression alternatives are pattern-only. +`Plan()` and `Outputs()` expose read-only recipe inspection. + +The standalone `Instantiate` API executes a recipe without a pattern, Match, +rule, or engine: + +```cpp +Replacement recipe("(x, perm) => Transpose(x, perm = perm)"); +ONNX_NAMESPACE::AttributeProto permutation; +permutation.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); +permutation.add_ints(1); +permutation.add_ints(0); +Replacement::Instantiation built; +ORT_RETURN_IF_ERROR(recipe.Instantiate( + graph, {{"x", input}}, {{"perm", permutation}}, built)); +// built.nodes are graph-owned NodeIndices; built.outputs follow return order. +graph.SetOutputs({built.outputs[0]}); +ORT_RETURN_IF_ERROR(graph.Resolve()); +``` + +Named tensors and attributes are separate explicit binding maps. Every tensor +pointer must be the canonical `NodeArg` in this graph; cross-graph pointers and +empty-name `NodeArg`s are errors. A null tensor binding means absent input, not +an unknown name: schema validation accepts it only in optional slots. Missing +referenced names and tensor/attribute name collisions return errors; extra +bindings are allowed for reusable environments. Replacements compiled internally +against pattern integer handles cannot be instantiated independently. +The final defaulted `absent_attributes` argument explicitly lists absent +attribute bindings without changing the existing `NodeAttributes` map API: + +```cpp +Replacement recipe("(x, perm) => Transpose(x, perm = perm)"); +Replacement::Instantiation built; +ORT_RETURN_IF_ERROR(recipe.Instantiate( + graph, {{"x", input}}, {}, built, nullptr, {"perm"})); +``` + +A name must not occur in both the absent set and the present attributes or +tensor bindings. Extra absent names are allowed like extra present bindings, +but a referenced name missing from both attribute environments remains an error. +Explicit absence omits the target keyword; target defaults/required validation +still apply. This differs from a null tensor input, which occupies an ONNX slot. + +All output names are fresh and generated by the API. Existing operators, edges, +graph inputs and outputs are not replaced or rewired. Result order matches recipe +return order, including returned intermediate outputs. `Instantiate` adds a +fragment to a resolved **or under-construction** graph. On success it deliberately +leaves the graph needing resolution: producer/consumer indexes, edges and inferred +type/shape metadata must not be consumed until the caller chooses graph outputs +and invokes `Graph::Resolve`. It does not infer types/shapes or run full Resolve +itself. Schema verification uses the graph's imported domain/opset versions, +requires those schemas and domains to exist, validates arities and attributes, +and populates default attributes. It is not a proof that final shape/type +inference will succeed; later Resolve errors are not rolled back. + +An optional `annotation_source` must be a live node in this graph. Created nodes +inherit its layering annotation and execution provider; without one they have +default layer context and are unassigned. The API does not accept user-chosen +destination outputs or reuse an existing producer's output. + +Schema/staging failure removes staged nodes and literal initializers and leaves +the caller's `Instantiation` unchanged. Existing graph topology is untouched. +Generated names, unused interned `NodeArg`s and graph-dirty flags may remain: +this is **not** byte-for-byte graph rollback. Returned node indices and tensors +remain graph-owned and invalid after their deletion. The engine reuses the same +staging/schema-validation code, but privately maps returned values to matched +outputs, commits deletion/rewiring and checks prospective cycles. + +In a rule, replacement tensor references must be pattern inputs, not matched +intermediates about to be deleted. A check-and-bind function can prepare attributes; +a replacement callback can select among pre-parsed recipes. Both features are +host C++ choices, not ONNX control flow. + +## The Match object + +`Match` is candidate-local evidence, not an independent graph or mutation API. +`nodes` follows pattern operator order (last is the root); `values` is indexed by +that selected pattern's handles. Use named/scoped lookups instead of transferring +integer handles between definitions. `valueOf`, `optionalValueOf`, `attributeOf`, +`optionalAttributeOf`, `dimensionOf`, `elementTypeOf` and `BindingView` provide read-only access to captures. +Unknown names are errors; explicitly declared optional unbound tensors are null. + +`update(name, AttributeProto)` prepares replacement attributes under fresh names +without changing the graph. It cannot shadow a captured attribute, tensor, +dimension, element-type or instance binding. `attributeOf` reads either captured +or prepared attributes; `ClearUpdates` drops only prepared attributes and the +rejection reason, leaving captures intact. `Reject(reason)` records diagnostic +text and returns false; it does not +throw, mutate the graph, stop `VisitMatches` enumeration, or veto a candidate by +itself if its caller ignores the boolean. + +Candidates from `MatchAt`, `VisitMatches` and engine callbacks retain their selected compiled +pattern. Copying a Match copies its binding/prepared/rejection state and shares +that immutable pattern ownership, so destroying the original Pattern/library or +Match does not invalidate the copied plan. Nodes, tensors and captured graph data +still borrow the graph: do not use a Match after graph mutation/deletion. A +`BindingView` additionally borrows its parent Match and cannot outlive or survive +moving that object. Engine callbacks should retain neither pointers nor views. + +The public low-level `Bind*` methods, raw vectors and `pattern` pointer are matcher +plumbing, not supported user construction/forgery APIs. Obtain evidence from +`MatchAt`, `VisitMatches` or engine callbacks; author checkers using lookups, `update` and +`Reject`. Replacement instantiation deliberately takes explicit bindings rather +than trusting a manually fabricated or stale Match. + +## The check-and-bind function + +`Rule::CheckAndBind` is `std::function`. +It checks semantic applicability and can prepare match-local bindings before +returning `true`. Return `false` (or `match.Reject(reason)`) to reject the rule. +Declare the attribute binding in the replacement's input list, set it with +`match.update`, and reference it as a keyword argument in a text replacement. +No replacement lambda is needed just to prepare attributes. + +It is an ordinary C++ function and can be unit-tested by calling it on candidates +obtained from `MatchAt`, without applying any rule: + +```cpp +std::optional candidate; +ORT_RETURN_IF_ERROR(pattern.MatchAt(graph, root_index, candidate)); +if (candidate && ComposeTransposePermutations(graph, *candidate)) { + const ONNX_NAMESPACE::AttributeProto& composed = candidate->attributeOf("composed"); + // Inspect the checked match and prepared permutation. + (void)composed; +} +``` + +For retry across disjunctions, see the +[ordered visitor example](#matching-alternatives-at-one-root). +In engine use, checker false discards candidate-local +updates and tries the next eligible rule/concrete plan. Exceptions and replacement +errors propagate; they are not ordinary rejection. The engine supplies provider +and opset eligibility before invoking the checker. The checker supplies semantic +equivalence/capability policy that pattern syntax alone cannot prove. + +The following rank-two specialization of the `TransposeMatMul` example rule +checks a node attribute rather than inspecting types in C++. The pattern's +annotations already require float32 matrices with compatible dimensions. +The callback accepts only `Transpose` with an explicit `perm = [1, 0]`, then +binds a new integer attribute variable, `rhs_transpose`, for the text replacement. +The resulting `FusedMatMul` uses `transB = 1` to perform the same transpose internally. +The Microsoft-domain opset must already be imported. + +```cpp +#include "core/optimizer/pattern_rewriter.h" + +using namespace onnxruntime; +using namespace onnxruntime::rewrite; + +Rule MakeTransposeMatMulRule() { + const auto* pattern = R"( + (x : float[M, N], y : float[K, N]) => { + transposed = Transpose(y) + return MatMul(x, transposed) + } + )"; + auto check_and_bind = [](const Graph&, Match& match) { + const auto& attributes = match.nodes[0]->GetAttributes(); + auto perm = attributes.find("perm"); + if (perm == attributes.end() || + perm->second.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_INTS || + perm->second.ints_size() != 2 || + perm->second.ints(0) != 1 || perm->second.ints(1) != 0) + return false; + ONNX_NAMESPACE::AttributeProto trans_b; + trans_b.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + trans_b.set_i(1); + match.update("rhs_transpose", std::move(trans_b)); + return true; + }; + const auto* replacement = "(x, y, rhs_transpose) => com.microsoft.FusedMatMul(x, y, transB = rhs_transpose)"; + return Rule{"TransposeMatMul", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; +} +``` + +`match.nodes[0]` is the first pattern operator, `Transpose`, not the root +`MatMul`. `rhs_transpose` is a new match-local name, not a graph value or an +attribute captured by the pattern. The replacement declares it in its header +and uses it in keyword position to set `transB`. No replacement lambda is needed. +This specialization deliberately rejects an omitted `perm`, even though the +rank-two default would also transpose the matrix. + +Only the selected rule's replacement runs, with the bindings prepared by its successful check-and-bind +callback. Replacement callbacks remain supported for choosing a parsed recipe; +they receive `const Match&` and cannot update it. Parse any such recipes outside +the callback to avoid per-match parsing. + +Each rule gets a separate copy of the structural match. Rejected rules and +other candidates cannot leak prepared attributes, captured bindings, or rejection +reasons. Match-only mode also runs check-and-bind callbacks, but never replacements +or graph mutations. Neither callback may mutate the graph or retain match +pointers; callbacks should have no externally visible side effects. + +`update` is non-const and accepts a concrete `AttributeProto`, including +list-valued and tensor-valued attributes. A tensor binding can be used as a +`ConstantOfShape` node's `value` attribute; it is not a new arbitrary tensor-input binding API. +It cannot overwrite a named pattern value or store a referenced/undefined +attribute. Prepared and captured attributes remain available to the replacement. +`attributeOf` retrieves either kind through a read-only lookup and errors if it is missing. A replacement header +name cannot be used as both a tensor and an attribute. Attribute-only names need +not appear in the pattern, but must be supplied by the callback. Reusing a parsed +plan does not reuse a previous match's attributes. `ClearUpdates` and `Reject` +are also non-const; there is no mutable storage behind a const match. + +Replacement tensor inputs may reference only pattern input variables, literal +constants, or earlier outputs in the replacement itself. A matched intermediate +is about to be deleted and cannot be a replacement input. Replacement sequences +must be connected to their final operator. + +## Constructing rules + +The experimental `onnxruntime::rewrite::PatternRewriter` replaces connected +operator patterns with semantically equivalent graphs. Include +`core/optimizer/pattern_rewriter.h`. Construct `Pattern` and `Replacement` from +text; programmatic builders and mutable replacement recipes are not public APIs. +Check-and-bind callbacks and replacement callbacks remain ordinary C++ functions. +An omitted check-and-bind callback (`{}`) accepts after built-in checks and +contributes zero to `check_and_bind_calls`. Use that placeholder when supplying +the fifth constructor argument without a checker. + +A `Rule` takes `(name, pattern, replacement, check_and_bind, opset_requirements)`. +The pattern may be text or a parsed `Pattern`; the replacement may be text, +a parsed `Replacement`, or a callback returning one. Text accepts string literals, +`std::string`, and `std::string_view` and is parsed once during rule construction; +syntax errors are reported immediately. Parsed definitions own their data, so +the original text need not outlive the rule. Rules remain copyable and movable. +The matcher establishes structural and deletion safety. **The rule author must establish +semantic equivalence, supported types, and kernel compatibility.** +The opset list defaults to empty, not an inferred version range: missing required +domain ranges are construction errors. The example's supported ranges are the +author's contract, not automatic proof about every schema version. + +## When a match succeeds or fails + +Before calling a rule's check-and-bind callback, the matcher checks: + +- Exact operator name/domain, ordered output counts and input slots, and repeated + binding identity. Omitted trailing inputs match only optional variables or + explicit `None`; variadic patterns are not supported. +- Literal and listed attribute constraints. It does not otherwise infer semantic + equivalence or check arbitrary shapes/types for the author. +- Distinct matched nodes with the same layer annotation. Matched operators cannot + own subgraphs or have implicit inputs. +- Pattern input variables are outside the matched deletion set. +- Every unreturned output stays internal. A graph output, an external consumer, or + an implicit capture by an external subgraph of an unreturned value prevents the rewrite. +- Execution-provider eligibility as described below. + +Returned values may have graph outputs or external consumers because their +identities are preserved. Even match-only mode applies deletion-boundary checks: +it reports safe rewrite eligibility, not arbitrary structural occurrences. + +### Types in the supplied examples + +The test-local `MakeExampleRules()` in +`onnxruntime/test/optimizer/pattern_rewriter_test.cc` constructs six examples +when contrib operators are enabled. These illustrate rule authoring and are +not a public factory or part of the default optimizer pipeline. Production +callers construct their own `Rules` and pass them to `PatternRewriter`. + +| Rule | Explicit float32 checks | +|---|---| +| HardSwish and HardSigmoidMul | `x` | +| MatMulDiv and FusedMatMulDiv | `x`, `y`; separate finite scalar initializer check on `divisor` | +| TransposeMatMul | `x`, `y` | +| RmsNormalization | `x`, `scale`; separate positive finite scalar initializer check on `epsilon` | + +There is no match-wide float32 restriction. Independently typed operands such +as shape/index tensors or a Pow exponent require their own semantic constraints. +The supplied RMS pattern specifically matches `Pow(x, 2.0)`; the float32 exponent +is a literal constraint, not a blanket callback check. + +The RMS example also requires last-axis reduction, keepdims, known positive +width and a matching one-dimensional scale, with pre-opset-18 attribute axes. +TransposeMatMul requires a known rank of at least two and exactly the last-two-axis +swap. Division examples require a finite scaling result. Consult the rule +definitions and correctness tests before extending their supported ranges. +Contrib-disabled tests construct only the two HardSwish examples. + +## Opset and execution-provider applicability + +Each `OpsetRequirement` is `{domain, min_version, max_version}`, with positive, +inclusive bounds. Declare every domain used by a pattern or replacement; +additional requirements are allowed. `""` and `"ai.onnx"` denote the same domain. +Duplicate domains, reversed ranges, and undeclared pattern domains are errors. + +All requirements are evaluated against each graph's `DomainToVersionMap()`, +including independently for subgraphs. Missing imports or out-of-range versions +make the rule inapplicable **before its check-and-bind callback runs**. The engine does not add +imports, compare individual node schema since-versions, or rewrite model-local +function bodies. Authors must establish correctness across their declared range. + +All matched nodes must have the same execution-provider assignment. Unassigned +nodes are eligible; assigned nodes must belong to +`Options::compatible_execution_providers`, which defaults to CPU only. +Mixed assignments are rejected. This policy applies before the rule callback and +also in match-only mode. To opt into another provider: + +```cpp +Options options; +options.compatible_execution_providers.insert(kCudaExecutionProvider); +PatternRewriter rewriter(Rules{MakeComposeTransposeRule()}, options); +``` + +Allowing a provider is the caller's assertion of kernel and precision +compatibility, not automatic kernel discovery. Clearing the allowed set permits +only unassigned nodes. Replacements inherit the root's provider and layer +annotation. + +## The rewriter engine: applying collections of rules + +Register the transformer through the normal session pipeline: + +```cpp +auto transformer = std::make_unique(Rules{MakeComposeTransposeRule()}); +ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer( + std::move(transformer), TransformerLevel::Level1)); +``` + +For one already-resolved graph, use `Run` and explicitly resolve afterwards: + +```cpp +PatternRewriter rewriter(Rules{MakeComposeTransposeRule()}); +Statistics statistics; +ORT_RETURN_IF_ERROR(rewriter.Run(graph, statistics)); +ORT_RETURN_IF_ERROR(graph.Resolve()); +``` + +`Run` handles one graph. The transformer's `Apply` path recurses into subgraphs. +After a successful modified `ApplyImpl`, the normal `GraphTransformer::Apply` +wrapper invokes `Graph::Resolve` (in non-minimal builds). It does not Resolve +before entering the transformer or after an error. Callers must not concurrently mutate a +graph. Options and compiled rules are immutable; match state and statistics are +per call. + +### Resolved metadata and mutation boundaries + +Call `Graph::Resolve` after constructing a source graph and before `Run` or +standalone pattern matching. These APIs reject `GraphResolveNeeded()`; they do +not silently repair input graphs. Resolve establishes schemas/default attributes, +edges and inferred type/shape metadata. Rules with annotations or C++ shape +checks may still reject when metadata remains unknown: Resolve cannot invent +missing dimension evidence. The Transpose checker needs known rank; repeated +dimension annotations require the selected known-equal evidence described above. + +`Run` does **not** invoke full Resolve between individual rewrites or batches, +or upon returning. It maintains edges and producer/consumer indexes for continued +matching and stages schemas/default attributes for new nodes, but does not run +general type/shape inference. Fresh replacement intermediates may lack metadata; +metadata-dependent later rules may only become applicable after the caller's +explicit Resolve and another Run (or a later transformer-manager pass). Never +describe schema verification as full inference. Existing output `NodeArg`s retain +their metadata and identities; checker authors remain responsible for semantic +and shape-preserving correctness. + +One replacement validates/stages before deleting, and ordinary staging/cycle +failures leave matched topology intact. A whole `Run` is **not transactional**: +an error in a later candidate does not undo earlier successful rewrites, and +`GraphTransformer::Apply` does not Resolve an error result. Inspect/handle the +status before deciding how to use that graph. Standalone `Instantiate` has the +separate fragment-only construction contract described above. + +Rules have deterministic **registration-order priority**, including across +different structural groups. At each eligible root, rewriting chooses the first +passing check-and-bind callback. With `Options{.match_only = true}`, all passing callbacks are counted +and replacements are not called. `indexed = false` selects the reference matching +mode with the same eligibility and priority semantics. + +Every pattern and replacement must contain at least one node, and their +returned-value counts must agree. Single-node patterns, same-size replacements, +and expansion are supported. Zero-node replacements and returning leaf inputs +remain unsupported; identity elimination is a separate transformation. + +`Options::max_iterations` is a positive upper bound, defaulting to 10. +Each iteration processes a fixed candidate batch, initially all live nodes in +topological order. Replacements and affected existing nodes are queued for the +next batch, never appended to the current batch. An existing later candidate +can still match an earlier replacement in the same batch. Processing stops +when a batch performs no rewrites, there are no pending candidates, or the bound +is reached. Match-only mode processes exactly one batch without mutation. + +```cpp +Options options; +options.max_iterations = 3; +PatternRewriter rewriter(rules, options); +``` + +The iteration cap bounds repeated self-rewrites and rule cycles, even when +node count grows. It does not prove convergence or semantic equivalence. +`Statistics::iterations` counts processed batches; +`iteration_limit_reached` is true if the final allowed batch rewrote nodes and +left candidates pending. Reaching the limit returns success with the current +valid rewrite result, not an error; the flag does not assert another match exists. +Calling `Run` again starts a fresh budget. External transformer-manager passes +also have their own budgets. Avoid expansion rules that multiply graph size +unnecessarily: the iteration limit is not a memory bound or a bound on callbacks. + +After a rewrite, the engine updates edges and producer/consumer indexes, +notifies node replacement, and requeues replacement nodes and affected +descendants through a bounded invalidation walk. This permits newly enabled +matches in subsequent batches without retaining pointers to deleted nodes. + +Malformed definitions throw `OnnxRuntimeException` with parse-location +diagnostics; exception-disabled builds follow ORT's abort behavior. Invalid rule +configuration raises an explicit error when the rewriter is constructed. +Direct parsed replacements are checked against their rule during construction; +callback-generated replacements also undergo runtime validation. Returning false from check-and-bind +is an ordinary nonmatch. + +Replacement schemas and attributes are validated before matched nodes are +deleted. Failed staging rolls back the staged nodes and literal initializers; +there is no whole-call transaction undoing earlier successful rewrites. +Errors propagate as statuses or ORT exceptions rather than being treated as +nonmatches. Final graph resolution checks types and inferred shapes; a rule +author must not rely on schema staging alone to prove type correctness. + +`Statistics` exposes live roots visited (including revisits), enumerated +candidates, structural checks, `check_and_bind_calls`, successful matches, and completed +rewrites. Match-only counts all passing check-and-bind callbacks; rewrite mode counts selected +winners. These counters do not alter rule semantics. + +### Opt-in failure diagnostics + +Set `Options::on_failure` to receive structured `MatchFailure` records containing +the rule name, candidate root index, `FailureStage`, and a reason string: + +```cpp +Options options; +options.on_failure = [](const MatchFailure& failure) { + std::cerr << failure.rule << ": " << failure.reason << '\n'; +}; +``` + +Stages distinguish opset filtering, structure, type/shape constraints, constants, attributes, deletion +boundaries, execution providers, check-and-bind callbacks, and replacement validation. An opset +rejection is reported once per rule/graph with root index +`std::numeric_limits::max()` because no candidate was inspected. +Other failures are reported for attempted candidates; indexed mode skips +unrelated root operators, so diagnostic counts need not match reference mode. +Shared structural failures are attributed to each affected rule. + +A check-and-bind callback can use `return match.Reject("unsupported element type");` instead of +bare `false` to supply a reason. This changes only diagnostics, not eligibility. +Records own their strings and contain no graph pointers. Callbacks run +synchronously, must not mutate the graph or reenter the rewriter, and may copy +records for later inspection. No diagnostic records are constructed when +diagnostics are disabled; `Reject` still stores an explicitly supplied callback +message. Replacement validation status errors are reported and +still propagated; exceptions from user callbacks or named lookups still propagate +normally and are not converted to nonmatches. + +## Component tests and rule integration tests + +`onnxruntime/test/optimizer/pattern_rewriter_components_test.cc` separates +`PatternComponentTests`, `ReplacementComponentTests`, `MatchComponentTests` +and `CheckAndBindComponentTests`. These invoke `MatchAt`, `VisitMatches`, `Instantiate`, +Match lookups/updates and an ordinary checker directly, without disguising +an engine run as a standalone component operation. The same file follows them +with the motivating Transpose rule's numerical/boundary integration test. +The existing `pattern_rewriter_test.cc` retains the broader parser, rule, +composition, replacement-rollback and engine regression coverage. + +From the build output directory, run all of them together: + +```text +onnxruntime_test_all --gtest_filter=PatternComponentTests.*:ReplacementComponentTests.*:MatchComponentTests.*:CheckAndBindComponentTests.*:PatternRewriterTests.* +``` + +The component split makes contracts independently testable; it does not bypass +the engine's integration tests or duplicate a second matcher/recipe executor. + +## Limits and unsupported forms + +Concrete entry plans are limited to 64 nodes, 256 values, and 64 inputs/outputs per +operator. They are connected acyclic graphs rooted at the final operator. +Named pattern alternatives are expanded as described below. There is no inline +operator/subgraph alternation, implicit commutation, variadic matching, +general tensor synthesis, domain-import editing, or model-local function +rewriting. Control-flow operators are not matched, although their subgraphs can +be rewritten through `Apply`. + +Tests in `onnxruntime/test/optimizer/pattern_rewriter_test.cc` cover parsing, +check-and-bind callbacks, boundaries, multi-output rewrites, rule priority, and numerical CPU +equivalence, including small decoder prefill/decode graphs. Add structural and +executed equivalence tests when writing or extending a rule. + +## Matching algorithm outline + +A single match is a **backward traversal from a candidate sink node**, without +structural backtracking. Starting with the pattern's final operator bound to that model +node, the matcher visits pattern operators in reverse definition order and +follows bound values to their producers through a value-to-producer dictionary. +Each input/output slot binds to one exact model value; repeated bindings must +agree. Independent DAG branches need not have the same node ordering in the +model, but operand order and data dependencies must match. A mismatch rejects +the candidate rather than trying other producers or operand permutations. +Constants, attributes, deletion boundaries, and EP eligibility are checked +before the rule's check-and-bind callback. Type unions use the separate bounded +annotation-choice search described above. + +Multiple rules are handled **node-first**, inside an outer bounded retry loop: + +```text +filter rules by this graph's opset imports +for each iteration, up to max_iterations (default 10): + for each still-live node in the current candidate batch: + find candidate rule groups by the node's (domain, operator name) + match their patterns and evaluate eligible check-and-bind callbacks + apply the passing rule with earliest registration priority + queue replacement nodes and affected nodes for the next batch + stop if no rewrite occurred or no candidates remain +``` + +The first batch contains all model nodes in topological order; later batches +contain only queued candidates. New nodes and revisits wait for the next batch, +although a later node already in the current batch sees earlier graph changes. +The retry loop is therefore **outside** the node loop, not a separate ten-retry +loop at each node. + +By default, a hash map indexed by `(domain, operator name)` selects candidate +groups using the model node as the prospective pattern root. Structurally +identical patterns can share a traversal, followed by their rule-specific check-and-bind callbacks; +patterns with attribute predicates or type/shape annotations are kept separate. The earliest passing +rule wins even across groups, and lower-priority work is skipped when possible. +Structural results are cached per group and candidate node, but callbacks +run in global registration order, not group order. Rewrite mode stops at +the first accepted rule without invoking lower-priority callbacks; match-only +mode visits all matching rules in that same order. Group signatures use +locale-independent serialization. +With `indexed = false`, the reference mode instead tries all opset-applicable +rules in registration order at each node. Match-only mode processes one batch, +counts all passing rules, and never applies replacements. + +## Instance binding views and handwritten C++ structs + +`match.instance("product")` returns a read-only `BindingView` for the activation +call in the [library example](#composable-pattern-libraries). +`view.instance("scaled")` selects its nested expression +occurrence. Its `valueOf`, `optionalValueOf`, `attributeOf`, `optionalAttributeOf`, `dimensionOf`, +and `elementTypeOf` use names relative to that occurrence. `definition()` identifies the called definition, +`interfaceName()` its contract, and `alternative()` the selected immediate member +(empty for a body call). `alternativeView()` enters a selected call/member's metadata; +calling it on a body or operator/value-alias leaf raises an error. Repeated calls traverse nested +alternatives/aliases without losing intermediate tags. This navigation does +not add tensor namespaces: all selected member locals use the same occurrence +path, and `instance("child")` addresses a nested call inside the selected body. +Entry metadata is available through `match.bindings()`. +Fully qualified lookups such as +`match.valueOf("product.alpha")` are equivalent. Missing instances and +unavailable member-local names raise errors; they are not nullable optional +fields. Explicitly declared optional tensors are the exception described above. +Only the selected member's concrete bindings exist. + +Views borrow their parent Match and graph. They cannot outlive or survive moving +that Match, and must not be retained across graph mutation; obtain a new view +from a copied Match instead. The following manual structs show what +a future generator could emit; **this API does not generate C++**: + +```cpp +struct GatedView { + const NodeArg* g; + const NodeArg* u; + const NodeArg* product; + explicit GatedView(BindingView v) + : g(v.valueOf("g")), u(v.valueOf("u")), product(v.valueOf("product")) {} +}; +struct PlainView { + GatedView ports; + const NodeArg* sigmoid; + explicit PlainView(BindingView v) : ports(v), sigmoid(v.valueOf("s")) {} +}; +struct ScaledView { + GatedView ports; + const NodeArg* alpha; + explicit ScaledView(BindingView v) : ports(v), alpha(v.optionalValueOf("alpha")) { + ORT_ENFORCE(alpha, "Selected scaling requires alpha."); + } +}; +struct ActivationView { + GatedView ports; + std::variant alternative; + static std::variant Select(BindingView v) { + const auto scaling = v.instance("scaled"); + if (scaling.alternative() == "g") return PlainView(v); + ORT_ENFORCE(scaling.alternative() == "Mul", "Unexpected scaling alternative."); + return ScaledView(v); + } + explicit ActivationView(BindingView v) : ports(v), alternative(Select(v)) {} +}; +struct FusionView { + ActivationView activation; + explicit FusionView(const Match& m) : activation(m.instance("product")) {} +}; + +bool CheckGatingShapes(const GatedView& v) { + if (!v.g->Shape() || !v.u->Shape() || + v.g->Shape()->dim_size() != 2 || v.u->Shape()->dim_size() != 2) + return false; + for (int i = 0; i < 2; ++i) { + const auto& g = v.g->Shape()->dim(i); + const auto& u = v.u->Shape()->dim(i); + if (!g.has_dim_value() || !u.has_dim_value() || + g.dim_value() <= 0 || g.dim_value() != u.dim_value()) + return false; + } + return true; +} + +Rule MakePatternFusion(const PatternLibrary& library) { + auto check = [](const Graph& graph, Match& match) { + FusionView view(match); + if (!CheckGatingShapes(view.activation.ports)) + return match.Reject("Unsupported gating shapes."); + float alpha = 1.0f; + if (const auto* scaled = std::get_if(&view.activation.alternative)) { + const auto* shape = scaled->alpha->Shape(); + if (!shape || !(shape->dim_size() == 0 || + (shape->dim_size() == 1 && shape->dim(0).has_dim_value() && + shape->dim(0).dim_value() == 1))) + return match.Reject("Alpha must have scalar or [1] shape."); + InlinedVector data; + bool available; + ORT_THROW_IF_ERROR(ReadConstantTensor(graph, scaled->alpha, data, available)); + if (!available || data.size() != 1 || !std::isfinite(data[0]) || data[0] <= 0) + return match.Reject("Unsupported alpha."); + alpha = data[0]; + } + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + attribute.set_f(alpha); + match.update("activation_alpha", std::move(attribute)); + return true; + }; + Replacement recipe(R"( + (g, u, activation_alpha) => { + a = com.microsoft.QuickGelu(g, alpha = activation_alpha) + return Mul(a, u) + } + )"); + return Rule{"PatternActivation", library.GetPattern("Fuse"), recipe, check, + {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; +} +``` + +Include `` for `std::isfinite`, in addition to the rewriter header. +This synthetic fusion recognizes `g * sigmoid(alpha*g) * u`, removes its +trailing Identity, and selects the QuickGelu alpha in C++. The tests execute both source and replacement with CPU EP for plain, +scalar-alpha and one-element-alpha cases, including the factored positional +form above. This is a small floating-point example, not a +general MoE/QMoE lowering or a bitwise-equivalence guarantee. + +Replacement callbacks may likewise select among pre-parsed recipes in C++; +that choice is host construction, **not** an ONNX If. Pre-parsed recipes can +reference scoped tensor leaves, for example `(x, p.alpha) => Mul(x, p.alpha)`. +Such names must exist in every concrete plan when the recipe is attached +directly to a rule; use a callback to choose a member-specific recipe. +Matched internal outputs remain forbidden replacement inputs, including +when accessed through an interface alias. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc new file mode 100644 index 0000000000000..8189169c51916 --- /dev/null +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -0,0 +1,2579 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "core/optimizer/pattern_rewriter.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/common/parse_string.h" +#include "core/graph/graph_utils.h" +#include "core/framework/tensorprotoutils.h" +#include "core/optimizer/initializer.h" +#include "onnx/checker.h" + +namespace onnxruntime { +namespace rewrite { +namespace { + +std::string Key(const std::string& op, const std::string& domain) { return domain + '\0' + op; } + +std::string Signature(const Pattern& pattern) { + std::ostringstream stream; + stream.imbue(std::locale::classic()); + stream << pattern.ValueCount() << ':'; + for (Value output : pattern.Outputs()) stream << 'r' << output.handle << ','; + for (int id = 0; id < static_cast(pattern.ValueCount()); ++id) + if (pattern.IsOptional(id)) stream << '?' << id << ','; + stream.precision(std::numeric_limits::max_digits10); + for (int id = 0; id < static_cast(pattern.ValueCount()); ++id) { + auto it = pattern.Constants().find(id); + if (it == pattern.Constants().end()) continue; + stream << 'c' << id << ':' << it->second.index() << ':'; + std::visit([&](const auto& value) { + using T = std::decay_t; + if constexpr (std::is_arithmetic_v) { + stream << value; + } else { + stream << value.size() << '['; + for (auto element : value) stream << element << ','; + stream << ']'; + } + }, + it->second); + stream << ';'; + } + for (const auto& node : pattern.Nodes()) { + stream << node.op.domain.size() << ':' << node.op.domain << node.op.name.size() << ':' << node.op.name; + stream << '/'; + for (Value output : node.outputs) stream << output << ','; + stream << '/'; + for (Value input : node.inputs) stream << ',' << input; + stream << ';'; + } + return stream.str(); +} + +using Producers = InlinedHashMap; + +bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const TypeConstraint& constraint, Match& match, + std::string& reason) { + if (pattern.IsAttribute(handle)) { + const auto* attribute = match.optionalAttributeOf(constraint.value_name); + return constraint.absent ? !attribute : attribute && attribute->type() == constraint.attribute_type; + } + const NodeArg* value = match.values[handle]; + if (constraint.absent) return !value; + if (!value) return false; + if (constraint.immutable && !graph_utils::GetConstantInitializer(graph, value->Name())) return false; + const auto* type = value->TypeAsProto(); + if (!type || !type->has_tensor_type()) return false; + const int32_t element = type->tensor_type().elem_type(); + if (constraint.element_type && constraint.element_type != element) return false; + if (!constraint.element_variable.empty() && !match.BindElementType(constraint.element_variable, element)) return false; + if (constraint.any_rank) return true; + const auto* shape = value->Shape(); + if (!shape || shape->dim_size() < static_cast(constraint.dimensions.size()) || + (!constraint.prefix && shape->dim_size() != static_cast(constraint.dimensions.size()))) + return false; + for (size_t i = 0; i < constraint.dimensions.size(); ++i) { + const auto& expected = constraint.dimensions[i]; + const auto& actual = shape->dim(static_cast(i)); + if (const auto* fixed = std::get_if(&expected)) { + if (!actual.has_dim_value() || actual.dim_value() != *fixed) return false; + } else if (const auto* name = std::get_if(&expected)) { + Dimension binding; + if (actual.has_dim_value()) + binding = actual.dim_value(); + else if (actual.has_dim_param() && !actual.dim_param().empty()) + binding = actual.dim_param(); + if (!match.BindDimension(*name, std::move(binding))) { + reason = "Conflicting bindings for repeated dimension '" + *name + "'."; + return false; + } + } + } + return true; +} + +bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match, std::string& reason) { + const auto& groups = pattern.OrderedTypeConstraints(); + for (const auto& [handle, alternatives] : groups) { + for (const auto& constraint : alternatives) { + if (!constraint.element_variable.empty()) match.BindElementType(constraint.element_variable, 0); + for (const auto& dim : constraint.dimensions) + if (const auto* name = std::get_if(&dim)) match.BindDimension(*name, std::monostate{}); + } + } + InlinedVector> selected; + // Count only a complete selection: a later alternative can turn equality into a singleton capture. + auto verify_dimensions = [&]() { + struct Evidence { + size_t occurrences = 0; + bool unknown = false; + }; + InlinedHashMap dimensions; + for (const auto& [handle, constraint] : selected) { + if (pattern.IsAttribute(handle) || constraint->absent || constraint->any_rank) continue; + const auto& shape = *match.values[handle]->Shape(); + for (size_t i = 0; i < constraint->dimensions.size(); ++i) { + const auto* name = std::get_if(&constraint->dimensions[i]); + if (!name) continue; + const auto& actual = shape.dim(static_cast(i)); + auto& evidence = dimensions[*name]; + ++evidence.occurrences; + evidence.unknown |= !actual.has_dim_value() && (!actual.has_dim_param() || actual.dim_param().empty()); + if (evidence.occurrences > 1 && evidence.unknown) { + reason = "Unknown occurrence of repeated dimension '" + *name + "'; equality cannot be verified."; + return false; + } + } + } + return true; + }; + std::function search = [&](size_t index, Match& state) { + if (index == groups.size()) return verify_dimensions(); + const auto& [handle, alternatives] = groups[index]; + for (const auto& constraint : alternatives) { + Match candidate = state; + if (MatchType(pattern, graph, handle, constraint, candidate, reason)) { + selected.emplace_back(handle, &constraint); + const bool matched = search(index + 1, candidate); + selected.pop_back(); + if (matched) { + state = std::move(candidate); + return true; + } + } + } + return false; + }; + return search(0, match); +} + +bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, + Match& match, FailureStage& stage, std::string& reason) { + auto fail = [&](FailureStage failure_stage, std::string_view message) { + stage = failure_stage; + reason = message; + return false; + }; + const auto& root_op = pattern.Nodes().back().op; + // Cheap root rejection precedes binding allocation in the reference implementation too. + if (root.OpType() != root_op.name || root.Domain() != root_op.domain) { + return fail(FailureStage::Structure, "Root operator/domain mismatch."); + } + match.values.assign(pattern.ValueCount(), nullptr); + match.nodes.assign(pattern.Nodes().size(), nullptr); + match.nodes.back() = &root; + InlinedVector bound(pattern.ValueCount(), false); + InlinedHashSet distinct; + for (size_t i = pattern.Nodes().size(); i-- > 0;) { + const auto& expected = pattern.Nodes()[i]; + const Node* node = match.nodes[i]; + if (!node) { + for (Value output : expected.outputs) { + if (!match.values[output]) continue; + auto it = producers.find(match.values[output]); + if (it != producers.end()) node = graph.GetNode(it->second); + break; + } + } + if (!node) return fail(FailureStage::Structure, "No producer for the expected node."); + if (!distinct.insert(node->Index()).second) + return fail(FailureStage::Structure, "Distinct pattern nodes mapped to one model node."); + if (node->OpType() != expected.op.name || node->Domain() != expected.op.domain) + return fail(FailureStage::Structure, "Producer operator/domain mismatch."); + if (node->InputDefs().size() > expected.inputs.size() || node->OutputDefs().size() != expected.outputs.size()) + return fail(FailureStage::Structure, "Operator input/output arity mismatch."); + if (!node->GetAttributeNameToSubgraphMap().empty() || !node->ImplicitInputDefs().empty()) + return fail(FailureStage::Structure, "Matched operators cannot own subgraphs or implicit inputs."); + if (node->GetLayeringAnnotation() != root.GetLayeringAnnotation()) + return fail(FailureStage::Structure, "Layer annotation mismatch."); + match.nodes[i] = node; + for (const auto& attribute : expected.attributes) { + auto found = node->GetAttributes().find(attribute.name); + if (found == node->GetAttributes().end()) { + if (!attribute.optional || attribute.capture.empty() || !match.BindAbsentAttribute(attribute.capture)) + return fail(FailureStage::Attribute, "Missing attribute or repeated capture presence mismatch."); + continue; + } + if (!attribute.Matches(found->second)) + return fail(FailureStage::Attribute, "Missing attribute or attribute constraint mismatch."); + if (!attribute.capture.empty() && !match.BindAttribute(attribute.capture, found->second)) + return fail(FailureStage::Attribute, "Repeated attribute capture mismatch."); + } + for (size_t j = 0; j < expected.outputs.size(); ++j) { + const NodeArg* output = node->OutputDefs()[j]; + const NodeArg*& binding = match.values[expected.outputs[j]]; + if (!output->Exists() || (binding && binding != output)) + return fail(FailureStage::Structure, "Output slot or repeated value mismatch."); + binding = output; + bound[expected.outputs[j].handle] = true; + } + for (size_t j = 0; j < expected.inputs.size(); ++j) { + const NodeArg* actual = j < node->InputDefs().size() && node->InputDefs()[j]->Exists() + ? node->InputDefs()[j] + : nullptr; + const Value value = expected.inputs[j]; + if (value.handle < 0) { + if (actual) return fail(FailureStage::Structure, "Input slot must be absent."); + continue; + } + const NodeArg*& binding = match.values[value]; + if ((!actual && !pattern.IsOptional(value)) || (bound[value.handle] && binding != actual)) + return fail(FailureStage::Structure, "Required input absent or repeated value mismatch."); + binding = actual; + bound[value.handle] = true; + } + } + // Input variables must be outside the deletion set, even when aliasing another binding. + if (!MatchTypes(pattern, graph, match, reason)) { + stage = FailureStage::TypeShape; + if (reason.empty()) reason = "Tensor type/rank/dimension or captured attribute type mismatch."; + return false; + } + for (int value = 0; value < static_cast(pattern.ValueCount()); ++value) { + if (!pattern.IsInput(value)) continue; + auto it = producers.find(match.values[value]); + if (it != producers.end() && distinct.count(it->second)) + return fail(FailureStage::Boundary, "Pattern input is produced inside the deletion set."); + } + for (const auto& [handle, expected] : pattern.Constants()) { + const auto* tensor = graph_utils::GetConstantInitializer(graph, match.values[handle]->Name()); + const bool vector = expected.index() >= 2; + if (!tensor || tensor->dims_size() != (vector ? 1 : 0) || + tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) + return fail(FailureStage::Constant, "Literal initializer rank, mutability or storage mismatch."); + const bool floating = std::holds_alternative(expected) || + std::holds_alternative>(expected); + if (tensor->data_type() != + (floating ? ONNX_NAMESPACE::TensorProto_DataType_FLOAT : ONNX_NAMESPACE::TensorProto_DataType_INT64)) + return fail(FailureStage::Constant, "Literal initializer dtype mismatch."); + Initializer initializer(graph, *tensor, graph.ModelPath()); + const bool equal = std::visit([&](const auto& value) { + using T = std::decay_t; + if constexpr (std::is_arithmetic_v) { + return initializer.size() == 1 && *initializer.data() == value; + } else { + using Element = typename T::value_type; + return tensor->dims(0) == static_cast(value.size()) && initializer.size() == value.size() && + std::equal(value.begin(), value.end(), initializer.data()); + } + }, + expected); + if (!equal) return fail(FailureStage::Constant, "Literal initializer values mismatch."); + } + InlinedHashSet exposed; + for (Value value : pattern.Outputs()) exposed.insert(match.values[value]); + for (const Node* matched : match.nodes) { + const Node& node = *matched; + for (const NodeArg* output : node.OutputDefs()) { + if (!exposed.count(output) && + std::find(graph.GetOutputs().begin(), graph.GetOutputs().end(), output) != graph.GetOutputs().end()) + return fail(FailureStage::Boundary, "Unreturned value is a graph output."); + } + for (auto edge = node.OutputEdgesBegin(); edge != node.OutputEdgesEnd(); ++edge) { + if (!distinct.count(edge->GetNode().Index()) && !exposed.count(node.OutputDefs()[edge->GetSrcArgIndex()])) + return fail(FailureStage::Boundary, "Unreturned value has an external consumer or capture."); + } + } + return true; +} + +bool CanRewritePartition(const Match& match, const Options& options) { + const auto& provider = match.nodes.back()->GetExecutionProviderType(); + if (!provider.empty() && !options.compatible_execution_providers.count(provider)) return false; + return std::all_of(match.nodes.begin(), match.nodes.end(), + [&](const Node* node) { return node->GetExecutionProviderType() == provider; }); +} + +void RollbackFragment(Graph& graph, InlinedVector& added, const InlinedVector& initializers) { + for (auto it = added.rbegin(); it != added.rend(); ++it) graph.RemoveNode(*it); + for (const auto& name : initializers) graph.RemoveInitializedTensor(name); + added.clear(); +} + +Status StageReplacement(Graph& graph, const Replacement& recipe, gsl::span returned_outputs, + const std::function& resolve_tensor, + const std::function&)>& resolve_attribute, + const Node* annotation_source, InlinedVector& added, + InlinedVector& initializers) { + ORT_RETURN_IF(!recipe.Plan(), "Replacement has no parsed plan."); + const auto& nodes = *recipe.Plan(); + ORT_RETURN_IF(nodes.empty(), "Replacements must contain at least one node."); + ORT_RETURN_IF(recipe.Outputs().size() != returned_outputs.size(), + "Replacement return count must equal the pattern return count."); + InlinedVector> exposed_outputs; + for (const auto& node : nodes) exposed_outputs.emplace_back(node.output_count, nullptr); + for (size_t i = 0; i < recipe.Outputs().size(); ++i) { + const auto& output = recipe.Outputs()[i]; + ORT_RETURN_IF(output.node >= nodes.size() || output.slot >= nodes[output.node].output_count, + "Invalid replacement return slot."); + exposed_outputs[output.node][output.slot] = returned_outputs[i]; + } + InlinedVector> external_inputs; + InlinedVector resolved_attributes; + for (size_t i = 0; i < nodes.size(); ++i) { + const auto& node = nodes[i]; + auto attributes = node.attributes; + for (const auto& [name, binding] : node.attribute_bindings) { + ORT_RETURN_IF(name.empty() || node.attributes.count(name), "Invalid or duplicate replacement attribute: ", name); + std::optional value; + ORT_RETURN_IF_ERROR(resolve_attribute(binding, value)); + if (value) { + value->set_name(name); + attributes.emplace(name, std::move(*value)); + } + } + resolved_attributes.push_back(std::move(attributes)); + const auto& domain = node.domain == kOnnxDomainAlias ? kOnnxDomain : node.domain; + ORT_RETURN_IF(node.op.empty() || node.inputs.size() > 64 || node.output_count == 0 || node.output_count > 64, + "Invalid replacement operator or arity."); + ORT_RETURN_IF(graph.DomainToVersionMap().find(domain) == graph.DomainToVersionMap().end(), + "Replacement domain is not imported: ", domain); + InlinedVector inputs; + for (const auto& input : node.inputs) { + if (std::holds_alternative(input)) { + inputs.push_back(&graph.GetOrCreateNodeArg("", nullptr)); + continue; + } + if (const auto* scalar = std::get_if(&input)) { + ORT_RETURN_IF(std::holds_alternative(scalar->value) && !std::isfinite(std::get(scalar->value)), + "Replacement float constants must be finite."); + inputs.push_back(nullptr); + continue; + } + if (const auto* output = std::get_if(&input)) { + ORT_RETURN_IF(output->node >= i || output->slot >= nodes[output->node].output_count, + "Replacement inputs must reference an earlier output slot."); + inputs.push_back(nullptr); + continue; + } + NodeArg* binding = nullptr; + ORT_RETURN_IF_ERROR(resolve_tensor(input, binding)); + inputs.push_back(binding ? binding : &graph.GetOrCreateNodeArg("", nullptr)); + } + external_inputs.push_back(std::move(inputs)); + } + // The unique sink is independent of the explicitly exposed output slots. + InlinedHashSet reachable{nodes.size() - 1}; + for (size_t i = nodes.size(); i-- > 0;) { + ORT_RETURN_IF(!reachable.count(i), "Disconnected replacement node."); + for (const auto& input : nodes[i].inputs) + if (const auto* output = std::get_if(&input)) reachable.insert(output->node); + } + bool staged = false; + auto cleanup = gsl::finally([&]() { + if (!staged) RollbackFragment(graph, added, initializers); + }); + // Stage and validate every schema before deleting any matched nodes or changing edges. + for (size_t i = 0; i < nodes.size(); ++i) { + auto& inputs = external_inputs[i]; + for (size_t j = 0; j < inputs.size(); ++j) { + if (const auto* output = std::get_if(&nodes[i].inputs[j])) + inputs[j] = graph.GetNode(added[output->node])->MutableOutputDefs()[output->slot]; + if (const auto* scalar = std::get_if(&nodes[i].inputs[j])) { + ONNX_NAMESPACE::TensorProto tensor; + tensor.set_name(graph.GenerateNodeArgName("rewrite_literal")); + std::visit([&](const auto& value) { + using T = std::decay_t; + if constexpr (std::is_arithmetic_v) { + if constexpr (std::is_same_v) { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor.add_float_data(value); + } else { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + tensor.add_int64_data(value); + } + } else { + tensor.add_dims(static_cast(value.size())); + if constexpr (std::is_same_v) { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + for (auto element : value) tensor.add_float_data(element); + } else { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + for (auto element : value) tensor.add_int64_data(element); + } + } + }, + scalar->value); + ONNX_NAMESPACE::TypeProto type; + type.mutable_tensor_type()->set_elem_type(tensor.data_type()); + auto* shape = type.mutable_tensor_type()->mutable_shape(); + if (tensor.dims_size()) shape->add_dim()->set_dim_value(tensor.dims(0)); + inputs[j] = &graph.GetOrCreateNodeArg(tensor.name(), &type); + initializers.push_back(tensor.name()); + graph.AddInitializedTensor(tensor); + } + } + InlinedVector outputs; + for (size_t j = 0; j < nodes[i].output_count; ++j) + outputs.push_back(exposed_outputs[i][j] ? exposed_outputs[i][j] + : &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("rewrite_value"), nullptr)); + const auto& node = nodes[i]; + const auto& domain = node.domain == kOnnxDomainAlias ? kOnnxDomain : node.domain; + const auto node_name = graph.GenerateNodeName("rewrite"); + Node& replacement = annotation_source + ? graph.AddNode(node_name, node.op, "PatternRewriter", inputs, outputs, + *annotation_source, &resolved_attributes[i], domain) + : graph.AddNode(node_name, node.op, "PatternRewriter", inputs, outputs, + &resolved_attributes[i], domain); + added.push_back(replacement.Index()); + if (!graph.SetOpSchemaFromRegistryForNode(replacement)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Replacement schema is unavailable: ", node.op); + } + ORT_TRY { + for (const auto& [name, attribute] : replacement.Op()->attributes()) { + if (!attribute.required && !attribute.default_value.name().empty() && + !replacement.GetAttributes().count(name)) + replacement.AddAttributeProto(attribute.default_value); + } + ONNX_NAMESPACE::NodeProto proto; + replacement.ToProto(proto); + replacement.Op()->Verify(proto); + } + ORT_CATCH(const ONNX_NAMESPACE::checker::ValidationError& error) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid replacement: ", error.what()); + } + if (annotation_source) replacement.SetExecutionProviderType(annotation_source->GetExecutionProviderType()); + } + staged = true; + return Status::OK(); +} + +Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const Replacement& recipe, + Producers& producers, InlinedVector& added) { + InlinedVector returned_outputs; + for (Value output : pattern.Outputs()) returned_outputs.push_back(graph.GetNodeArg(match.values[output]->Name())); + auto tensor = [&](const Replacement::Input& input, NodeArg*& result) -> Status { + const Value value = std::holds_alternative(input) ? std::get(input) + : pattern.NamedValue(std::get(input)); + ORT_RETURN_IF(value.handle < 0 || static_cast(value.handle) >= match.values.size() || + !pattern.IsInput(value) || pattern.IsAttribute(value), + "Replacement inputs must be pattern input variables."); + result = match.values[value] ? graph.GetNodeArg(match.values[value]->Name()) : nullptr; + return Status::OK(); + }; + auto attribute = [&](const std::string& name, std::optional& result) -> Status { + const auto* value = match.optionalAttributeOf(name); + if (value) + result = *value; + else + result.reset(); + return Status::OK(); + }; + InlinedVector initializers; + ORT_RETURN_IF_ERROR(StageReplacement(graph, recipe, returned_outputs, tensor, attribute, match.nodes.back(), added, initializers)); + auto rollback = [&]() { RollbackFragment(graph, added, initializers); }; + // External consumers of a returned intermediate may feed a pattern input. + // Validate the prospective producer mapping before replacing their edges. + if (std::any_of(pattern.Outputs().begin(), pattern.Outputs().end(), [&](Value value) { + const auto& root_outputs = pattern.Nodes().back().outputs; + return std::find(root_outputs.begin(), root_outputs.end(), value) == root_outputs.end(); + })) { + InlinedHashMap staged_producers; + for (NodeIndex index : added) + for (const NodeArg* output : graph.GetNode(index)->OutputDefs()) staged_producers[output] = index; + InlinedHashMap colors; + InlinedVector> pending; + for (NodeIndex index : added) { + pending.emplace_back(index, false); + while (!pending.empty()) { + auto [current, exiting] = pending.back(); + pending.pop_back(); + if (exiting) { + colors[current] = 2; + continue; + } + if (colors[current] == 2) continue; + if (colors[current] == 1) { + rollback(); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Replacement output mapping introduces a cycle."); + } + colors[current] = 1; + pending.emplace_back(current, true); + const Node& node = *graph.GetNode(current); + auto visit_input = [&](const NodeArg* input) { + auto staged = staged_producers.find(input); + if (staged != staged_producers.end()) { + pending.emplace_back(staged->second, false); + } else { + auto existing = producers.find(input); + if (existing != producers.end()) pending.emplace_back(existing->second, false); + } + }; + for (const NodeArg* input : node.InputDefs()) visit_input(input); + for (const NodeArg* input : node.ImplicitInputDefs()) visit_input(input); + } + } + } + struct Edge { + const NodeArg* value; + NodeIndex destination; + int destination_arg; + }; + InlinedVector outgoing; + InlinedHashSet deletion_set; + for (const Node* node : match.nodes) deletion_set.insert(node->Index()); + for (const Node* node : match.nodes) { + for (auto it = node->OutputEdgesBegin(); it != node->OutputEdgesEnd(); ++it) { + if (!deletion_set.count(it->GetNode().Index())) + outgoing.push_back({node->OutputDefs()[it->GetSrcArgIndex()], it->GetNode().Index(), it->GetDstArgIndex()}); + } + } + InlinedVector removed; + for (const Node* node : match.nodes) { + removed.push_back(node->Index()); + for (const NodeArg* output : node->OutputDefs()) producers.erase(output); + for (const NodeArg* input : node->InputDefs()) { + if (!input->Exists()) continue; + graph.RemoveConsumerNode(input->Name(), graph.GetNode(node->Index())); + } + graph_utils::RemoveNodeOutputEdges(graph, *graph.GetNode(node->Index())); + } + // Remove in reverse order so producers cannot be left with dangling edges. + for (auto it = removed.rbegin(); it != removed.rend(); ++it) { + ORT_ENFORCE(graph.RemoveNode(*it), "Failed to remove matched node."); + } + for (NodeIndex index : added) { + Node& replacement = *graph.GetNode(index); + for (size_t i = 0; i < replacement.InputDefs().size(); ++i) { + const NodeArg* input = replacement.InputDefs()[i]; + if (!input->Exists()) continue; + graph.AddConsumerNode(input->Name(), &replacement); + auto it = producers.find(input); + if (it == producers.end()) continue; + const Node& producer = *graph.GetNode(it->second); + auto output_it = std::find(producer.OutputDefs().begin(), producer.OutputDefs().end(), input); + graph.AddEdge(producer.Index(), index, static_cast(std::distance(producer.OutputDefs().begin(), output_it)), + static_cast(i)); + } + for (const NodeArg* output : replacement.OutputDefs()) { + producers[output] = index; + graph.UpdateProducerNode(output->Name(), index); + } + } + for (const auto& edge : outgoing) { + const NodeIndex source = producers.at(edge.value); + const auto& outputs = graph.GetNode(source)->OutputDefs(); + const auto slot = std::distance(outputs.begin(), std::find(outputs.begin(), outputs.end(), edge.value)); + graph.AddEdge(source, edge.destination, static_cast(slot), edge.destination_arg); + } + graph.NotifyNodeReplacement(removed, added.back()); + return Status::OK(); +} +} // namespace + +AttributeMatcher AttributeMatcher::Equal(std::string name, int64_t value) { + return {std::move(name), [value](const auto& attribute) { + return attribute.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT && attribute.i() == value; + }}; +} + +AttributeMatcher AttributeMatcher::Equal(std::string name, float value) { + ORT_ENFORCE(std::isfinite(value), "Attribute float constants must be finite."); + return {std::move(name), [value](const auto& attribute) { + return attribute.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT && attribute.f() == value; + }}; +} + +AttributeMatcher AttributeMatcher::Equal(std::string name, std::string value) { + return {std::move(name), [value = std::move(value)](const auto& attribute) { + return attribute.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_STRING && attribute.s() == value; + }}; +} + +Value Pattern::AddValue(std::string_view name) { + ORT_ENFORCE(value_count_ < 256, "Patterns are limited to 256 values."); + if (!name.empty()) { + ORT_ENFORCE(std::none_of(values_.begin(), values_.end(), [&](const auto& value) { return value.name == name; }), + "Duplicate pattern value name: ", name); + } + Value value{std::string(name), static_cast(value_count_++)}; + values_.push_back(value); + return value; +} + +Value Pattern::Input(std::string_view name) { return AddValue(name); } + +Value Pattern::Constant(float constant) { + ORT_ENFORCE(std::isfinite(constant), "Pattern float constants must be finite."); + auto value = AddValue(""); + constants_.emplace(value.handle, constant); + return value; +} + +Value Pattern::Constant(int64_t constant) { + auto value = AddValue(""); + constants_.emplace(value.handle, constant); + return value; +} + +Value Pattern::Constant(Literal constant) { + auto value = AddValue(""); + constants_.emplace(value.handle, std::move(constant)); + return value; +} + +Value Pattern::NamedValue(std::string_view name) const { + auto alias = aliases_.find(std::string(name)); + if (alias != aliases_.end()) return alias->second; + auto it = std::find_if(values_.begin(), values_.end(), [&](const auto& value) { return value.name == name; }); + ORT_ENFORCE(!name.empty() && it != values_.end(), "Unknown pattern value: ", name); + return *it; +} + +void Pattern::Name(Value value, std::string_view name) { + ORT_ENFORCE(value.handle >= 0 && static_cast(value.handle) < value_count_ && !name.empty(), + "Invalid pattern value or empty name."); + ORT_ENFORCE(std::none_of(values_.begin(), values_.end(), + [&](const auto& other) { return other.handle != value.handle && other.name == name; }), + "Duplicate pattern value name: ", name); + values_[value.handle].name = name; + for (auto& node : nodes_) { + for (auto& output : node.outputs) + if (output.handle == value.handle) output.name = name; + for (auto& input : node.inputs) + if (input.handle == value.handle) input.name = name; + } +} + +void Pattern::Attributes(InlinedVector matchers) { + ORT_ENFORCE(!nodes_.empty(), "Attribute matchers require a pattern node."); + InlinedHashSet names; + for (const auto& attr : nodes_.back().attributes) names.insert(attr.name); + for (const auto& attr : matchers) { + ORT_ENFORCE(!attr.name.empty() && attr.predicate && names.insert(attr.name).second, + "Invalid or duplicate attribute matcher: ", attr.name); + } + auto& target = nodes_.back().attributes; + target.insert(target.end(), std::make_move_iterator(matchers.begin()), std::make_move_iterator(matchers.end())); +} + +InlinedVector Pattern::MultiOp(Operator op, gsl::span inputs, size_t output_count) { + ORT_ENFORCE(!op.name.empty(), "Pattern operators need a name."); + ORT_ENFORCE(nodes_.size() < 64, "Patterns are limited to 64 nodes."); + ORT_ENFORCE(output_count > 0 && output_count <= 64, "Pattern operators need between 1 and 64 outputs."); + ORT_ENFORCE(value_count_ <= 256 - output_count && inputs.size() <= 64, + "Patterns are limited to 256 values and 64 inputs per operator."); + if (op.domain == kOnnxDomainAlias) op.domain = kOnnxDomain; + for (Value input : inputs) + ORT_ENFORCE(input.handle == -1 || (input.handle >= 0 && static_cast(input.handle) < value_count_), + "Pattern inputs must already be defined."); + InlinedVector outputs; + for (size_t i = 0; i < output_count; ++i) outputs.push_back(AddValue("")); + InlinedVector canonical_inputs; + for (Value input : inputs) canonical_inputs.push_back(input.handle < 0 ? input : values_[input.handle]); + nodes_.push_back({std::move(op), std::move(canonical_inputs), outputs, {}}); + return outputs; +} + +bool Pattern::IsInput(Value value) const { + if (value.handle < 0 || static_cast(value.handle) >= value_count_) return false; + return std::none_of(nodes_.begin(), nodes_.end(), [value](const auto& node) { + return std::find(node.outputs.begin(), node.outputs.end(), value) != node.outputs.end(); + }); +} + +namespace detail { +struct Port { + std::string name; + InlinedVector constraints; +}; + +struct Definition { + enum class Kind { Interface, + Body, + Alternatives }; + Kind kind; + std::string name; + std::string interface_name; + std::string signature; + std::string body; + InlinedVector inputs; + InlinedVector outputs; + InlinedVector members; + InlinedVector, 1> plans; + size_t dependency_depth = 1; + bool explicit_outputs = false; + bool compiling = false; +}; + +using ResolveDefinition = std::function; + +bool CompatiblePorts(const InlinedVector& a, const InlinedVector& b) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) { + if (a[i].name != b[i].name || a[i].constraints.size() != b[i].constraints.size()) return false; + for (size_t j = 0; j < a[i].constraints.size(); ++j) { + const auto& x = a[i].constraints[j]; + const auto& y = b[i].constraints[j]; + if (x.element_type != y.element_type || x.element_variable != y.element_variable || + x.dimensions != y.dimensions || x.any_rank != y.any_rank || (!x.any_rank && x.prefix != y.prefix) || + x.attribute_type != y.attribute_type || x.tensor_type != y.tensor_type || + x.absent != y.absent || x.immutable != y.immutable) + return false; + } + } + return true; +} + +// Reconcile declaration metadata without adding constraints to any concrete plan. +bool MergeResultPorts(InlinedVector& target, const InlinedVector& results, bool require_names) { + if (target.size() != results.size()) return false; + for (size_t i = 0; i < results.size(); ++i) { + if (require_names && target[i].name != results[i].name) return false; + if (target[i].constraints.empty()) { + target[i].constraints = results[i].constraints; + } else if (!results[i].constraints.empty()) { + auto a = target[i], b = results[i]; + a.name = b.name = ""; + if (!CompatiblePorts({a}, {b})) return false; + } + } + return true; +} + +class PatternParser { + public: + explicit PatternParser(std::string_view text, InlinedVector* attributes = nullptr, + InlinedVector, 1>* bindings = nullptr) + : text_(text), replacement_attributes_(attributes), replacement_bindings_(bindings) {} + + PatternParser(std::string_view text, ResolveDefinition resolve, gsl::span choices, + const Definition* definition = nullptr, size_t* compile_depth = nullptr) + : text_(text), resolve_(std::move(resolve)), choices_(choices), definition_(definition), compile_depth_(compile_depth) {} + + const auto& Radices() const { return radices_; } + const auto& ReturnedNames() const { return returned_names_; } + const auto& ExpressionResults() const { return expression_results_; } + Status Declarations(InlinedVector& definitions); + static Status ValidateTypes(const Pattern& pattern); + static Status AppendAlternative(Definition& target, const Definition& member); + + Status Parse(Pattern& destination) { + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Expect("(")); + if (token_ != ")") { + do { + std::string name; + ORT_RETURN_IF_ERROR(Identifier(name, replacement_attributes_ != nullptr)); + ORT_RETURN_IF(symbols_.count(name), Error("duplicate input name")); + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); + Value input = pattern_.Input(name); + symbols_.emplace(name, input); + if (token_ == ":") { + InlinedVector alternatives; + ORT_RETURN_IF_ERROR(Annotation(name, alternatives)); + if (std::any_of(alternatives.begin(), alternatives.end(), [](const auto& type) { return type.absent; })) + pattern_.optional_inputs_.insert(input.handle); + pattern_.ordered_type_constraints_.emplace_back(input.handle, alternatives); + pattern_.type_constraints_.emplace(input.handle, std::move(alternatives)); + } + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + } + ORT_RETURN_IF_ERROR(Expect(")")); + ORT_RETURN_IF_ERROR(Expect("=>")); + const bool braced = token_ == "{"; + if (braced) ORT_RETURN_IF_ERROR(Next()); + bool returned = false; + while (kind_ != Kind::End && token_ != "}") { + if (kind_ == Kind::Symbol && token_ == ";") { + ORT_RETURN_IF_ERROR(Next()); + continue; + } + InlinedVector names; + InlinedHashMap> annotations; + if (kind_ == Kind::Identifier && token_ == "return") { + returned = true; + ORT_RETURN_IF_ERROR(Next()); + } + bool returned_values = false; + if (returned && definition_ && token_ == "(") { + ORT_RETURN_IF_ERROR(Next()); + InlinedHashMap ports; + do { + std::string port, local; + ORT_RETURN_IF_ERROR(Identifier(port)); + ORT_RETURN_IF_ERROR(Expect("=")); + ORT_RETURN_IF_ERROR(Identifier(local, true)); + auto found = symbols_.find(local); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined return value")); + ORT_RETURN_IF(!ports.emplace(port, found->second).second, Error("duplicate return port")); + returned_names_.push_back(port); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF_ERROR(Expect(")")); + if (definition_->interface_name.empty()) { + for (const auto& name : returned_names_) pattern_.outputs_.push_back(ports.at(name)); + } else { + ORT_RETURN_IF(ports.size() != definition_->outputs.size(), Error("interface return ports differ")); + returned_names_.clear(); + for (const auto& port : definition_->outputs) { + auto found = ports.find(port.name); + ORT_RETURN_IF(found == ports.end(), Error("missing interface return port")); + pattern_.outputs_.push_back(found->second); + returned_names_.push_back(port.name); + } + } + returned_values = true; + } else { + const size_t start = token_start_; + do { + std::string name; + ORT_RETURN_IF_ERROR(Identifier(name, names.empty() || static_cast(resolve_))); + ORT_RETURN_IF(std::find(names.begin(), names.end(), name) != names.end(), Error("duplicate assignment name")); + names.push_back(std::move(name)); + if (token_ == ":") { + ORT_RETURN_IF(returned, Error("annotate the definition, not the return")); + InlinedVector alternatives; + ORT_RETURN_IF_ERROR(Annotation(names.back(), alternatives)); + ORT_RETURN_IF(std::any_of(alternatives.begin(), alternatives.end(), [](const auto& type) { return type.absent; }), + Error("operator outputs cannot be optional")); + annotations.emplace(names.back(), std::move(alternatives)); + } + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + if (names.size() == 1 && (token_ == "(" || token_ == "::")) { + ORT_RETURN_IF(!returned && !replacement_attributes_, Error("expected assignment or return")); + position_ = start; + ORT_RETURN_IF_ERROR(Next()); + names.clear(); + returned = true; + } else { + ORT_RETURN_IF(!resolve_ && std::any_of(names.begin(), names.end(), + [](const auto& name) { return name.find('.') != std::string::npos; }), + Error("expected value identifier")); + if (returned) { + ORT_RETURN_IF(token_ == "=", Error("assign outputs before returning them")); + ORT_RETURN_IF(pattern_.Nodes().empty() && !definition_, Error("return must identify operator outputs")); + for (size_t i = 0; i < names.size(); ++i) { + auto found = symbols_.find(names[i]); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined return value")); + ORT_RETURN_IF(pattern_.IsInput(found->second) && !definition_, Error("cannot return a leaf input")); + pattern_.outputs_.push_back(found->second); + returned_names_.push_back(names[i].substr(names[i].rfind('.') == std::string::npos + ? 0 + : names[i].rfind('.') + 1)); + } + returned_values = true; + } else { + for (const auto& name : names) { + ORT_RETURN_IF(name.find('.') != std::string::npos, Error("expected assignment identifier")); + ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name) || pattern_.instances_.count(name), + Error("duplicate assignment or capture name")); + } + ORT_RETURN_IF_ERROR(Expect("=")); + } + } + } + if (!returned_values) { + ORT_RETURN_IF_ERROR(resolve_ ? Expression(names) : Call(names)); + for (const auto& name : names) { + auto annotation = annotations.find(name); + if (annotation == annotations.end()) continue; + auto& constraint = annotation->second; + auto found = symbols_.find(name); + ORT_RETURN_IF(found == symbols_.end(), Error("cannot annotate an instance-only binding")); + pattern_.ordered_type_constraints_.emplace_back(found->second.handle, constraint); + pattern_.type_constraints_.emplace(found->second.handle, std::move(constraint)); + } + } + if (returned) { + if (!returned_values) { + pattern_.outputs_ = last_outputs_; + for (size_t i = 0; i < last_outputs_.size(); ++i) + returned_names_.push_back(last_outputs_.size() == 1 ? "output" : "output" + std::to_string(i)); + } + if (kind_ == Kind::Symbol && token_ == ";") ORT_RETURN_IF_ERROR(Next()); + if (braced) ORT_RETURN_IF_ERROR(Expect("}")); + ORT_RETURN_IF(kind_ != Kind::End, Error("unexpected text after return")); + break; + } + } + ORT_RETURN_IF(!returned, Error("missing return")); + ORT_RETURN_IF_ERROR(ValidateTypes(pattern_)); + for (const auto& node : pattern_.Nodes()) + for (Value input : node.inputs) + ORT_RETURN_IF(pattern_.IsAttribute(input), Error("attribute variable used as a tensor input")); + InlinedHashSet reachable; + if (!pattern_.Nodes().empty() && !definition_) + for (Value output : pattern_.Nodes().back().outputs) reachable.insert(output.handle); + else + for (Value output : pattern_.Outputs()) reachable.insert(output.handle); + for (int handle : attribute_inputs_) reachable.insert(handle); + for (auto it = pattern_.Nodes().rbegin(); it != pattern_.Nodes().rend(); ++it) { + ORT_RETURN_IF(std::none_of(it->outputs.begin(), it->outputs.end(), + [&](Value output) { return reachable.count(output.handle); }), + Error("disconnected pattern node")); + for (Value output : it->outputs) reachable.insert(output.handle); + for (Value input : it->inputs) + if (input.handle >= 0) reachable.insert(input.handle); + } + for (int handle : pattern_.optional_inputs_) reachable.insert(handle); + ORT_RETURN_IF(reachable.size() != pattern_.ValueCount(), Error("unused input")); + destination = std::move(pattern_); + return Status::OK(); + } + + private: + Status Error(const char* message) const { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Pattern parse error at byte ", token_start_, ": ", message); + } + + static bool Alpha(char c) { return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_'; } + static bool Digit(char c) { return c >= '0' && c <= '9'; } + + Status Next() { + while (position_ < text_.size()) { + if (std::isspace(static_cast(text_[position_]))) { + ++position_; + } else if (text_[position_] == '#') { + while (position_ < text_.size() && text_[position_] != '\n' && text_[position_] != '\r') ++position_; + } else { + break; + } + } + token_start_ = position_; + token_.clear(); + kind_ = Kind::End; + if (position_ == text_.size()) return Status::OK(); + kind_ = Kind::Symbol; + char c = text_[position_++]; + if (c == '.' && position_ + 1 < text_.size() && text_.substr(token_start_, 3) == "...") { + position_ += 2; + } else if (Alpha(c)) { + kind_ = Kind::Identifier; + while (position_ < text_.size() && + (Alpha(text_[position_]) || Digit(text_[position_]) || text_[position_] == '.')) + ++position_; + } else if (c == '-' && position_ < text_.size() && text_[position_] == '>') { + ++position_; + } else if (Digit(c) || c == '-' || c == '+' || c == '.') { + kind_ = Kind::Number; + while (position_ < text_.size()) { + char d = text_[position_]; + if (!Digit(d) && d != '.' && d != 'e' && d != 'E' && d != '+' && d != '-') break; + ++position_; + } + } else if (c == '\'' || c == '"') { + kind_ = Kind::String; + bool closed = false; + while (position_ < text_.size()) { + char d = text_[position_++]; + if (d == c) { + closed = true; + break; + } + if (d == '\\') { + ORT_RETURN_IF(position_ == text_.size(), Error("unterminated string escape")); + d = text_[position_++]; + ORT_RETURN_IF(d != '\\' && d != '\'' && d != '"', Error("unsupported string escape")); + } + token_.push_back(d); + } + ORT_RETURN_IF(!closed, Error("unterminated string")); + return Status::OK(); + } else if ((c == '=' && position_ < text_.size() && text_[position_] == '>') || + (c == ':' && position_ < text_.size() && text_[position_] == ':')) { + ++position_; + } else { + ORT_RETURN_IF(std::string_view("(),=;{}[]:<>|").find(c) == std::string_view::npos, Error("unexpected character")); + } + token_ = text_.substr(token_start_, position_ - token_start_); + return Status::OK(); + } + + Status Expect(std::string_view token) { + ORT_RETURN_IF(kind_ != Kind::Symbol || token_ != token, Error("unexpected token")); + return Next(); + } + Status Identifier(std::string& name, bool allow_domain = false) { + ORT_RETURN_IF( + kind_ != Kind::Identifier || token_ == "return" || token_ == "None" || + (!allow_domain && token_.find('.') != std::string::npos), + Error("expected value identifier")); + name = token_; + return Next(); + } + + Status Annotation(const std::string& name, InlinedVector& alternatives) { + ORT_RETURN_IF(replacement_attributes_, Error("type annotations are supported in patterns only")); + ORT_RETURN_IF_ERROR(Expect(":")); + do { + TypeConstraint constraint; + constraint.value_name = name; + if (token_ == "None") { + constraint.absent = true; + ORT_RETURN_IF_ERROR(Next()); + } else { + ORT_RETURN_IF_ERROR(AnnotationType(constraint)); + } + auto equivalent = [&](const TypeConstraint& other) { + return other.absent == constraint.absent && other.element_type == constraint.element_type && + other.element_variable == constraint.element_variable && other.dimensions == constraint.dimensions && + other.any_rank == constraint.any_rank && (constraint.any_rank || other.prefix == constraint.prefix) && + other.attribute_type == constraint.attribute_type && other.tensor_type == constraint.tensor_type && + other.immutable == constraint.immutable; + }; + ORT_RETURN_IF(std::any_of(alternatives.begin(), alternatives.end(), equivalent), Error("duplicate type alternative")); + alternatives.push_back(std::move(constraint)); + ORT_RETURN_IF(alternatives.size() > 64, Error("too many type alternatives")); + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF(alternatives.size() == 1 && alternatives.front().absent, Error("None requires a concrete type alternative")); + return Status::OK(); + } + + Status AnnotationType(TypeConstraint& constraint) { + std::string type; + ORT_RETURN_IF_ERROR(Identifier(type)); + if (type == "const") { + constraint.immutable = true; + ORT_RETURN_IF_ERROR(Identifier(type)); + } + if (type == "int" || type == "ints" || type == "floats" || type == "strings") { + constraint.tensor_type = false; + if (type == "int") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_INT; + if (type == "ints") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_INTS; + if (type == "floats") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_FLOATS; + if (type == "strings") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_STRINGS; + } else if (type == "tensor") { + constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR; + } else { + std::string upper = type; + std::transform(upper.begin(), upper.end(), upper.begin(), + [](unsigned char c) { return static_cast(std::toupper(c)); }); + ONNX_NAMESPACE::TensorProto_DataType element; + ORT_RETURN_IF(!ONNX_NAMESPACE::TensorProto_DataType_Parse(upper, &element) || + element == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED, + Error("unknown tensor element type")); + constraint.element_type = element; + if (type == "float") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT; + if (type == "string") constraint.attribute_type = ONNX_NAMESPACE::AttributeProto_AttributeType_STRING; + } + if (token_ == "<") { + ORT_RETURN_IF(type != "tensor", Error("only tensor permits an element type variable")); + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Identifier(constraint.element_variable)); + ORT_RETURN_IF_ERROR(Expect(">")); + constraint.attribute_type = 0; + } + if (token_ == "[") { + ORT_RETURN_IF(!constraint.tensor_type, Error("attribute list types do not accept a shape")); + constraint.attribute_type = 0; + ORT_RETURN_IF_ERROR(Next()); + if (token_ == "]") { + constraint.any_rank = true; + } else { + while (true) { + if (token_ == "...") { + constraint.any_rank = constraint.dimensions.empty(); + constraint.prefix = true; + ORT_RETURN_IF_ERROR(Next()); + break; + } + ORT_RETURN_IF(constraint.dimensions.size() >= 64, Error("annotations are limited to 64 dimensions")); + if (kind_ == Kind::Identifier) { + std::string dimension; + ORT_RETURN_IF_ERROR(Identifier(dimension)); + if (dimension == "_") + constraint.dimensions.emplace_back(std::monostate{}); + else + constraint.dimensions.emplace_back(std::move(dimension)); + } else { + std::variant number; + ORT_RETURN_IF_ERROR(Number(number)); + const auto* fixed = std::get_if(&number); + ORT_RETURN_IF(!fixed || *fixed < 0, Error("dimensions must be nonnegative integers or names")); + constraint.dimensions.emplace_back(*fixed); + } + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(token_ == "]", Error("missing dimension after comma")); + } + } + ORT_RETURN_IF_ERROR(Expect("]")); + } + ORT_RETURN_IF(constraint.immutable && !constraint.tensor_type, Error("const requires a tensor type")); + if (constraint.immutable) constraint.attribute_type = 0; + return Status::OK(); + } + + Status Number(std::variant& result) { + ORT_RETURN_IF(kind_ != Kind::Number, Error("expected numeric literal")); + std::string_view number = token_; + if (!number.empty() && number.front() == '+') number.remove_prefix(1); + if (number.find_first_of(".eE") != std::string_view::npos) { + float value; + ORT_RETURN_IF(!TryParseStringWithClassicLocale(number, value) || !std::isfinite(value), + Error("invalid or nonfinite float literal")); + result = value; + } else { + int64_t value; + auto parsed = std::from_chars(number.data(), number.data() + number.size(), value); + ORT_RETURN_IF(parsed.ec != std::errc{} || parsed.ptr != number.data() + number.size(), + Error("invalid integer literal")); + result = value; + } + return Next(); + } + + Status List(Literal& result) { + ORT_RETURN_IF_ERROR(Expect("[")); + InlinedVector integers; + InlinedVector floats; + std::optional floating; + while (token_ != "]") { + std::variant number; + ORT_RETURN_IF_ERROR(Number(number)); + const bool is_float = std::holds_alternative(number); + ORT_RETURN_IF(floating && *floating != is_float, Error("list elements must have the same numeric type")); + floating = is_float; + if (is_float) + floats.push_back(std::get(number)); + else + integers.push_back(std::get(number)); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(token_ == "]", Error("expected list element after comma")); + } + ORT_RETURN_IF_ERROR(Expect("]")); + if (floating.value_or(false)) + result = std::move(floats); + else + result = std::move(integers); + return Status::OK(); + } + + Status Call(const InlinedVector& names) { + if (resolve_ && token_ == "(") { + Definition anonymous; + anonymous.kind = Definition::Kind::Alternatives; + anonymous.name = ""; + InlinedHashSet members; + ORT_RETURN_IF_ERROR(Next()); + do { + std::string member; + ORT_RETURN_IF_ERROR(Identifier(member)); + ORT_RETURN_IF(!members.insert(member).second, Error("duplicate alternative member")); + const auto* definition = resolve_(member); + ORT_RETURN_IF(!definition, Error("unknown alternative pattern")); + ORT_RETURN_IF_ERROR(AppendAlternative(anonymous, *definition)); + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF_ERROR(Expect(")")); + return PatternCall(names, anonymous); + } + ORT_RETURN_IF(kind_ != Kind::Identifier || token_ == "return", Error("expected operator name")); + Operator op{token_, ""}; + ORT_RETURN_IF_ERROR(Next()); + if (token_ == "::") { + op.domain = std::move(op.name); + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(kind_ != Kind::Identifier, Error("expected qualified operator name")); + op.name = token_; + ORT_RETURN_IF_ERROR(Next()); + } else { + auto dot = op.name.rfind('.'); + if (dot != std::string::npos) { + op.domain = op.name.substr(0, dot); + op.name.erase(0, dot + 1); + ORT_RETURN_IF(op.name.empty(), Error("expected qualified operator name")); + } + } + if (resolve_ && op.domain.empty()) { + if (const auto* definition = resolve_(op.name)) + return PatternCall(names, *definition); + ORT_RETURN_IF(!ONNX_NAMESPACE::OpSchemaRegistry::Schema(op.name), + Error("unknown pattern or unqualified ONNX operator")); + } + ORT_RETURN_IF_ERROR(Expect("(")); + InlinedVector inputs; + InlinedVector attributes; + NodeAttributes literal_attributes; + InlinedHashMap bindings; + InlinedHashSet attribute_names; + bool keyword = false; + if (token_ != ")") { + do { + if (kind_ == Kind::Identifier) { + std::string name = token_; + ORT_RETURN_IF_ERROR(Next()); + if (token_ == "=") { + keyword = true; + ORT_RETURN_IF(name == "None" || name == "return", Error("reserved attribute name")); + ORT_RETURN_IF(!attribute_names.insert(name).second, Error("duplicate attribute")); + ORT_RETURN_IF_ERROR(Next()); + if (kind_ == Kind::Identifier) { + std::string capture; + ORT_RETURN_IF_ERROR(Identifier(capture, replacement_bindings_ != nullptr)); + ORT_RETURN_IF(pattern_.instances_.count(capture), Error("attribute capture conflicts with a pattern instance")); + if (replacement_bindings_) { + auto found = symbols_.find(capture); + ORT_RETURN_IF(found == symbols_.end() || !pattern_.IsInput(found->second), + Error("attribute binding must reference a declared replacement input")); + bindings.emplace(name, capture); + attribute_inputs_.insert(found->second.handle); + } else { + auto declared = symbols_.find(capture); + if (declared != symbols_.end()) { + ORT_RETURN_IF(!pattern_.IsInput(declared->second), + Error("attribute capture conflicts with an operator output")); + pattern_.attribute_inputs_.insert(declared->second.handle); + attribute_inputs_.insert(declared->second.handle); + } + pattern_.attribute_names_.insert(capture); + const bool optional = declared != symbols_.end() && pattern_.IsOptional(declared->second); + attributes.push_back({name, [](const auto&) { return true; }, capture, optional}); + } + } else if (kind_ == Kind::String) { + attributes.push_back(AttributeMatcher::Equal(name, token_)); + auto& attr = literal_attributes[name]; + attr.set_name(name); + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + attr.set_s(token_); + ORT_RETURN_IF_ERROR(Next()); + } else if (token_ == "[") { + Literal list; + ORT_RETURN_IF_ERROR(List(list)); + ONNX_NAMESPACE::AttributeProto attr; + attr.set_name(name); + if (const auto* values = std::get_if>(&list)) { + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOATS); + for (auto value : *values) attr.add_floats(value); + } else { + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); + for (auto value : std::get>(list)) attr.add_ints(value); + } + attributes.push_back({name, [attr](const auto& actual) { + if (actual.type() != attr.type()) return false; + if (attr.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_FLOATS) { + return actual.floats_size() == attr.floats_size() && + std::equal(attr.floats().begin(), attr.floats().end(), + actual.floats().begin()); + } + return actual.ints_size() == attr.ints_size() && + std::equal(attr.ints().begin(), attr.ints().end(), actual.ints().begin()); + }}); + literal_attributes.emplace(name, std::move(attr)); + } else { + std::variant number; + ORT_RETURN_IF_ERROR(Number(number)); + attributes.push_back( + std::visit([&](auto value) { return AttributeMatcher::Equal(name, value); }, number)); + auto& attr = literal_attributes[name]; + attr.set_name(name); + if (const auto* floating = std::get_if(&number)) { + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + attr.set_f(*floating); + } else { + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attr.set_i(std::get(number)); + } + } + } else { + ORT_RETURN_IF(keyword, Error("positional input after attribute")); + if (name == "_" && resolve_) { + inputs.push_back(pattern_.AddValue("")); + } else if (name == "None") { + inputs.push_back(Value{-1}); + } else { + auto found = symbols_.find(name); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined input value")); + inputs.push_back(found->second); + } + } + } else { + ORT_RETURN_IF(keyword, Error("positional input after attribute")); + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); + if (token_ == "[") { + Literal list; + ORT_RETURN_IF_ERROR(List(list)); + inputs.push_back(pattern_.Constant(std::move(list))); + } else { + std::variant number; + ORT_RETURN_IF_ERROR(Number(number)); + inputs.push_back(std::visit([&](auto value) { return pattern_.Constant(value); }, number)); + } + } + ORT_RETURN_IF(inputs.size() > 64, Error("too many operator inputs")); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + } + ORT_RETURN_IF_ERROR(Expect(")")); + size_t outputs = names.empty() ? 1 : names.size(); + ORT_RETURN_IF(outputs > 64 || pattern_.Nodes().size() >= 64 || pattern_.ValueCount() + outputs > 256, + Error("pattern limits exceeded")); + auto values = pattern_.MultiOp(std::move(op), gsl::span(inputs), outputs); + last_outputs_ = values; + pattern_.Attributes(std::move(attributes)); + if (replacement_attributes_) replacement_attributes_->push_back(std::move(literal_attributes)); + if (replacement_bindings_) replacement_bindings_->push_back(std::move(bindings)); + for (size_t i = 0; i < names.size(); ++i) { + ORT_RETURN_IF(pattern_.attribute_names_.count(names[i]), Error("output name conflicts with an attribute capture")); + pattern_.Name(values[i], names[i]); + symbols_.emplace(names[i], pattern_.NamedValue(names[i])); + } + return Status::OK(); + } + + enum class Kind { + End, + Symbol, + Identifier, + Number, + String + }; + std::string_view text_; + size_t position_ = 0; + size_t token_start_ = 0; + Kind kind_ = Kind::Symbol; + std::string token_; + Pattern pattern_; + InlinedHashMap symbols_; + InlinedVector* replacement_attributes_ = nullptr; + InlinedVector, 1>* replacement_bindings_ = nullptr; + InlinedHashSet attribute_inputs_; + ResolveDefinition resolve_; + gsl::span choices_; + const Definition* definition_ = nullptr; + InlinedVector radices_; + InlinedVector last_outputs_; + InlinedVector returned_names_; + InlinedHashMap> expression_results_; + InlinedVector last_result_ports_; + size_t expression_depth_ = 0; + size_t* compile_depth_ = nullptr; + Status Header(InlinedVector& ports, std::string& signature); + Status PatternCall(const InlinedVector& names, const Definition& definition); + Status Expression(const InlinedVector& names); + Status ScanAtom(); +}; + +Status PatternParser::AppendAlternative(Definition& target, const Definition& member) { + ORT_RETURN_IF(member.kind == Definition::Kind::Interface, "Interfaces are not alternative patterns: ", member.name); + if (target.plans.empty()) { + target.interface_name = member.interface_name; + target.inputs = member.inputs; + target.outputs = member.outputs; + } else { + ORT_RETURN_IF(target.interface_name != member.interface_name || + !CompatiblePorts(target.inputs, member.inputs) || !CompatiblePorts(target.outputs, member.outputs), + "Incompatible alternative interfaces/ports: ", target.name, " / ", member.name); + } + for (const auto& plan : member.plans) { + Pattern copy = *plan; + auto selected = std::make_shared(copy.instances_.at("")); + copy.instances_[""] = {target.name, target.interface_name, member.name, std::move(selected)}; + target.plans.push_back(std::make_shared(std::move(copy))); + ORT_RETURN_IF(target.plans.size() > 64, "Graph alternatives exceed 64 concrete plans."); + } + return Status::OK(); +} + +Status PatternParser::ValidateTypes(const Pattern& pattern) { + for (int handle : pattern.required_tensor_bindings_) + ORT_RETURN_IF(pattern.IsAttribute(handle), "Declared tensor binding cannot become an attribute."); + for (int handle : pattern.composed_tensor_annotations_) + ORT_RETURN_IF(pattern.IsAttribute(handle), "Composed tensor annotation cannot become an attribute."); + InlinedHashSet dimensions, elements, names; + for (const auto& value : pattern.values_) + if (!value.name.empty()) names.insert(value.name); + for (const auto& [name, value] : pattern.aliases_) names.insert(name); + for (const auto& name : pattern.attribute_names_) names.insert(name); + for (const auto& [name, instance] : pattern.instances_) names.insert(name); + size_t combinations = 1; + auto validate = [&](int handle, const auto& alternatives) -> Status { + ORT_RETURN_IF(alternatives.empty() || combinations > 64 / alternatives.size(), + "Type alternatives exceed 64 combinations."); + combinations *= alternatives.size(); + for (const auto& type : alternatives) { + if (pattern.IsAttribute(handle)) { + const bool optional = alternatives.size() == 2 && + std::count_if(alternatives.begin(), alternatives.end(), + [](const auto& alternative) { return alternative.absent; }) == 1; + ORT_RETURN_IF((alternatives.size() != 1 && !optional) || + (!type.absent && (!type.attribute_type || type.immutable)), + "Attribute annotations require one concrete attribute type, optionally combined with None."); + } else { + ORT_RETURN_IF(!type.tensor_type, "Annotation requires an attribute variable."); + } + if (!type.element_variable.empty()) elements.insert(type.element_variable); + for (const auto& dim : type.dimensions) + if (const auto* name = std::get_if(&dim)) dimensions.insert(*name); + } + return Status::OK(); + }; + for (const auto& [handle, alternatives] : pattern.ordered_type_constraints_) + ORT_RETURN_IF_ERROR(validate(handle, alternatives)); + ORT_RETURN_IF(dimensions.size() + elements.size() > 256, "Too many annotation variables."); + for (const auto& name : dimensions) + ORT_RETURN_IF(names.count(name) || elements.count(name), "Dimension name conflicts with another variable: ", name); + for (const auto& name : elements) + ORT_RETURN_IF(names.count(name), "Element type name conflicts with another variable: ", name); + return Status::OK(); +} + +Status PatternParser::Header(InlinedVector& ports, std::string& signature) { + const size_t start = token_start_; + ORT_RETURN_IF_ERROR(Expect("(")); + InlinedHashSet names; + while (token_ != ")") { + Port port; + ORT_RETURN_IF_ERROR(Identifier(port.name)); + ORT_RETURN_IF(!names.insert(port.name).second, Error("duplicate port")); + if (token_ == ":") ORT_RETURN_IF_ERROR(Annotation(port.name, port.constraints)); + ports.push_back(std::move(port)); + ORT_RETURN_IF(ports.size() > 64, Error("too many ports")); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } + const size_t end = token_start_ + 1; + ORT_RETURN_IF_ERROR(Expect(")")); + signature = text_.substr(start, end - start); + return Status::OK(); +} + +Status PatternParser::ScanAtom() { + ORT_RETURN_IF(kind_ != Kind::Identifier && token_ != "(", Error("expected call or value expression")); + const bool grouped = token_ == "("; + if (!grouped) { + ORT_RETURN_IF_ERROR(Next()); + if (token_ == "::") { + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(kind_ != Kind::Identifier, Error("expected operator name")); + ORT_RETURN_IF_ERROR(Next()); + } + } + while (token_ == "(") { + int depth = 0; + do { + if (kind_ == Kind::Symbol && token_ == "(") ++depth; + if (kind_ == Kind::Symbol && token_ == ")") --depth; + ORT_RETURN_IF(depth > 64, Error("expression nesting limit of 64 exceeded")); + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(kind_ == Kind::End && depth, Error("unterminated expression")); + } while (depth); + // Legacy (First | Second)(arguments) syntax. + if (!grouped) break; + } + return Status::OK(); +} + +Status PatternParser::Expression(const InlinedVector& names) { + ORT_RETURN_IF(expression_depth_ >= 64, Error("expression nesting limit of 64 exceeded")); + const bool nested = compile_depth_ && expression_depth_ > 0; + ORT_RETURN_IF(nested && *compile_depth_ >= 64, Error("combined expression/dependency nesting limit of 64 exceeded")); + if (nested) ++*compile_depth_; + ++expression_depth_; + const size_t start = token_start_; + InlinedVector branches; + do { + branches.push_back(token_start_); + ORT_RETURN_IF(branches.size() > 64, Error("expression alternatives exceed 64")); + ORT_RETURN_IF_ERROR(ScanAtom()); + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + const size_t end = token_start_; + size_t choice = 0; + if (branches.size() > 1) { + const size_t slot = radices_.size(); + radices_.push_back(branches.size()); + choice = slot < choices_.size() ? choices_[slot] : 0; + ORT_RETURN_IF(choice >= branches.size(), Error("invalid expression alternative")); + } + position_ = branches[choice]; + ORT_RETURN_IF_ERROR(Next()); + const std::string label = token_ == "(" ? "" : token_; + last_result_ports_.clear(); + if (token_ == "(") { + // Distinguish a grouped expression from the supported legacy anonymous call. + ORT_RETURN_IF_ERROR(ScanAtom()); + const size_t after = token_start_; + position_ = branches[choice]; + ORT_RETURN_IF_ERROR(Next()); + int depth = 0; + do { + if (kind_ == Kind::Symbol && token_ == "(") ++depth; + if (kind_ == Kind::Symbol && token_ == ")") --depth; + ORT_RETURN_IF_ERROR(Next()); + } while (depth); + const bool legacy = token_ == "("; + position_ = branches[choice]; + ORT_RETURN_IF_ERROR(Next()); + if (legacy) { + ORT_RETURN_IF_ERROR(Call(names)); + } else { + ORT_RETURN_IF_ERROR(Expect("(")); + ORT_RETURN_IF_ERROR(Expression(names)); + ORT_RETURN_IF_ERROR(Expect(")")); + } + ORT_RETURN_IF(token_start_ != after, Error("invalid grouped expression")); + } else { + const size_t atom = token_start_; + std::string value; + ORT_RETURN_IF_ERROR(Identifier(value, true)); + if (token_ == "(" || token_ == "::") { + position_ = atom; + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Call(names)); + } else { + ORT_RETURN_IF(names.size() != 1, Error("value aliases require one assignment target")); + auto found = symbols_.find(value); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined value alias")); + pattern_.required_tensor_bindings_.insert(found->second.handle); + symbols_.emplace(names.front(), found->second); + pattern_.aliases_.emplace(names.front(), found->second); + last_outputs_ = {found->second}; + auto constraints = pattern_.type_constraints_.find(found->second.handle); + last_result_ports_.push_back({names.front(), constraints == pattern_.type_constraints_.end() + ? InlinedVector{} + : constraints->second}); + } + } + if (last_result_ports_.empty()) + for (size_t i = 0; i < last_outputs_.size(); ++i) last_result_ports_.push_back({"", {}}); + if (branches.size() > 1) { + expression_results_.emplace(start, last_result_ports_); + if (!names.empty()) { + const std::string& path = names.front(); + auto existing = pattern_.instances_.find(path); + std::shared_ptr selected; + if (existing != pattern_.instances_.end()) + selected = std::make_shared(existing->second); + pattern_.instances_.insert_or_assign(path, Pattern::Instance{"", "", label, std::move(selected)}); + } + } + position_ = end; + ORT_RETURN_IF_ERROR(Next()); + --expression_depth_; + if (nested) --*compile_depth_; + return Status::OK(); +} + +Status PatternParser::Declarations(InlinedVector& definitions) { + ORT_RETURN_IF_ERROR(Next()); + InlinedHashSet names; + while (kind_ != Kind::End) { + if (token_ == ";") { + ORT_RETURN_IF_ERROR(Next()); + continue; + } + Definition definition; + if (token_ == "interface") + definition.kind = Definition::Kind::Interface; + else if (token_ == "pattern") + definition.kind = Definition::Kind::Body; + else + return Error("expected interface or pattern"); + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Identifier(definition.name)); + ORT_RETURN_IF(!names.insert(definition.name).second, Error("duplicate definition")); + if (definition.kind != Definition::Kind::Interface && token_ == "=") { + definition.kind = Definition::Kind::Alternatives; + ORT_RETURN_IF_ERROR(Expect("=")); + do { + std::string member; + ORT_RETURN_IF_ERROR(Identifier(member)); + ORT_RETURN_IF(std::find(definition.members.begin(), definition.members.end(), member) != definition.members.end(), + Error("duplicate alternative member")); + definition.members.push_back(std::move(member)); + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF(definition.members.size() > 64, Error("alternatives are limited to 64 members")); + } else { + if (definition.kind == Definition::Kind::Body && token_ == "implements") { + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Identifier(definition.interface_name)); + definition.signature = "()"; + if (token_ == ":") { + ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF_ERROR(Header(definition.inputs, definition.signature)); + ORT_RETURN_IF_ERROR(Expect("->")); + } + } else { + ORT_RETURN_IF_ERROR(Header(definition.inputs, definition.signature)); + } + if (definition.kind == Definition::Kind::Interface || + (definition.kind == Definition::Kind::Body && token_ == "->")) { + std::string output_signature; + ORT_RETURN_IF_ERROR(Expect("->")); + ORT_RETURN_IF_ERROR(Header(definition.outputs, output_signature)); + definition.explicit_outputs = true; + ORT_RETURN_IF(definition.outputs.empty(), Error("output signature requires at least one output")); + InlinedHashSet ports; + for (const auto& input : definition.inputs) ports.insert(input.name); + for (const auto& output : definition.outputs) { + ORT_RETURN_IF(!ports.insert(output.name).second, Error("input/output port name collision")); + ORT_RETURN_IF(std::any_of(output.constraints.begin(), output.constraints.end(), + [](const auto& constraint) { return constraint.absent; }), + Error("interface outputs cannot be optional")); + } + Pattern contract; + for (const auto* port_list : {&definition.inputs, &definition.outputs}) { + if (definition.kind != Definition::Kind::Interface && port_list == &definition.inputs) continue; + for (const auto& port : *port_list) { + for (const auto& type : port.constraints) + ORT_RETURN_IF(!type.tensor_type, Error("interface ports must be tensors")); + const Value value = contract.Input(port.name); + if (!port.constraints.empty()) contract.ordered_type_constraints_.emplace_back(value.handle, port.constraints); + } + } + ORT_RETURN_IF_ERROR(ValidateTypes(contract)); + } + if (definition.kind != Definition::Kind::Interface) { + const size_t start = token_start_; + ORT_RETURN_IF_ERROR(Expect("{")); + while (!(kind_ == Kind::Symbol && token_ == "}") && kind_ != Kind::End) ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(kind_ == Kind::End, Error("unterminated definition body")); + definition.body = text_.substr(start, token_start_ + 1 - start); + ORT_RETURN_IF_ERROR(Expect("}")); + } + } + definitions.push_back(std::move(definition)); + ORT_RETURN_IF(definitions.size() > 256, Error("too many definitions")); + } + return Status::OK(); +} + +Status PatternParser::PatternCall(const InlinedVector& names, const Definition& definition) { + ORT_RETURN_IF(definition.kind == Definition::Kind::Interface, Error("interfaces are not callable")); + ORT_RETURN_IF(names.empty(), Error("pattern calls require assigned outputs")); + ORT_RETURN_IF(names.size() != 1 && names.size() != definition.outputs.size(), + Error("pattern output assignment arity mismatch")); + const std::string& path = names.front(); + ORT_RETURN_IF(pattern_.instances_.count(path), Error("duplicate pattern instance")); + ORT_RETURN_IF_ERROR(Expect("(")); + InlinedVector inputs; + while (token_ != ")") { + std::string actual; + ORT_RETURN_IF_ERROR(Identifier(actual, true)); + if (actual == "_") { + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("expanded value limit exceeded")); + inputs.push_back(pattern_.AddValue("")); + } else { + auto found = symbols_.find(actual); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined pattern argument")); + inputs.push_back(found->second); + } + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } + ORT_RETURN_IF_ERROR(Expect(")")); + ORT_RETURN_IF(inputs.size() != definition.inputs.size(), Error("pattern argument count mismatch")); + ORT_RETURN_IF(definition.plans.empty(), Error("pattern has no concrete plans")); + const size_t call = radices_.size(); + radices_.push_back(definition.plans.size()); + const size_t choice = call < choices_.size() ? choices_[call] : 0; + ORT_RETURN_IF(choice >= definition.plans.size(), Error("invalid alternative selection")); + const Pattern& source = *definition.plans[choice]; + // Identity-only calls still expand scoped aliases and instance metadata. + ORT_RETURN_IF(source.expanded_calls_ >= 256 || pattern_.expanded_calls_ >= 256 - source.expanded_calls_, + Error("expanded call limit of 256 exceeded")); + pattern_.expanded_calls_ += source.expanded_calls_ + 1; + ORT_RETURN_IF(pattern_.Nodes().size() + source.Nodes().size() > 64, Error("expanded node limit exceeded")); + InlinedVector map(source.ValueCount(), Value{-1}); + InlinedHashMap attributes; + for (size_t i = 0; i < inputs.size(); ++i) { + const Value formal = source.NamedValue(definition.inputs[i].name); + if (source.IsAttribute(formal)) { + ORT_RETURN_IF(inputs[i].name.empty(), Error("wildcards require tensor arguments")); + ORT_RETURN_IF(!pattern_.IsInput(inputs[i]), Error("attribute arguments must be external inputs")); + pattern_.attribute_inputs_.insert(inputs[i].handle); + attribute_inputs_.insert(inputs[i].handle); + attributes.emplace(formal.name, inputs[i].name); + } + map[formal.handle] = inputs[i]; + if (!source.IsOptional(formal) || inputs[i].name.empty()) pattern_.optional_inputs_.erase(inputs[i].handle); + } + auto scoped = [&](std::string_view name) { return path + "." + std::string(name); }; + auto attribute_name = [&](const std::string& name) { + auto found = attributes.find(name); + return found == attributes.end() ? scoped(name) : found->second; + }; + for (const auto& value : source.values_) { + if (map[value.handle].handle < 0) { + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("expanded value limit exceeded")); + map[value.handle] = pattern_.AddValue(value.name.empty() ? "" : scoped(value.name)); + if (source.IsOptional(value)) pattern_.optional_inputs_.insert(map[value.handle].handle); + } + if (!value.name.empty()) { + pattern_.aliases_.emplace(scoped(value.name), map[value.handle]); + symbols_.emplace(scoped(value.name), map[value.handle]); + } + } + auto remap = [&](Value value) { return value.handle < 0 ? value : map[value.handle]; }; + for (const auto& [name, value] : source.aliases_) { + pattern_.aliases_.insert_or_assign(scoped(name), remap(value)); + symbols_.insert_or_assign(scoped(name), remap(value)); + } + for (const auto& [handle, literal] : source.constants_) + pattern_.constants_.emplace(map[handle].handle, literal); + for (const auto& node : source.nodes_) { + Pattern::NodePattern copy = node; + for (auto& input : copy.inputs) input = remap(input); + for (auto& output : copy.outputs) output = remap(output); + for (auto& attribute : copy.attributes) + if (!attribute.capture.empty()) attribute.capture = attribute_name(attribute.capture); + pattern_.nodes_.push_back(std::move(copy)); + } + for (const auto& name : source.attribute_names_) { + const auto binding = attribute_name(name); + pattern_.attribute_names_.insert(binding); + pattern_.attribute_aliases_.emplace(scoped(name), binding); + } + for (const auto& [name, binding] : source.attribute_aliases_) + pattern_.attribute_aliases_.emplace(scoped(name), attribute_name(binding)); + for (int handle : source.attribute_inputs_) { + pattern_.attribute_inputs_.insert(map[handle].handle); + attribute_inputs_.insert(map[handle].handle); + } + for (int handle : source.required_tensor_bindings_) + pattern_.required_tensor_bindings_.insert(map[handle].handle); + for (int handle : source.composed_tensor_annotations_) + pattern_.composed_tensor_annotations_.insert(map[handle].handle); + auto constraints = [&](int handle, const auto& alternatives) { + if (!source.IsAttribute(handle)) pattern_.composed_tensor_annotations_.insert(map[handle].handle); + auto copy = alternatives; + for (auto& type : copy) { + type.value_name = source.IsAttribute(handle) ? attribute_name(type.value_name) : scoped(type.value_name); + if (!type.element_variable.empty()) type.element_variable = scoped(type.element_variable); + for (auto& dim : type.dimensions) + if (auto* name = std::get_if(&dim)) *name = scoped(*name); + } + pattern_.ordered_type_constraints_.emplace_back(map[handle].handle, std::move(copy)); + }; + for (const auto& [handle, alternatives] : source.ordered_type_constraints_) constraints(handle, alternatives); + for (const auto& [name, instance] : source.instances_) + pattern_.instances_.emplace(name.empty() ? path : scoped(name), instance); + last_outputs_.clear(); + for (size_t i = 0; i < source.outputs_.size(); ++i) { + const Value output = remap(source.outputs_[i]); + const auto output_name = scoped(definition.outputs[i].name); + pattern_.aliases_.insert_or_assign(output_name, output); + symbols_.insert_or_assign(output_name, output); + last_outputs_.push_back(output); + if (names.size() == source.outputs_.size()) { + ORT_RETURN_IF(pattern_.attribute_names_.count(names[i]), Error("output conflicts with attribute capture")); + symbols_.insert_or_assign(names[i], output); + pattern_.aliases_.insert_or_assign(names[i], output); + } + } + if (last_outputs_.size() == 1) { + symbols_.emplace(path, last_outputs_.front()); + pattern_.aliases_.emplace(path, last_outputs_.front()); + } + last_result_ports_ = definition.outputs; + return Status::OK(); +} + +class PatternModule { + public: + explicit PatternModule(std::string_view text) { + ORT_THROW_IF_ERROR(PatternParser(text).Declarations(definitions_)); + for (size_t i = 0; i < definitions_.size(); ++i) by_name_.emplace(definitions_[i].name, i); + for (auto& definition : definitions_) Compile(definition); + } + + Pattern GetPattern(std::string_view name) const { + auto found = by_name_.find(std::string(name)); + ORT_ENFORCE(found != by_name_.end(), "Unknown library pattern: ", name); + const auto& definition = definitions_[found->second]; + ORT_ENFORCE(definition.kind != Definition::Kind::Interface, "Interfaces cannot be rule entries: ", name); + for (const auto& plan : definition.plans) ValidateEntry(*plan, name); + Pattern result = *definition.plans.front(); + for (size_t i = 1; i < definition.plans.size(); ++i) result.alternatives_.push_back(definition.plans[i]); + return result; + } + + static void ValidateEntry(const Pattern& plan, std::string_view name) { + ORT_ENFORCE(!plan.Nodes().empty(), "Pattern ", name, " cannot be a rule entry: no operator root."); + InlinedHashSet returned; + for (Value output : plan.Outputs()) { + ORT_ENFORCE(!plan.IsInput(output), "Pattern ", name, + " cannot be a rule entry: returned leaf/identity alias; compose it into a rooted pattern."); + ORT_ENFORCE(returned.insert(output.handle).second, "Pattern ", name, + " cannot be a rule entry: duplicate returned value after alias expansion."); + } + InlinedHashSet reachable; + for (Value value : plan.Nodes().back().outputs) reachable.insert(value.handle); + for (auto it = plan.Nodes().rbegin(); it != plan.Nodes().rend(); ++it) { + ORT_ENFORCE(std::any_of(it->outputs.begin(), it->outputs.end(), + [&](Value value) { return reachable.count(value.handle); }), + "Pattern ", name, " cannot be a rule entry: independent branches do not reach one operator root."); + for (Value value : it->outputs) reachable.insert(value.handle); + for (Value value : it->inputs) + if (value.handle >= 0) reachable.insert(value.handle); + } + for (int handle : plan.attribute_inputs_) reachable.insert(handle); + for (int handle : plan.optional_inputs_) reachable.insert(handle); + ORT_ENFORCE(reachable.size() == plan.ValueCount(), "Pattern ", name, " cannot be a rule entry: unused inputs."); + } + + private: + const Definition* Resolve(std::string_view name) { + auto found = by_name_.find(std::string(name)); + if (found == by_name_.end()) return nullptr; // Unqualified ONNX operator. + auto& definition = definitions_[found->second]; + Compile(definition); + if (active_definition_ && definition.kind != Definition::Kind::Interface) { + active_definition_->dependency_depth = std::max(active_definition_->dependency_depth, definition.dependency_depth + 1); + ORT_ENFORCE(active_definition_->dependency_depth <= 64, + "Pattern dependency nesting limit of 64 exceeded: ", active_definition_->name); + } + return &definition; + } + + void Compile(Definition& definition) { + ORT_ENFORCE(!definition.compiling, "Recursive pattern definition: ", definition.name); + if (definition.kind == Definition::Kind::Interface || !definition.plans.empty()) return; + ORT_ENFORCE(active_compile_depth_ < 64, "Pattern dependency nesting limit of 64 exceeded: ", definition.name); + ++active_compile_depth_; + auto* enclosing = active_definition_; + active_definition_ = &definition; + definition.compiling = true; + if (definition.kind == Definition::Kind::Alternatives) { + for (const auto& member : definition.members) { + const auto* concrete = Resolve(member); + ORT_ENFORCE(concrete, "Unknown alternative pattern: ", member); + ORT_THROW_IF_ERROR(PatternParser::AppendAlternative(definition, *concrete)); + } + } else { + if (!definition.interface_name.empty()) { + const auto* contract = Resolve(definition.interface_name); + ORT_ENFORCE(contract && contract->kind == Definition::Kind::Interface, + "Unknown pattern interface: ", definition.interface_name); + // Extra external inputs are part of the body signature, not the call signature. + std::string extra = definition.signature.substr(1, definition.signature.size() - 2); + definition.signature = contract->signature; + if (!extra.empty()) { + definition.signature.pop_back(); + if (!contract->inputs.empty()) definition.signature += ","; + definition.signature += extra + ")"; + } + definition.inputs = contract->inputs; + definition.outputs = contract->outputs; + } + const std::string text = definition.signature + " => " + definition.body; + InlinedVector choices; + InlinedHashMap> expression_results; + do { + Pattern plan; + PatternParser parser(text, [&](std::string_view name) { return Resolve(name); }, choices, &definition, &active_compile_depth_); + const auto status = parser.Parse(plan); + ORT_ENFORCE(status.IsOK(), "In definition ", definition.name, ": ", status.ErrorMessage()); + for (const auto& [position, results] : parser.ExpressionResults()) { + auto [it, inserted] = expression_results.emplace(position, results); + if (!inserted) { + ORT_ENFORCE(it->second.size() == results.size(), "Expression alternative output arity mismatch in ", definition.name); + ORT_ENFORCE(MergeResultPorts(it->second, results, false), + "Expression alternative output type mismatch in ", definition.name); + } + } + if (!definition.interface_name.empty()) { + for (const auto& port : definition.inputs) { + ORT_ENFORCE(!plan.IsAttribute(plan.NamedValue(port.name)), + "Interface input must remain a tensor in pattern ", definition.name, ": ", port.name); + plan.required_tensor_bindings_.insert(plan.NamedValue(port.name).handle); + } + for (Value output : plan.outputs_) { + ORT_ENFORCE(!plan.IsAttribute(output), + "Interface output must remain a tensor in pattern ", definition.name); + plan.required_tensor_bindings_.insert(output.handle); + } + for (Value output : plan.outputs_) plan.optional_inputs_.erase(output.handle); + } + if (definition.explicit_outputs) { + for (Value output : plan.outputs_) { + ORT_ENFORCE(!plan.IsAttribute(output), "Declared pattern output must remain a tensor in ", definition.name); + plan.required_tensor_bindings_.insert(output.handle); + plan.optional_inputs_.erase(output.handle); + } + } + if (definition.interface_name.empty() && !definition.explicit_outputs) { + InlinedVector outputs; + InlinedHashSet names; + for (size_t i = 0; i < parser.ReturnedNames().size(); ++i) { + const auto& name = parser.ReturnedNames()[i]; + ORT_ENFORCE(names.insert(name).second, "Duplicate inferred output port in pattern ", definition.name, ": ", name); + auto constraints = plan.type_constraints_.find(plan.outputs_[i].handle); + outputs.push_back({name, constraints == plan.type_constraints_.end() + ? InlinedVector{} + : constraints->second}); + } + if (definition.outputs.empty()) + definition.outputs = std::move(outputs); + else + ORT_ENFORCE(MergeResultPorts(definition.outputs, outputs, true), + "Pattern alternatives change output ports: ", definition.name); + } + { + ORT_ENFORCE(plan.outputs_.size() == definition.outputs.size(), + "Pattern return count differs from its ports: ", definition.name); + for (size_t i = 0; i < definition.outputs.size(); ++i) { + const auto& port = definition.outputs[i]; + const Value output = plan.outputs_[i]; + const auto local = std::find_if(plan.values_.begin(), plan.values_.end(), + [&](const auto& value) { return value.name == port.name; }); + ORT_ENFORCE(local == plan.values_.end() || local->handle == output.handle, + "Output port shadows a different local value in pattern ", definition.name, ": ", port.name); + const auto alias = plan.aliases_.find(port.name); + ORT_ENFORCE(alias == plan.aliases_.end() || alias->second.handle == output.handle, + "Output port shadows a different local value in pattern ", definition.name, ": ", port.name); + plan.aliases_.insert_or_assign(port.name, output); + if ((!definition.interface_name.empty() || definition.explicit_outputs) && !port.constraints.empty()) { + plan.ordered_type_constraints_.emplace_back(output.handle, port.constraints); + plan.type_constraints_.emplace(output.handle, port.constraints); + } + } + } + ORT_THROW_IF_ERROR(PatternParser::ValidateTypes(plan)); + plan.instances_[""] = {definition.name, definition.interface_name, ""}; + definition.plans.push_back(std::make_shared(std::move(plan))); + ORT_ENFORCE(definition.plans.size() <= 64, "Graph alternatives exceed 64 concrete plans."); + const auto& radices = parser.Radices(); + choices.resize(radices.size(), 0); + size_t index = choices.size(); + while (index > 0) { + --index; + if (++choices[index] < radices[index]) break; + choices[index] = 0; + } + if (index == 0 && (choices.empty() || choices[0] == 0)) break; + } while (true); + } + definition.compiling = false; + --active_compile_depth_; + active_definition_ = enclosing; + } + + InlinedVector definitions_; + InlinedHashMap by_name_; + size_t active_compile_depth_ = 0; + Definition* active_definition_ = nullptr; +}; +} // namespace detail + +PatternLibrary::PatternLibrary(std::string_view text) : module_(std::make_shared(text)) {} +Pattern PatternLibrary::GetPattern(std::string_view name) const { return module_->GetPattern(name); } + +Rule::Rule(std::string name, Pattern pattern, ReplacementFunction replacement, CheckAndBind check_and_bind, + InlinedVector opsets) + : name(std::move(name)), + pattern(std::move(pattern)), + replacement(std::move(replacement)), + check_and_bind(std::move(check_and_bind)), + opsets(std::move(opsets)) {} + +Rule::Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, CheckAndBind check_and_bind, + InlinedVector opsets) + : Rule(std::move(name), Pattern(pattern), std::move(replacement), std::move(check_and_bind), std::move(opsets)) {} + +Rule::Rule(std::string name, Pattern pattern, std::string_view replacement, CheckAndBind check_and_bind, + InlinedVector opsets) + : Rule(std::move(name), std::move(pattern), Replacement(replacement), std::move(check_and_bind), std::move(opsets)) {} + +Rule::Rule(std::string name, std::string_view pattern, std::string_view replacement, CheckAndBind check_and_bind, + InlinedVector opsets) + : Rule(std::move(name), Pattern(pattern), Replacement(replacement), std::move(check_and_bind), std::move(opsets)) {} + +Status Pattern::MatchAt(const Graph& graph, NodeIndex root, std::optional& result) const { + std::optional candidate; + ORT_RETURN_IF_ERROR(VisitMatches(graph, root, [&](Match& match) { + candidate = std::move(match); + return false; + })); + result = std::move(candidate); + return Status::OK(); +} + +Status Pattern::VisitMatches(const Graph& graph, NodeIndex root_index, + const std::function& visitor) const { + ORT_RETURN_IF(!visitor, "Pattern visitor must be callable."); + ORT_RETURN_IF(graph.GraphResolveNeeded(), "Pattern matching requires an already resolved graph."); + ORT_RETURN_IF(root_index >= static_cast(graph.MaxNodeIndex()), + "Pattern root does not exist in the supplied graph."); + const Node* root = graph.GetNode(root_index); + ORT_RETURN_IF(!root, "Pattern root does not exist in the supplied graph."); + detail::PatternModule::ValidateEntry(*this, ""); + for (const auto& alternative : alternatives_) detail::PatternModule::ValidateEntry(*alternative, ""); + Producers producers; + for (const auto& node : graph.Nodes()) + for (const NodeArg* output : node.OutputDefs()) + if (output->Exists()) producers[output] = node.Index(); + auto visit = [&](const Pattern& plan) { + Match match; + match.pattern = &plan; + FailureStage stage = FailureStage::Structure; + std::string reason; + if (!MatchPattern(plan, graph, *root, producers, match, stage, reason)) return true; + match.owned_pattern_ = std::make_shared(plan); + match.pattern = match.owned_pattern_.get(); + return visitor(match); + }; + if (!visit(*this)) return Status::OK(); + for (const auto& alternative : alternatives_) + if (!visit(*alternative)) break; + return Status::OK(); +} + +Status Replacement::Instantiate(Graph& graph, const TensorBindings& tensors, const NodeAttributes& attributes, + Instantiation& result, const Node* annotation_source, + const InlinedHashSet& absent_attributes) const { + ORT_RETURN_IF(annotation_source && + (annotation_source->Index() >= static_cast(graph.MaxNodeIndex()) || + graph.GetNode(annotation_source->Index()) != annotation_source), + "Replacement annotation source must belong to the supplied graph."); + for (const auto& [name, value] : tensors) { + ORT_RETURN_IF(attributes.count(name) || absent_attributes.count(name), + "Replacement binding is both tensor and attribute: ", name); + ORT_RETURN_IF(value && (!value->Exists() || graph.GetNodeArg(value->Name()) != value), + "Replacement tensor binding must belong to the supplied graph: ", name); + } + for (const auto& name : absent_attributes) + ORT_RETURN_IF(name.empty() || attributes.count(name), "Invalid or conflicting absent attribute binding: ", name); + auto tensor = [&](const Input& input, NodeArg*& value) -> Status { + const auto* name = std::get_if(&input); + ORT_RETURN_IF(!name, "Standalone replacement cannot use compiled pattern handles."); + auto found = tensors.find(*name); + ORT_RETURN_IF(found == tensors.end(), "Missing replacement tensor binding: ", *name); + value = found->second ? graph.GetNodeArg(found->second->Name()) : nullptr; + return Status::OK(); + }; + auto attribute = [&](const std::string& name, std::optional& value) -> Status { + if (absent_attributes.count(name)) { + value.reset(); + return Status::OK(); + } + auto found = attributes.find(name); + ORT_RETURN_IF(found == attributes.end(), "Missing replacement attribute binding: ", name); + value = found->second; + return Status::OK(); + }; + Instantiation staged; + for (size_t i = 0; i < outputs_.size(); ++i) + staged.outputs.push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("rewrite_output"), nullptr)); + InlinedVector initializers; + ORT_RETURN_IF_ERROR(StageReplacement(graph, *this, staged.outputs, tensor, attribute, annotation_source, + staged.nodes, initializers)); + result = std::move(staged); + return Status::OK(); +} + +Pattern::Pattern(std::string_view text) { ORT_THROW_IF_ERROR(detail::PatternParser(text).Parse(*this)); } + +Replacement::Replacement(std::string_view text) { + Pattern parsed; + InlinedVector parsed_attributes; + InlinedVector, 1> parsed_bindings; + ORT_THROW_IF_ERROR(detail::PatternParser(text, &parsed_attributes, &parsed_bindings).Parse(parsed)); + InlinedHashMap outputs; + InlinedVector nodes; + InlinedHashSet attribute_names; + InlinedHashSet tensor_names; + for (size_t i = 0; i < parsed.Nodes().size(); ++i) { + const auto& node = parsed.Nodes()[i]; + NodeRecipe recipe{node.op.name, node.op.domain, {}, std::move(parsed_attributes[i]), node.outputs.size()}; + recipe.attribute_bindings = std::move(parsed_bindings[i]); + for (const auto& [name, binding] : recipe.attribute_bindings) attribute_names.insert(binding); + for (const Value& input : node.inputs) { + if (input.handle < 0) { + recipe.inputs.push_back(std::monostate{}); + continue; + } + auto literal = parsed.Constants().find(input.handle); + if (literal != parsed.Constants().end()) { + recipe.inputs.push_back(Scalar{literal->second}); + continue; + } + auto found = outputs.find(input.handle); + if (found == outputs.end()) { + tensor_names.insert(input.name); + recipe.inputs.push_back(input.name); + } else + recipe.inputs.push_back(found->second); + } + for (const auto& name : attribute_names) + ORT_ENFORCE(!tensor_names.count(name), "Replacement input cannot be both a tensor and an attribute: ", name); + for (size_t slot = 0; slot < node.outputs.size(); ++slot) + outputs.emplace(node.outputs[slot].handle, Output{nodes.size(), slot}); + nodes.push_back(std::move(recipe)); + } + for (Value value : parsed.Outputs()) outputs_.push_back(outputs.at(value.handle)); + plan_ = std::make_shared>(std::move(nodes)); +} + +const NodeArg* Match::valueOf(std::string_view name) const { + ORT_ENFORCE(pattern, "Named lookup requires a matched pattern."); + const Value value = pattern->NamedValue(name); + ORT_ENFORCE(!pattern->IsAttribute(value), "Pattern binding is an attribute: ", name); + ORT_ENFORCE(value.handle >= 0 && static_cast(value.handle) < values.size() && + (values[value.handle] || pattern->IsOptional(value)), + "Pattern value is not bound: ", name); + return values[value.handle]; +} + +const NodeArg* Match::optionalValueOf(std::string_view name) const { + ORT_ENFORCE(pattern, "Named lookup requires a matched pattern."); + const Value value = pattern->NamedValue(name); + ORT_ENFORCE(!pattern->IsAttribute(value), "Pattern binding is an attribute: ", name); + return valueOf(name); +} + +std::string BindingView::Qualify(std::string_view name) const { + return path_.empty() ? std::string(name) : path_ + "." + std::string(name); +} + +const NodeArg* BindingView::valueOf(std::string_view name) const { return match_->valueOf(Qualify(name)); } +const NodeArg* BindingView::optionalValueOf(std::string_view name) const { + return match_->optionalValueOf(Qualify(name)); +} +const ONNX_NAMESPACE::AttributeProto& BindingView::attributeOf(std::string_view name) const { + return match_->attributeOf(Qualify(name)); +} + +const ONNX_NAMESPACE::AttributeProto* BindingView::optionalAttributeOf(std::string_view name) const { + return match_->optionalAttributeOf(Qualify(name)); +} +const Dimension& BindingView::dimensionOf(std::string_view name) const { return match_->dimensionOf(Qualify(name)); } +std::optional BindingView::elementTypeOf(std::string_view name) const { + return match_->elementTypeOf(Qualify(name)); +} +BindingView BindingView::instance(std::string_view name) const { + ORT_ENFORCE(match_->pattern && !name.empty(), "Instance lookup requires a matched pattern and a name."); + const std::string path = Qualify(name); + ORT_ENFORCE(match_->pattern->instances_.count(path), "Unknown pattern instance: ", path); + return BindingView(*match_, path); +} +const Pattern::Instance& BindingView::Selection() const { + ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); + const auto* selection = &match_->pattern->instances_.at(path_); + for (size_t depth = 0; depth < alternative_depth_; ++depth) { + ORT_ENFORCE(selection->selected, "No selected alternative at binding path: ", path_); + selection = selection->selected.get(); + } + return *selection; +} +std::string_view BindingView::definition() const { return Selection().definition; } +std::string_view BindingView::interfaceName() const { return Selection().interface_name; } +std::string_view BindingView::alternative() const { return Selection().alternative; } +BindingView BindingView::alternativeView() const { + ORT_ENFORCE(Selection().selected, "Pattern has no selected alternative: ", path_); + BindingView result = *this; + ++result.alternative_depth_; + return result; +} + +namespace { +template +Status ReadConstant(const Graph& graph, const NodeArg* value, int32_t type, InlinedVector& data, bool& available) { + const auto* tensor = value ? graph_utils::GetConstantInitializer(graph, value->Name()) : nullptr; + if (!tensor || tensor->data_type() != type || + tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) { + available = false; + return Status::OK(); + } + size_t size = 1; + for (int64_t dimension : tensor->dims()) { + ORT_RETURN_IF(dimension < 0 || static_cast(dimension) > std::numeric_limits::max(), + "Invalid constant tensor dimension."); + const auto dim = static_cast(dimension); + ORT_RETURN_IF(dim != 0 && size > std::numeric_limits::max() / dim, "Constant tensor size overflow."); + size *= dim; + } + InlinedVector decoded(size); + ORT_RETURN_IF_ERROR(utils::UnpackTensor(*tensor, tensor->has_raw_data() ? tensor->raw_data().data() : nullptr, + tensor->raw_data().size(), + decoded.data(), decoded.size())); + data = std::move(decoded); + available = true; + return Status::OK(); +} +} // namespace + +Status ReadConstantTensor(const Graph& graph, const NodeArg* value, InlinedVector& data, bool& available) { + return ReadConstant(graph, value, ONNX_NAMESPACE::TensorProto_DataType_FLOAT, data, available); +} +Status ReadConstantTensor(const Graph& graph, const NodeArg* value, InlinedVector& data, bool& available) { + return ReadConstant(graph, value, ONNX_NAMESPACE::TensorProto_DataType_INT64, data, available); +} + +void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) { + ORT_ENFORCE(!name.empty() && value.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_UNDEFINED && + !value.has_ref_attr_name(), + "Extra bindings require a name and a concrete attribute value."); + ORT_ENFORCE(!dimensions_.count(std::string(name)) && !element_types_.count(std::string(name)), + "Extra attribute binding cannot overwrite a dimension or element type: ", name); + if (pattern) { + ORT_ENFORCE(!pattern->aliases_.count(std::string(name)) && !pattern->instances_.count(std::string(name)), + "Extra attribute binding cannot overwrite a pattern binding: ", name); + ORT_ENFORCE(std::none_of(pattern->values_.begin(), pattern->values_.end(), + [&](const Value& binding) { return binding.name == name; }), + "Extra attribute binding cannot overwrite a pattern value: ", name); + ORT_ENFORCE(!pattern->AttributeNames().count(std::string(name)) && + !pattern->attribute_aliases_.count(std::string(name)), + "Extra attribute binding cannot overwrite a captured attribute: ", name); + } + value.set_name(std::string(name)); + extra_attributes_.insert_or_assign(std::string(name), std::move(value)); +} + +const ONNX_NAMESPACE::AttributeProto& Match::attributeOf(std::string_view name) const { + const auto* value = optionalAttributeOf(name); + ORT_ENFORCE(value, "Attribute capture is absent: ", name); + return *value; +} + +const ONNX_NAMESPACE::AttributeProto* Match::optionalAttributeOf(std::string_view name) const { + if (pattern) { + auto alias = pattern->attribute_aliases_.find(std::string(name)); + if (alias != pattern->attribute_aliases_.end()) name = alias->second; + } + auto captured = captured_attributes_.find(std::string(name)); + if (captured != captured_attributes_.end()) return &captured->second; + if (absent_attributes_.count(std::string(name))) return nullptr; + auto found = extra_attributes_.find(std::string(name)); + ORT_ENFORCE(!name.empty() && found != extra_attributes_.end(), "Missing replacement attribute binding: ", name); + return &found->second; +} + +bool Match::BindAttribute(std::string_view name, const ONNX_NAMESPACE::AttributeProto& value) { + if (absent_attributes_.count(std::string(name))) return false; + if (value.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_UNDEFINED || value.has_ref_attr_name()) return false; + auto normalized = value; + normalized.set_name(std::string(name)); + normalized.clear_doc_string(); + auto [it, inserted] = captured_attributes_.emplace(std::string(name), normalized); + return inserted || it->second.SerializeAsString() == normalized.SerializeAsString(); +} + +bool Match::BindAbsentAttribute(std::string_view name) { + if (captured_attributes_.count(std::string(name))) return false; + absent_attributes_.insert(std::string(name)); + return true; +} + +const Dimension& Match::dimensionOf(std::string_view name) const { + auto found = dimensions_.find(std::string(name)); + ORT_ENFORCE(found != dimensions_.end(), "Unknown pattern dimension: ", name); + return found->second; +} + +std::optional Match::elementTypeOf(std::string_view name) const { + auto found = element_types_.find(std::string(name)); + ORT_ENFORCE(found != element_types_.end(), "Unknown pattern element type: ", name); + if (!found->second) return std::nullopt; + return found->second; +} + +bool Match::BindDimension(std::string_view name, Dimension value) { + auto [it, inserted] = dimensions_.emplace(std::string(name), value); + if (inserted || std::holds_alternative(value)) return true; + if (std::holds_alternative(it->second)) { + it->second = std::move(value); + return true; + } + return it->second == value; +} + +bool Match::BindElementType(std::string_view name, int32_t value) { + auto [it, inserted] = element_types_.emplace(std::string(name), value); + if (inserted || !value) return true; + if (!it->second) { + it->second = value; + return true; + } + return it->second == value; +} + +PatternRewriter::PatternRewriter(Rules rules, Options options) + : GraphTransformer("PatternRewriter"), options_(options) { + // Expand in rule/branch order before indexing. Every attempt gets a fresh Match, + // including check-and-bind updates; replacement failures are never search failures. + for (auto& rule : rules) { + auto alternatives = std::move(rule.pattern.alternatives_); + rule.pattern.alternatives_.clear(); + rules_.push_back(rule); + for (const auto& alternative : alternatives) { + Rule concrete = rule; + concrete.pattern = *alternative; + rules_.push_back(std::move(concrete)); + } + } + ORT_ENFORCE(options_.max_iterations > 0, "PatternRewriter max_iterations must be positive."); + InlinedHashMap signatures; + for (size_t i = 0; i < rules_.size(); ++i) { + Rule& rule = rules_[i]; + ORT_ENFORCE(!rule.name.empty() && !rule.pattern.Nodes().empty() && rule.replacement, + "Rules need a name, at least one node and a recipe."); + InlinedHashSet domains; + for (auto& requirement : rule.opsets) { + if (requirement.domain == kOnnxDomainAlias) requirement.domain = kOnnxDomain; + ORT_ENFORCE(requirement.min_version > 0 && requirement.max_version >= requirement.min_version, + "Rule opset requirements need a positive, nonempty inclusive range."); + ORT_ENFORCE(domains.insert(requirement.domain).second, "Duplicate rule opset domain: ", requirement.domain); + } + for (const auto& node : rule.pattern.Nodes()) { + ORT_ENFORCE(domains.count(node.op.domain), "Missing rule opset requirement for domain: ", node.op.domain); + } + if (auto* recipe = rule.replacement.target(); recipe && recipe->plan_) { + auto nodes = *recipe->plan_; + for (auto& node : nodes) { + if (node.domain == kOnnxDomainAlias) node.domain = kOnnxDomain; + ORT_ENFORCE(domains.count(node.domain), "Missing rule opset requirement for replacement domain: ", node.domain); + for (const auto& [name, binding] : node.attribute_bindings) { + auto alias = rule.pattern.aliases_.find(binding); + ORT_ENFORCE(alias == rule.pattern.aliases_.end() || rule.pattern.IsAttribute(alias->second), + "Replacement attribute binding conflicts with a pattern value: ", binding); + ORT_ENFORCE(std::none_of(rule.pattern.values_.begin(), rule.pattern.values_.end(), + [&](const Value& value) { + return value.name == binding && !rule.pattern.IsAttribute(value); + }), + "Replacement attribute binding conflicts with a pattern value: ", binding); + } + for (auto& input : node.inputs) { + if (const auto* name = std::get_if(&input)) { + auto value = rule.pattern.NamedValue(*name); + ORT_ENFORCE(rule.pattern.IsInput(value) && !rule.pattern.IsAttribute(value), + "Replacement inputs must be pattern tensor input variables: ", *name); + input = std::move(value); + } + } + } + recipe->plan_ = std::make_shared>(std::move(nodes)); + } + // Every node must reach the last node. Disconnected definitions are not patterns. + const auto& root_outputs = rule.pattern.Nodes().back().outputs; + InlinedHashSet reachable; + for (Value output : root_outputs) reachable.insert(output.handle); + for (auto it = rule.pattern.Nodes().rbegin(); it != rule.pattern.Nodes().rend(); ++it) { + ORT_ENFORCE( + std::any_of(it->outputs.begin(), it->outputs.end(), [&](Value output) { return reachable.count(output); }), + "Disconnected pattern node."); + for (Value output : it->outputs) reachable.insert(output.handle); + for (Value input : it->inputs) + if (input.handle >= 0) reachable.insert(input.handle); + } + for (int handle : rule.pattern.attribute_inputs_) reachable.insert(handle); + for (int handle : rule.pattern.optional_inputs_) reachable.insert(handle); + ORT_ENFORCE(reachable.size() == rule.pattern.ValueCount(), "Unused pattern input."); + invalidation_depth_ = std::max(invalidation_depth_, rule.pattern.Nodes().size()); + if (!options_.indexed) continue; + std::string signature = Signature(rule.pattern); + // Opaque predicates cannot be compared safely; keep their plans distinct. + if (std::any_of(rule.pattern.Nodes().begin(), rule.pattern.Nodes().end(), + [](const auto& node) { return !node.attributes.empty(); }) || + !rule.pattern.OrderedTypeConstraints().empty()) + signature += "#" + std::to_string(i); + auto [it, inserted] = signatures.emplace(signature, groups_.size()); + if (inserted) { + const auto& root = rule.pattern.Nodes().back().op; + roots_[Key(root.name, root.domain)].push_back(groups_.size()); + groups_.push_back({rule.pattern, {}}); + } + groups_[it->second].rules.push_back(i); + } +} + +Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { + statistics = {}; + ORT_RETURN_IF(graph.GraphResolveNeeded(), "PatternRewriter requires an already resolved graph."); + auto report = [&](size_t id, NodeIndex root, FailureStage stage, std::string_view reason) { + if (options_.on_failure) + options_.on_failure({rules_[id].name, root, stage, std::string(reason)}); + }; + // Import applicability is evaluated once per rule and graph, never per matched node. + InlinedVector applicable; + InlinedVector applicable_rules; + for (size_t id = 0; id < rules_.size(); ++id) { + bool supported = std::all_of(rules_[id].opsets.begin(), rules_[id].opsets.end(), [&](const auto& requirement) { + auto imported = graph.DomainToVersionMap().find(requirement.domain); + return imported != graph.DomainToVersionMap().end() && imported->second >= requirement.min_version && + imported->second <= requirement.max_version; + }); + applicable.push_back(supported); + if (supported) + applicable_rules.push_back(id); + else + report(id, std::numeric_limits::max(), FailureStage::Opset, "Missing or unsupported domain import."); + } + if (applicable_rules.empty()) return Status::OK(); + InlinedVector, 1> applicable_groups; + for (const auto& group : groups_) { + InlinedVector ids; + for (size_t id : group.rules) { + if (applicable[id]) ids.push_back(id); + } + applicable_groups.push_back(std::move(ids)); + } + Producers producers; + InlinedVector worklist; + InlinedHashSet queued; + GraphViewer viewer(graph); + for (NodeIndex index : viewer.GetNodesInTopologicalOrder()) { + const Node& node = *graph.GetNode(index); + worklist.push_back(index); + queued.insert(index); + for (const NodeArg* output : node.OutputDefs()) { + if (output->Exists()) producers[output] = index; + } + } + for (size_t iteration = 0; iteration < options_.max_iterations && !worklist.empty(); ++iteration) { + ++statistics.iterations; + InlinedVector next_worklist; + queued.clear(); + const size_t rewrites_before = statistics.rewrites; + // Each batch is immutable: revisits and new nodes wait until the next iteration. + for (size_t cursor = 0; cursor < worklist.size(); ++cursor) { + const Node* root = graph.GetNode(worklist[cursor]); + if (!root) continue; + ++statistics.visited; + size_t winner = std::numeric_limits::max(); + Match winning_match; + const Pattern* winning_pattern = nullptr; + struct CandidateMatch { + Match match; + FailureStage stage = FailureStage::Structure; + std::string reason; + bool matched = false; + }; + auto match_group = [&](const Pattern& pattern, CandidateMatch& candidate) { + ++statistics.structural_checks; + candidate.matched = MatchPattern(pattern, graph, *root, producers, candidate.match, + candidate.stage, candidate.reason); + if (candidate.matched && !CanRewritePartition(candidate.match, options_)) { + candidate.matched = false; + candidate.stage = FailureStage::ExecutionProvider; + candidate.reason = "Mixed or unsupported execution providers."; + } + }; + auto consider = [&](size_t id, const CandidateMatch& candidate) { + ++statistics.candidates; + if (!candidate.matched) { + report(id, root->Index(), candidate.stage, candidate.reason); + return; + } + Match rule_match = candidate.match; + rule_match.owned_pattern_ = std::make_shared(rules_[id].pattern); + rule_match.pattern = rule_match.owned_pattern_.get(); + if (rules_[id].check_and_bind) { + ++statistics.check_and_bind_calls; + if (!rules_[id].check_and_bind(graph, rule_match)) { + report(id, root->Index(), FailureStage::CheckAndBind, + rule_match.RejectionReason().empty() ? std::string_view("Check-and-bind rejected match.") + : rule_match.RejectionReason()); + return; + } + } + if (options_.match_only) ++statistics.matches; + if (id < winner) { + winner = id; + winning_match = std::move(rule_match); + winning_pattern = winning_match.pattern; + } + }; + if (options_.indexed) { + auto it = roots_.find(Key(root->OpType(), root->Domain())); + if (it != roots_.end()) { + InlinedVector> candidates; + for (size_t group : it->second) + for (size_t id : applicable_groups[group]) candidates.emplace_back(id, group); + std::sort(candidates.begin(), candidates.end()); + InlinedHashMap cache; + // Share structural work, but never run a lower-priority callback before an earlier rule. + for (const auto& [id, group] : candidates) { + auto [cached, inserted] = cache.try_emplace(group); + if (inserted) match_group(groups_[group].pattern, cached->second); + consider(id, cached->second); + if (!options_.match_only && winning_pattern) break; + } + } + } else { + for (size_t id : applicable_rules) { + CandidateMatch candidate; + match_group(rules_[id].pattern, candidate); + consider(id, candidate); + if (!options_.match_only && winning_pattern) break; + } + } + if (options_.match_only || !winning_pattern) continue; + ++statistics.matches; + InlinedVector affected; + for (int value = 0; value < static_cast(winning_pattern->ValueCount()); ++value) { + if (!winning_pattern->IsInput(value)) continue; + auto it = producers.find(winning_match.values[value]); + if (it != producers.end()) affected.push_back(it->second); + } + InlinedVector added; + auto recipe = rules_[winner].replacement(graph, winning_match); + auto declared = [&](const std::string& node_domain) { + const auto& domain = node_domain == kOnnxDomainAlias ? kOnnxDomain : node_domain; + return std::any_of(rules_[winner].opsets.begin(), rules_[winner].opsets.end(), + [&](const auto& requirement) { return requirement.domain == domain; }); + }; + auto replace = [&]() -> Status { + ORT_RETURN_IF(!recipe.Plan(), "Replacement has no parsed plan."); + for (const auto& node : *recipe.Plan()) + ORT_RETURN_IF(!declared(node.domain), "Replacement domain has no rule opset requirement: ", node.domain); + return Replace(graph, *winning_pattern, winning_match, recipe, producers, added); + }; + auto status = replace(); + if (!status.IsOK()) { + report(winner, root->Index(), FailureStage::Replacement, status.ErrorMessage()); + return status; + } + ++statistics.rewrites; + affected.insert(affected.end(), added.begin(), added.end()); + // A changed producer or use count can affect a deeper pattern root, not just a direct consumer. + InlinedHashSet seen; + for (size_t depth = 0; depth <= invalidation_depth_ && !affected.empty(); ++depth) { + std::sort(affected.begin(), affected.end()); + InlinedVector next; + for (NodeIndex index : affected) { + if (!seen.insert(index).second) continue; + const Node* node = graph.GetNode(index); + if (!node) continue; + if (queued.insert(index).second) next_worklist.push_back(index); + for (auto it = node->OutputEdgesBegin(); it != node->OutputEdgesEnd(); ++it) { + next.push_back(it->GetNode().Index()); + } + } + affected = std::move(next); + } + } + if (options_.match_only || statistics.rewrites == rewrites_before) break; + worklist = std::move(next_worklist); + if (iteration + 1 == options_.max_iterations && !worklist.empty()) + statistics.iteration_limit_reached = true; + } + return Status::OK(); +} + +Status PatternRewriter::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { + for (auto& node : graph.Nodes()) ORT_RETURN_IF_ERROR(Recurse(node, modified, graph_level, logger)); + Statistics statistics; + ORT_RETURN_IF_ERROR(Run(graph, statistics)); + modified |= statistics.rewrites != 0; + return Status::OK(); +} + +} // namespace rewrite +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h new file mode 100644 index 0000000000000..df8be925cdbe4 --- /dev/null +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -0,0 +1,355 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { +namespace rewrite { + +namespace detail { +class PatternParser; +class PatternModule; +} // namespace detail + +struct Value { + std::string name; + int handle = -1; + Value() = default; + Value(int id) : handle(id) {} + Value(std::string label, int id) : name(std::move(label)), handle(id) {} + // Handles index Match::values; names are diagnostic and do not affect identity. + operator int() const { return handle; } +}; + +struct AttributeMatcher { + std::string name; + std::function predicate; + std::string capture; + bool optional = false; + bool Matches(const ONNX_NAMESPACE::AttributeProto& value) const { return predicate(value); } + static AttributeMatcher Equal(std::string name, int64_t value); + static AttributeMatcher Equal(std::string name, int value) { return Equal(std::move(name), int64_t{value}); } + static AttributeMatcher Equal(std::string name, float value); + static AttributeMatcher Equal(std::string name, double value) { + return Equal(std::move(name), static_cast(value)); + } + static AttributeMatcher Equal(std::string name, std::string value); +}; + +struct Operator { + std::string name; + std::string domain; +}; + +using Literal = std::variant, InlinedVector>; + +using Dimension = std::variant; + +struct TypeConstraint { + std::string value_name; + int32_t element_type = 0; // Zero means unconstrained. + std::string element_variable; + InlinedVector dimensions; + bool any_rank = false; + bool prefix = false; + int32_t attribute_type = 0; // Scalar/list attribute interpretation, when applicable. + bool tensor_type = true; + bool absent = false; + bool immutable = false; +}; + +struct Match; + +struct OpsetRequirement { + std::string domain; + int min_version; + int max_version; // Inclusive model/function import range, not schema since-versions. +}; + +// A straight-line definition describes a DAG, not necessarily a chain. +// Reusing a Value requires exactly the same NodeArg, including for input variables. +class Pattern { + public: + explicit Pattern(std::string_view text); + Value NamedValue(std::string_view name) const; + struct NodePattern { + Operator op; + InlinedVector inputs; + InlinedVector outputs; + InlinedVector attributes; + }; + const auto& Nodes() const { return nodes_; } + size_t ValueCount() const { return value_count_; } + bool IsInput(Value value) const; + const auto& Constants() const { return constants_; } + const auto& Outputs() const { return outputs_; } + bool IsOptional(Value value) const { return optional_inputs_.count(value.handle) != 0; } + const auto& AttributeNames() const { return attribute_names_; } + const auto& TypeConstraints() const { return type_constraints_; } + bool IsAttribute(Value value) const { return attribute_inputs_.count(value.handle) != 0; } + const auto& OrderedTypeConstraints() const { return ordered_type_constraints_; } + // First rewrite-safe match at this root, or nullopt. On error, result is unchanged. + Status MatchAt(const Graph& graph, NodeIndex root, std::optional& result) const; + // Visit rewrite-safe candidates at one root in expansion order. true continues; false stops. + // Requires a resolved graph. The visitor must not mutate it. Does not filter opsets or providers. + Status VisitMatches(const Graph& graph, NodeIndex root, const std::function& visitor) const; + + private: + friend class detail::PatternParser; + friend class detail::PatternModule; + friend class PatternLibrary; + friend class BindingView; + friend struct Replacement; + friend struct Match; + friend class PatternRewriter; + Pattern() = default; + Value Input(std::string_view name); + Value Constant(float value); + Value Constant(int64_t value); + Value Constant(Literal value); + void Name(Value value, std::string_view name); + void Attributes(InlinedVector matchers); + InlinedVector MultiOp(Operator op, gsl::span inputs, size_t output_count); + + size_t value_count_ = 0; + size_t expanded_calls_ = 0; + InlinedVector nodes_; + InlinedVector values_; + InlinedHashMap constants_; + InlinedVector outputs_; + InlinedHashSet optional_inputs_; + InlinedHashSet attribute_names_; + InlinedHashMap> type_constraints_; + InlinedHashSet attribute_inputs_; + InlinedHashSet required_tensor_bindings_; + InlinedHashSet composed_tensor_annotations_; + InlinedHashMap aliases_; + InlinedHashMap attribute_aliases_; + InlinedVector>> ordered_type_constraints_; + struct Instance { + std::string definition; + std::string interface_name; + std::string alternative; + std::shared_ptr selected; + }; + InlinedHashMap instances_; + InlinedVector, 1> alternatives_; + Value AddValue(std::string_view name); +}; + +// A library owns interfaces and named patterns (graph bodies or alternatives). +class PatternLibrary { + public: + explicit PatternLibrary(std::string_view text); + Pattern GetPattern(std::string_view name) const; + + private: + std::shared_ptr module_; +}; + +// Non-owning access; the parent Match and its unchanged graph must outlive this view. +class BindingView { + public: + const NodeArg* valueOf(std::string_view name) const; + const NodeArg* optionalValueOf(std::string_view name) const; + const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; + // Null only for a captured absence; unknown or wrong-kind names are errors. + const ONNX_NAMESPACE::AttributeProto* optionalAttributeOf(std::string_view name) const; + const Dimension& dimensionOf(std::string_view name) const; + std::optional elementTypeOf(std::string_view name) const; + BindingView instance(std::string_view name) const; + std::string_view definition() const; + std::string_view interfaceName() const; + // Empty for a body call; otherwise the selected immediate member name. + std::string_view alternative() const; + // Descend one named alternative without losing the enclosing selection tag. + BindingView alternativeView() const; + + private: + friend struct Match; + BindingView(const Match& match, std::string path) : match_(&match), path_(std::move(path)) {} + std::string Qualify(std::string_view name) const; + const Pattern::Instance& Selection() const; + const Match* match_; + std::string path_; + size_t alternative_depth_ = 0; +}; + +// A missing/nonconstant/wrongly typed or externally stored payload is unavailable. +// Malformed available payloads return an error. Outputs change only on success. +Status ReadConstantTensor(const Graph& graph, const NodeArg* value, InlinedVector& data, bool& available); +Status ReadConstantTensor(const Graph& graph, const NodeArg* value, InlinedVector& data, bool& available); + +struct Match { + // Matcher plumbing, not a supported interface for constructing or modifying matches. + InlinedVector nodes; + InlinedVector values; + const NodeArg* valueOf(std::string_view name) const; + const NodeArg* optionalValueOf(std::string_view name) const; + void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value); + const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; + const ONNX_NAMESPACE::AttributeProto* optionalAttributeOf(std::string_view name) const; + const Dimension& dimensionOf(std::string_view name) const; + std::optional elementTypeOf(std::string_view name) const; + BindingView bindings() const { return BindingView(*this, ""); } + BindingView instance(std::string_view name) const { return bindings().instance(name); } + void ClearUpdates() { + extra_attributes_.clear(); + rejection_reason_.clear(); + } + bool BindAttribute(std::string_view name, const ONNX_NAMESPACE::AttributeProto& value); + bool BindAbsentAttribute(std::string_view name); + bool BindDimension(std::string_view name, Dimension value); + bool BindElementType(std::string_view name, int32_t value); + bool Reject(std::string_view reason) { + rejection_reason_ = reason; + return false; + } + const std::string& RejectionReason() const { return rejection_reason_; } + const Pattern* pattern = nullptr; + + private: + friend class Pattern; + friend class PatternRewriter; + std::shared_ptr owned_pattern_; // Copies own the plan, but still borrow graph objects. + NodeAttributes captured_attributes_; + InlinedHashSet absent_attributes_; + InlinedHashMap dimensions_; + InlinedHashMap element_types_; + NodeAttributes extra_attributes_; + std::string rejection_reason_; +}; + +struct Replacement { + struct Instantiation { + InlinedVector nodes; + InlinedVector outputs; + }; + using TensorBindings = InlinedHashMap; + struct Output { + size_t node; + size_t slot = 0; + }; + struct Scalar { + Literal value; + }; + using Input = std::variant; + struct NodeRecipe { + std::string op; + std::string domain; + InlinedVector inputs; + NodeAttributes attributes; + size_t output_count = 1; + InlinedHashMap attribute_bindings; + }; + explicit Replacement(std::string_view text); + // A parsed replacement is a reusable recipe callback; copying it shares the immutable plan. + Replacement operator()(const Graph&, const Match&) const { return *this; } + const auto& Plan() const { return plan_; } + const auto& Outputs() const { return outputs_; } + // Build a fresh fragment without matching, deleting, rewiring, or changing graph outputs. + // result changes only on success. Call Graph::Resolve before using the constructed graph. + // absent_attributes explicitly omits keyword bindings; unresolved names remain errors. + Status Instantiate(Graph& graph, const TensorBindings& tensors, const NodeAttributes& attributes, + Instantiation& result, const Node* annotation_source = nullptr, + const InlinedHashSet& absent_attributes = {}) const; + + private: + friend class PatternRewriter; + std::shared_ptr> plan_; + InlinedVector outputs_; +}; + +struct Rule { + using CheckAndBind = std::function; + using ReplacementFunction = std::function; + + Rule(std::string name, Pattern pattern, ReplacementFunction replacement, CheckAndBind check_and_bind = {}, + InlinedVector opsets = {}); + Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, CheckAndBind check_and_bind = {}, + InlinedVector opsets = {}); + Rule(std::string name, Pattern pattern, std::string_view replacement, CheckAndBind check_and_bind = {}, + InlinedVector opsets = {}); + Rule(std::string name, std::string_view pattern, std::string_view replacement, CheckAndBind check_and_bind = {}, + InlinedVector opsets = {}); + + std::string name; + Pattern pattern; + // Check-and-bind callbacks may prepare match-local attributes. Neither callback may mutate the graph + // or retain pointers into a Match; replacement callbacks receive a read-only match. + ReplacementFunction replacement; + CheckAndBind check_and_bind; // Empty accepts a structurally valid match without additional checks. + InlinedVector opsets; +}; + +using Rules = InlinedVector; + +struct Statistics { + size_t visited = 0; + size_t candidates = 0; + size_t structural_checks = 0; + size_t check_and_bind_calls = 0; + size_t matches = 0; + size_t rewrites = 0; + size_t iterations = 0; + bool iteration_limit_reached = false; +}; + +enum class FailureStage { Opset, + Structure, + TypeShape, + Constant, + Attribute, + Boundary, + ExecutionProvider, + CheckAndBind, + Replacement }; + +struct MatchFailure { + std::string rule; + NodeIndex root; + FailureStage stage; + std::string reason; +}; + +struct Options { + bool match_only = false; + bool indexed = true; // false is the rule-by-rule reference implementation. + // Unassigned nodes are always eligible. Assigned nodes must share an allowed provider. + InlinedHashSet compatible_execution_providers{kCpuExecutionProvider}; + // Called synchronously. Records own their strings and contain no graph/match pointers. + std::function on_failure; + size_t max_iterations = 10; // Positive upper bound on candidate batches; match-only uses one batch. +}; + +// Immutable compiled rules; statistics belong to a call, not this transformer. +// Not registered among ORT's default optimizers. Call Apply or Run explicitly. +class PatternRewriter final : public GraphTransformer { + public: + explicit PatternRewriter(Rules rules, Options options = {}); + Status Run(Graph& graph, Statistics& statistics) const; + + private: + struct Group { + Pattern pattern; + InlinedVector rules; + }; + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; + Rules rules_; + InlinedVector groups_; + InlinedHashMap> roots_; + size_t invalidation_depth_ = 0; + Options options_; +}; + +} // namespace rewrite +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc new file mode 100644 index 0000000000000..6cb78e86dc988 --- /dev/null +++ b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc @@ -0,0 +1,765 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include +#include + +#include "core/optimizer/pattern_rewriter.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { +namespace { +using namespace rewrite; + +struct TestGraph { + TestGraph() + : model{"rewrite-components", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 17}}, {}, DefaultLoggingManager().DefaultLogger()} {} + Model model; + Graph& graph = model.MainGraph(); + ModelTestBuilder builder{graph}; +}; + +size_t Count(const Graph& graph, const std::string& op) { + size_t count = 0; + for (const auto& node : graph.Nodes()) count += node.OpType() == op; + return count; +} + +bool ComposeTransposePermutations(const Graph&, Match& match) { + const auto* shape = match.valueOf("x")->Shape(); + if (!shape) return match.Reject("Transpose composition requires known rank."); + const int rank = shape->dim_size(); + const auto& p = match.attributeOf("p"); + const auto& q = match.attributeOf("q"); + auto valid = [rank](const ONNX_NAMESPACE::AttributeProto& permutation) { + if (permutation.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_INTS || + permutation.ints_size() != rank) + return false; + InlinedHashSet seen; + for (int64_t axis : permutation.ints()) + if (axis < 0 || axis >= rank || !seen.insert(axis).second) return false; + return true; + }; + if (!valid(p) || !valid(q)) return match.Reject("Invalid transpose permutation."); + ONNX_NAMESPACE::AttributeProto composed; + composed.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); + for (int64_t axis : q.ints()) composed.add_ints(p.ints(static_cast(axis))); + match.update("composed", std::move(composed)); + return true; +} + +Pattern TransposePairPattern() { + return Pattern("(x,p:ints,q:ints)=>{a=Transpose(x,perm=p);return Transpose(a,perm=q)}"); +} + +Rule ComposeTransposeRule() { + return Rule{"ComposeTranspose", TransposePairPattern(), "(x,composed)=>Transpose(x,perm=composed)", ComposeTransposePermutations, {{kOnnxDomain, 13, 25}}}; +} +} // namespace + +TEST(PatternComponentTests, AnchoredMatchingDoesNotRewriteOrRequireARule) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {y}); + auto& root = b.AddNode("Sigmoid", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const auto before = test.model.ToProto().SerializeAsString(); + int calls = 0; + ASSERT_STATUS_OK(Pattern("(x:float[2])=>{a=Identity(x);return Sigmoid(a)}") + .VisitMatches(test.graph, root.Index(), [&](Match& match) { + ++calls; + EXPECT_EQ(match.valueOf("x"), x); + EXPECT_EQ(match.valueOf("a"), y); + return true; + })); + EXPECT_EQ(calls, 1); + EXPECT_EQ(test.model.ToProto().SerializeAsString(), before); + ASSERT_STATUS_OK(Pattern("(x:int64[2])=>return Sigmoid(x)") + .VisitMatches(test.graph, root.Index(), [&](Match&) { ADD_FAILURE(); return true; })); +} + +TEST(PatternComponentTests, OrderedCandidatesOwnPlansAndIsolatePreparedState) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {y}); + auto& root = b.AddNode("Sigmoid", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Match saved; + int calls = 0; + Rule::CheckAndBind check = [&](const Graph&, Match& match) { + ++calls; + if (calls == 1) { + ONNX_NAMESPACE::AttributeProto value; + value.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + value.set_i(1); + match.update("rejected", value); + return match.Reject("try next"); + } + EXPECT_THROW(match.attributeOf("rejected"), OnnxRuntimeException); + EXPECT_TRUE(match.RejectionReason().empty()); + ONNX_NAMESPACE::AttributeProto accepted; + accepted.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + accepted.set_i(2); + match.update("accepted", accepted); + return true; + }; + ASSERT_STATUS_OK(PatternLibrary(R"( + pattern E(x){v=Identity(x)|x|x;return Sigmoid(v)} + )") + .GetPattern("E") + .VisitMatches(test.graph, root.Index(), [&](Match& match) { + if (!check(test.graph, match)) return true; + saved = match; + return false; + })); + EXPECT_EQ(calls, 2); + EXPECT_EQ(saved.valueOf("x"), y); + EXPECT_EQ(saved.attributeOf("accepted").i(), 2); + EXPECT_EQ(saved.instance("v").alternative(), "x"); + Match copied = saved; + saved = {}; + EXPECT_EQ(copied.valueOf("x"), y); + EXPECT_EQ(copied.attributeOf("accepted").i(), 2); +} + +TEST(PatternComponentTests, InvalidContextAndCallbackErrorsAreNotNonmatches) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto& root = b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + Pattern pattern("(x)=>return Identity(x)"); + auto visitor = [](Match&) { return true; }; + EXPECT_FALSE(pattern.VisitMatches(test.graph, root.Index(), visitor).IsOK()); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_FALSE(pattern.VisitMatches(test.graph, std::numeric_limits::max(), visitor).IsOK()); + EXPECT_FALSE(pattern.VisitMatches(test.graph, root.Index(), {}).IsOK()); + EXPECT_THROW((void)pattern.VisitMatches(test.graph, root.Index(), [](Match&) -> bool { + ORT_THROW("callback error"); + }), + OnnxRuntimeException); +} + +TEST(PatternComponentTests, MatchAtReturnsFirstCandidateAndClearsOrdinaryNonmatch) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + auto& inner = b.AddNode("Identity", {x}, {y}); + auto& root = b.AddNode("Sigmoid", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const auto before = test.model.ToProto().SerializeAsString(); + std::optional result; + ASSERT_STATUS_OK(PatternLibrary("pattern E(x){v=Identity(x)|x;return Sigmoid(v)}") + .GetPattern("E") + .MatchAt(test.graph, root.Index(), result)); + ASSERT_TRUE(result); + EXPECT_EQ(result->valueOf("x"), x); + EXPECT_EQ(result->nodes.size(), 2u); + Match copy = *result; + EXPECT_FALSE(Pattern("(x)=>return Identity(x)").MatchAt(test.graph, std::numeric_limits::max(), result).IsOK()); + ASSERT_TRUE(result); + EXPECT_EQ(result->valueOf("x"), x); + EXPECT_THROW((void)Pattern("(x)=>return x").MatchAt(test.graph, root.Index(), result), OnnxRuntimeException); + ASSERT_TRUE(result); + EXPECT_EQ(result->valueOf("x"), x); + ASSERT_STATUS_OK(Pattern("(x)=>return Identity(x)").MatchAt(test.graph, root.Index(), result)); + EXPECT_FALSE(result); + EXPECT_EQ(copy.valueOf("x"), x); + ASSERT_STATUS_OK(Pattern("(x)=>return Identity(x)").MatchAt(test.graph, inner.Index(), result)); + ASSERT_TRUE(result); + EXPECT_EQ(result->valueOf("x"), x); + EXPECT_EQ(test.model.ToProto().SerializeAsString(), before); + test.graph.SetGraphResolveNeeded(); + EXPECT_FALSE(Pattern("(x)=>return Identity(x)").MatchAt(test.graph, inner.Index(), result).IsOK()); + ASSERT_TRUE(result); +} + +TEST(PatternComponentTests, OptionalAttributeCapturesDistinguishAbsenceFromUnknownNames) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto& root = b.AddNode("Transpose", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + std::optional result; + ASSERT_STATUS_OK(Pattern("(x,p:ints)=>return Transpose(x,perm=p)").MatchAt(test.graph, root.Index(), result)); + EXPECT_FALSE(result); + ASSERT_STATUS_OK(Pattern("(x)=>return Transpose(x)").MatchAt(test.graph, root.Index(), result)); + ASSERT_TRUE(result); + EXPECT_EQ(result->nodes[0]->GetAttributes().find("perm"), result->nodes[0]->GetAttributes().end()); + EXPECT_THROW(result->attributeOf("p"), OnnxRuntimeException); + ASSERT_STATUS_OK(Pattern("(x,p:ints|None)=>return Transpose(x,perm=p)").MatchAt(test.graph, root.Index(), result)); + ASSERT_TRUE(result); + EXPECT_EQ(result->optionalAttributeOf("p"), nullptr); + EXPECT_EQ(result->bindings().optionalAttributeOf("p"), nullptr); + EXPECT_THROW(result->attributeOf("p"), OnnxRuntimeException); + EXPECT_THROW(result->optionalAttributeOf("missing"), OnnxRuntimeException); + EXPECT_THROW(result->optionalAttributeOf("x"), OnnxRuntimeException); + EXPECT_THROW(result->valueOf("p"), OnnxRuntimeException); + EXPECT_THROW(result->optionalValueOf("p"), OnnxRuntimeException); + Match copy = *result; + ONNX_NAMESPACE::AttributeProto prepared; + prepared.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + prepared.set_i(1); + EXPECT_THROW(result->update("p", prepared), OnnxRuntimeException); + result->update("prepared", prepared); + EXPECT_EQ(result->optionalAttributeOf("prepared")->i(), 1); + result->Reject("rejected"); + result->ClearUpdates(); + EXPECT_EQ(result->optionalAttributeOf("p"), nullptr); + EXPECT_EQ(copy.optionalAttributeOf("p"), nullptr); + EXPECT_THROW(result->optionalAttributeOf("prepared"), OnnxRuntimeException); +} + +TEST(PatternComponentTests, OptionalAttributeTypesAndResolvedDefaults) { + for (const char* type : {"int", "ints", "float", "floats", "string", "strings", "tensor"}) { + const auto text = MakeString("(x,p:", type, "|None)=>return Identity(x,unused=p)"); + EXPECT_NO_THROW((void)Pattern(text)); + } + for (const char* text : { + "(x,p:None)=>return Identity(x,unused=p)", + "(x,p:ints|float|None)=>return Identity(x,unused=p)", + "(x,p:const int64[1]|None)=>return Identity(x,unused=p)", + "(x,p:float[1]|None)=>return Identity(x,unused=p)", + "(x,p:ints|None)=>return Add(x,p,unused=p)"}) { + EXPECT_THROW((void)Pattern(text), OnnxRuntimeException) << text; + } + for (bool present : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({2, 3}, -1.0f, 1.0f); + auto& root = test.builder.AddNode("Transpose", {x}, {test.builder.MakeOutput()}); + if (present) root.AddAttribute("perm", std::vector{1, 0}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + for (const char* annotation : {"ints|None", "None|ints", "float|None"}) { + std::optional match; + ASSERT_STATUS_OK(Pattern(MakeString("(x,p:", annotation, ")=>return Transpose(x,perm=p)")) + .MatchAt(test.graph, root.Index(), match)); + EXPECT_EQ(match.has_value(), !present || std::string_view(annotation) != "float|None"); + if (match) { + const auto* attribute = match->optionalAttributeOf("p"); + EXPECT_EQ(attribute != nullptr, present); + if (attribute) EXPECT_EQ(attribute->ints_size(), 2); + } + } + } + TestGraph defaults; + auto* x = defaults.builder.MakeInput({2}, -1.0f, 1.0f); + auto& root = defaults.builder.AddNode("LeakyRelu", {x}, {defaults.builder.MakeOutput()}); + defaults.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(defaults.graph.Resolve()); + std::optional match; + ASSERT_STATUS_OK(Pattern("(x,p:float|None)=>return LeakyRelu(x,alpha=p)") + .MatchAt(defaults.graph, root.Index(), match)); + ASSERT_TRUE(match); + ASSERT_NE(match->optionalAttributeOf("p"), nullptr); + EXPECT_FLOAT_EQ(match->attributeOf("p").f(), 0.01f); + TestGraph integers; + auto* input = integers.builder.MakeInput({2, 3}, -1.0f, 1.0f); + auto& softmax = integers.builder.AddNode("Softmax", {input}, {integers.builder.MakeOutput()}); + integers.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(integers.graph.Resolve()); + ASSERT_STATUS_OK(PatternLibrary("pattern E(x,p:int|None)->(y){y=Softmax(x,axis=p);return y}") + .GetPattern("E") + .MatchAt(integers.graph, softmax.Index(), match)); + ASSERT_TRUE(match); + EXPECT_EQ(match->attributeOf("p").i(), -1); +} + +TEST(PatternComponentTests, RepeatedOptionalAttributesRequireEqualPresenceAndPayload) { + for (int inner = 0; inner < 3; ++inner) { + for (int outer = 0; outer < 3; ++outer) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + auto& a = b.AddNode("Transpose", {x}, {y}); + auto& root = b.AddNode("Transpose", {y}, {b.MakeOutput()}); + if (inner) a.AddAttribute("perm", inner == 1 ? std::vector{1, 0} : std::vector{0, 1}); + if (outer) root.AddAttribute("perm", outer == 1 ? std::vector{1, 0} : std::vector{0, 1}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + std::optional match; + ASSERT_STATUS_OK(Pattern("(x,p:ints|None)=>{a=Transpose(x,perm=p);return Transpose(a,perm=p)}") + .MatchAt(test.graph, root.Index(), match)); + EXPECT_EQ(match.has_value(), inner == outer) << inner << "," << outer; + } + } + TestGraph scalar; + auto* x = scalar.builder.MakeInput({}, -1.0f, 1.0f); + auto& root = scalar.builder.AddNode("Transpose", {x}, {scalar.builder.MakeOutput()}); + root.AddAttribute("perm", std::vector{}); + scalar.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(scalar.graph.Resolve()); + std::optional match; + ASSERT_STATUS_OK(Pattern("(x,p:ints|None)=>return Transpose(x,perm=p)") + .MatchAt(scalar.graph, root.Index(), match)); + ASSERT_TRUE(match); + ASSERT_NE(match->optionalAttributeOf("p"), nullptr); + EXPECT_EQ(match->attributeOf("p").ints_size(), 0); +} + +TEST(PatternComponentTests, OptionalAttributeCompositionAndCheckerRetryAreIsolated) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {y}).AddAttribute("perm", std::vector{1, 0}); + auto& root = b.AddNode("Transpose", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + PatternLibrary library(R"( + pattern Short(x,p:ints|None) { return Transpose(x,perm=p) } + pattern Long(x,p:ints|None) { a=Transpose(x,perm=p); return Transpose(a) } + pattern Nested(x,p:ints|None) { a=Short(x,p); return a } + pattern E(x,p:ints|None) { a=Nested(x,p)|Long(x,p); return a } + )"); + int calls = 0; + Match saved; + auto visitor = [&](Match& match) { + if (++calls == 1) { + EXPECT_EQ(match.optionalAttributeOf("p"), nullptr); + EXPECT_EQ(match.instance("a").optionalAttributeOf("p"), nullptr); + return true; + } + EXPECT_NE(match.optionalAttributeOf("p"), nullptr); + EXPECT_EQ(match.attributeOf("p").ints(0), 1); + saved = match; + return false; + }; + ASSERT_STATUS_OK(library.GetPattern("E").VisitMatches(test.graph, root.Index(), visitor)); + EXPECT_EQ(calls, 2); + EXPECT_EQ(saved.attributeOf("p").ints_size(), 2); + EXPECT_EQ(saved.instance("a").optionalAttributeOf("p"), &saved.attributeOf("p")); +} + +TEST(ReplacementComponentTests, ExplicitAbsentAttributesOmitKeywordsAndPreserveMissingErrors) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + Replacement::Instantiation result; + ASSERT_STATUS_OK(Replacement("(x,p)=>LeakyRelu(x,alpha=p)") + .Instantiate(test.graph, {{"x", x}}, {}, result, nullptr, {"p", "unused"})); + EXPECT_FLOAT_EQ(test.graph.GetNode(result.nodes.front())->GetAttributes().at("alpha").f(), 0.01f); + Replacement::Instantiation transpose; + ASSERT_STATUS_OK(Replacement("(x,p)=>Transpose(x,perm=p)") + .Instantiate(test.graph, {{"x", x}}, {}, transpose, nullptr, {"p"})); + EXPECT_TRUE(test.graph.GetNode(transpose.nodes.front())->GetAttributes().empty()); + const auto count = test.graph.NumberOfNodes(); + const auto initializers = test.graph.GetAllInitializedTensors().size(); + for (const char* text : {"(x,p)=>LeakyRelu(x,alpha=p)", "(x,p)=>{a=Add(x,1.0);return Cast(a,to=p)}"}) { + Replacement::Instantiation unchanged = result; + const bool required = std::string_view(text).find("Cast") != std::string_view::npos; + const InlinedHashSet absent = required ? InlinedHashSet{"p"} : InlinedHashSet{"typo"}; + EXPECT_FALSE(Replacement(text).Instantiate(test.graph, {{"x", x}}, {}, unchanged, nullptr, absent).IsOK()); + EXPECT_EQ(unchanged.nodes, result.nodes); + EXPECT_EQ(unchanged.outputs, result.outputs); + EXPECT_EQ(test.graph.NumberOfNodes(), count); + EXPECT_EQ(test.graph.GetAllInitializedTensors().size(), initializers); + } + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + attribute.set_f(0.5f); + EXPECT_FALSE(Replacement("(x,p)=>LeakyRelu(x,alpha=p)") + .Instantiate(test.graph, {{"x", x}}, {{"p", attribute}}, result, nullptr, {"p"}) + .IsOK()); + EXPECT_FALSE(Replacement("(x)=>Identity(x)") + .Instantiate(test.graph, {{"x", x}}, {}, result, nullptr, {"x"}) + .IsOK()); +} + +TEST(PatternRewriterTests, OptionalAttributeReplacementAndRequiredRuleIsolation) { + for (bool indexed : {false, true}) { + for (bool present : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({2, 3}, -1.0f, 1.0f); + auto& root = test.builder.AddNode("Transpose", {x}, {test.builder.MakeOutput()}); + if (present) root.AddAttribute("perm", std::vector{1, 0}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int required_calls = 0, optional_calls = 0; + Rules rules{ + Rule{"Required", "(x,p:ints)=>return Transpose(x,perm=p)", "(x,p)=>Transpose(x,perm=p)", [&](const Graph&, Match&) { ++required_calls; return false; }, {{kOnnxDomain, 13, 25}}}, + Rule{"Optional", "(x,p:ints|None)=>return Transpose(x,perm=p)", "(x,p)=>Transpose(x,perm=p)", [&](const Graph&, Match& match) { + ++optional_calls; + EXPECT_EQ(match.optionalAttributeOf("p") != nullptr, present); + return true; + }, + {{kOnnxDomain, 13, 25}}}}; + Options options; + options.indexed = indexed; + options.max_iterations = 1; + PatternRewriter rewriter(std::move(rules), options); + Statistics statistics; + ASSERT_STATUS_OK(rewriter.Run(test.graph, statistics)); + EXPECT_EQ(statistics.rewrites, 1u); + EXPECT_EQ(required_calls, present ? 1 : 0); + EXPECT_EQ(optional_calls, 1); + for (const auto& node : test.graph.Nodes()) { + EXPECT_EQ(node.GetAttributes().count("perm"), present ? 1u : 0u); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, OptionalAttributeCheckerRejectionRetriesAndRequiredOmissionRollsBack) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {y}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Transpose", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + PatternLibrary library(R"( + pattern Short(x,p:ints|None) { return Transpose(x,perm=p) } + pattern Long(x,p:ints|None) { a=Transpose(x,perm=p);return Transpose(a) } + pattern E(x,p:ints|None) { a=Short(x,p)|Long(x,p);return a } + )"); + int rejected = 0, accepted = 0; + Rule rule{"OptionalRetry", library.GetPattern("E"), "(x,p)=>Transpose(x,perm=p)", [&](const Graph&, Match& match) { + if (!match.optionalAttributeOf("p")) { + ++rejected; + return match.Reject("Need explicit perm"); + } + ++accepted; + return true; + }, + {{kOnnxDomain, 13, 25}}}; + Options options; + options.indexed = indexed; + options.match_only = true; + Statistics statistics; + ASSERT_STATUS_OK(PatternRewriter({std::move(rule)}, options).Run(test.graph, statistics)); + EXPECT_EQ(rejected, 1); + EXPECT_EQ(accepted, 2); // The inner short match and the outer long match. + + const auto before = test.model.ToProto().SerializeAsString(); + Rule invalid{"MissingRequired", "(x,p:ints|None)=>return Transpose(x,perm=p)", "(x,p)=>{a=Add(x,1.0);return Cast(a,to=p)}", [](const Graph&, Match& match) { return !match.optionalAttributeOf("p"); }, {{kOnnxDomain, 13, 25}}}; + options.match_only = false; + EXPECT_FALSE(PatternRewriter({std::move(invalid)}, options).Run(test.graph, statistics).IsOK()); + EXPECT_EQ(test.model.ToProto().SerializeAsString(), before); + EXPECT_EQ(test.graph.NumberOfNodes(), 2); + EXPECT_TRUE(test.graph.GetAllInitializedTensors().empty()); + } +} + +TEST(PatternComponentTests, DeletionBoundaryButNotProviderPolicyIsIncluded) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + auto& inner = b.AddNode("Identity", {x}, {y}); + auto& root = b.AddNode("Sigmoid", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + inner.SetExecutionProviderType(kCudaExecutionProvider); + root.SetExecutionProviderType(kCpuExecutionProvider); + if (scenario == 2) inner.SetLayeringAnnotation("different-layer"); + int calls = 0; + ASSERT_STATUS_OK(Pattern("(x)=>{a=Identity(x);return Sigmoid(a)}") + .VisitMatches(test.graph, root.Index(), [&](Match&) { ++calls; return true; })); + EXPECT_EQ(calls, scenario == 0 ? 1 : 0); + } +} + +TEST(ReplacementComponentTests, InstantiateFreshFragmentWithoutPatternOrMatch) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + Replacement recipe("(x)=>{a=Add(x,1.0);return LeakyRelu(a)}"); + Replacement::Instantiation result; + ONNX_NAMESPACE::AttributeProto unused; + unused.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + unused.set_i(1); + ASSERT_STATUS_OK(recipe.Instantiate(test.graph, {{"x", x}, {"unused", x}}, {{"unused_attribute", unused}}, result)); + ASSERT_EQ(result.nodes.size(), 2u); + ASSERT_EQ(result.outputs.size(), 1u); + EXPECT_TRUE(test.graph.GraphResolveNeeded()); + EXPECT_TRUE(test.graph.GetOutputs().empty()); + EXPECT_TRUE(test.graph.GetNode(result.nodes.back())->GetExecutionProviderType().empty()); + EXPECT_EQ(test.graph.GetNode(result.nodes.back())->GetAttributes().at("alpha").f(), 0.01f); + test.graph.SetOutputs({result.outputs.front()}); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(result.outputs.front()->Shape()->dim(0).dim_value(), 2); + EXPECT_EQ(test.graph.GetNode(result.nodes.back())->GetInputEdgesCount(), 1u); +} + +TEST(ReplacementComponentTests, PositionalOutputsAttributesOptionalInputsAndSource) { + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto& source = test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + source.SetExecutionProviderType(kCpuExecutionProvider); + source.SetLayeringAnnotation("layer0"); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + ONNX_NAMESPACE::AttributeProto axis; + axis.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis.set_i(0); + Replacement::Instantiation result; + ASSERT_STATUS_OK(Replacement("(x,axis)=>{a,b=Split(x,axis=axis);c=Concat(a,b,axis=axis);return b,c}") + .Instantiate(test.graph, {{"x", x}}, {{"axis", axis}}, result, &source)); + ASSERT_EQ(result.outputs.size(), 2u); + EXPECT_NE(result.outputs[0], result.outputs[1]); + EXPECT_EQ(test.graph.GetNode(result.nodes.front())->GetExecutionProviderType(), kCpuExecutionProvider); + EXPECT_EQ(test.graph.GetNode(result.nodes.front())->GetLayeringAnnotation(), "layer0"); + EXPECT_EQ(test.graph.GetOutputs().front(), source.OutputDefs()[0]); + test.graph.SetOutputs({result.outputs[0], result.outputs[1]}); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(result.outputs[0]->Shape()->dim(0).dim_value(), 2); + Replacement::Instantiation clip; + ASSERT_STATUS_OK(Replacement("(x,min)=>Clip(x,min)").Instantiate(test.graph, {{"x", x}, {"min", nullptr}}, {}, clip)); + EXPECT_FALSE(test.graph.GetNode(clip.nodes.front())->InputDefs()[1]->Exists()); + Replacement::Instantiation another_clip; + ASSERT_STATUS_OK(Replacement("(x)=>Clip(x)").Instantiate(test.graph, {{"x", x}}, {}, another_clip)); + EXPECT_NE(clip.outputs.front(), another_clip.outputs.front()); + EXPECT_EQ(test.graph.GetOutputs()[0], result.outputs[0]); +} + +TEST(ReplacementComponentTests, SchemaVerificationIsNotShapeInferenceOrResolveTransaction) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + Replacement::Instantiation result; + ASSERT_STATUS_OK(Replacement("(x)=>Transpose(x,perm=[0,1])").Instantiate(test.graph, {{"x", x}}, {}, result)); + test.graph.SetOutputs({result.outputs.front()}); + EXPECT_FALSE(test.graph.Resolve().IsOK()); + EXPECT_NE(test.graph.GetNode(result.nodes.front()), nullptr); + + TestGraph wrong_attribute; + auto* input = wrong_attribute.builder.MakeInput({2}, -1.0f, 1.0f); + ONNX_NAMESPACE::AttributeProto permutation; + permutation.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + permutation.set_i(0); + Replacement::Instantiation unchanged; + EXPECT_FALSE(Replacement("(x,p)=>Transpose(x,perm=p)") + .Instantiate(wrong_attribute.graph, {{"x", input}}, {{"p", permutation}}, unchanged) + .IsOK()); + EXPECT_EQ(wrong_attribute.graph.NumberOfNodes(), 0); + EXPECT_TRUE(unchanged.nodes.empty()); +} + +TEST(ReplacementComponentTests, ValidationRollsBackAndPreservesCallerResult) { + for (const char* text : { + "(x)=>{a=Add(x,1.0);return Identity(a,unexpected=1)}", + "(x)=>{a=Add(x,1.0);return missing.Op(a)}", + "(x)=>{a=Add(x,1.0);return MissingStandard(a)}", + "(x)=>Identity(x)", "(a,missing)=>Transpose(a,perm=missing)"}) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + Replacement::Instantiation result; + result.nodes.push_back(123); + result.outputs.push_back(x); + const auto initializers = test.graph.GetAllInitializedTensors().size(); + Replacement::TensorBindings bindings; + if (std::string_view(text).find("unexpected") != std::string_view::npos || + std::string_view(text).find("missing.Op") != std::string_view::npos || + std::string_view(text).find("MissingStandard") != std::string_view::npos) + bindings.emplace("x", x); + if (std::string_view(text).find("perm=missing") != std::string_view::npos) bindings.emplace("a", x); + EXPECT_FALSE(Replacement(text).Instantiate(test.graph, bindings, {}, result).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 0); + EXPECT_EQ(test.graph.GetAllInitializedTensors().size(), initializers); + EXPECT_EQ(result.nodes, (InlinedVector{123})); + EXPECT_EQ(result.outputs, (InlinedVector{x})); + } + TestGraph test; + Replacement::Instantiation result; + EXPECT_FALSE(Replacement("(x)=>Identity(x)").Instantiate(test.graph, {{"x", nullptr}}, {}, result).IsOK()); + EXPECT_TRUE(result.nodes.empty()); +} + +TEST(ReplacementComponentTests, CrossGraphBindingsAndAnnotationSourcesAreRejected) { + TestGraph test, other; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* foreign = other.builder.MakeInput({2}, -1.0f, 1.0f); + auto& foreign_source = other.builder.AddNode("Identity", {foreign}, {other.builder.MakeOutput()}); + Replacement recipe("(x)=>Identity(x)"); + Replacement::Instantiation result; + EXPECT_FALSE(recipe.Instantiate(test.graph, {{"x", foreign}}, {}, result).IsOK()); + EXPECT_FALSE(recipe.Instantiate(test.graph, {{"x", x}}, {}, result, &foreign_source).IsOK()); + ONNX_NAMESPACE::AttributeProto attr; + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attr.set_i(1); + EXPECT_FALSE(recipe.Instantiate(test.graph, {{"x", x}}, {{"x", attr}}, result).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 0); + auto& local_source = test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + ASSERT_EQ(local_source.Index(), foreign_source.Index()); + EXPECT_FALSE(recipe.Instantiate(test.graph, {{"x", x}}, {}, result, &foreign_source).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 1); + EXPECT_TRUE(result.nodes.empty()); + EXPECT_TRUE(result.outputs.empty()); +} + +TEST(MatchComponentTests, CapturedPreparedBindingsAndCopyAreSeparate) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto& root = b.AddNode("LeakyRelu", {x}, {b.MakeOutput()}); + root.AddAttribute("alpha", 0.25f); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + bool visited = false; + auto visitor = [&](Match& match) { + visited = true; + auto prepared = match.attributeOf("a"); + prepared.set_f(0.5f); + EXPECT_THROW(match.update("a", prepared), OnnxRuntimeException); + EXPECT_THROW(match.update("x", prepared), OnnxRuntimeException); + match.update("prepared", prepared); + Match copy = match; + prepared.set_f(0.75f); + match.update("prepared", prepared); + EXPECT_FLOAT_EQ(match.attributeOf("prepared").f(), 0.75f); + EXPECT_FLOAT_EQ(copy.attributeOf("prepared").f(), 0.5f); + EXPECT_FLOAT_EQ(copy.attributeOf("a").f(), 0.25f); + EXPECT_FALSE(match.Reject("unsupported")); + match.ClearUpdates(); + EXPECT_TRUE(match.RejectionReason().empty()); + EXPECT_FLOAT_EQ(match.attributeOf("a").f(), 0.25f); + EXPECT_THROW(match.attributeOf("prepared"), OnnxRuntimeException); + EXPECT_EQ(std::get(match.dimensionOf("N")), 2); + EXPECT_THROW(match.optionalValueOf("missing"), OnnxRuntimeException); + EXPECT_THROW(match.valueOf("a"), OnnxRuntimeException); + return false; + }; + ASSERT_STATUS_OK(Pattern("(x:float[N],a:float)=>return LeakyRelu(x,alpha=a)") + .VisitMatches(test.graph, root.Index(), visitor)); + EXPECT_TRUE(visited); + Match retained; + { + Rule rule{"Evidence", "(x,a:float)=>return LeakyRelu(x,alpha=a)", "(x)=>Identity(x)", [&](const Graph&, Match& match) { retained = match; return true; }, {{kOnnxDomain, 13, 25}}}; + PatternRewriter rewriter({std::move(rule)}, Options{.match_only = true}); + Statistics statistics; + ASSERT_STATUS_OK(rewriter.Run(test.graph, statistics)); + EXPECT_EQ(statistics.matches, 1u); + } + EXPECT_EQ(retained.valueOf("x"), x); + EXPECT_FLOAT_EQ(retained.attributeOf("a").f(), 0.25f); +} + +TEST(CheckAndBindComponentTests, TransposeCheckerChecksAndBindsWithoutEngine) { + for (int scenario = 0; scenario < 6; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput(scenario == 4 ? std::nullopt + : std::optional>{{scenario == 5 ? -1 : 2, 3, 4}}); + auto* intermediate = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {intermediate}).AddAttribute("perm", std::vector{1, 0, 2}); + auto& root = b.AddNode("Transpose", {intermediate}, {b.MakeOutput()}); + root.AddAttribute("perm", std::vector{0, 2, 1}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + bool visited = false; + ASSERT_STATUS_OK(Pattern("(x)=>{a=Transpose(x);return Transpose(a)}") + .VisitMatches(test.graph, root.Index(), [&](Match& match) { + visited = true; + auto p = match.nodes[0]->GetAttributes().at("perm"); + auto q = match.nodes[1]->GetAttributes().at("perm"); + if (scenario == 1) p.set_ints(1, 1); + if (scenario == 2) q.set_ints(1, 3); + if (scenario == 3) q.mutable_ints()->RemoveLast(); + match.update("p", p); + match.update("q", q); + const bool valid = scenario == 0 || scenario == 5; + EXPECT_EQ(ComposeTransposePermutations(test.graph, match), valid); + if (valid) { + const auto& perm = match.attributeOf("composed"); + EXPECT_EQ(perm.ints(0), 1); + EXPECT_EQ(perm.ints(1), 2); + EXPECT_EQ(perm.ints(2), 0); + } else { + EXPECT_FALSE(match.RejectionReason().empty()); + EXPECT_THROW(match.attributeOf("composed"), OnnxRuntimeException); + } + return false; + })); + EXPECT_TRUE(visited); + } +} + +TEST(PatternRewriterTests, RunAndTransformerApplyHaveDifferentResolveBoundaries) { + for (bool use_apply : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {y}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Transpose", {y}, {b.MakeOutput()}).AddAttribute("perm", std::vector{1, 0}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + PatternRewriter rewriter({ComposeTransposeRule()}); + if (use_apply) { + bool modified = false; + ASSERT_STATUS_OK(rewriter.Apply(test.graph, modified, DefaultLoggingManager().DefaultLogger())); + EXPECT_TRUE(modified); + EXPECT_FALSE(test.graph.GraphResolveNeeded()); + } else { + Statistics statistics; + ASSERT_STATUS_OK(rewriter.Run(test.graph, statistics)); + EXPECT_EQ(statistics.rewrites, 1u); + EXPECT_TRUE(test.graph.GraphResolveNeeded()); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + EXPECT_EQ(Count(test.graph, "Transpose"), 1u); + } +} + +TEST(PatternRewriterTests, ComposeTransposeRuleIntegrationAndNumerics) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3, 4}, -1.0f, 1.0f); + auto* middle = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {middle}).AddAttribute("perm", std::vector{1, 0, 2}); + auto& root = b.AddNode("Transpose", {middle}, {b.MakeOutput()}); + if (scenario != 3) root.AddAttribute("perm", std::vector{0, 2, 1}); + if (scenario == 2) b.AddNode("Identity", {middle}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const auto original = test.model.ToProto().SerializeAsString(); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({ComposeTransposeRule()}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario != 0) continue; + EXPECT_EQ(Count(test.graph, "Transpose"), 1u); + const auto rewritten = test.model.ToProto().SerializeAsString(); + auto execute = [&](const std::string& serialized, std::vector& outputs) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + InferenceSessionWrapper session{options, GetEnvironment()}; + std::istringstream stream(serialized); + ORT_RETURN_IF_ERROR(session.Load(stream)); + ORT_RETURN_IF_ERROR(session.Initialize()); + return session.Run(RunOptions{}, b.feeds_, b.output_names_, &outputs); + }; + std::vector before, after; + ASSERT_STATUS_OK(execute(original, before)); + ASSERT_STATUS_OK(execute(rewritten, after)); + ASSERT_EQ(before.size(), 1u); + ASSERT_EQ(after.size(), 1u); + const auto& expected = before[0].Get(); + const auto& actual = after[0].Get(); + EXPECT_EQ(actual.Shape(), expected.Shape()); + ASSERT_EQ(actual.Shape().Size(), expected.Shape().Size()); + for (int64_t i = 0; i < expected.Shape().Size(); ++i) + EXPECT_EQ(actual.Data()[i], expected.Data()[i]); + } + } +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc new file mode 100644 index 0000000000000..c363d2809c64d --- /dev/null +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -0,0 +1,6202 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/optimizer/pattern_rewriter.h" +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" +#include "test/optimizer/graph_transform_test_fixture.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { +namespace { +using namespace rewrite; + +#ifndef DISABLE_CONTRIB_OPS +ONNX_NAMESPACE::AttributeProto FloatAttribute(const char* name, float value) { + ONNX_NAMESPACE::AttributeProto attr; + attr.set_name(name); + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + attr.set_f(value); + return attr; +} +#endif + +bool IsFloat32(const NodeArg& value) { + const auto* type = value.TypeAsProto(); + return type && type->has_tensor_type() && + type->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT; +} + +#ifndef DISABLE_CONTRIB_OPS +std::optional Scalar(const Graph& graph, const NodeArg& value) { + const auto* tensor = graph_utils::GetConstantInitializer(graph, value.Name()); + if (!tensor || tensor->data_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL || tensor->dims_size() != 0) + return std::nullopt; + Initializer initializer(graph, *tensor, graph.ModelPath()); + if (initializer.size() != 1) return std::nullopt; + float scalar = *initializer.data(); + if (!std::isfinite(scalar)) return std::nullopt; + return scalar; +} + +bool Integer(const Node& node, const char* name, int64_t value) { + auto it = node.GetAttributes().find(name); + return it != node.GetAttributes().end() && it->second.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT && + it->second.i() == value; +} +#endif + +Rules MakeExampleRules() { + Rules rules; + { + const auto* pattern = R"( + (x) => { + add = Add(x, 3.0) + clip = Clip(add, 0.0, 6.0) + mul = Mul(clip, x) + return Div(mul, 6.0) + } + )"; + auto check_and_bind = [](const Graph&, Match& m) { return IsFloat32(*m.valueOf("x")); }; + const auto* replacement = "(x) => HardSwish(x)"; + Rule rule{"HardSwish", pattern, replacement, check_and_bind, {{kOnnxDomain, 14, 21}}}; + rules.push_back(std::move(rule)); + } + { + const auto* pattern = R"( + (x) => { + sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) + return Mul(sigmoid, x) + } + )"; + auto check_and_bind = [](const Graph&, Match& m) { return IsFloat32(*m.valueOf("x")); }; + const auto* replacement = "(x) => HardSwish(x)"; + Rule rule{"HardSigmoidMul", pattern, replacement, check_and_bind, {{kOnnxDomain, 14, 21}}}; + rules.push_back(std::move(rule)); + } +#ifndef DISABLE_CONTRIB_OPS + { + const auto* pattern = "(x, y, divisor) => { mm = MatMul(x, y); return Div(mm, divisor); }"; + auto check_and_bind = [](const Graph& graph, Match& m) { + auto scalar = Scalar(graph, *m.valueOf("divisor")); + if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("y")) || + !scalar || *scalar == 0.0f || !std::isfinite(1.0f / *scalar)) + return false; + m.update("alpha", FloatAttribute("alpha", 1.0f / *scalar)); + return true; + }; + const auto* replacement = "(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"; + Rule rule{"MatMulDiv", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + rules.push_back(std::move(rule)); + } + { + const auto* pattern = "(x, y) => { transpose = Transpose(y); return MatMul(x, transpose); }"; + auto check_and_bind = [](const Graph&, Match& m) { + if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("y"))) return false; + const auto* shape = m.valueOf("y")->Shape(); + if (!shape || shape->dim_size() < 2) return false; + auto it = m.nodes[0]->GetAttributes().find("perm"); + if (it == m.nodes[0]->GetAttributes().end()) return shape->dim_size() == 2; + const auto& perm = it->second; + if (perm.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_INTS || + perm.ints_size() != shape->dim_size()) + return false; + int rank = shape->dim_size(); + for (int i = 0; i < rank; ++i) { + int expected = i < rank - 2 ? i : (i == rank - 2 ? rank - 1 : rank - 2); + if (perm.ints(i) != expected) return false; + } + return true; + }; + const auto* replacement = "(x,y) => com.microsoft.FusedMatMul(x,y,transB=1)"; + Rule rule{"TransposeMatMul", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + rules.push_back(std::move(rule)); + } + { + const auto* pattern = "(x, y, divisor) => { mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor); }"; + auto check_and_bind = [](const Graph& graph, Match& m) { + auto scalar = Scalar(graph, *m.valueOf("divisor")); + const auto& node = *m.nodes[0]; + const auto& attributes = node.GetAttributes(); + auto it = attributes.find("alpha"); + float alpha = it == attributes.end() ? 1.0f : it->second.f(); + if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("y")) || + !scalar || *scalar == 0.0f || !std::isfinite(alpha / *scalar)) + return false; + m.update("alpha", FloatAttribute("alpha", alpha / *scalar)); + for (const char* name : {"transA", "transB", "transBatchA", "transBatchB"}) { + auto attribute = attributes.find(name); + m.update(name, attribute == attributes.end() ? node.Op()->attributes().at(name).default_value + : attribute->second); + } + return true; + }; + const auto* replacement = R"( + (x, y, alpha, transA, transB, transBatchA, transBatchB) => + com.microsoft.FusedMatMul(x, y, alpha=alpha, transA=transA, transB=transB, + transBatchA=transBatchA, transBatchB=transBatchB) + )"; + Rule rule{"FusedMatMulDiv", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + rules.push_back(std::move(rule)); + } + { + const auto* pattern = R"( + (x, scale, epsilon) => { + square = Pow(x, 2.0) + mean = ReduceMean(square) + add = Add(mean, epsilon) + sqrt = Sqrt(add) + reciprocal = Reciprocal(sqrt) + normalize = Mul(x, reciprocal) + return Mul(normalize, scale) + } + )"; + auto check_and_bind = [](const Graph& graph, Match& m) { + if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("scale"))) return false; + auto eps = Scalar(graph, *m.valueOf("epsilon")); + const auto* shape = m.valueOf("x")->Shape(); + const auto* scale_shape = m.valueOf("scale")->Shape(); + if (!eps || *eps <= 0 || !shape || shape->dim_size() < 1 || !scale_shape || scale_shape->dim_size() != 1) + return false; + const auto& dim = shape->dim(shape->dim_size() - 1); + const auto& scale_dim = scale_shape->dim(0); + if (!dim.has_dim_value() || !scale_dim.has_dim_value() || dim.dim_value() <= 0 || + dim.dim_value() != scale_dim.dim_value()) + return false; + const auto& mean_node = *m.nodes[1]; + auto axes = mean_node.GetAttributes().find("axes"); + if (!Integer(mean_node, "keepdims", 1) || axes == mean_node.GetAttributes().end() || + axes->second.type() != ONNX_NAMESPACE::AttributeProto_AttributeType_INTS || + axes->second.ints_size() != 1 || axes->second.ints(0) != -1) + return false; + m.update("fused_epsilon", FloatAttribute("epsilon", *eps)); + return true; + }; + const auto* replacement = R"( + (x, scale, fused_epsilon) => + SimplifiedLayerNormalization(x, scale, axis=-1, epsilon=fused_epsilon, stash_type=1) + )"; + Rule rule{"RmsNormalization", pattern, replacement, check_and_bind, {{kOnnxDomain, 7, 17}}}; + rules.push_back(std::move(rule)); + } +#endif + return rules; +} + +struct TestGraph { + explicit TestGraph(const std::unordered_map& imports = {{kOnnxDomain, 17}, {kMSDomain, 1}}) + : model{"rewrite", + false, + ModelMetaData(), + PathString(), + IOnnxRuntimeOpSchemaRegistryList(), + imports, + {}, + DefaultLoggingManager().DefaultLogger()} {} + Model model; + Graph& graph = model.MainGraph(); + ModelTestBuilder builder{graph}; +}; + +void BuildHardSwish(ModelTestBuilder& b, bool mismatch = false, bool intermediate_output = false, float divisor = 6.0f, + bool external_consumer = false) { + auto* x = b.MakeInput({2, 4}, -5.0f, 5.0f); + auto* y = mismatch ? b.MakeInput({2, 4}, -5.0f, 5.0f) : x; + auto* add = b.MakeIntermediate(); + auto* clip = intermediate_output ? b.MakeOutput() : b.MakeIntermediate(); + auto* mul = b.MakeIntermediate(); + b.AddNode("Add", {x, b.MakeScalarInitializer(3.0f)}, {add}); + b.AddNode("Clip", {add, b.MakeScalarInitializer(0.0f), b.MakeScalarInitializer(6.0f)}, {clip}); + b.AddNode("Mul", {clip, y}, {mul}); + b.AddNode("Div", {mul, b.MakeScalarInitializer(divisor)}, {b.MakeOutput()}); + if (external_consumer) b.AddNode("Identity", {clip}, {b.MakeOutput()}); +} + +void BuildRms(ModelTestBuilder& b, NodeArg* x, NodeArg* output, int64_t width, float epsilon = 1e-5f, + bool wrong_axis = false) { + auto* square = b.MakeIntermediate(); + auto* mean = b.MakeIntermediate(); + auto* add = b.MakeIntermediate(); + auto* sqrt = b.MakeIntermediate(); + auto* reciprocal = b.MakeIntermediate(); + auto* normalized = b.MakeIntermediate(); + b.AddNode("Pow", {x, b.MakeScalarInitializer(2.0f)}, {square}); + auto& reduce = b.AddNode("ReduceMean", {square}, {mean}); + reduce.AddAttribute("axes", std::vector{wrong_axis ? -2 : -1}); + reduce.AddAttribute("keepdims", int64_t{1}); + b.AddNode("Add", {mean, b.MakeScalarInitializer(epsilon)}, {add}); + b.AddNode("Sqrt", {add}, {sqrt}); + b.AddNode("Reciprocal", {sqrt}, {reciprocal}); + b.AddNode("Mul", {x, reciprocal}, {normalized}); + b.AddNode("Mul", {normalized, b.MakeInitializer({width}, 0.8f, 1.2f)}, {output}); +} + +size_t Count(const Graph& graph, const std::string& op) { + size_t count = 0; + for (const auto& node : graph.Nodes()) count += node.OpType() == op; + return count; +} + +Rule IdentityPair(const char* name = "identity", bool accept = true) { + const auto* pattern = "(x) => a = Identity(x); return Identity(a)"; + auto check_and_bind = [accept](const Graph&, Match&) { return accept; }; + const auto* replacement = "(x) => Identity(x)"; + Rule rule{name, pattern, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}; + return rule; +} + +Rule IdentitySplit() { + const auto* pattern = "(x) => { first = Identity(x); a, b = Split(first); return a, b; }"; + auto check_and_bind = [](const Graph&, Match& match) { + const auto& node = *match.nodes.back(); + auto axis = node.GetAttributes().find("axis"); + match.update("axis", axis == node.GetAttributes().end() ? node.Op()->attributes().at("axis").default_value + : axis->second); + return true; + }; + const auto* replacement = "(x, axis) => { a, b = Split(x, axis=axis); return a, b; }"; + Rule rule{"IdentitySplit", pattern, replacement, check_and_bind, {{kOnnxDomain, 13, 17}}}; + return rule; +} + +Rule SplitConcat() { + const auto* pattern = "(x) => a, b = Split(x); return Concat(a, b)"; + auto check_and_bind = [](const Graph&, Match& match) { + const auto& split_attributes = match.nodes[0]->GetAttributes(); + auto axis = split_attributes.find("axis"); + return (axis == split_attributes.end() || axis->second.i() == 0) && + match.nodes[1]->GetAttributes().at("axis").i() == 0; + }; + const auto* replacement = "(x) => Identity(x)"; + Rule rule{"SplitConcat", pattern, replacement, check_and_bind, {{kOnnxDomain, 13, 17}}}; + return rule; +} + +// Multi-output adaptation of ONNXScript's LayerNormBiasFusion; unused statistics may be removed. +Rule LayerNormBias() { + const auto* pattern = "(x, scale, bias) => y, mean, inv_std = LayerNormalization(x, scale); return Add(y, bias)"; + auto check_and_bind = [](const Graph&, Match& match) { + const auto* shape = match.valueOf("x")->Shape(); + const auto* scale_shape = match.valueOf("scale")->Shape(); + const auto* bias_shape = match.valueOf("bias")->Shape(); + if (!shape || shape->dim_size() < 1 || !scale_shape || !bias_shape || scale_shape->dim_size() != 1 || + bias_shape->dim_size() != 1) + return false; + const auto& width = shape->dim(shape->dim_size() - 1); + const auto& scale_width = scale_shape->dim(0); + const auto& bias_width = bias_shape->dim(0); + auto axis = match.nodes[0]->GetAttributes().find("axis"); + if (!width.has_dim_value() || width.dim_value() <= 0 || !scale_width.has_dim_value() || + !bias_width.has_dim_value() || width.dim_value() != scale_width.dim_value() || + width.dim_value() != bias_width.dim_value() || + (axis != match.nodes[0]->GetAttributes().end() && axis->second.i() != -1)) + return false; + const auto& node = *match.nodes[0]; + for (const char* name : {"axis", "epsilon", "stash_type"}) { + auto attribute = node.GetAttributes().find(name); + match.update(name, attribute == node.GetAttributes().end() ? node.Op()->attributes().at(name).default_value + : attribute->second); + } + return true; + }; + const auto* replacement = R"( + (x, scale, bias, axis, epsilon, stash_type) => + LayerNormalization(x, scale, bias, axis=axis, epsilon=epsilon, stash_type=stash_type) + )"; + Rule rule{"LayerNormBias", pattern, replacement, check_and_bind, {{kOnnxDomain, 17, 17}}}; + return rule; +} + +// Small-weight eager decoder: B=1, Hq=4, Hkv=2, D=4. Cache length is 3. +void BuildDecoder(ModelTestBuilder& b, size_t layers, int64_t sequence = 1) { + constexpr int64_t width = 16, kv_width = 8, head = 4, past = 3; + auto* x = b.MakeInput({1, sequence, width}, -0.5f, 0.5f); + auto unary = [&](const char* op, NodeArg* input) { + auto* output = b.MakeIntermediate(); + b.AddNode(op, {input}, {output}); + return output; + }; + auto binary = [&](const char* op, NodeArg* a, NodeArg* c) { + auto* output = b.MakeIntermediate(); + b.AddNode(op, {a, c}, {output}); + return output; + }; + auto reshape = [&](NodeArg* input, std::vector shape) { + auto* output = b.MakeIntermediate(); + auto rank = static_cast(shape.size()); + b.AddNode("Reshape", {input, b.MakeInitializer({rank}, shape)}, {output}); + return output; + }; + auto transpose = [&](NodeArg* input, std::vector perm) { + auto* output = b.MakeIntermediate(); + b.AddNode("Transpose", {input}, {output}).AddAttribute("perm", perm); + return output; + }; + auto projection = [&](NodeArg* input, int64_t in, int64_t out) { + return binary("MatMul", input, b.MakeInitializer({in, out}, -0.08f, 0.08f)); + }; + auto slice = [&](NodeArg* input, int64_t start, int64_t end) { + auto* output = b.MakeIntermediate(); + b.AddNode("Slice", + {input, b.MakeInitializer({1}, {start}), b.MakeInitializer({1}, {end}), + b.MakeInitializer({1}, {3})}, + {output}); + return output; + }; + auto concat = [&](NodeArg* a, NodeArg* c, int64_t axis, NodeArg* output) { + b.AddNode("Concat", {a, c}, {output}).AddAttribute("axis", axis); + return output; + }; + auto rope = [&](NodeArg* input, int64_t heads) { + auto* shaped = transpose(reshape(input, {1, sequence, heads, head}), {0, 2, 1, 3}); + auto* first = slice(shaped, 0, head / 2); + auto* second = slice(shaped, head / 2, head); + auto* rotated = concat(unary("Neg", second), first, 3, b.MakeIntermediate()); + // Nontrivial position-dependent cache entries, gathered at decode positions. + std::vector positions; + for (int64_t s = 0; s < sequence; ++s) positions.push_back(past + s); + auto* indices = b.MakeInitializer({sequence}, positions); + std::vector cosine(32 * head), sine(32 * head); + for (int64_t position = 0; position < 32; ++position) { + for (int64_t d = 0; d < head; ++d) { + float angle = static_cast(position) * std::pow(10000.0f, -2.0f * (d % (head / 2)) / head); + cosine[static_cast(position * head + d)] = std::cos(angle); + sine[static_cast(position * head + d)] = std::sin(angle); + } + } + auto* cos = binary("Gather", b.MakeInitializer({32, head}, cosine), indices); + auto* sin = binary("Gather", b.MakeInitializer({32, head}, sine), indices); + return binary("Add", binary("Mul", shaped, reshape(cos, {1, 1, sequence, head})), + binary("Mul", rotated, reshape(sin, {1, 1, sequence, head}))); + }; + auto repeat_kv = [&](NodeArg* input) { + auto* expanded = b.MakeIntermediate(); + b.AddNode("Expand", + {reshape(input, {1, 2, 1, past + sequence, head}), + b.MakeInitializer({5}, {1, 2, 2, past + sequence, head})}, + {expanded}); + return reshape(expanded, {1, 4, past + sequence, head}); + }; + for (size_t layer = 0; layer < layers; ++layer) { + auto* normalized = b.MakeIntermediate(); + BuildRms(b, x, normalized, width); + auto* q = rope(projection(normalized, width, width), 4); + auto* k = rope(projection(normalized, width, kv_width), 2); + auto* v = transpose(reshape(projection(normalized, width, kv_width), {1, sequence, 2, head}), {0, 2, 1, 3}); + auto* present_k = concat(b.MakeInput({1, 2, past, head}, -0.1f, 0.1f), k, 2, b.MakeOutput()); + auto* present_v = concat(b.MakeInput({1, 2, past, head}, -0.1f, 0.1f), v, 2, b.MakeOutput()); + auto* scores = binary("MatMul", q, transpose(repeat_kv(present_k), {0, 1, 3, 2})); + auto* scaled = binary("Div", scores, b.MakeScalarInitializer(2.0f)); + std::vector mask(static_cast(sequence * (past + sequence)), 0.0f); + for (int64_t s = 0; s < sequence; ++s) { + for (int64_t t = past + s + 1; t < past + sequence; ++t) { + mask[static_cast(s * (past + sequence) + t)] = -10000.0f; + } + } + auto* masked = binary("Add", scaled, b.MakeInitializer({1, 1, sequence, past + sequence}, mask)); + auto* probabilities = b.MakeIntermediate(); + b.AddNode("Softmax", {masked}, {probabilities}).AddAttribute("axis", int64_t{-1}); + auto* attention = binary("MatMul", probabilities, repeat_kv(present_v)); + auto* context = reshape(transpose(attention, {0, 2, 1, 3}), {1, sequence, width}); + auto* residual = binary("Add", x, projection(context, width, width)); + auto* norm2 = b.MakeIntermediate(); + BuildRms(b, residual, norm2, width); + auto* gate = projection(norm2, width, width * 4); + auto* silu = binary("Mul", gate, unary("Sigmoid", gate)); + auto* mlp = binary("Mul", silu, projection(norm2, width, width * 4)); + x = binary("Add", residual, projection(mlp, width * 4, width)); + } + b.AddNode("Identity", {x}, {b.MakeOutput()}); +} +} // namespace + +TEST(PatternRewriterTests, HardSwishPositiveAndBoundaries) { + for (int scenario = 0; scenario < 6; ++scenario) { + TestGraph test; + BuildHardSwish(test.builder, scenario == 1, scenario == 2, scenario == 3 ? 7.0f : 6.0f, scenario == 4); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 5) test.graph.Nodes().begin()->SetExecutionProviderType(kCudaExecutionProvider); + PatternRewriter rewriter(MakeExampleRules()); + Statistics statistics; + ASSERT_STATUS_OK(rewriter.Run(test.graph, statistics)); + EXPECT_EQ(statistics.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(Count(test.graph, "HardSwish"), scenario == 0 ? 1u : 0u); + } +} + +TEST(PatternRewriterTests, NumericalHardSwish) { + TransformerTester([](ModelTestBuilder& b) { BuildHardSwish(b); }, + [](InferenceSessionWrapper& s) { + // CPU executes HardSwish's ONNX function body (HardSigmoid + Mul). + EXPECT_EQ(Count(s.GetGraph(), "HardSigmoid"), 1u); + EXPECT_EQ(Count(s.GetGraph(), "Clip"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(MakeExampleRules())); +} + +TEST(PatternRewriterTests, HardSigmoidAttributeGuard) { + for (bool correct : {true, false}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 4}, -5.0f, 5.0f); + auto* h = b.MakeIntermediate(); + auto& sigmoid = b.AddNode("HardSigmoid", {x}, {h}); + sigmoid.AddAttribute("alpha", correct ? 1.0f / 6.0f : 0.2f); + sigmoid.AddAttribute("beta", 0.5f); + b.AddNode("Mul", {h, x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules()).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, correct ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, NamedValuesAndAttributeGuards) { + Pattern p("(x) => sigmoid = HardSigmoid(x, alpha=0.25); return Mul(sigmoid, x)"); + auto x = p.NamedValue("x"); + auto first = p.NamedValue("sigmoid"); + EXPECT_EQ(x.name, "x"); + EXPECT_EQ(x.handle, 0); + EXPECT_EQ(p.NamedValue("sigmoid").handle, first.handle); + EXPECT_EQ(p.Nodes()[0].outputs[0].name, "sigmoid"); + EXPECT_EQ(p.Nodes()[1].inputs[0].name, "sigmoid"); + EXPECT_THROW(Pattern("(x,x) => return Identity(x)"), OnnxRuntimeException); + EXPECT_THROW(p.NamedValue("missing"), OnnxRuntimeException); + EXPECT_THROW(Pattern("(x) => x = Identity(x); return Identity(x)"), OnnxRuntimeException); + EXPECT_THROW(Pattern("(x) => return Identity(missing)"), OnnxRuntimeException); + auto integer = AttributeMatcher::Equal("axis", int64_t{1}); + auto string = AttributeMatcher::Equal("mode", std::string("linear")); + ONNX_NAMESPACE::AttributeProto attr; + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attr.set_i(1); + EXPECT_TRUE(integer.Matches(attr)); + EXPECT_FALSE(string.Matches(attr)); + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + attr.set_s("linear"); + EXPECT_TRUE(string.Matches(attr)); + EXPECT_FALSE(integer.Matches(attr)); + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto* input = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* output = test.builder.MakeIntermediate(); + auto& sigmoid = test.builder.AddNode("HardSigmoid", {input}, {output}); + if (scenario != 3) sigmoid.AddAttribute("alpha", scenario == 1 ? 0.2f : 0.25f); + sigmoid.AddAttribute("beta", scenario == 2 ? 0.3f : 0.5f); + test.builder.AddNode("Mul", {output, input}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = p; + rule.check_and_bind = [](const Graph&, Match& match) { + const auto& beta = match.nodes[0]->GetAttributes().at("beta"); + return beta.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT && beta.f() > 0.4f; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 0 ? 1u : 0u); + } + } +} + +TEST(PatternRewriterTests, NamedMatchLookupAndReplacementSafety) { + Match empty; + EXPECT_THROW(empty.valueOf("x"), OnnxRuntimeException); + Pattern incomplete("(x) => return Identity(x)"); + empty.pattern = &incomplete; + EXPECT_THROW(empty.valueOf("x"), OnnxRuntimeException); + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* intermediate = test.builder.MakeIntermediate(); + auto* output = test.builder.MakeOutput(); + test.builder.AddNode("Identity", {x}, {intermediate}); + test.builder.AddNode("Identity", {intermediate}, {output}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto first = IdentityPair(), second = IdentityPair(); + second.pattern = Pattern("(input) => { middle = Identity(input); root = Identity(middle); return root; }"); + first.check_and_bind = [=](const Graph&, Match& m) { + EXPECT_EQ(m.valueOf("x"), x); + EXPECT_EQ(m.valueOf("a"), intermediate); + EXPECT_THROW(m.valueOf("input"), OnnxRuntimeException); + EXPECT_THROW(m.valueOf(""), OnnxRuntimeException); + return false; + }; + second.check_and_bind = [=](const Graph&, Match& m) { + EXPECT_EQ(m.valueOf("input"), x); + EXPECT_EQ(m.valueOf("middle"), intermediate); + EXPECT_THROW(m.valueOf("x"), OnnxRuntimeException); + return true; + }; + second.replacement = [scenario](const Graph&, const Match&) { + if (scenario == 1) return Replacement("(middle) => Identity(middle)"); + if (scenario == 2) return Replacement("(root) => Identity(root)"); + if (scenario == 3) return Replacement("(input) => Identity(input, invalid=0)"); + return Replacement("(input) => Identity(input)"); + }; + PatternRewriter rewriter({first, second}, {.indexed = indexed}); + Statistics stats; + auto status = rewriter.Run(test.graph, stats); + EXPECT_EQ(status.IsOK(), scenario == 0); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + EXPECT_EQ(test.graph.NumberOfNodes(), scenario == 0 ? 1u : 2u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ParsedReplacementSequenceAndNumerics) { + Replacement replacement(R"( + (x) => + h = HardSigmoid(x, alpha=0.16666667, beta=0.5) + return Mul(h, x) + )"); + const auto* plan = replacement.Plan().get(); + auto copy = replacement; + EXPECT_EQ(copy.Plan().get(), plan); + EXPECT_EQ(replacement.Plan()->size(), 2u); + EXPECT_FLOAT_EQ(replacement.Plan()->front().attributes.at("alpha").f(), 1.0f / 6.0f); + for (bool indexed : {false, true}) { + TestGraph test; + EXPECT_EQ(replacement(test.graph, {}).Plan().get(), plan); + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto outputs = test.graph.GetOutputs(); + auto rule = MakeExampleRules()[0]; + rule.replacement = replacement; + PatternRewriter rewriter({rule}, {.indexed = indexed}); + Statistics stats; + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 2u); + EXPECT_EQ(test.graph.GetOutputs(), outputs); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(Count(test.graph, "HardSigmoid"), 1u); + EXPECT_EQ(Count(test.graph, "Mul"), 1u); + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); + } + auto rule = MakeExampleRules()[0]; + rule.replacement = replacement; + TransformerTester([](ModelTestBuilder& b) { BuildHardSwish(b); }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "HardSigmoid"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Mul"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Clip"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(Rules{rule})); +} + +TEST(PatternRewriterTests, ReplacementBranchesAndMultiOutputRoot) { + auto branched = IdentityPair(); + branched.pattern = Pattern(R"( + (x) => a = Identity(x); b = Identity(a); c = Identity(b); return Identity(c) + )"); + branched.replacement = Replacement("(x) => a, b = Split(x, axis=0); return Concat(a, b, axis=0)"); + auto build_branch = [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* upstream = b.MakeIntermediate(); + b.AddNode("Add", {x, x}, {upstream}); + auto* current = upstream; + for (int i = 0; i < 4; ++i) { + auto* next = i == 3 ? b.MakeOutput() : b.MakeIntermediate(); + b.AddNode("Identity", {current}, {next}); + current = next; + } + b.AddNode("Add", {current, current}, {b.MakeOutput()}); + }; + for (bool indexed : {false, true}) { + TestGraph test; + build_branch(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + for (auto& node : test.graph.Nodes()) node.SetExecutionProviderType(kCpuExecutionProvider); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({branched}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 4u); + ASSERT_STATUS_OK(test.graph.Resolve()); + for (const auto& node : test.graph.Nodes()) { + EXPECT_EQ(node.GetExecutionProviderType(), kCpuExecutionProvider); + if (node.OpType() == "Concat") EXPECT_EQ(node.GetOutputEdgesCount(), 2u); + if (node.OpType() == "Split") { + EXPECT_EQ(node.OutputDefs().size(), 2u); + EXPECT_EQ(node.GetOutputEdgesCount(), 2u); + } + } + } + TransformerTester( + build_branch, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Split"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Concat"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, + std::make_unique(Rules{branched})); + + auto split = IdentitySplit(); + split.pattern = Pattern("(x) => { a = Identity(x); b = Identity(a); left, right = Split(b); return left, right; }"); + split.replacement = Replacement("(x) => { a=Identity(x); left, right = Split(a, axis=0); return left, right; }"); + auto build_split = [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* intermediate = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {intermediate}); + b.AddNode("Split", {intermediate}, {b.MakeOutput(), b.MakeOutput()}).AddAttribute("axis", int64_t{0}); + }; + TransformerTester( + build_split, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Split"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, std::make_unique(Rules{split})); +} + +TEST(PatternRewriterTests, ReplacementParsingCompilationAndRollback) { + for (auto text : {"(x)=>Unknown(y)", "(x)=>a=Identity(x); Identity(x)", "(x)=>Add(x,1e999)", + "(x)=>a=Identity(x);return Identity(a,axis=[0,])", "(x)=>Identity(x);more", + "(x)=>Identity(x,axis=0,axis=1)", "(x)=>Identity(x,mode='bad)", + "(x)=>a=Identity(a);return Identity(a)", "(x)=>a=Split(x);return Identity(b)", + "(x)=>a=Identity(x);return Identity(x)", "(x)=>return = Identity(x)", + "(x)=>return a,a = Split(x)"}) + EXPECT_THROW(Replacement{text}, OnnxRuntimeException) << text; + auto invalid = IdentityPair(); + invalid.replacement = Replacement("(missing)=>Identity(missing)"); + EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); + invalid.replacement = Replacement("(a)=>Identity(a)"); + EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); + invalid.replacement = Replacement("(x)=>a=Identity(x);return Identity(a)"); + EXPECT_NO_THROW(PatternRewriter({invalid})); + invalid.replacement = Replacement("(x)=>com.microsoft::FusedMatMul(x,x)"); + EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); + EXPECT_NO_THROW(Replacement("(Identity)=>Identity(Identity)")); + for (bool indexed : {false, true}) { + for (auto text : {"(x)=>a=Identity(x);return NoSuchReplacementOperator(a)", + "(x)=>a=Identity(x);return Softmax(a,axis='wrong')", + "(x)=>a=Identity(x);return com.microsoft::HardSwish(a)", + "(x)=>a=Add(x,0.0);return NoSuchReplacementOperator(a)", + "(x)=>HardSwish(x,invalid=0)"}) { + SCOPED_TRACE(text); + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto initializer_count = test.graph.GetAllInitializedTensors().size(); + auto rule = MakeExampleRules()[0]; + rule.replacement = [recipe = Replacement(text)](const Graph&, const Match&) { return recipe; }; + Statistics stats; + ASSERT_FALSE(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(test.graph.NumberOfNodes(), 4u); + EXPECT_EQ(test.graph.GetAllInitializedTensors().size(), initializer_count); + ASSERT_STATUS_OK(test.graph.Resolve()); + rule.replacement = Replacement("(x)=>a=Identity(x);return HardSwish(a)"); + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ReplacementScalarConstantsAndParsedSequence) { + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); + for (bool integer : {false, true}) { + rule.replacement = Replacement(integer ? "(x)=>Add(x,0)" : "(x)=>Add(x,0.0)"); + auto* compiled = rule.replacement.target(); + ASSERT_NE(compiled, nullptr); + ASSERT_EQ(compiled->Plan()->size(), 1u); + const auto& scalar = std::get(compiled->Plan()->front().inputs[1]); + EXPECT_EQ(std::holds_alternative(scalar.value), integer); + auto build = [integer](ModelTestBuilder& b) { + auto* x = integer ? b.MakeInput({4}, -5, 5) : b.MakeInput({4}, -1.0f, 1.0f); + auto* current = x; + for (int i = 0; i < 3; ++i) { + auto* next = i == 2 ? b.MakeOutput() : b.MakeIntermediate(); + b.AddNode("Identity", {current}, {next}); + current = next; + } + }; + TransformerTester( + build, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Add"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, std::make_unique(Rules{rule})); + } + rule.replacement = Replacement("(x)=>a=Identity(x);return Identity(a)"); + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* current = x; + for (int i = 0; i < 3; ++i) { + auto* next = i == 2 ? test.builder.MakeOutput() : test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {current}, {next}); + current = next; + } + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 2u); + ASSERT_STATUS_OK(test.graph.Resolve()); +} + +TEST(PatternRewriterTests, ReplacementPlanBindingAndSequenceInvalidation) { + Replacement shared("(x,y)=>Sub(x,y)"); + for (auto text : {"(x,y)=>a=Identity(x);b=Identity(y);return Sub(a,b)", + "(y,x)=>a=Identity(x);b=Identity(y);return Sub(a,b)"}) { + auto rule = IdentityPair(); + rule.pattern = Pattern(text); + rule.replacement = shared; + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2}, 1.0f, 2.0f); + auto* y = b.MakeInput({2}, -2.0f, -1.0f); + auto* a = b.MakeIntermediate(); + auto* intermediate = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {y}, {intermediate}); + b.AddNode("Sub", {a, intermediate}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Sub"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, + std::make_unique(Rules{rule})); + } + for (bool indexed : {false, true}) { + auto first = IdentityPair("first", false); + first.pattern = Pattern("(first)=>a=Identity(first);b=Identity(a);c=Identity(b);return Identity(c)"); + first.replacement = Replacement("(first)=>Identity(first)"); + auto second = first; + second.pattern = Pattern("(second)=>a=Identity(second);b=Identity(a);c=Identity(b);return Identity(c)"); + second.check_and_bind = [](const Graph&, Match&) { return true; }; + second.replacement = Replacement("(second)=>a=Identity(second);return Identity(a)"); + auto followup = IdentityPair(); + followup.check_and_bind = [](const Graph&, Match& m) { return m.nodes.back()->Name().find("rewrite") == 0; }; + PatternRewriter rewriter({first, second, followup}, {.indexed = indexed}); + for (int repeat = 0; repeat < 2; ++repeat) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* current = x; + for (int i = 0; i < 4; ++i) { + auto* next = i == 3 ? test.builder.MakeOutput() : test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {current}, {next}); + current = next; + } + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 2u); + EXPECT_EQ(test.graph.NumberOfNodes(), 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto invalid = second; + invalid.replacement = [](const Graph&, const Match&) { + return Replacement("(second)=>a=Identity(second);b=Identity(a);c=Identity(b);return Identity(c,invalid=1)"); + }; + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* current = x; + for (int i = 0; i < 4; ++i) { + auto* next = i == 3 ? test.builder.MakeOutput() : test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {current}, {next}); + current = next; + } + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + EXPECT_FALSE(PatternRewriter({invalid}, {.indexed = indexed}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(test.graph.NumberOfNodes(), 4u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, ReplacementExtraAttributeBindings) { + static_assert(std::is_invocable_v); + static_assert(!std::is_invocable_v); + static_assert(!std::is_invocable_v); + static_assert(!std::is_invocable_v); + static_assert(!std::is_invocable_v); + static_assert(std::is_invocable_v); + auto make_attribute = [](ONNX_NAMESPACE::AttributeProto_AttributeType type) { + ONNX_NAMESPACE::AttributeProto attr; + attr.set_name("ignored_source_name"); + attr.set_type(type); + return attr; + }; + Match empty; + EXPECT_THROW(empty.attributeOf("missing"), OnnxRuntimeException); + auto alpha = make_attribute(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + alpha.set_f(0.25f); + EXPECT_THROW(empty.update("", alpha), OnnxRuntimeException); + EXPECT_THROW(empty.update("alpha", {}), OnnxRuntimeException); + auto reference = alpha; + reference.set_ref_attr_name("other"); + EXPECT_THROW(empty.update("alpha", reference), OnnxRuntimeException); + empty.update("alpha", alpha); + EXPECT_FLOAT_EQ(empty.attributeOf("alpha").f(), 0.25f); + alpha.set_f(0.5f); + empty.update("alpha", alpha); + EXPECT_FLOAT_EQ(empty.attributeOf("alpha").f(), 0.5f); + auto text_attribute = make_attribute(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + text_attribute.set_s("none"); + empty.update("approximate", text_attribute); + EXPECT_EQ(empty.attributeOf("approximate").s(), "none"); + empty.ClearUpdates(); + EXPECT_THROW(empty.attributeOf("alpha"), OnnxRuntimeException); + + for (auto text : {"(x)=>HardSigmoid(x,alpha=unknown)", + "(x,a)=>HardSigmoid(x,alpha=x)", + "(x)=>a=Identity(x);return HardSigmoid(a,alpha=a)", + "(x,a)=>HardSigmoid(x,alpha=a,alpha=0.5)"}) + EXPECT_THROW(Replacement{text}, OnnxRuntimeException); + EXPECT_NO_THROW(Pattern{"(x,a)=>return HardSigmoid(x,alpha=a)"}); + Replacement dotted("(x,alpha)=>com.microsoft.FusedMatMul(x,x,alpha=alpha)"); + EXPECT_EQ(dotted.Plan()->front().domain, kMSDomain); + EXPECT_EQ(dotted.Plan()->front().op, "FusedMatMul"); + EXPECT_EQ(dotted.Plan()->front().attribute_bindings.at("alpha"), "alpha"); + + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto first = MakeExampleRules()[0]; + auto second = first; + // Rejected rules cannot leak prepared bindings into the next rule. + first.check_and_bind = [alpha](const Graph&, Match& m) { + m.update("alpha", alpha); + return false; + }; + second.replacement = Replacement(scenario == 4 + ? "(x,alpha,beta)=>h=HardSigmoid(x,alpha=alpha,beta=beta,extra=0.0);return Mul(h,x)" + : "(x,alpha,beta)=>h=HardSigmoid(x,alpha=alpha,beta=beta);return Mul(h,x)"); + second.check_and_bind = [scenario, make_attribute](const Graph&, Match& m) { + auto beta = make_attribute(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + beta.set_f(0.5f); + m.update("beta", beta); + if (scenario == 1) return true; // alpha is missing; must not leak from the rejected rule. + if (scenario == 2) m.update("x", beta); + auto alpha_value = make_attribute(scenario == 3 ? ONNX_NAMESPACE::AttributeProto_AttributeType_INT + : ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + alpha_value.set_f(1.0f / 6.0f); + m.update("alpha", alpha_value); + return true; + }; + PatternRewriter rewriter({first, second}, {.indexed = indexed}); + Statistics stats; + if (scenario == 1 || scenario == 2) { + EXPECT_THROW((void)rewriter.Run(test.graph, stats), OnnxRuntimeException); + } else { + EXPECT_EQ(rewriter.Run(test.graph, stats).IsOK(), scenario == 0); + } + EXPECT_EQ(test.graph.NumberOfNodes(), scenario == 0 ? 2u : 4u); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 0) { + for (const auto& node : test.graph.Nodes()) { + if (node.OpType() != "HardSigmoid") continue; + EXPECT_EQ(node.GetAttributes().at("alpha").name(), "alpha"); + EXPECT_FLOAT_EQ(node.GetAttributes().at("alpha").f(), 1.0f / 6.0f); + EXPECT_FLOAT_EQ(node.GetAttributes().at("beta").f(), 0.5f); + } + } + } + } +} + +TEST(PatternRewriterTests, ReplacementExtraBindingsVaryPerMatchAndAttributeTypes) { + for (bool indexed : {false, true}) { + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); + rule.replacement = Replacement("(x,axis)=>Softmax(x,axis=axis)"); + int calls = 0; + rule.check_and_bind = [&calls](const Graph&, Match& m) { + ONNX_NAMESPACE::AttributeProto axis; + axis.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis.set_i((calls++ / 2) % 2); + m.update("axis", axis); + return true; + }; + PatternRewriter matcher({rule}, {.match_only = true, .indexed = indexed}); + PatternRewriter rewriter({rule}, {.indexed = indexed}); + for (int repeat = 0; repeat < 2; ++repeat) { + TestGraph test; + auto* x = test.builder.MakeInput({2, 2}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + auto* b = test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {x}, {a}); + test.builder.AddNode("Identity", {a}, {b}); + test.builder.AddNode("Identity", {b}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(matcher.Run(test.graph, stats)); + EXPECT_EQ(calls, 2 * repeat + 1); + EXPECT_EQ(test.graph.NumberOfNodes(), 3u); + EXPECT_EQ(stats.matches, 1u); + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(calls, 2 * repeat + 2); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.Nodes().begin()->GetAttributes().at("axis").i(), repeat); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); + rule.replacement = Replacement("(x,perm)=>Transpose(x,perm=perm)"); + auto check_and_bind = [](const Graph&, Match& m) { + ONNX_NAMESPACE::AttributeProto perm; + perm.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); + perm.add_ints(0); + perm.add_ints(1); + m.update("perm", perm); + return true; + }; + rule.check_and_bind = check_and_bind; + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* intermediate = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {intermediate}); + b.AddNode("Identity", {intermediate}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Transpose"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, + std::make_unique(Rules{rule})); +} + +TEST(PatternRewriterTests, CheckAndBindTensorAttributeToTextReplacement) { + for (bool indexed : {false, true}) { + auto check_and_bind = [](const Graph&, Match& match) { + ONNX_NAMESPACE::AttributeProto zero; + zero.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* tensor = zero.mutable_t(); + tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor->add_dims(1); + tensor->add_float_data(0.0f); + match.update("zero", std::move(zero)); + return true; + }; + Rule rule{"prepared_tensor", + "(x:float[2,3])=>{a=Identity(x);return Identity(a);}", + "(x,zero)=>{z=ConstantOfShape([2,3],value=zero);return Add(x,z);}", + check_and_bind, + {{kOnnxDomain, 13, 17}}}; + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + const auto& graph = session.GetGraph(); + EXPECT_EQ(Count(graph, "Identity"), 0u); + EXPECT_EQ(Count(graph, "Add"), 1u); + EXPECT_EQ(Count(graph, "ConstantOfShape"), 1u); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "ConstantOfShape") { + const auto& tensor = node.GetAttributes().at("value").t(); + EXPECT_EQ(tensor.data_type(), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + ASSERT_EQ(tensor.dims_size(), 1); + EXPECT_EQ(tensor.dims(0), 1); + ASSERT_EQ(tensor.float_data_size(), 1); + EXPECT_FLOAT_EQ(tensor.float_data(0), 0.0f); + } + } + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, CheckAndBindRejectedStateIsolation) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + TestGraph test; + auto& b = test.builder; + for (int i = 0; i < 2; ++i) { + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rejected = IdentityPair("rejected"); + rejected.check_and_bind = [](const Graph&, Match& match) { + ONNX_NAMESPACE::AttributeProto value; + value.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + value.set_f(0.25f); + match.update("prepared", value); + EXPECT_TRUE(match.BindAttribute("captured", value)); + EXPECT_TRUE(match.BindDimension("scratch_dimension", int64_t{2})); + EXPECT_TRUE(match.BindElementType("scratch_type", ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + return match.Reject("rejected after preparation"); + }; + auto accepted = IdentityPair("accepted"); + size_t calls = 0; + accepted.check_and_bind = [&calls](const Graph&, Match& match) { + EXPECT_THROW(match.attributeOf("prepared"), OnnxRuntimeException); + EXPECT_THROW(match.attributeOf("captured"), OnnxRuntimeException); + EXPECT_THROW(match.dimensionOf("scratch_dimension"), OnnxRuntimeException); + EXPECT_THROW(match.elementTypeOf("scratch_type"), OnnxRuntimeException); + EXPECT_TRUE(match.RejectionReason().empty()); + ++calls; + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rejected, accepted}, {.match_only = match_only, .indexed = indexed}) + .Run(test.graph, stats)); + EXPECT_EQ(calls, 2u); + EXPECT_EQ(stats.check_and_bind_calls, 4u); + EXPECT_EQ(stats.matches, 2u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 2u); + EXPECT_EQ(test.graph.NumberOfNodes(), match_only ? 4 : 2); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ParsedHardSwishConstantsAndExecution) { + constexpr auto text = R"( + (x, low, high, divisor) => + add = Add(x, 3.0) + clip = Clip(add, low, high) + mul = Mul(clip, x) + return Div(mul, divisor) + )"; + Pattern parsed(text); + EXPECT_EQ(parsed.NamedValue("x").name, "x"); + EXPECT_EQ(parsed.NamedValue("low").handle, 1); + EXPECT_EQ(parsed.NamedValue("divisor").handle, 3); + EXPECT_EQ(parsed.Nodes()[0].outputs[0].name, "add"); + EXPECT_EQ(parsed.Constants().size(), 1u); + auto semantic_guard = [type_guard = MakeExampleRules()[0].check_and_bind](const Graph& graph, Match& match) { + if (!type_guard(graph, match)) return false; + for (const auto& [name, expected] : + {std::pair{"low", 0.0f}, std::pair{"high", 6.0f}, std::pair{"divisor", 6.0f}}) { + const auto* tensor = graph_utils::GetConstantInitializer(graph, match.valueOf(name)->Name()); + if (!tensor || tensor->dims_size() != 0 || tensor->data_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) + return false; + Initializer initializer(graph, *tensor, graph.ModelPath()); + if (initializer.size() != 1 || *initializer.data() != expected) return false; + } + return true; + }; + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + NodeArg* bias = nullptr; + for (auto& node : test.graph.Nodes()) { + if (node.OpType() == "Add") bias = node.MutableInputDefs()[1]; + } + ASSERT_NE(bias, nullptr); + if (scenario == 1) { + const auto* tensor = test.graph.GetInitializer(bias->Name(), false); + ASSERT_NE(tensor, nullptr); + auto updated = *tensor; + updated.clear_raw_data(); + updated.clear_float_data(); + updated.add_float_data(4.0f); + test.graph.RemoveInitializedTensor(bias->Name()); + test.graph.AddInitializedTensor(updated); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + if (scenario == 2) { + auto inputs = test.graph.GetInputsIncludingInitializers(); + inputs.push_back(bias); + test.graph.SetInputs(inputs); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + if (scenario == 3) { + for (auto& node : test.graph.Nodes()) { + if (node.OpType() == "Mul") node.MutableInputDefs()[1] = node.MutableInputDefs()[0]; + } + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto rule = MakeExampleRules()[0]; + rule.pattern = parsed; + rule.check_and_bind = semantic_guard; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + auto rule = MakeExampleRules()[0]; + rule.pattern = parsed; + rule.check_and_bind = semantic_guard; + TransformerTester([](ModelTestBuilder& b) { BuildHardSwish(b); }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "HardSigmoid"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Clip"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(Rules{rule})); +} + +TEST(PatternRewriterTests, ExampleFixedConstantsArePatternLiterals) { + const auto rule = MakeExampleRules()[0]; + EXPECT_EQ(rule.pattern.Constants().size(), 4u); + EXPECT_EQ(rule.pattern.NamedValue("x").handle, 0); + for (auto name : {"bias", "low", "high", "divisor"}) + EXPECT_THROW(rule.pattern.NamedValue(name), OnnxRuntimeException); + for (bool indexed : {false, true}) { + for (auto op : {"Add", "Clip", "Div"}) { + for (int slot = op == std::string_view("Clip") ? 1 : 0; + slot < (op == std::string_view("Clip") ? 3 : 1); ++slot) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const NodeArg* argument = nullptr; + for (const auto& node : test.graph.Nodes()) { + if (node.OpType() == op) argument = node.InputDefs()[op == std::string_view("Clip") ? slot : 1]; + } + ASSERT_NE(argument, nullptr); + const auto* original = test.graph.GetInitializer(argument->Name(), false); + ASSERT_NE(original, nullptr); + auto tensor = *original; + if (scenario == 0) { + tensor.clear_raw_data(); + tensor.clear_float_data(); + tensor.add_float_data(9.0f); + } + if (scenario == 1) { + tensor.add_dims(1); + test.graph.GetNodeArg(argument->Name())->ClearShape(); + } + if (scenario == 2) { + tensor.clear_raw_data(); + tensor.clear_float_data(); + tensor.add_float_data(std::numeric_limits::infinity()); + } + if (scenario != 3) { + test.graph.RemoveInitializedTensor(argument->Name()); + test.graph.AddInitializedTensor(tensor); + } else { + auto inputs = test.graph.GetInputsIncludingInitializers(); + inputs.push_back(argument); + test.graph.SetInputs(inputs); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u) << op << " input " << slot << " scenario " << scenario; + EXPECT_EQ(test.graph.NumberOfNodes(), 4u); + } + } + } + } +#ifndef DISABLE_CONTRIB_OPS + const auto rms = MakeExampleRules().back(); + ASSERT_EQ(rms.pattern.Constants().size(), 1u); + EXPECT_FLOAT_EQ(std::get(rms.pattern.Constants().begin()->second), 2.0f); + EXPECT_THROW(rms.pattern.NamedValue("exponent"), OnnxRuntimeException); + EXPECT_NO_THROW(rms.pattern.NamedValue("epsilon")); +#endif +} + +TEST(PatternRewriterTests, ParsedAttributesMultiOutputsAndDomains) { + for (bool indexed : {false, true}) { + Pattern split("(x) => a, b = Split(x, axis = 0); return Concat(a, b, axis = 0)"); + auto rule = SplitConcat(); + rule.pattern = split; + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + auto* b = test.builder.MakeIntermediate(); + test.builder.AddNode("Split", {x}, {a, b}).AddAttribute("axis", int64_t{0}); + test.builder.AddNode("Concat", {a, b}, {test.builder.MakeOutput()}).AddAttribute("axis", int64_t{0}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + Pattern p("(x, y) => a = com.microsoft::FusedMatMul(x, y, alpha=+1.0e-1); return Identity(a)"); + EXPECT_EQ(p.Nodes()[0].op.domain, kMSDomain); + EXPECT_EQ(p.Nodes()[0].attributes.size(), 1u); + p = Pattern("(x) => a = Identity(x, mode='linear'); return Identity(a)"); + ONNX_NAMESPACE::AttributeProto attr; + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + attr.set_s("linear"); + EXPECT_TRUE(p.Nodes()[0].attributes[0].Matches(attr)); + Pattern original("(x) => return Identity(x)"); + for (auto text : + {"(x,x)=>return Identity(x)", "(x)=>return Identity(y)", "(x)=>a=Identity(x)", + "(x)=>a=Identity(x);return Identity(a);extra", "(x)=>a=Identity(x);return Add(a, 1e999)", + "(x)=>a=Identity(x);return Add(a, 9223372036854775808)", "(x)=>a=Identity(x);return Add(a, 1.2.3)", + "(x)=>a=Identity(x);return Identity(a, mode='bad)", "(x)=>a=Identity(x);return Identity(a, axis=1, axis=2)", + "(x)=>a=Identity(x);return Add(axis=1, a)", "(x)=>a=Identity(x);return Identity(a,) ", + "(x)=>a=Identity(x);return Identity(a, axis=[1,])", "(x,y)=>return Identity(x)", + "(x)=>a=Identity(x);return Identity(x)", "(x)=>x=Identity(x);return Identity(x)", + "(x)=>a=Identity(x);return Identity(a) ''", "(x)=>'return' Identity(x)", "(x)=>return Identity(x ')'"}) { + try { + original = Pattern(text); + FAIL() << "Expected malformed pattern to throw: " << text; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("byte"), std::string::npos); + } + EXPECT_EQ(original.Nodes().size(), 1u); + EXPECT_EQ(original.NamedValue("x").handle, 0); + } + p = Pattern("(x)=>a=Identity(x, mode='');return Identity(a)"); + attr.set_s(""); + EXPECT_TRUE(p.Nodes()[0].attributes[0].Matches(attr)); + std::string excessive = "(x)=>"; + for (int i = 0; i < 65; ++i) + excessive += "v" + std::to_string(i) + "=Identity(" + (i == 0 ? "x" : "v" + std::to_string(i - 1)) + ");"; + excessive += "return Identity(v64)"; + EXPECT_THROW(original = Pattern(excessive), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, LiteralTypeAndRankMatching) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, 0, 5); + auto* intermediate = test.builder.MakeIntermediate(); + auto* literal = scenario == 1 ? test.builder.MakeInitializer({1}, {3}) + : test.builder.MakeScalarInitializer(scenario == 2 ? 4 : 3); + test.builder.AddNode("Add", {x, literal}, {intermediate}); + test.builder.AddNode("Identity", {intermediate}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>a=Add(x,3);return Identity(a)"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 0 ? 1u : 0u); + rule.pattern = Pattern("(x)=>a=Add(x,3.0);return Identity(a)"); + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 0u); + } + } + Pattern p("(x)=>a=Add(x,3);return Add(a,3.0)"); + auto integer = p.Nodes()[0].inputs[1]; + auto floating = p.Nodes()[1].inputs[1]; + EXPECT_TRUE(std::holds_alternative(p.Constants().at(integer))); + EXPECT_TRUE(std::holds_alternative(p.Constants().at(floating))); +} + +TEST(PatternRewriterTests, VectorLiteralParsingAndErrors) { + Pattern pattern("(x)=>{a=Add(x,[1,2]);return Identity(a);}"); + const auto& integers = std::get>(pattern.Constants().begin()->second); + ASSERT_EQ(integers.size(), 2u); + EXPECT_EQ(integers[0], 1); + EXPECT_EQ(integers[1], 2); + Replacement plan("(x)=>Add(x,[1.0,2e0])"); + const auto& values = std::get>( + std::get(plan.Plan()->front().inputs[1]).value); + ASSERT_EQ(values.size(), 2u); + EXPECT_FLOAT_EQ(values[1], 2.0f); + for (auto text : {"(x)=>return Add(x,[1,2.0])", "(x)=>return Add(x,[[1]])", + "(x)=>return Add(x,[1,])", "(x)=>return Add(x,[1)", "(x)=>return Add(x,[1e999])", + "(x)=>return Add(x,[9223372036854775808])", "(x)=>return Add(x,[1 2])", + "(x)=>return Identity(x,axes=[1,2.0])", "(x)=>return Identity(x,axes=[1,])"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + } + Pattern empty("(x)=>{a=Add(x,[]);return Identity(a);}"); + EXPECT_TRUE(std::get>(empty.Constants().begin()->second).empty()); +} + +TEST(PatternRewriterTests, VectorLiteralMatchingAndNumerics) { + for (bool indexed : {false, true}) { + SCOPED_TRACE(indexed); + for (int scenario = 0; scenario < 5; ++scenario) { + SCOPED_TRACE(scenario); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0, 5); + NodeArg* literal = scenario == 1 ? b.MakeInitializer({1, 2}, {1, 2}) + : scenario == 2 ? b.MakeInitializer({2}, {1.0f, 2.0f}) + : b.MakeInitializer({2}, {1, scenario == 3 ? 3 : 2}); + if (scenario == 2) { + x = b.MakeInput({2}, 0.0f, 5.0f); + } + auto* a = b.MakeIntermediate(); + b.AddNode("Add", {x, literal}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 4) { + auto inputs = test.graph.GetInputsIncludingInitializers(); + inputs.push_back(literal); + test.graph.SetInputs(inputs); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{a=Add(x,[1,2]);return Identity(a);}"); + rule.replacement = Replacement("(x)=>Add(x,[1,2])"); + auto rejected = rule; + rejected.pattern = Pattern("(x)=>{a=Add(x,[4,5]);return Identity(a);}"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rejected, rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + for (bool floating : {false, true}) { + auto rule = IdentityPair(); + rule.replacement = Replacement(floating ? "(x)=>Add(x,[0.0,0.0])" : "(x)=>Add(x,[0,0])"); + TransformerTester( + [floating](ModelTestBuilder& b) { + auto* x = floating ? b.MakeInput({2}, -1.0f, 1.0f) : b.MakeInput({2}, -5, 5); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "Add"), 1u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } + for (auto text : {"(x)=>{a=Add(x,[0.0,0.0]);return Identity(a);}", + "(x)=>{a=Add(x,[]);return Identity(a);}"}) { + TestGraph test; + const bool empty = std::string_view(text).find("[]") != std::string_view::npos; + auto* x = empty ? test.builder.MakeInput({0}, std::vector{}) + : test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* literal = empty ? test.builder.MakeInitializer({0}, std::vector{}) + : test.builder.MakeInitializer({2}, {0.0f, 0.0f}); + auto* a = test.builder.MakeIntermediate(); + test.builder.AddNode("Add", {x, literal}, {a}); + test.builder.AddNode("Identity", {a}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern(text); + rule.replacement = Replacement(empty ? "(x)=>Add(x,[])" : "(x)=>Add(x,[0.0,0.0])"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + for (auto text : {"(x)=>{a=Add(x,[0,0]);return NoSuchOperator(a);}", + "(x)=>{a=Add(x,[0.0,0.0]);return NoSuchOperator(a);}"}) { + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto count = test.graph.GetAllInitializedTensors().size(); + auto rule = MakeExampleRules()[0]; + rule.replacement = Replacement(text); + Statistics stats; + EXPECT_FALSE(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 4u); + EXPECT_EQ(test.graph.GetAllInitializedTensors().size(), count); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ListAttributesAndDottedOperators) { + for (auto text : {"(x,y)=>{a=com.microsoft.FusedMatMul(x,y);return Identity(a);}", + "(x,y)=>{a=com.microsoft::FusedMatMul(x,y);return Identity(a);}"}) { + Pattern pattern(text); + EXPECT_EQ(pattern.Nodes()[0].op.domain, kMSDomain); + EXPECT_EQ(pattern.Nodes()[0].op.name, "FusedMatMul"); + } + for (auto text : {"(x)=>return ai.onnx.Identity(x)", "(x)=>return ai.onnx::Identity(x)"}) { + EXPECT_EQ(Pattern(text).Nodes()[0].op.domain, kOnnxDomain); + EXPECT_EQ(Replacement(text).Plan()->front().domain, kOnnxDomain); + } + for (bool indexed : {false, true}) { + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{a=ai.onnx.Transpose(x,perm=[1,0]);return Identity(a);}"); + rule.replacement = Replacement("(x)=>ai.onnx.Transpose(x,perm=[1,0])"); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {a}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + for (auto text : {"(x)=>{a=Transpose(x,perm=[0,1]);return Identity(a);}", + "(x)=>{a=Transpose(x,perm=[1]);return Identity(a);}", + "(x)=>{a=Transpose(x,perm=[1.0,0.0]);return Identity(a);}"}) { + TestGraph test; + auto* x = test.builder.MakeInput({2, 3}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + test.builder.AddNode("Transpose", {x}, {a}).AddAttribute("perm", std::vector{1, 0}); + test.builder.AddNode("Identity", {a}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + rule.pattern = Pattern(text); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 0u); + } + } + Pattern floats("(x)=>{a=Identity(x,coefficients=[1.0,2.0]);return Identity(a);}"); + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOATS); + attribute.add_floats(1.0f); + attribute.add_floats(2.0f); + EXPECT_TRUE(floats.Nodes()[0].attributes[0].Matches(attribute)); + attribute.set_floats(1, 3.0f); + EXPECT_FALSE(floats.Nodes()[0].attributes[0].Matches(attribute)); + Replacement float_plan("(x)=>Identity(x,coefficients=[1.0,2.0])"); + EXPECT_EQ(float_plan.Plan()->front().attributes.at("coefficients").type(), + ONNX_NAMESPACE::AttributeProto_AttributeType_FLOATS); +} + +TEST(PatternRewriterTests, AttributeAndConstantPlansCannotAlias) { + for (bool indexed : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + test.builder.AddNode("HardSigmoid", {x}, {a}).AddAttribute("alpha", 0.25f); + test.builder.AddNode("Mul", {a, x}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rejected = IdentityPair(), accepted = IdentityPair(); + rejected.pattern = Pattern("(x)=>a=HardSigmoid(x, alpha=0.2);return Mul(a,x)"); + accepted.pattern = Pattern("(x)=>a=HardSigmoid(x, alpha=0.25);return Mul(a,x)"); + Statistics stats; + ASSERT_STATUS_OK( + PatternRewriter({rejected, accepted}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + Pattern first("(x)=>a=Add(x,3.0);return Identity(a)"); + Pattern second("(x)=>a=Add(x,4.0);return Identity(a)"); + TestGraph constants; + auto* input = constants.builder.MakeInput({2}, -1.0f, 1.0f); + auto* intermediate = constants.builder.MakeIntermediate(); + constants.builder.AddNode("Add", {input, constants.builder.MakeScalarInitializer(4.0f)}, {intermediate}); + constants.builder.AddNode("Identity", {intermediate}, {constants.builder.MakeOutput()}); + constants.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(constants.graph.Resolve()); + rejected.pattern = first; + accepted.pattern = second; + ASSERT_STATUS_OK( + PatternRewriter({rejected, accepted}, {.match_only = true, .indexed = indexed}).Run(constants.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + } +} + +TEST(PatternRewriterTests, GenericExecutionProviderPolicy) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* first = test.builder.MakeIntermediate(); + auto& a = test.builder.AddNode("Identity", {x}, {first}); + auto& b = test.builder.AddNode("Identity", {first}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 1) { + a.SetExecutionProviderType(kCpuExecutionProvider); + b.SetExecutionProviderType(kCpuExecutionProvider); + } + if (scenario >= 2) { + a.SetExecutionProviderType(scenario == 4 ? kCpuExecutionProvider : kCudaExecutionProvider); + b.SetExecutionProviderType(kCudaExecutionProvider); + } + Options options{.match_only = match_only, .indexed = indexed}; + if (scenario >= 3) options.compatible_execution_providers.insert(kCudaExecutionProvider); + auto rule = IdentityPair(); + int guards = 0; + rule.check_and_bind = [&guards](const Graph&, Match&) { + ++guards; + return true; + }; + const auto expected_provider = b.GetExecutionProviderType(); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + const bool supported = scenario == 0 || scenario == 1 || scenario == 3; + EXPECT_EQ(guards, supported ? 1 : 0); + EXPECT_EQ(stats.matches, supported ? 1u : 0u); + EXPECT_EQ(stats.rewrites, supported && !match_only ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (stats.rewrites) { + EXPECT_EQ(test.graph.Nodes().begin()->GetExecutionProviderType(), expected_provider); + } + } + } + } +} + +TEST(PatternRewriterTests, ExampleRulesDoNotCheckExecutionProvider) { + TestGraph test; + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + for (auto& node : test.graph.Nodes()) node.SetExecutionProviderType(kCudaExecutionProvider); + Options options; + options.compatible_execution_providers.insert(kCudaExecutionProvider); + Statistics stats; + // Graph-only validation: permitting a provider is not evidence of GPU kernel execution. + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules(), options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(test.graph.Nodes().begin()->GetExecutionProviderType(), kCudaExecutionProvider); +} + +TEST(PatternRewriterTests, PriorityOverlapTerminationAndMatchOnly) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + for (int i = 0; i < 8; ++i) { + auto* output = i == 7 ? b.MakeOutput() : b.MakeIntermediate(); + b.AddNode("Identity", {x}, {output}); + x = output; + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int first = 0, second = 0; + Rule a = IdentityPair("first"), c = IdentityPair("second"); + a.check_and_bind = [&first](const Graph&, Match&) { + ++first; + return true; + }; + c.check_and_bind = [&second](const Graph&, Match&) { + ++second; + return true; + }; + PatternRewriter rewriter({a, c}, {.indexed = indexed}); + Statistics match_stats; + ASSERT_STATUS_OK(PatternRewriter({a, c}, {.match_only = true, .indexed = indexed}).Run(test.graph, match_stats)); + EXPECT_EQ(test.graph.NumberOfNodes(), 8); + EXPECT_EQ(match_stats.matches, 14u); + first = second = 0; + Statistics stats; + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 7u); + EXPECT_EQ(first, 7); + EXPECT_EQ(second, 0); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics repeated; + ASSERT_STATUS_OK(rewriter.Run(test.graph, repeated)); + EXPECT_EQ(repeated.rewrites, 0u); + } +} + +TEST(PatternRewriterTests, SingleNodeSameSizeAndExpansionNumerics) { + for (bool indexed : {false, true}) { + for (bool expand : {false, true}) { + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>return Neg(x)"); + rule.replacement = Replacement(expand ? "(x)=>{a=Identity(x);return Neg(a);}" : "(x)=>Neg(x)"); + auto numerical_rule = rule; + if (!expand) numerical_rule.replacement = Replacement("(x)=>Mul(x,-1.0)"); + numerical_rule.check_and_bind = [](const Graph& graph, Match& match) { + const Node* producer = graph.GetProducerNode(match.valueOf("x")->Name()); + return !producer || producer->OpType() != "Identity"; + }; + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + b.AddNode("Neg", {x}, {b.MakeOutput()}); + }, + [expand](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Neg"), expand ? 1u : 0u); + EXPECT_EQ(Count(session.GetGraph(), "Mul"), expand ? 0u : 1u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), expand ? 1u : 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{numerical_rule}, Options{.indexed = indexed, .max_iterations = 1})); + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto* output = test.builder.MakeOutput(); + test.builder.AddNode("Neg", {x}, {output}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed, .max_iterations = 3}).Run(test.graph, stats)); + EXPECT_EQ(stats.iterations, 3u); + EXPECT_TRUE(stats.iteration_limit_reached); + EXPECT_EQ(stats.rewrites, 3u); + EXPECT_EQ(Count(test.graph, "Identity"), expand ? 3u : 0u); + EXPECT_EQ(test.graph.GetProducerNode(output->Name())->OpType(), "Neg"); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + EXPECT_THROW((void)Replacement{"(x)=>return x"}, OnnxRuntimeException); +} + +TEST(PatternRewriterTests, BoundedCyclesPriorityAndMatchOnly) { + for (bool indexed : {false, true}) { + for (size_t limit : {1u, 2u, 5u}) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto forward = IdentityPair("forward"); + forward.pattern = Pattern("(x)=>return Identity(x)"); + forward.replacement = Replacement("(x)=>Transpose(x,perm=[0])"); + auto backward = IdentityPair("backward"); + backward.pattern = Pattern("(x)=>return Transpose(x,perm=[0])"); + backward.replacement = Replacement("(x)=>Identity(x)"); + auto later = forward; + later.name = "later"; + later.check_and_bind = [](const Graph&, Match&) { + ADD_FAILURE() << "Lower priority rule selected."; + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({forward, later, backward}, + {.indexed = indexed, .max_iterations = limit}) + .Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, limit); + EXPECT_EQ(stats.iterations, limit); + EXPECT_TRUE(stats.iteration_limit_reached); + EXPECT_EQ(Count(test.graph, "Transpose"), limit % 2); + ASSERT_STATUS_OK(test.graph.Resolve()); + ASSERT_STATUS_OK(PatternRewriter({forward, backward}, + {.match_only = true, .indexed = indexed, .max_iterations = limit}) + .Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(stats.iterations, 1u); + EXPECT_FALSE(stats.iteration_limit_reached); + } + } +} + +TEST(PatternRewriterTests, BoundedIterationsStopOnConvergence) { + for (bool indexed : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>return Identity(x)"); + rule.replacement = Replacement("(x)=>Transpose(x,perm=[0])"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed, .max_iterations = 5}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(stats.iterations, 2u); + EXPECT_FALSE(stats.iteration_limit_reached); + ASSERT_STATUS_OK(test.graph.Resolve()); + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(stats.iterations, 1u); + EXPECT_FALSE(stats.iteration_limit_reached); + } +} + +TEST(PatternRewriterTests, InvalidRuleAndRecipe) { + EXPECT_THROW(Pattern("(x)=>a=Identity(x);return Identity(x)"), OnnxRuntimeException); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>return Identity(x)"); + EXPECT_NO_THROW(PatternRewriter({rule})); + EXPECT_THROW((PatternRewriter({rule}, {.max_iterations = 0})), OnnxRuntimeException); + TestGraph test; + auto* input = test.builder.MakeInput({1}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {input}, {a}); + test.builder.AddNode("Identity", {a}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + rule = IdentityPair(); + for (auto text : {"(x)=>NoSuchReplacementOperator(x)", "(x)=>not.imported::Identity(x)", + "(a)=>Identity(a)", "(x)=>Identity(x,invalid=1)", "(x)=>com.microsoft::Identity(x)"}) { + rule.replacement = [recipe = Replacement(text)](const Graph&, const Match&) { return recipe; }; + Statistics stats; + EXPECT_FALSE(PatternRewriter({rule}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 2); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, InvalidOpsetRequirements) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + auto rule = IdentityPair(); + if (scenario == 0) rule.opsets.clear(); + if (scenario == 1) rule.opsets = {{kOnnxDomain, 0, 17}}; + if (scenario == 2) rule.opsets = {{kOnnxDomain, 18, 17}}; + if (scenario == 3) rule.opsets = {{kOnnxDomain, 1, 17}, {kOnnxDomainAlias, 1, 17}}; + if (scenario == 4) rule.opsets = {{kMSDomain, 1, 1}}; + EXPECT_THROW((PatternRewriter({rule}, {.indexed = indexed})), OnnxRuntimeException); + } + } +} + +TEST(PatternRewriterTests, RuleImportRangesAndAliases) { + for (bool indexed : {false, true}) { + auto rule = IdentityPair(); + rule.opsets = {{kOnnxDomainAlias, 14, 17}}; + rule.pattern = Pattern("(x)=>a=ai.onnx::Identity(x);return ai.onnx::Identity(a)"); + int guards = 0; + rule.check_and_bind = [&guards](const Graph&, Match&) { + ++guards; + return true; + }; + rule.replacement = Replacement("(x)=>ai.onnx::Identity(x)"); + // Reusing one compiled rule across models must not retain the first model's imports. + PatternRewriter rewriter({rule}, {.indexed = indexed}); + for (int opset : {13, 14, 15, 16, 17, 18}) { + TestGraph test({{kOnnxDomain, opset}}); + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* middle = test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {x}, {middle}); + test.builder.AddNode("Identity", {middle}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const bool supported = opset >= 14 && opset <= 17; + guards = 0; + Statistics stats; + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(guards, supported ? 1 : 0); + EXPECT_EQ(stats.rewrites, supported ? 1u : 0u); + if (!supported) { + EXPECT_EQ(stats.visited, 0u); + EXPECT_EQ(stats.candidates, 0u); + EXPECT_EQ(stats.structural_checks, 0u); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(test.graph.NumberOfNodes(), supported ? 1 : 2); + } + } +} + +TEST(PatternRewriterTests, ImportFilteringWithinSharedStructuralGroups) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* middle = test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {x}, {middle}); + test.builder.AddNode("Identity", {middle}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto inactive = IdentityPair("inactive"), active = IdentityPair("active"), later = IdentityPair("later"); + inactive.opsets = {{kOnnxDomain, 18, 21}}; + int inactive_guards = 0, active_guards = 0, later_guards = 0; + inactive.check_and_bind = [&inactive_guards](const Graph&, Match&) { + ++inactive_guards; + return true; + }; + active.check_and_bind = [&active_guards](const Graph&, Match&) { + ++active_guards; + return true; + }; + later.check_and_bind = [&later_guards](const Graph&, Match&) { + ++later_guards; + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({inactive, active, later}, {.match_only = match_only, .indexed = indexed}) + .Run(test.graph, stats)); + EXPECT_EQ(inactive_guards, 0); + EXPECT_EQ(active_guards, 1); + EXPECT_EQ(later_guards, match_only ? 1 : 0); + EXPECT_EQ(stats.matches, match_only ? 2u : 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, MultipleDomainImportRequirements) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test(scenario == 0 ? std::unordered_map{{kOnnxDomain, 17}} + : std::unordered_map{{kOnnxDomain, 17}, {kMSDomain, 1}}); + auto* x = test.builder.MakeInput({2}, -1.0f, 1.0f); + auto* middle = test.builder.MakeIntermediate(); + test.builder.AddNode("Identity", {x}, {middle}); + test.builder.AddNode("Identity", {middle}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.opsets.push_back({kMSDomain, scenario == 1 ? 2 : 1, 2}); + int guards = 0; + rule.check_and_bind = [&guards](const Graph&, Match&) { + ++guards; + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(guards, scenario == 2 ? 1 : 0); + EXPECT_EQ(stats.rewrites, scenario == 2 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ExampleRuleImportBoundaries) { + for (bool indexed : {false, true}) { + PatternRewriter rewriter(MakeExampleRules(), {.indexed = indexed}); + for (int opset : {13, 14, 17, 21, 22}) { + TestGraph test({{kOnnxDomain, opset}}); + BuildHardSwish(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + const bool supported = opset >= 14 && opset <= 21; + EXPECT_EQ(stats.rewrites, supported ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(Count(test.graph, "HardSwish"), supported ? 1u : 0u); + } + } +} + +TEST(PatternRewriterTests, RuleConstructorTextAndParsedForms) { + static_assert(std::is_copy_constructible_v && std::is_copy_assignable_v); + static_assert(std::is_move_constructible_v && std::is_move_assignable_v); + const auto* pattern_text = "(x)=>{a=Identity(x);return Identity(a);}"; + const auto* replacement_text = "(x)=>Identity(x)"; + auto check = [](const auto& pattern_text, const auto& replacement_text) { + Pattern parsed_pattern(pattern_text); + Replacement parsed_replacement(replacement_text); + auto check_and_bind = [](const Graph&, Match&) { return true; }; + size_t calls = 0; + auto replacement = [parsed_replacement, &calls](const Graph&, const Match&) { + ++calls; + return parsed_replacement; + }; + Rules rules{ + {"literal", "(x)=>{a=Identity(x);return Identity(a);}", "(x)=>Identity(x)", check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"text", pattern_text, replacement_text, check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"parsed", parsed_pattern, parsed_replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"parsed_pattern", parsed_pattern, replacement_text, check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"parsed_replacement", pattern_text, parsed_replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"callback", pattern_text, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}, + {"parsed_callback", parsed_pattern, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}}; + for (bool indexed : {false, true}) { + for (const auto& rule : rules) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* middle = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {middle}); + b.AddNode("Identity", {middle}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 1); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + EXPECT_EQ(calls, 4u); + }; + check(pattern_text, replacement_text); + check(std::string(pattern_text), std::string(replacement_text)); + check(std::string_view(pattern_text), std::string_view(replacement_text)); +} + +TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { + std::string pattern = "(x)=>{a=Identity(x);return Identity(a);}"; + std::string replacement = "(x)=>Identity(x)"; + auto check_and_bind = [](const Graph&, Match&) { return true; }; + Rule rule{"owned", pattern, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}; + pattern.assign("not a pattern"); + replacement.assign("not a replacement"); + EXPECT_EQ(rule.pattern.Nodes().size(), 2u); + const auto* compiled = rule.replacement.target(); + ASSERT_NE(compiled, nullptr); + const auto* plan = compiled->Plan().get(); + auto copied = rule; + auto moved = std::move(copied); + copied = moved; + copied = std::move(moved); + TestGraph test; + Match match; + for (int i = 0; i < 3; ++i) { + auto result = copied.replacement(test.graph, match); + EXPECT_EQ(result.Plan().get(), plan); + } + size_t calls = 0; + auto callback = [recipe = *compiled, &calls](const Graph&, const Match&) { + ++calls; + return recipe; + }; + Rule computed{"computed", rule.pattern, callback, check_and_bind, rule.opsets}; + EXPECT_EQ(calls, 0u); + for (int i = 0; i < 3; ++i) { + auto result = computed.replacement(test.graph, match); + EXPECT_EQ(result.Plan().get(), plan); + } + EXPECT_EQ(calls, 3u); +} + +TEST(PatternRewriterTests, DefaultCheckAndBindAcceptsAnnotatedPatterns) { + const auto* pattern = "(x:float[])=>{a=Identity(x);return Identity(a);}"; + const auto* replacement = "(x)=>Identity(x)"; + Pattern parsed_pattern(pattern); + Replacement parsed_replacement(replacement); + Rules rules{ + {"text", pattern, replacement}, + {"parsed", parsed_pattern, parsed_replacement}, + {"parsed_pattern", parsed_pattern, replacement}, + {"parsed_replacement", pattern, parsed_replacement}, + {"explicit_empty", pattern, replacement, {}, {{kOnnxDomain, 13, 17}}}}; + for (auto& rule : rules) { + EXPECT_FALSE(rule.check_and_bind); + if (rule.opsets.empty()) { + EXPECT_THROW((void)PatternRewriter{Rules{rule}}, OnnxRuntimeException); + rule.opsets = {{kOnnxDomain, 13, 17}}; + } + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool floating : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = floating ? b.MakeInput({2, 3}, -1.0f, 1.0f) + : b.MakeInput({2, 3}, -1, 1); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed}) + .Run(test.graph, stats)); + EXPECT_EQ(stats.matches, floating ? 1u : 0u); + EXPECT_EQ(stats.rewrites, floating && !match_only ? 1u : 0u); + EXPECT_EQ(stats.check_and_bind_calls, 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } + } + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "Identity"), 1u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rules.front()})); +} + +TEST(PatternRewriterTests, RuleConstructorRejectsMalformedText) { + const auto* pattern = "(x)=>{a=Identity(x);return Identity(a);}"; + const auto* replacement = "(x)=>Identity(x)"; + auto check_and_bind = [](const Graph&, Match&) { return true; }; + auto make_rule = [&](auto pattern, auto replacement) { + return Rule{"bad", pattern, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}; + }; + EXPECT_THROW(make_rule("not a pattern", replacement), OnnxRuntimeException); + EXPECT_THROW(make_rule(pattern, "not a replacement"), OnnxRuntimeException); + EXPECT_THROW(make_rule(Pattern(pattern), "not a replacement"), OnnxRuntimeException); + EXPECT_THROW(make_rule("not a pattern", Replacement(replacement)), OnnxRuntimeException); + auto callback = [plan = Replacement(replacement)](const Graph&, const Match&) { return plan; }; + EXPECT_THROW(make_rule("not a pattern", callback), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, PriorityAcrossInterleavedStructuralGroups) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + for (int i = 0; i < 3; ++i) { + auto* next = i == 2 ? b.MakeOutput() : b.MakeIntermediate(); + b.AddNode("Identity", {x}, {next}); + x = next; + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rejected = IdentityPair("rejected", false); + auto lower_priority = IdentityPair("lower_priority"); + lower_priority.check_and_bind = [](const Graph& graph, Match& m) { + return graph.NodeProducesGraphOutput(*m.nodes.back()); + }; + const auto* pattern = "(x)=>a=Identity(x);b=Identity(a);return Identity(b)"; + auto check_and_bind = [](const Graph&, Match&) { return true; }; + const auto* replacement = "(x)=>Identity(x)"; + Rule higher_priority{"higher_priority", pattern, replacement, check_and_bind, {{kOnnxDomain, 1, 21}}}; + Statistics stats; + ASSERT_STATUS_OK( + PatternRewriter({rejected, higher_priority, lower_priority}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 1); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, LayerAnnotationsAndRootConsumers) { + for (bool same_layer : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* first = b.MakeIntermediate(); + auto* root = b.MakeOutput(); + auto& a = b.AddNode("Identity", {x}, {first}); + auto& c = b.AddNode("Identity", {first}, {root}); + a.SetLayeringAnnotation(same_layer ? "layer0" : "layer1"); + c.SetLayeringAnnotation("layer0"); + b.AddNode("Neg", {root}, {b.MakeOutput()}); + b.AddNode("Abs", {root}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({IdentityPair()}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, same_layer ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + const Node* producer = test.graph.GetProducerNode(root->Name()); + ASSERT_NE(producer, nullptr); + EXPECT_EQ(producer->GetLayeringAnnotation(), "layer0"); + EXPECT_EQ(producer->GetOutputEdgesCount(), 2); + } +} + +TEST(PatternRewriterTests, OutputCountMismatch) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Dropout", {a}, {b.MakeOutput(), b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Pattern p("(x)=>a=Identity(x);return Dropout(a)"); + auto rule = IdentityPair(); + rule.pattern = p; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); +} + +TEST(PatternRewriterTests, LineCommentParsingAndQuotedHashes) { + for (const auto* newline : {"\n", "\r\n", "\r"}) { + SCOPED_TRACE(newline); + std::string text = "# before header"; + for (const auto* line : { + "(# input", "x# name", ")# header", "=># arrow", "{# body", + "a# output", "=# assignment", "Identity# operator", "(# argument", "x# input", ")# call", + "return# result", "Identity(a)# final call", "}# EOF comment"}) { + text += newline; + text += line; + } + Pattern pattern(text); + Replacement replacement(text); + EXPECT_EQ(pattern.Nodes().size(), 2u); + ASSERT_TRUE(replacement.Plan()); + EXPECT_EQ(replacement.Plan()->size(), 2u); + } + Pattern annotated("(x:# type\nfloat[# dimension\nM] |# alternative\nfloat[M,N])=>return Identity(x)#EOF"); + EXPECT_EQ(annotated.TypeConstraints().at(0).size(), 2u); + Replacement lists("(x)=>Transpose(x,perm=[1,# next element\n0])#EOF"); + EXPECT_EQ(lists.Plan()->front().attributes.at("perm").ints(1), 0); + for (auto text : {R"((x)=>return Identity(x,label='left#right')# trailing)", + R"((x)=>return Identity(x,label="left#right")# trailing)", + R"((x)=>return Identity(x,label="left\"#right")# trailing)"}) { + Replacement replacement(text); + const auto& value = replacement.Plan()->front().attributes.at("label"); + EXPECT_NE(value.s().find('#'), std::string::npos); + Pattern pattern(text); + ASSERT_EQ(pattern.Nodes().front().attributes.size(), 1u); + EXPECT_TRUE(pattern.Nodes().front().attributes.front().Matches(value)); + auto different = value; + different.set_s("left"); + EXPECT_FALSE(pattern.Nodes().front().attributes.front().Matches(different)); + } +} + +TEST(PatternRewriterTests, LineCommentsDoNotRepairMalformedTokens) { + for (auto text : {"# no definition", + "(x)=>return Iden# split identifier\ntity(x)", + "(x)=# split arrow\n>return Identity(x)", + "(x)=>return Add(x,1# split number\n.0)", + "(x)=>return Add(x,1e# split exponent\n2)", + "(x)=>return com.microsoft:# split qualifier\n:FusedMatMul(x,x)", + "(x)=>return Identity(x# closing delimiter swallowed)", + "(x)=>{return Identity(x)# closing brace swallowed}", + "(x)=>return Identity(x,label='unterminated# not a comment)", + "(x)=>return Identity(x)# comment\nunexpected"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + } +} + +TEST(PatternRewriterTests, CommentedRuleNumerics) { + for (bool indexed : {false, true}) { + auto check_and_bind = [](const Graph&, Match&) { return true; }; + Rule rule{"commented", + R"(# pattern + (x:float[# dimensions + M,N]) => { a=Identity(x) # intermediate + return Identity(a) # root + } # end)", + R"(# replacement + (x) => { return Identity(# argument + x) } # end)", + check_and_bind, + {{kOnnxDomain, 13, 17}}}; + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "Identity"), 1u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, ParserOnlyConstruction) { + static_assert(!std::is_default_constructible_v); + static_assert(!std::is_default_constructible_v); + static_assert(!std::is_constructible_v>); + static_assert(std::is_const_v().Nodes())>>); + EXPECT_NO_THROW(Pattern("(x)=>{a=Identity(x);result=Identity(a);return result;}")); + EXPECT_NO_THROW(Replacement("(x)=>{result=Identity(x);return result;}")); +} + +TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { + for (auto text : {"(x)=>{a=Identity(x);return a;}", "(x)=>{a=Identity(x);return a}", + "(x)=>a=Identity(x);return a", "(x)=>{return Identity(x);}", + "(x)=>return Identity(x)", " ( x ) => { \n a = Identity(x); \n return a; \n } \n", + "(x)=>{left,right=Split(x);return left,right;}", + "(x)=>left,right=Split(x);return left,right", + "(x)=>{left,right=ai.onnx::Split(x);return left,right;}", + "(x)=>{a=Identity(x);b=Identity(a);return a,b;}", + "(x)=>{a=Identity(x);b=Identity(a);return a;}", + "(x)=>{left,right=Split(x);return left;}", + "(x)=>{left,right=Split(x);return right,left;}"}) { + SCOPED_TRACE(text); + EXPECT_NO_THROW((void)Pattern{text}); + EXPECT_NO_THROW((void)Replacement{text}); + } + EXPECT_NO_THROW(Replacement("(x)=>{Identity(x);}")); + EXPECT_NO_THROW(Replacement("(x)=>Identity(x)")); + for (auto text : {"(x)=>{a=Identity(x);return a", "(x)=>a=Identity(x);return a}", + "(x)=>{{return Identity(x);}}", "(x)=>{return Identity(x);}}", + "(x)=>{return Identity(x);} extra", "(x)=>{return Identity(x);Identity(x);}", + "(x)=>{}", "(x)=>{a=Identity(x);}", "(x)=>{return x;}", + "(x)=>{a=Identity(x);return missing;}", + "(x)=>{a=Identity(x);b=Identity(x);return a,b;}", + "(x)=>{left,right=Split(x);return left,left;}", + "(x)=>{left,right=Split(x);return left,right,missing;}", + "(x)=>{return a=Identity(x);}", "(x)=>return left,right=Split(x)", + "(x)=>{a=Identity(x);return a;;}", "(x)=>{return Identity(x);};"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + } + Pattern direct("(x)=>{a=Identity(x);return Identity(a);}"); + Pattern named("(x)=>{a=Identity(x);result=Identity(a);return result;}"); + EXPECT_EQ(direct.ValueCount(), named.ValueCount()); + EXPECT_EQ(direct.Nodes().size(), named.Nodes().size()); + EXPECT_EQ(named.Nodes().back().outputs[0].name, "result"); +} + +TEST(PatternRewriterTests, ExplicitBoundaryOutputsAndNumerics) { + for (bool indexed : {false, true}) { + for (bool reversed : {false, true}) { + auto rule = IdentityPair("boundary"); + rule.pattern = Pattern(reversed + ? "(x)=>{a=Identity(x);y=Add(a,a);z=Mul(y,y);return z,y;}" + : "(x)=>{a=Identity(x);y=Add(a,a);z=Mul(y,y);return y,z;}"); + rule.replacement = Replacement(reversed + ? "(x)=>{y=Add(x,x);z=Mul(y,y);return z,y;}" + : "(x)=>{y=Add(x,x);z=Mul(y,y);return y,z;}"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* y = b.MakeOutput(); + auto* z = b.MakeOutput(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Add", {a, a}, {y}); + b.AddNode("Mul", {y, y}, {z}); + b.AddNode("Sub", {y, z}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(test.graph.NumberOfNodes(), 3); + const Node* sum = test.graph.GetProducerNode(y->Name()); + const Node* product = test.graph.GetProducerNode(z->Name()); + ASSERT_NE(sum, nullptr); + ASSERT_NE(product, nullptr); + EXPECT_EQ(sum->OpType(), "Add"); + EXPECT_EQ(product->OpType(), "Mul"); + EXPECT_EQ(sum->GetOutputEdgesCount(), 3); + EXPECT_EQ(product->GetOutputEdgesCount(), 1); + EXPECT_EQ(product->InputDefs()[0], y); + EXPECT_EQ(product->InputDefs()[1], y); + ASSERT_STATUS_OK(test.graph.Resolve()); + TransformerTester( + [](ModelTestBuilder& builder) { + auto* input = builder.MakeInput({2}, -1.0f, 1.0f); + auto* copied = builder.MakeIntermediate(); + auto* sum_output = builder.MakeOutput(); + auto* product_output = builder.MakeOutput(); + builder.AddNode("Identity", {input}, {copied}); + builder.AddNode("Add", {copied, copied}, {sum_output}); + builder.AddNode("Mul", {sum_output, sum_output}, {product_output}); + builder.AddNode("Sub", {sum_output, product_output}, {builder.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Add"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Mul"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } + } +} + +TEST(PatternRewriterTests, ReturnedOutputSelectionAndBoundaries) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + SCOPED_TRACE(scenario); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{a=Identity(x);left,right=Split(a,axis=0);return right;}"); + rule.replacement = Replacement("(x)=>{left,right=Split(x,axis=0);return right;}"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* a = scenario == 3 ? b.MakeOutput() : b.MakeIntermediate(); + auto* left = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + auto* right = b.MakeOutput(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Split", {a}, {left, right}).AddAttribute("axis", int64_t{0}); + if (scenario == 2) b.AddNode("Neg", {left}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + if (scenario == 0) { + const Node* producer = test.graph.GetProducerNode(right->Name()); + ASSERT_NE(producer, nullptr); + EXPECT_EQ(producer->OutputDefs()[1], right); + EXPECT_NE(producer->OutputDefs()[0], left); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{a=Identity(x);left,right=Split(a,axis=0);return right,left;}"); + rule.replacement = Replacement("(x)=>{left,right=Split(x,axis=0);return right,left;}"); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* copied = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {copied}); + b.AddNode("Split", {copied}, {b.MakeOutput(), b.MakeOutput()}).AddAttribute("axis", int64_t{0}); + }, + [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, BoundaryOutputMappingCycleChecks) { + for (bool indexed : {false, true}) { + for (bool valid : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* y = b.MakeOutput(); + auto* external = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Add", {a, a}, {y}); + b.AddNode("Neg", {y}, {external}); + b.AddNode("Mul", {y, external}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x,e)=>{a=Identity(x);y=Add(a,a);z=Mul(y,e);return y,z;}"); + rule.replacement = Replacement(valid + ? "(x,e)=>{y=Add(x,x);z=Mul(y,e);return y,z;}" + : "(e)=>{left,right=Split(e,axis=0);return left,right;}"); + Statistics stats; + auto status = PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats); + EXPECT_EQ(status.IsOK(), valid); + EXPECT_EQ(stats.rewrites, valid ? 1u : 0u); + if (!valid) EXPECT_NE(status.ErrorMessage().find("cycle"), std::string::npos); + EXPECT_EQ(test.graph.NumberOfNodes(), valid ? 3 : 4); + EXPECT_EQ(test.graph.GetProducerNode(y->Name())->OpType(), "Add"); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ReturnedIntermediateImplicitCapture) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* copied = b.MakeIntermediate(); + auto* left = b.MakeIntermediate(); + auto* right = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {copied}); + b.AddNode("Split", {copied}, {left, right}).AddAttribute("axis", int64_t{0}); + b.AddNode("Concat", {left, right}, {b.MakeOutput()}).AddAttribute("axis", int64_t{0}); + ONNX_NAMESPACE::GraphProto branch; + branch.set_name("returned_intermediate"); + auto* identity = branch.add_node(); + identity->set_op_type("Identity"); + identity->add_input(right->Name()); + identity->add_output("branch_out"); + auto* output = branch.add_output(); + output->set_name("branch_out"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output->mutable_type()->mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + auto& if_node = b.AddNode("If", {b.MakeInput({}, std::vector{true})}, {b.MakeOutput()}); + if_node.AddAttribute("then_branch", branch); + if_node.AddAttribute("else_branch", branch); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{a=Identity(x);l,r=Split(a,axis=0);y=Concat(l,r,axis=0);return r,y;}"); + rule.replacement = Replacement("(x)=>{l,r=Split(x,axis=0);y=Concat(l,r,axis=0);return r,y;}"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + const Node* producer = test.graph.GetProducerNode(right->Name()); + ASSERT_NE(producer, nullptr); + EXPECT_EQ(producer->OpType(), "Split"); + EXPECT_EQ(producer->OutputDefs()[1], right); + EXPECT_EQ(Count(test.graph, "Identity"), 0u); + } +} + +TEST(PatternRewriterTests, NamedAttributeCaptures) { + for (auto text : {"(x)=>{t=Transpose(x,perm=p);return Identity(t);}", + "(x)=>{t=Transpose(x,perm=p);return Transpose(t,perm=p);}", + "(x)=>{h=HardSigmoid(x,alpha=a,beta=b);return Mul(h,x);}"}) + EXPECT_NO_THROW((void)Pattern{text}); + for (auto text : {"(x)=>{t=Transpose(x,perm=x);return Identity(t);}", + "(x)=>{t=Transpose(x,perm=p);p=Identity(t);return p;}", + "(x)=>{t=Transpose(x,perm=None);return Identity(t);}"}) + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + for (bool indexed : {false, true}) { + for (bool equal : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 2}, -1.0f, 1.0f); + auto* t = b.MakeIntermediate(); + auto& first = b.AddNode("Transpose", {x}, {t}); + first.AddAttribute("perm", std::vector{1, 0}); + auto& second = b.AddNode("Transpose", {t}, {b.MakeOutput()}); + second.AddAttribute("perm", equal ? std::vector{1, 0} : std::vector{0, 1}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{t=Transpose(x,perm=p);return Transpose(t,perm=p);}"); + int guards = 0; + rule.check_and_bind = [&guards](const Graph&, Match& match) { + ++guards; + EXPECT_EQ(match.attributeOf("p").type(), ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); + EXPECT_EQ(match.attributeOf("p").ints(0), 1); + EXPECT_THROW(match.update("p", match.attributeOf("p")), OnnxRuntimeException); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(guards, equal ? 1 : 0); + EXPECT_EQ(stats.rewrites, equal ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{t=Transpose(x,perm=axes);return Identity(t);}"); + rule.replacement = Replacement("(x,axes)=>Transpose(x,perm=axes)"); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* t = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {t}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Identity", {t}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(session.GetGraph().Nodes().begin()->GetAttributes().at("perm").ints(0), 1); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, OptionalInputSlotsAndNumerics) { + for (auto text : {"(x,low:tensor[...] | None,high:tensor[...] | None)=>{a=Identity(x);return Clip(a,low,high);}", + "(x,high)=>{a=Identity(x);return Clip(a,None,high);}", + "(x)=>{a=Identity(x);return Clip(a,None,None);}"}) { + EXPECT_NO_THROW((void)Pattern{text}); + if (std::string_view(text).find(':') != std::string_view::npos) + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + else + EXPECT_NO_THROW((void)Replacement{text}); + } + for (auto text : {"(None)=>return Identity(None)", "(x??)=>return Identity(x)", + "(x)=>{a=Identity(x);return Clip(a,missing);}", + "(x?)=>return x", "(x)=>{a=Identity(x);return Clip(a,None=0);}"}) + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + SCOPED_TRACE(scenario); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x,low:tensor[...] | None,high:tensor[...] | None)=>{a=Identity(x);return Clip(a,low,high);}"); + rule.replacement = Replacement("(x,low,high)=>Clip(x,low,high)"); + auto build = [scenario](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -2.0f, 2.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + std::vector inputs{x}; + inputs[0] = a; + if (scenario > 0) { + auto* empty = b.MakeOptionalTensor(); + inputs.push_back(scenario == 2 ? empty : b.MakeScalarInitializer(-1.0f)); + if (scenario >= 2) + inputs.push_back(scenario == 4 ? empty : b.MakeScalarInitializer(1.0f)); + } + b.AddNode("Clip", inputs, {b.MakeOutput()}); + }; + TestGraph test; + build(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 0) { + auto invalid = rule; + invalid.replacement = Replacement("(low)=>Identity(low)"); + Statistics rejected; + EXPECT_FALSE(PatternRewriter({invalid}, {.indexed = indexed}).Run(test.graph, rejected).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 2); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + int guards = 0; + rule.check_and_bind = [scenario, &guards](const Graph&, Match& match) { + ++guards; + EXPECT_EQ(match.valueOf("low") == nullptr, scenario == 0 || scenario == 2); + EXPECT_EQ(match.valueOf("high") == nullptr, scenario < 2 || scenario == 4); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(guards, 1); + ASSERT_STATUS_OK(test.graph.Resolve()); + // Numerical check_and_bind must not retain references to a previous graph's match. + rule.check_and_bind = [](const Graph&, Match&) { return true; }; + TransformerTester( + build, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Clip"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } + } +} + +TEST(PatternRewriterTests, CapturedAttributesVaryAcrossMatches) { + for (bool indexed : {false, true}) { + auto rule = IdentityPair(); + rule.pattern = Pattern("(x)=>{h=HardSigmoid(x,alpha=slope,beta=bias);return Identity(h);}"); + rule.replacement = Replacement("(x,slope,bias)=>HardSigmoid(x,alpha=slope,beta=bias)"); + auto build = [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -2.0f, 2.0f); + for (float alpha : {0.25f, 0.5f}) { + auto* h = b.MakeIntermediate(); + auto& sigmoid = b.AddNode("HardSigmoid", {x}, {h}); + sigmoid.AddAttribute("alpha", alpha); + sigmoid.AddAttribute("beta", 0.5f); + b.AddNode("Identity", {h}, {b.MakeOutput()}); + } + }; + TestGraph test; + build(test.builder); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int guards = 0; + rule.check_and_bind = [&guards](const Graph&, Match& match) { + ++guards; + EXPECT_EQ(match.attributeOf("slope").type(), ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + EXPECT_FLOAT_EQ(match.attributeOf("bias").f(), 0.5f); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 2u); + EXPECT_EQ(guards, 2); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedHashSet slopes; + for (const auto& node : test.graph.Nodes()) slopes.insert(node.GetAttributes().at("alpha").f()); + EXPECT_EQ(slopes.size(), 2u); + rule.check_and_bind = [](const Graph&, Match&) { return true; }; + TransformerTester( + build, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "HardSigmoid"), 2u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, OptionalInputConstraintsAndAbsentRepetition) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + auto* empty = &test.graph.GetOrCreateNodeArg("", nullptr); + auto* low = scenario == 0 ? empty : b.MakeScalarInitializer(0.0f); + auto* high = scenario < 2 ? empty : (scenario == 2 ? low : b.MakeScalarInitializer(1.0f)); + b.AddNode("Clip", {a, low, high}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + for (int constraint = 0; constraint < 3; ++constraint) { + auto rule = IdentityPair(); + rule.pattern = Pattern(constraint == 0 ? "(x,v:tensor[...] | None)=>{a=Identity(x);return Clip(a,v,v);}" + : constraint == 1 ? "(x)=>{a=Identity(x);return Clip(a,None,None);}" + : "(x,lo,hi)=>{a=Identity(x);return Clip(a,lo,hi);}"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + const bool expected = constraint == 0 ? scenario == 0 || scenario == 2 + : constraint == 1 ? scenario == 0 + : scenario >= 2; + EXPECT_EQ(stats.matches, expected ? 1u : 0u); + } + } + } +} + +TEST(PatternRewriterTests, StructuredFailureDiagnostics) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 6; ++scenario) { + SCOPED_TRACE(scenario); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + auto& first = b.AddNode("Identity", {x}, {a}); + auto& root = b.AddNode("Identity", {a}, {b.MakeOutput()}); + const NodeIndex root_index = root.Index(); + if (scenario == 2) first.SetExecutionProviderType(kCudaExecutionProvider); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair("diagnostic"); + if (scenario == 0) rule.check_and_bind = [](const Graph&, Match& match) { return match.Reject("unsupported dtype"); }; + if (scenario == 3) rule.opsets = {{kOnnxDomain, 18, 21}}; + if (scenario == 4) rule.pattern = Pattern("(x)=>{a=Identity(x,absent=captured);return Identity(a);}"); + if (scenario == 5) { + rule.replacement = [](const Graph&, const Match&) { return Replacement("(x)=>Identity(x,invalid=1)"); }; + } + InlinedVector failures; + Options options{.indexed = indexed}; + options.on_failure = [&failures](const MatchFailure& failure) { failures.push_back(failure); }; + Statistics stats; + auto status = PatternRewriter({rule}, options).Run(test.graph, stats); + EXPECT_EQ(status.IsOK(), scenario != 5); + EXPECT_EQ(stats.rewrites, 0u); + const auto expected_stage = scenario == 0 ? FailureStage::CheckAndBind + : scenario == 1 ? FailureStage::Boundary + : scenario == 2 ? FailureStage::ExecutionProvider + : scenario == 3 ? FailureStage::Opset + : scenario == 4 ? FailureStage::Attribute + : FailureStage::Replacement; + auto found = std::find_if(failures.begin(), failures.end(), + [expected_stage](const auto& failure) { return failure.stage == expected_stage; }); + ASSERT_NE(found, failures.end()); + EXPECT_EQ(found->rule, "diagnostic"); + EXPECT_EQ(found->root, scenario == 3 ? std::numeric_limits::max() : root_index); + EXPECT_FALSE(found->reason.empty()); + if (scenario == 0) EXPECT_EQ(found->reason, "unsupported dtype"); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(test.graph.NumberOfNodes(), 2); + } + } +} + +TEST(PatternRewriterTests, TypeShapeAnnotationParsing) { + for (auto text : { + "(x:float[M,N])=>return Identity(x)", + "(x:float[])=>return Identity(x)", + "(x:float[...])=>return Identity(x)", + "(x:float)=>return Identity(x)", + "(x:tensor[M,0])=>return Identity(x)", + "(x:tensor[M,...],y:tensor[M,_])=>return CastLike(x,y)", + "(x:int64[...] | None)=>return Identity(x)", + "(x:float[M,N])=>{y:float[M,N]=Identity(x);return y;}", + "(x:float[M,N])=>{a:float[M,N],b:float[M,N]=Split(x);return a,b;}", + "(x:float[],axis:int)=>return Softmax(x,axis=axis)", + "(x:float[],p:ints)=>return Transpose(x,perm=p)", + "(x,a:float,b:float)=>return HardSigmoid(x,alpha=a,beta=b)", + "(x,s:string)=>return Identity(x,mode=s)", + "(x,t:tensor)=>return Identity(x,value=t)", + "(x,f:floats,s:strings)=>return Identity(x,f=f,s=s)"}) { + SCOPED_TRACE(text); + EXPECT_NO_THROW((void)Pattern{text}); + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + } + for (auto text : { + "(x:bogus)=>return Identity(x)", "(x:undefined)=>return Identity(x)", + "(x:float[-1])=>return Identity(x)", "(x:float[1.0])=>return Identity(x)", + "(x:float[M,])=>return Identity(x)", "(x:float[...,M])=>return Identity(x)", + "(x:float[M,...,N])=>return Identity(x)", "(x:float[M)=>return Identity(x)", + "(x:float[M])=>return Identity(x)", "(x:tensor<>[M])=>return Identity(x)", + "(x:tensorreturn Identity(x)", "(x:float[M,N])=>return x", + "(M,x:float[M])=>return Identity(x)", "(x:float[x])=>return Identity(x)", + "(x:tensor[M])=>return Identity(x)", "(x:tensor[T])=>return Identity(x)", + "(x:float[M],M:int)=>return Softmax(x,axis=M)", + "(x:float[...],p:float[M])=>return Transpose(x,perm=p)", + "(x:int)=>return Identity(x)", "(x,axis?:int)=>return Softmax(x,axis=axis)", + "(x,axis:int)=>return Add(x,axis)", "(x,axis:int)=>return Identity(x)", + "(x)=>{y:float=Identity(x);return y:float;}"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + } + Match empty; + EXPECT_THROW(empty.dimensionOf("M"), OnnxRuntimeException); + EXPECT_THROW(empty.elementTypeOf("T"), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, TypeShapeAnnotationRankAndDtype) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + SCOPED_TRACE(scenario); + TestGraph test; + auto& b = test.builder; + auto* x = scenario == 0 ? b.MakeInput({}, -1.0f, 1.0f) + : scenario == 1 ? b.MakeInput({2}, -1.0f, 1.0f) + : scenario == 2 ? b.MakeInput({2, 3}, -1.0f, 1.0f) + : scenario == 3 ? b.MakeInput({2, 3}, -1, 1) + : b.MakeInput({2, 3, 4}, -1.0f, 1.0f); + auto* output = b.MakeOutput(); + b.AddNode("Identity", {x}, {output}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const std::pair cases[] = { + {"float", scenario == 0}, + {"float[]", scenario != 3}, + {"float[...]", scenario != 3}, + {"float[M,N]", scenario == 2}, + {"float[2,3]", scenario == 2}, + {"float[M,...]", scenario != 0 && scenario != 3}, + {"tensor[M,N]", scenario == 2 || scenario == 3}, + {"tensor[...]", true}, + {"int64[2,3]", scenario == 3}, + {"float[_,_]", scenario == 2}}; + for (const auto& [type, expected] : cases) { + SCOPED_TRACE(type); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:" + std::string(type) + ")=>return Identity(x)"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, expected ? 1u : 0u); + } + x->ClearShape(); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:float[])=>return Identity(x)"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 3 ? 0u : 1u); + rule.pattern = Pattern("(x:float[M,...])=>return Identity(x)"); + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 0u); + } + } +} + +TEST(PatternRewriterTests, SharedDimensionAndElementBindings) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 7; ++scenario) { + SCOPED_TRACE(scenario); + TestGraph test; + auto& b = test.builder; + using Shape = std::vector>; + Shape sx{int64_t{2}, int64_t{3}}, sy = sx; + if (scenario == 1) sy[0] = int64_t{4}; + if (scenario >= 2 && scenario <= 4) { + sx[0] = std::string("batch"); + sy[0] = scenario == 3 ? std::string("other") : std::string("batch"); + if (scenario == 4) sy[0] = int64_t{2}; + } + if (scenario == 5) sx[0] = sy[0] = std::string{}; + auto* x = b.MakeSymbolicInput(sx); + auto* y = scenario == 6 ? b.MakeSymbolicInput(sy) : b.MakeSymbolicInput(sy); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("CastLike", {a, y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:tensor[M,N],y:tensor[M,N])=>{a=Identity(x);return CastLike(a,y);}"); + int guards = 0; + rule.check_and_bind = [scenario, &guards](const Graph&, Match& match) { + ++guards; + EXPECT_EQ(match.elementTypeOf("T"), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + EXPECT_EQ(std::get(match.dimensionOf("N")), 3); + if (scenario == 0) EXPECT_EQ(std::get(match.dimensionOf("M")), 2); + if (scenario == 2) EXPECT_EQ(std::get(match.dimensionOf("M")), "batch"); + EXPECT_THROW(match.dimensionOf("missing"), OnnxRuntimeException); + EXPECT_THROW(match.elementTypeOf("missing"), OnnxRuntimeException); + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attribute.set_i(1); + EXPECT_THROW(match.update("M", attribute), OnnxRuntimeException); + EXPECT_THROW(match.update("T", attribute), OnnxRuntimeException); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + const bool expected = scenario == 0 || scenario == 2; + EXPECT_EQ(stats.matches, expected ? 1u : 0u); + EXPECT_EQ(guards, expected ? 1 : 0); + } + } +} + +TEST(PatternRewriterTests, TypedAttributeCapturesAndNumerics) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + const bool accepted = scenario == 0; + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* t = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {t}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Identity", {t}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern(accepted + ? "(x:float[M,N],axes:ints)=>{t:float[N,M]=Transpose(x,perm=axes);return Identity(t);}" + : scenario == 1 + ? "(x:float[M,N],axes:floats)=>{t=Transpose(x,perm=axes);return Identity(t);}" + : "(x:float[M,N],axes:ints)=>{t:float[M,N]=Transpose(x,perm=axes);return Identity(t);}"); + rule.replacement = Replacement("(x,axes)=>Transpose(x,perm=axes)"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, accepted ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:float[M,N],axes:ints)=>{t:float[N,M]=Transpose(x,perm=axes);return Identity(t);}"); + rule.replacement = Replacement("(x,axes)=>Transpose(x,perm=axes)"); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* t = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {t}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Identity", {t}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Transpose"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{rule}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* clipped = b.MakeIntermediate(); + b.AddNode("Clip", {x}, {clipped}); + b.AddNode("Identity", {clipped}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto first = IdentityPair("reject"), second = IdentityPair("accept"); + first.pattern = Pattern("(x:float[2,4],lo:float | None,hi:float | None)=>{a=Clip(x,lo,hi);return Identity(a);}"); + second.pattern = Pattern( + "(input:tensor[M,N],low:tensor | None,high:float | None)=>" + "{a:tensor[M,N]=Clip(input,low,high);return Identity(a);}"); + second.replacement = Replacement("(input)=>Clip(input)"); + second.check_and_bind = [](const Graph&, Match& match) { + EXPECT_EQ(match.valueOf("low"), nullptr); + EXPECT_FALSE(match.elementTypeOf("U").has_value()); + EXPECT_EQ(std::get(match.dimensionOf("M")), 2); + return true; + }; + if (scenario == 1) first.pattern = Pattern( + "(x:float[2,3],lo:float | None,hi:float | None)=>" + "{a=Clip(x,lo,hi);return Identity(a);}"); + if (scenario == 1) first.check_and_bind = [](const Graph&, Match&) { return false; }; + if (scenario == 2) second.pattern = Pattern( + "(input:tensor[M,M],low:tensor | None,high:float | None)=>" + "{a=Clip(input,low,high);return Identity(a);}"); + InlinedVector failures; + Options options{.indexed = indexed}; + options.on_failure = [&failures](const MatchFailure& failure) { failures.push_back(failure); }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({first, second}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 2 ? 0u : 1u); + EXPECT_TRUE(std::any_of(failures.begin(), failures.end(), + [](const auto& failure) { return failure.stage == FailureStage::TypeShape; }) || + scenario == 1); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, TypeUnionParsingAndBounds) { + for (auto text : {"(x:float[M] | float[M,N])=>return Identity(x)", + "(x:float[M] | None)=>return Identity(x)", + "(x:None | tensor[...])=>return Identity(x)", + "(x)=>{y:float[M] | float[M,N]=Identity(x);return y;}"}) { + EXPECT_NO_THROW((void)Pattern{text}); + EXPECT_THROW((void)Replacement{text}, OnnxRuntimeException); + } + for (auto text : {"(x:|float[M])=>return Identity(x)", "(x:float[M]|)=>return Identity(x)", + "(x:float[M]||None)=>return Identity(x)", "(x:None)=>return Identity(x)", + "(x:None|None)=>return Identity(x)", "(x:float[M]|float[M])=>return Identity(x)", + "(x:float[]|float[...])=>return Identity(x)", + "(x?)=>return Identity(x)", "(x?:float[])=>return Identity(x)", + "(x:float[]?)=>return Identity(x)", + "(x,p:ints|floats)=>return Identity(x,perm=p)", + "(x)=>{y:float[]|None=Identity(x);return y;}"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); + } + auto make_text = [](int count) { + std::string text = "("; + for (int i = 0; i < count; ++i) { + if (i) text += ','; + text += "x" + std::to_string(i) + ":float[]|int64[]"; + } + text += ")=>return Concat("; + for (int i = 0; i < count; ++i) { + if (i) text += ','; + text += "x" + std::to_string(i); + } + return text + ",axis=0)"; + }; + EXPECT_NO_THROW((void)Pattern{make_text(6)}); + EXPECT_THROW((void)Pattern{make_text(7)}, OnnxRuntimeException); +} + +TEST(PatternRewriterTests, TypeUnionRanksAndNone) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = scenario == 0 ? b.MakeInput({2}, -1.0f, 1.0f) + : scenario == 1 ? b.MakeInput({2, 3}, -1.0f, 1.0f) + : scenario == 2 ? b.MakeInput({}, -1.0f, 1.0f) + : scenario == 3 ? b.MakeInput({2, 3, 4}, -1.0f, 1.0f) + : b.MakeInput({2}, -1, 1); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:float[M] | float[M,N])=>return Identity(x)"); + rule.check_and_bind = [scenario](const Graph&, Match& match) { + EXPECT_EQ(std::get(match.dimensionOf("M")), 2); + if (scenario == 0) + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("N"))); + else + EXPECT_EQ(std::get(match.dimensionOf("N")), 3); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario < 2 ? 1u : 0u); + } + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + std::vector inputs{x}; + if (scenario == 1) inputs.push_back(b.MakeOptionalTensor()); + if (scenario == 2) inputs.push_back(b.MakeScalarInitializer(0.0f)); + if (scenario == 3) inputs.push_back(b.MakeInput({2}, -1.0f, 1.0f)); + b.AddNode("Clip", inputs, {b.MakeOutput()}); + b.SetGraphOutputs(); + // A vector bound is schema-valid but cannot satisfy the scalar annotation. + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:float[M],low:float | None)=>return Clip(x,low)"); + rule.check_and_bind = [scenario](const Graph&, Match& match) { + EXPECT_EQ(match.valueOf("low") == nullptr, scenario < 2); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario < 3 ? 1u : 0u); + } + } +} + +TEST(PatternRewriterTests, TypeUnionGlobalBindingsAndRollback) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = b.MakeInput({scenario == 2 ? 4 : 3}, -1.0f, 1.0f); + b.AddNode("CastLike", {x, y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern(scenario == 0 + ? "(x:float[M,9] | float[2,M],y:float[M])=>return CastLike(x,y)" + : scenario == 1 || scenario == 2 + ? "(x:float[M,_] | float[_,M],y:float[M])=>return CastLike(x,y)" + : "(x:tensor[M,_] | tensor[_,M],y:tensor[M])=>return CastLike(x,y)"); + rule.check_and_bind = [scenario](const Graph&, Match& match) { + EXPECT_EQ(std::get(match.dimensionOf("M")), 3); + if (scenario == 3) { + EXPECT_FALSE(match.elementTypeOf("T").has_value()); + EXPECT_EQ(match.elementTypeOf("U"), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + } + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 2 ? 0u : 1u); + } + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + b.AddNode("Clip", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x,lo:tensor[D] | None)=>return Clip(x,lo)"); + rule.check_and_bind = [](const Graph&, Match& match) { + EXPECT_EQ(match.valueOf("lo"), nullptr); + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("D"))); + EXPECT_FALSE(match.elementTypeOf("T").has_value()); + return true; + }; + rule.replacement = Replacement("(x,lo)=>Clip(x,lo)"); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed, .max_iterations = 1}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, ParsedMultiOutputDefinitionsAndLimits) { + Pattern p("(x)=>{a=Identity(x);b=ai.onnx::Identity(a);left,right=ai.onnx::Split(b);return left,right;}"); + auto x = p.NamedValue("x"); + const auto& outputs = p.Nodes().back().outputs; + EXPECT_EQ(outputs.size(), 2u); + EXPECT_NE(outputs[0], outputs[1]); + EXPECT_TRUE(p.IsInput(x)); + EXPECT_FALSE(p.IsInput(outputs[0])); + EXPECT_FALSE(p.IsInput(outputs[1])); + EXPECT_EQ(p.Nodes().back().op.domain, kOnnxDomain); + const auto count = p.ValueCount(); + EXPECT_THROW(p = Pattern("(x)=>return = Split(x)"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>a,b=Split(missing);return a,b"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>left,left=Split(x);return left,left"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>x,right=Split(x);return x,right"), OnnxRuntimeException); + EXPECT_NO_THROW(Pattern("(x)=>a,b=Split(x);return a,b")); + EXPECT_EQ(p.ValueCount(), count); + auto outputs_text = [](int node, int count) { + std::string text; + for (int slot = 0; slot < count; ++slot) { + if (slot) text += ','; + text += "v" + std::to_string(node) + "_" + std::to_string(slot); + } + return text; + }; + EXPECT_THROW(Pattern("(x)=>{" + outputs_text(0, 65) + "=Split(x);return " + outputs_text(0, 65) + ";}"), + OnnxRuntimeException); + std::string many_values = "(x)=>"; + for (int node = 0; node < 3; ++node) { + const auto input = node == 0 ? "x" : "v" + std::to_string(node - 1) + "_0"; + many_values += outputs_text(node, 64) + "=Split(" + input + ");"; + } + EXPECT_NO_THROW(Pattern(many_values + "return Identity(v2_0)")); + EXPECT_THROW(Pattern(many_values + outputs_text(3, 64) + "=Split(v2_0);return " + outputs_text(3, 64)), + OnnxRuntimeException); +} + +TEST(PatternRewriterTests, MultiOutputIntermediateBoundariesAndBindings) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 5; ++scenario) { + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + auto* b = scenario == 1 ? test.builder.MakeOutput() : test.builder.MakeIntermediate(); + test.builder.AddNode("Split", {x}, {a, b}); + auto& concat = test.builder.AddNode("Concat", {a, scenario == 3 ? a : b}, {test.builder.MakeOutput()}); + concat.AddAttribute("axis", int64_t{0}); + if (scenario == 2) test.builder.AddNode("Identity", {b}, {test.builder.MakeOutput()}); + if (scenario == 4) { + // A repeated output binding must refer to the same node, not a second Split. + auto* c = test.builder.MakeIntermediate(); + auto* d = test.builder.MakeIntermediate(); + test.builder.AddNode("Split", {x}, {c, d}); + concat.MutableInputDefs()[1] = d; + } + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({SplitConcat()}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 0) EXPECT_EQ(test.graph.NumberOfNodes(), 1); + } + } +} + +TEST(PatternRewriterTests, MultiOutputProducerReachedThroughSecondOutput) { + for (bool indexed : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto* a = test.builder.MakeIntermediate(); + auto* b = test.builder.MakeIntermediate(); + test.builder.AddNode("Split", {x}, {a, b}); + test.builder.AddNode("Identity", {b}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Pattern pattern("(x)=>a,b=Split(x);return Identity(b)"); + auto rule = SplitConcat(); + rule.pattern = pattern; + rule.check_and_bind = [](const Graph&, Match&) { return true; }; + Statistics stats; + // Only match: an Identity replacement would not preserve slicing semantics. + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, 0u); + } +} + +TEST(PatternRewriterTests, MultiOutputRootEdgesAndInvalidRecipe) { + for (bool indexed : {false, true}) { + TestGraph test; + auto* x = test.builder.MakeInput({4}, -1.0f, 1.0f); + auto* middle = test.builder.MakeIntermediate(); + auto* a = test.builder.MakeOutput(); + auto* b = test.builder.MakeOutput(); + test.builder.AddNode("Identity", {x}, {middle}); + test.builder.AddNode("Split", {middle}, {a, b}); + test.builder.AddNode("Neg", {a}, {test.builder.MakeOutput()}); + test.builder.AddNode("Add", {b, b}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + for (auto text : {"(x)=>Identity(x)", "(x)=>Split(x)", + "(x)=>{left,middle,right=Split(x);return left,middle,right;}"}) { + auto invalid = IdentitySplit(); + invalid.replacement = Replacement(text); + EXPECT_FALSE(PatternRewriter({invalid}, {.indexed = indexed}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(test.graph.NumberOfNodes(), 4); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + auto wrong_count = IdentityPair("single_output_split"); + wrong_count.pattern = Pattern("(x)=>a=Identity(x);return Split(a)"); + auto valid = IdentitySplit(); + valid.check_and_bind = [a, b, prepare = valid.check_and_bind](const Graph& graph, Match& match) { + EXPECT_EQ(match.valueOf("a"), a); + EXPECT_EQ(match.valueOf("b"), b); + return prepare(graph, match); + }; + ASSERT_STATUS_OK(PatternRewriter({wrong_count, valid}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + const Node* producer = test.graph.GetProducerNode(a->Name()); + ASSERT_NE(producer, nullptr); + EXPECT_EQ(producer, test.graph.GetProducerNode(b->Name())); + EXPECT_EQ(producer->OutputDefs()[0], a); + EXPECT_EQ(producer->OutputDefs()[1], b); + EXPECT_EQ(producer->GetOutputEdgesCount(), 3); + size_t slot0 = 0, slot1 = 0; + for (auto edge = producer->OutputEdgesBegin(); edge != producer->OutputEdgesEnd(); ++edge) { + if (edge->GetSrcArgIndex() == 0) ++slot0; + if (edge->GetSrcArgIndex() == 1) ++slot1; + } + EXPECT_EQ(slot0, 1u); + EXPECT_EQ(slot1, 2u); + ASSERT_STATUS_OK(test.graph.Resolve()); + ASSERT_STATUS_OK(PatternRewriter({IdentitySplit()}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); + } +} + +TEST(PatternRewriterTests, MultiOutputImplicitCaptureBoundary) { + for (bool indexed : {false, true}) { + for (bool root : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* split_input = x; + if (root) { + split_input = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {split_input}); + } + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Split", {split_input}, {a, c}); + b.AddNode("Concat", {a, c}, {b.MakeOutput()}).AddAttribute("axis", int64_t{0}); + ONNX_NAMESPACE::GraphProto branch; + branch.set_name("captured_split_output"); + auto* identity = branch.add_node(); + identity->set_op_type("Identity"); + identity->add_input(c->Name()); + identity->add_output("branch_out"); + auto* output = branch.add_output(); + output->set_name("branch_out"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output->mutable_type()->mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + auto& if_node = b.AddNode("If", {b.MakeInput({}, std::vector{true})}, {b.MakeOutput()}); + if_node.AddAttribute("then_branch", branch); + if_node.AddAttribute("else_branch", branch); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({IdentitySplit(), SplitConcat()}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, root ? 1u : 0u); + const Node* producer = test.graph.GetProducerNode(c->Name()); + ASSERT_NE(producer, nullptr); + EXPECT_EQ(producer->OutputDefs()[1], c); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, NumericalMultiOutputPatterns) { + for (bool indexed : {false, true}) { + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Split", {x}, {a, c}); + b.AddNode("Concat", {a, c}, {b.MakeOutput()}).AddAttribute("axis", int64_t{0}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Split"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Concat"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{SplitConcat()}, Options{.indexed = indexed})); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* middle = b.MakeIntermediate(); + auto* a = b.MakeOutput(); + auto* c = b.MakeOutput(); + b.AddNode("Identity", {x}, {middle}); + b.AddNode("Split", {middle}, {a, c}); + b.AddNode("Neg", {a}, {b.MakeOutput()}); + b.AddNode("Add", {c, c}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 0u); + EXPECT_EQ(Count(session.GetGraph(), "Split"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{IdentitySplit()}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, OnnxscriptMultiOutputLayerNormBias) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + auto* mean = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + auto* inverse_std = b.MakeIntermediate(); + b.AddNode("LayerNormalization", {x, b.MakeInitializer({4}, 0.8f, 1.2f)}, {y, mean, inverse_std}); + b.AddNode("Add", {y, b.MakeInitializer({4}, -0.1f, 0.1f)}, {b.MakeOutput()}); + if (scenario == 2) b.AddNode("Identity", {inverse_std}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({LayerNormBias()}, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("LayerNormalization", {x, b.MakeInitializer({4}, 0.8f, 1.2f)}, + {y, b.MakeIntermediate(), b.MakeIntermediate()}); + b.AddNode("Add", {y, b.MakeInitializer({4}, -0.1f, 0.1f)}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "LayerNormalization"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Add"), 0u); + for (const auto& node : session.GetGraph().Nodes()) { + if (node.OpType() == "LayerNormalization") { + EXPECT_EQ(node.InputDefs().size(), 3u); + EXPECT_EQ(node.OutputDefs().size(), 1u); + } + } + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(Rules{LayerNormBias()}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, NumericalMultiOutputLeafPreservation) { + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* dropout = b.MakeIntermediate(); + auto* first = b.MakeIntermediate(); + b.AddNode("Dropout", {x}, {dropout, b.MakeOutput()}); + b.AddNode("Identity", {dropout}, {first}); + b.AddNode("Identity", {first}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "Dropout"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Identity"), 1u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 0.0, 0.0, + std::make_unique(Rules{IdentityPair()})); +} + +TEST(PatternRewriterTests, DoubleTypeRejection) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0, 1.0); + auto* h = b.MakeIntermediate(); + auto& sigmoid = b.AddNode("HardSigmoid", {x}, {h}); + sigmoid.AddAttribute("alpha", 1.0f / 6.0f); + sigmoid.AddAttribute("beta", 0.5f); + b.AddNode("Mul", {h, x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules()).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); +} + +TEST(PatternRewriterTests, ExampleGuardsCheckRelevantBindings) { + auto examples = MakeExampleRules(); + for (size_t index = 0; index < examples.size(); ++index) { + SCOPED_TRACE(examples[index].name); + ASSERT_NE(examples[index].replacement.target(), nullptr); + TestGraph test; + auto& b = test.builder; + auto* non_float = b.MakeScalarInitializer(2.0); + auto* integer = b.MakeScalarInitializer(2); + // Resolve removes unused untyped graph NodeArgs. + NodeArg unknown{"guard_unknown", nullptr}; + b.AddNode("Identity", {non_float}, {b.MakeOutput()}); + b.AddNode("Identity", {integer}, {b.MakeOutput()}); + InlinedVector bindings{"x"}; + if (index == 0) { + BuildHardSwish(b); + } else if (index == 1) { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* sigmoid = b.MakeIntermediate(); + auto& node = b.AddNode("HardSigmoid", {x}, {sigmoid}); + node.AddAttribute("alpha", 1.0f / 6.0f); + node.AddAttribute("beta", 0.5f); + b.AddNode("Mul", {sigmoid, x}, {b.MakeOutput()}); + } else if (index == 5) { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + BuildRms(b, x, b.MakeOutput(), 4); + bindings.push_back("scale"); + bindings.push_back("epsilon"); + } else { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* y = b.MakeInput(index == 3 ? std::vector{3, 4} + : std::vector{4, 3}, + -1.0f, 1.0f); + bindings.push_back("y"); + if (index == 3) { + auto* transpose = b.MakeIntermediate(); + b.AddNode("Transpose", {y}, {transpose}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("MatMul", {x, transpose}, {b.MakeOutput()}); + } else { + auto* mm = b.MakeIntermediate(); + b.AddNode(index == 4 ? "FusedMatMul" : "MatMul", {x, y}, {mm}, + index == 4 ? kMSDomain : kOnnxDomain); + b.AddNode("Div", {mm, b.MakeScalarInitializer(2.0f)}, {b.MakeOutput()}); + bindings.push_back("divisor"); + } + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = examples[index]; + auto check_and_bind = rule.check_and_bind; + size_t checked = 0; + rule.check_and_bind = [&](const Graph& graph, Match& match) { + EXPECT_TRUE(check_and_bind(graph, match)); + for (auto binding : bindings) { + SCOPED_TRACE(binding); + // Exercise the check_and_bind independently without making the resolved graph ill-typed. + for (const NodeArg* value : {non_float, integer, &unknown}) { + Match changed = match; + changed.values[match.pattern->NamedValue(binding).handle] = value; + EXPECT_FALSE(check_and_bind(graph, changed)); + } + } + ++checked; + return check_and_bind(graph, match); + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({std::move(rule)}).Run(test.graph, stats)); + EXPECT_EQ(checked, 1u); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, BranchedSharedNodeBindings) { + for (bool external : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeInput({2}, -1.0f, 1.0f); + auto* sum = b.MakeIntermediate(); + auto* product = b.MakeIntermediate(); + b.AddNode("Add", {x, y}, {sum}); + b.AddNode("Mul", {sum, x}, {product}); + b.AddNode("Add", {sum, product}, {b.MakeOutput()}); + if (external) b.AddNode("Identity", {sum}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const auto* pattern = "(x,y)=>sum=Add(x,y);product=Mul(sum,x);return Add(sum,product)"; + auto check_and_bind = [](const Graph&, Match&) { return true; }; + const auto* replacement = "(x)=>Identity(x)"; + Rule rule{"diamond", pattern, replacement, check_and_bind, {{kOnnxDomain, 7, 21}}}; + Statistics stats; + // Structural-only test: this deliberately has no semantically equivalent replacement. + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, external ? 0u : 1u); + EXPECT_EQ(stats.rewrites, 0u); + } +} + +TEST(PatternRewriterTests, SubgraphsAndImplicitCaptureBoundary) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* captured = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {captured}); + b.AddNode("Identity", {captured}, {b.MakeOutput()}); + ONNX_NAMESPACE::GraphProto branch; + branch.set_name("branch"); + auto* first = branch.add_node(); + first->set_op_type("Identity"); + first->add_input(captured->Name()); + first->add_output("branch_mid"); + auto* second = branch.add_node(); + second->set_op_type("Identity"); + second->add_input("branch_mid"); + second->add_output("branch_out"); + auto* output = branch.add_output(); + output->set_name("branch_out"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output->mutable_type()->mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + auto* condition = b.MakeInput({}, std::vector{true}); + auto& if_node = b.AddNode("If", {condition}, {b.MakeOutput()}); + if_node.AddAttribute("then_branch", branch); + if_node.AddAttribute("else_branch", branch); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + bool modified = false; + ASSERT_STATUS_OK( + PatternRewriter({IdentityPair()}).Apply(test.graph, modified, DefaultLoggingManager().DefaultLogger())); + EXPECT_TRUE(modified); + // The captured intermediate keeps both outer Identity nodes alive. + EXPECT_EQ(Count(test.graph, "Identity"), 2u); + for (const auto& entry : if_node.GetAttributeNameToSubgraphMap()) { + EXPECT_EQ(Count(*entry.second, "Identity"), 1u); + } +} + +#ifndef DISABLE_CONTRIB_OPS +TEST(PatternRewriterTests, RmsGuardAllowsIntegerExponent) { + for (bool indexed : {false, true}) { + Rule rule = MakeExampleRules()[5]; + rule.pattern = Pattern(R"( + (x, scale, epsilon) => + square = Pow(x, 2) + mean = ReduceMean(square) + add = Add(mean, epsilon) + sqrt = Sqrt(add) + reciprocal = Reciprocal(sqrt) + normalize = Mul(x, reciprocal) + return Mul(normalize, scale) + )"); + TransformerTester( + [](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 4}, -1.0f, 1.0f); + BuildRms(b, x, b.MakeOutput(), 4); + auto* exponent = b.MakeScalarInitializer(2); + for (auto& node : b.graph_.Nodes()) { + if (node.OpType() == "Pow") node.MutableInputDefs()[1] = exponent; + } + }, + [](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "SimplifiedLayerNormalization"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Pow"), 0u); + }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(Rules{std::move(rule)}, Options{.indexed = indexed})); + } +} + +TEST(PatternRewriterTests, MatMulConstantGuardsAndBroadcastRank) { + for (int scenario = 0; scenario < 6; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeInput({2}, -1.0f, 1.0f); + auto* mm = b.MakeIntermediate(); + b.AddNode("MatMul", {x, y}, {mm}); + float value = 2.0f; + if (scenario == 1) value = 0.0f; + if (scenario == 2) value = std::numeric_limits::infinity(); + NodeArg* divisor; + if (scenario == 3) { + divisor = b.MakeInitializer({1}, {value}); + } else if (scenario == 4) { + divisor = b.MakeInput({}, std::vector{value}); + } else { + divisor = b.MakeScalarInitializer(value); + } + b.AddNode("Div", {mm, divisor}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 5) { + auto inputs = test.graph.GetInputsIncludingInitializers(); + inputs.push_back(divisor); + test.graph.SetInputs(inputs); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules()).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, NumericalMatMulDivAndTranspose) { + for (int scenario = 0; scenario < 3; ++scenario) { + bool transpose = scenario != 0; + bool divide = scenario != 1; + TransformerTester( + [transpose, divide](ModelTestBuilder& b) { + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* y = + b.MakeInput(transpose ? std::vector{4, 3} : std::vector{3, 4}, -1.0f, 1.0f); + auto* mm = divide ? b.MakeIntermediate() : b.MakeOutput(); + if (transpose) { + auto* t = b.MakeIntermediate(); + b.AddNode("Transpose", {y}, {t}).AddAttribute("perm", std::vector{1, 0}); + y = t; + } + b.AddNode("MatMul", {x, y}, {mm}); + if (divide) b.AddNode("Div", {mm, b.MakeScalarInitializer(2.0f)}, {b.MakeOutput()}); + }, + [](InferenceSessionWrapper& s) { EXPECT_EQ(Count(s.GetGraph(), "FusedMatMul"), 1u); }, + TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(MakeExampleRules())); + } +} + +TEST(PatternRewriterTests, NumericalFusedMatMulDivPreservesAttributes) { + for (bool indexed : {false, true}) { + for (int scenario = -1; scenario < 16; ++scenario) { + SCOPED_TRACE(scenario); + const int flags = std::max(scenario, 0); + auto build = [scenario, flags](ModelTestBuilder& b) { + auto input_shape = [flags](std::vector shape, int transpose, int transpose_batch) { + if (flags & transpose) std::swap(shape[1], shape[2]); + if (flags & transpose_batch) std::swap(shape[0], shape[1]); + return shape; + }; + auto* x = b.MakeInput(input_shape({2, 3, 4}, 1, 4), -1.0f, 1.0f); + auto* y = b.MakeInput(input_shape({2, 4, 5}, 2, 8), -1.0f, 1.0f); + auto* mm = b.MakeIntermediate(); + auto& node = b.AddNode("FusedMatMul", {x, y}, {mm}, kMSDomain); + if (scenario >= 0) { + node.AddAttribute("alpha", -1.5f); + node.AddAttribute("transA", int64_t{(flags & 1) != 0}); + node.AddAttribute("transB", int64_t{(flags & 2) != 0}); + node.AddAttribute("transBatchA", int64_t{(flags & 4) != 0}); + node.AddAttribute("transBatchB", int64_t{(flags & 8) != 0}); + } + b.AddNode("Div", {mm, b.MakeScalarInitializer(2.0f)}, {b.MakeOutput()}); + }; + auto check = [scenario, flags](InferenceSessionWrapper& session) { + EXPECT_EQ(Count(session.GetGraph(), "FusedMatMul"), 1u); + EXPECT_EQ(Count(session.GetGraph(), "Div"), 0u); + for (const auto& node : session.GetGraph().Nodes()) { + if (node.OpType() != "FusedMatMul") continue; + const auto& attributes = node.GetAttributes(); + EXPECT_FLOAT_EQ(attributes.at("alpha").f(), scenario < 0 ? 0.5f : -0.75f); + EXPECT_EQ(attributes.at("transA").i(), (flags & 1) != 0); + EXPECT_EQ(attributes.at("transB").i(), (flags & 2) != 0); + EXPECT_EQ(attributes.at("transBatchA").i(), (flags & 4) != 0); + EXPECT_EQ(attributes.at("transBatchB").i(), (flags & 8) != 0); + } + }; + Rule rule = MakeExampleRules()[4]; + rule.pattern = Pattern("(x,y,divisor)=>{mm=com.microsoft.FusedMatMul(x,y);return Div(mm,divisor);}"); + TransformerTester(build, check, TransformerLevel::Default, TransformerLevel::Level1, 17, 1e-6, 1e-5, + std::make_unique(Rules{std::move(rule)}, Options{.indexed = indexed})); + } + } +} + +TEST(PatternRewriterTests, RmsConstantAndAttributeGuards) { + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto* x = test.builder.MakeInput({2, 4}, -1.0f, 1.0f); + BuildRms(test.builder, x, test.builder.MakeOutput(), 4, scenario == 1 ? -1e-5f : 1e-5f, scenario == 2); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 3) { + for (auto& node : test.graph.Nodes()) { + if (node.OpType() == "Pow") node.MutableInputDefs()[1] = x; + } + ASSERT_STATUS_OK(test.graph.Resolve()); + } + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules()).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, NumericalDecoderPrefillAndDecode) { + for (int64_t sequence : {1, 4}) { + TestGraph test; + BuildDecoder(test.builder, 2, sequence); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + std::string original; + test.model.ToProto().SerializeToString(&original); + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(MakeExampleRules()).Run(test.graph, stats)); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(Count(test.graph, "SimplifiedLayerNormalization"), 4u); + EXPECT_EQ(Count(test.graph, "FusedMatMul"), 2u); + EXPECT_EQ(stats.rewrites, 8u); + std::string rewritten; + test.model.ToProto().SerializeToString(&rewritten); + auto execute = [&](const std::string& serialized, std::vector& outputs) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + InferenceSessionWrapper session{options, GetEnvironment()}; + std::istringstream stream(serialized); + ORT_RETURN_IF_ERROR(session.Load(stream)); + ORT_RETURN_IF_ERROR(session.Initialize()); + return session.Run(RunOptions{}, test.builder.feeds_, test.builder.output_names_, &outputs); + }; + std::vector baseline, fused; + ASSERT_STATUS_OK(execute(original, baseline)); + ASSERT_STATUS_OK(execute(rewritten, fused)); + ASSERT_EQ(baseline.size(), fused.size()); + float max_absolute = 0.0f, max_relative = 0.0f; + for (size_t output = 0; output < baseline.size(); ++output) { + const auto& a = baseline[output].Get(); + const auto& c = fused[output].Get(); + ASSERT_EQ(a.Shape(), c.Shape()); + for (int64_t i = 0; i < a.Shape().Size(); ++i) { + float difference = std::abs(a.Data()[i] - c.Data()[i]); + float magnitude = std::abs(a.Data()[i]); + max_absolute = std::max(max_absolute, difference); + max_relative = std::max(max_relative, difference / std::max(magnitude, 1e-6f)); + ASSERT_LE(difference, 2e-5f + 2e-4f * magnitude); + } + } + std::cout << "decoder sequence=" << sequence << " outputs=" << fused.size() << " rewrites=" << stats.rewrites + << " max_absolute=" << max_absolute << " max_relative_floor_1e-6=" << max_relative << '\n'; + } +} +#endif + +#ifndef DISABLE_CONTRIB_OPS +namespace { + +const char* ActivationLibrary = R"( + interface GatedActivation(g:float[M,N], u:float[M,N]) -> (product:float[M,N]) + pattern plain implements GatedActivation { + s = Sigmoid(g) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + pattern scaled implements GatedActivation: (alpha:float[...]) -> { + scaled = Mul(g,alpha) + s = Sigmoid(scaled) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + pattern Forms = plain | scaled + pattern wrapped implements GatedActivation { + act = Forms(g,u) + return (product=act.product) + } + pattern Fuse(g,u) { + outer = wrapped(g,u) + return Identity(outer.product) + } +)"; + +// Handwritten views illustrate a future generator's common/concrete/nested API. +struct GatedView { + const NodeArg* g; + const NodeArg* u; + const NodeArg* product; + explicit GatedView(BindingView view) + : g(view.valueOf("g")), u(view.valueOf("u")), product(view.valueOf("product")) { + ORT_ENFORCE(view.interfaceName() == "GatedActivation"); + } +}; + +struct PlainView { + GatedView ports; + const NodeArg* sigmoid; + explicit PlainView(BindingView view) : ports(view), sigmoid(view.valueOf("s")) {} +}; + +struct ScaledView { + GatedView ports; + const NodeArg* alpha; + const NodeArg* scaled; + explicit ScaledView(BindingView view) : ports(view), alpha(view.valueOf("alpha")), scaled(view.valueOf("scaled")) {} +}; + +struct ActivationView { + GatedView ports; + std::variant alternative; + static std::variant Select(BindingView view) { + if (view.alternative() == "plain") return PlainView(view.alternativeView()); + ORT_ENFORCE(view.alternative() == "scaled", "Unexpected activation member."); + return ScaledView(view.alternativeView()); + } + explicit ActivationView(BindingView view) : ports(view), alternative(Select(view)) {} +}; + +struct FusionView { + GatedView outer; + ActivationView act; + explicit FusionView(const Match& match) + : outer(match.instance("outer")), act(match.instance("outer").instance("act")) {} +}; + +bool CheckCommon(const GatedView& view) { + if (!view.g || !view.u || !view.product || !view.g->Shape() || !view.u->Shape() || + view.g->Shape()->dim_size() != 2 || view.u->Shape()->dim_size() != 2) + return false; + for (int i = 0; i < 2; ++i) { + const auto& g = view.g->Shape()->dim(i); + const auto& u = view.u->Shape()->dim(i); + if (!g.has_dim_value() || !u.has_dim_value() || g.dim_value() <= 0 || g.dim_value() != u.dim_value()) return false; + } + return true; +} + +Rule ActivationRule(const PatternLibrary& library, std::string_view entry = "Fuse") { + auto check = [](const Graph& graph, Match& match) { + const FusionView view(match); + if (!CheckCommon(view.outer) || !CheckCommon(view.act.ports)) return match.Reject("Incompatible gating shapes."); + float alpha = 1.0f; + if (const auto* scaled = std::get_if(&view.act.alternative)) { + const auto* shape = scaled->alpha->Shape(); + if (!shape || !(shape->dim_size() == 0 || + (shape->dim_size() == 1 && shape->dim(0).has_dim_value() && shape->dim(0).dim_value() == 1))) + return match.Reject("Alpha must have scalar or [1] shape."); + InlinedVector values; + bool available; + ORT_THROW_IF_ERROR(ReadConstantTensor(graph, scaled->alpha, values, available)); + if (!available || values.size() != 1 || !std::isfinite(values[0]) || values[0] <= 0) + return match.Reject("Alpha must be a supported immutable float."); + alpha = values[0]; + } + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + attribute.set_f(alpha); + match.update("activation_alpha", std::move(attribute)); + return true; + }; + Replacement replacement(R"( + (g,u,activation_alpha) => { + a = com.microsoft.QuickGelu(g,alpha=activation_alpha) + return Mul(a,u) + } + )"); + return Rule{"PatternCompositionActivation", library.GetPattern(entry), replacement, check, {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; +} + +} // namespace + +TEST(PatternRewriterTests, PatternCompositionManualViewsAndSyntheticFusion) { + const PatternLibrary library(ActivationLibrary); + for (bool indexed : {false, true}) { + for (int variant = 0; variant < 5; ++variant) { + TestGraph test; + auto& b = test.builder; + auto* g = b.MakeInput({2, 4}, -1.0f, 1.0f); + auto* u = b.MakeInput({2, 4}, -1.0f, 1.0f); + NodeArg* scaled = g; + if (variant != 0) { + auto* alpha = variant == 1 ? b.MakeScalarInitializer(1.5f) + : variant == 2 ? b.MakeInitializer({1}, {0.5f}) + : variant == 3 ? b.MakeInput({1}, 1.0f, 2.0f) + : b.MakeInitializer({2, 4}, 0.5f, 1.0f); + scaled = b.MakeIntermediate(); + b.AddNode("Mul", {g, alpha}, {scaled}); + } + auto* sigmoid = b.MakeIntermediate(); + auto* activated = b.MakeIntermediate(); + auto* product = b.MakeIntermediate(); + b.AddNode("Sigmoid", {scaled}, {sigmoid}); + b.AddNode("Mul", {g, sigmoid}, {activated}); + b.AddNode("Mul", {activated, u}, {product}); + b.AddNode("Identity", {product}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + std::string original; + test.model.ToProto().SerializeToString(&original); + const bool accepted = variant < 3; + Options options; + options.indexed = indexed; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({ActivationRule(library)}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, accepted ? 1u : 0u); + EXPECT_EQ(Count(test.graph, "QuickGelu"), accepted ? 1u : 0u); + if (accepted) { + for (const auto& node : test.graph.Nodes()) + if (node.OpType() == "QuickGelu") + EXPECT_EQ(node.GetAttributes().at("alpha").f(), variant == 0 ? 1.0f : variant == 1 ? 1.5f + : 0.5f); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + if (accepted) { + std::string rewritten; + test.model.ToProto().SerializeToString(&rewritten); + auto execute = [&](const std::string& serialized, std::vector& outputs) { + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + session_options.intra_op_param.thread_pool_size = 1; + InferenceSessionWrapper session{session_options, GetEnvironment()}; + std::istringstream stream(serialized); + ORT_RETURN_IF_ERROR(session.Load(stream)); + ORT_RETURN_IF_ERROR(session.Initialize()); + return session.Run(RunOptions{}, b.feeds_, b.output_names_, &outputs); + }; + std::vector baseline, fused; + ASSERT_STATUS_OK(execute(original, baseline)); + ASSERT_STATUS_OK(execute(rewritten, fused)); + ASSERT_EQ(baseline.size(), 1u); + ASSERT_EQ(fused.size(), 1u); + const auto& expected = baseline[0].Get(); + const auto& actual = fused[0].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + for (int64_t i = 0; i < expected.Shape().Size(); ++i) + EXPECT_NEAR(expected.Data()[i], actual.Data()[i], 1e-6f); + } + } + } +} +#endif + +TEST(PatternRewriterTests, PatternCompositionNamedAndAnonymousAlternativesAreOccurrenceLocal) { + for (const char* call : {"Forms", "(first | second)"}) { + const std::string text = std::string(R"( + interface Unary(x:tensor[N]) -> (y:tensor[N]) + pattern first implements Unary { local=Identity(x); return (y=local) } + pattern second implements Unary { local=Identity(x); return (y=local) } + pattern Forms = first | second + pattern P(x,z) { a=)") + + call + "(x); b=" + call + R"((z); return Add(a.y,b.y) } + )"; + const PatternLibrary library(text); + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* z = b.MakeInput({2}, 0.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {z}, {c}); + b.AddNode("Add", {a, c}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector attempts; + auto check = [&](const Graph&, Match& match) { + auto av = match.instance("a"), bv = match.instance("b"); + attempts.push_back(std::string(av.alternative()) + "/" + std::string(bv.alternative())); + EXPECT_EQ(av.valueOf("x"), x); + EXPECT_EQ(bv.valueOf("x"), z); + EXPECT_EQ(av.valueOf("y"), av.valueOf("local")); + EXPECT_EQ(bv.valueOf("y"), bv.valueOf("local")); + EXPECT_NE(av.valueOf("local"), bv.valueOf("local")); + EXPECT_EQ(std::get(av.dimensionOf("N")), 2); + EXPECT_EQ(av.elementTypeOf("T"), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + EXPECT_THROW(match.attributeOf("leaked"), OnnxRuntimeException); + EXPECT_THROW(match.instance("absent"), OnnxRuntimeException); + if (av.alternative() == "first" && bv.alternative() == "second") return true; + ONNX_NAMESPACE::AttributeProto leaked; + leaked.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + leaked.set_i(99); + match.update("leaked", std::move(leaked)); + return match.Reject("Try the next complete candidate."); + }; + Options options; + options.indexed = indexed; + Statistics stats; + Rule rule{"Alternatives", library.GetPattern("P"), "(x,z) => Add(x,z)", check, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(attempts, (InlinedVector{"first/first", "first/second"})); + EXPECT_EQ(Count(test.graph, "Identity"), 0u); + } + } +} + +TEST(PatternRewriterTests, PatternCompositionHygienicDimensionsAttributesAndCommonOutputs) { + PatternLibrary library(R"( + interface Permute(x:float[N,M]) -> (y:float[M,N]) + pattern transpose implements Permute: (axes:ints) -> { + local=Transpose(x,perm=axes) + return (y=local) + } + pattern P(x,z) { + a=transpose(x) + b=transpose(z) + return MatMul(a.y,b.y) + } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4, 3}, 0.0f, 1.0f); + auto* z = b.MakeInput({4, 4}, 0.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {a}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Transpose", {z}, {c}).AddAttribute("perm", std::vector{0, 1}); + b.AddNode("MatMul", {a, c}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(std::get(match.instance("a").dimensionOf("M")), 3); + EXPECT_EQ(std::get(match.instance("b").dimensionOf("M")), 4); + EXPECT_EQ(match.instance("a").attributeOf("axes").ints(0), 1); + EXPECT_EQ(match.instance("b").attributeOf("axes").ints(1), 1); + EXPECT_EQ(match.instance("a").valueOf("y"), a); + EXPECT_THROW(match.update("a.y", {}), OnnxRuntimeException); + return true; + }; + Statistics stats; + Options options; + options.max_iterations = 1; + Rule rule{"Hygiene", library.GetPattern("P"), "(x,z,a.axes,b.axes)=>{a=Transpose(x,perm=a.axes);b=Transpose(z,perm=b.axes);return MatMul(a,b)}", check, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); +} + +TEST(PatternRewriterTests, PatternCompositionIdentityPortsRemapAndInterfaceConstraintsConjoin) { + PatternLibrary library(R"( + interface Pass(x:float[N]) -> (y:float[N]) + pattern pass implements Pass { return (y=x) } + pattern P(x:float[2]) { p=pass(x); return Identity(p.y) } + )"); + for (int64_t size : {2, 3}) { + TestGraph test; + auto* x = test.builder.MakeInput({size}, 0.0f, 1.0f); + test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.instance("p").valueOf("x"), x); + EXPECT_EQ(match.instance("p").valueOf("y"), x); + return true; + }; + Options options; + options.match_only = true; + Statistics stats; + Rule rule{"Pass", library.GetPattern("P"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, size == 2 ? 1u : 0u); + } +} + +TEST(PatternRewriterTests, PatternCompositionOutputAnnotationsAreEnforcedOnAliases) { + const PatternLibrary library(R"( + interface Pass(x:float[]) -> (y:float[3]) + pattern pass implements Pass { return (y=x) } + pattern P(x) { p=pass(x); return Identity(p.y) } + )"); + for (int64_t size : {2, 3}) { + TestGraph test; + auto* x = test.builder.MakeInput({size}, 0.0f, 1.0f); + test.builder.AddNode("Identity", {x}, {test.builder.MakeOutput()}); + test.builder.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Options options; + options.match_only = true; + Statistics stats; + Rule rule{"OutputConstraint", library.GetPattern("P"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, size == 3 ? 1u : 0u); + } +} + +TEST(PatternRewriterTests, PatternCompositionUnionTypeAndConstantRejectionRetry) { + const PatternLibrary library(R"( + interface Unary(x:float[...]) -> (y:float[...]) + pattern fixed implements Unary: (a:const float[1]) -> { y=Mul(x,a); return (y=y) } + pattern dynamic implements Unary: (a:float[]) -> { y=Mul(x,a); return (y=y) } + pattern Forms = fixed | dynamic + pattern P(x) { f=Forms(x); return Identity(f.y) } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 4}, 0.0f, 1.0f); + auto* a = b.MakeInput({1}, 0.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Mul", {x, a}, {y}); + b.AddNode("Identity", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.instance("f").alternative(), "dynamic"); + EXPECT_EQ(match.instance("f").valueOf("a"), a); + return true; + }; + Options options; + options.match_only = true; + Statistics stats; + Rule rule{"ConstantRetry", library.GetPattern("P"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.check_and_bind_calls, 1u); +} + +TEST(PatternRewriterTests, PatternCompositionDeletionProviderAndLayerBoundaries) { + const PatternLibrary library(R"( + interface Unary(x) -> (y) + pattern pair implements Unary { a=Identity(x); y=Identity(a); return (y=y) } + pattern P(x) { p=pair(x); return p.y } + )"); + for (int scenario = 0; scenario < 6; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* a = scenario == 1 ? b.MakeOutput() : b.MakeIntermediate(); + auto& first = b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + if (scenario == 2) b.AddNode("Neg", {a}, {b.MakeOutput()}); + if (scenario == 3) first.SetExecutionProviderType(kCudaExecutionProvider); + if (scenario == 4) first.SetLayeringAnnotation("other"); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Statistics stats; + Rule rule{"Safety", library.GetPattern("P"), "(x)=>Identity(x)", {}, {{kOnnxDomain, scenario == 5 ? 18 : 17, 21}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, scenario == 0 ? 1u : 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, PatternCompositionReplacementFailureDoesNotRetry) { + PatternLibrary library(R"( + interface Unary(x) -> (y) + pattern first implements Unary { a=Identity(x); y=Identity(a); return (y=y) } + pattern second implements Unary { a=Identity(x); y=Identity(a); return (y=y) } + pattern Forms=first|second + pattern P(x) { p=Forms(x); return p.y } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {a}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const auto before = test.graph.ToGraphProto().SerializeAsString(); + size_t checks = 0; + auto check = [&](const Graph&, Match&) { ++checks; return true; }; + Rule rule{"BadRecipe", library.GetPattern("P"), "(x)=>UndefinedOperator(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + EXPECT_FALSE(PatternRewriter({rule}).Run(test.graph, stats).IsOK()); + EXPECT_EQ(checks, 1u); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(test.graph.ToGraphProto().SerializeAsString(), before); + ASSERT_STATUS_OK(test.graph.Resolve()); +} + +TEST(PatternRewriterTests, PatternCompositionExtraInputsAreUsableByParsedReplacement) { + PatternLibrary library(R"( + interface Unary(x:float[]) -> (y:float[]) + pattern scale implements Unary: (alpha:float[]) -> { + y=Mul(x,alpha) + return (y=y) + } + pattern P(x) { p=scale(x); return Identity(p.y) } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* alpha = b.MakeInitializer({1}, {2.0f}); + auto* y = b.MakeIntermediate(); + b.AddNode("Mul", {x, alpha}, {y}); + b.AddNode("Identity", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"Scale", library.GetPattern("P"), "(x,p.alpha)=>Mul(x,p.alpha)", {}, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(Count(test.graph, "Identity"), 0u); + ASSERT_EQ(test.graph.Nodes().begin()->InputDefs()[1], alpha); + EXPECT_THROW((PatternRewriter({Rule{"Unsafe", library.GetPattern("P"), "(p.y)=>Identity(p.y)", {}, {{kOnnxDomain, 17, 17}}}})), + OnnxRuntimeException); +} + +TEST(PatternRewriterTests, PatternCompositionDifferentRootsAndHostRecipeSelection) { + // The returned pattern owns its linked plans after both text and library die. + const Pattern pattern = [] { + std::string text = R"( + interface Unary(x) -> (y) + pattern positive implements Unary { y=Relu(x); return (y=y) } + pattern negative implements Unary { y=Neg(x); return (y=y) } + pattern Forms=positive|negative + pattern P(x) { p=Forms(x); return p.y } + )"; + return PatternLibrary(text).GetPattern("P"); + }(); + for (bool indexed : {false, true}) { + for (const char* op : {"Relu", "Neg"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + b.AddNode(op, {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto choose = [positive = Replacement("(x)=>Relu(x)"), negative = Replacement("(x)=>Neg(x)")]( + const Graph&, const Match& match) { + return match.instance("p").alternative() == "positive" ? positive : negative; + }; + Rule rule{"SelectRecipe", pattern, choose, {}, {{kOnnxDomain, 17, 17}}}; + Options options; + options.indexed = indexed; + options.max_iterations = 1; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(Count(test.graph, op), 1u); + EXPECT_EQ(Count(test.graph, "If"), 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, PatternCompositionCommonMultiOutputsPreserveGraphOutputs) { + PatternLibrary library(R"( + interface SplitPair(x:float[4]) -> (left:float[2],right:float[2]) + pattern split implements SplitPair { + a,b=Split(x,axis=0) + return (right=b,left=a) + } + pattern P(x) { p=split(x); y=Concat(p.left,p.right,axis=0); return p.left,y } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, 0.0f, 1.0f); + auto* a = b.MakeOutput(); + auto* c = b.MakeIntermediate(); + auto* output = b.MakeOutput(); + b.AddNode("Split", {x}, {a, c}).AddAttribute("axis", int64_t{0}); + b.AddNode("Concat", {a, c}, {output}).AddAttribute("axis", int64_t{0}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.instance("p").valueOf("left"), a); + EXPECT_EQ(match.instance("p").valueOf("a"), a); + EXPECT_EQ(match.instance("p").valueOf("right"), c); + return true; + }; + Rule rule{"MultiOutput", library.GetPattern("P"), "(x)=>{a,b=Split(x,axis=0); y=Concat(a,b,axis=0); return a,y}", check, {{kOnnxDomain, 17, 17}}}; + Options options; + options.max_iterations = 1; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + EXPECT_EQ(test.graph.GetOutputs()[0], a); + EXPECT_EQ(test.graph.GetOutputs()[1], output); +} + +TEST(PatternRewriterTests, PatternCompositionLibraryErrorsAndBounds) { + const char* invalid[] = { + "interface I(x)->(y) pattern P(x){a=I(x);return a.y}", + "pattern f implements Missing {return (y=x)}", + "interface I(x)->(y) pattern f implements I {a=f(x);return (y=a.y)}", + "interface I(x)->(y) pattern f implements I {a=g(x);return (y=a.y)} " + "pattern g implements I {a=f(x);return (y=a.y)}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (wrong=y)}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y,y=y)}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=z)}", + "interface I(x)->(x)", + "interface I(x:float[x])->(y)", + "interface I(x:float[N])->(y:tensor[])", + "interface I(x)->(y) pattern f implements I:(unused:float[])->{y=Identity(x);return (y=y)}", + "interface I(x)->(y) pattern f implements I {a=Identity(x);y=Identity(a);return (y=a)}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a=f(x,x);return a.y}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a,b=f(x);return a}", + "interface I(x)->(y) interface J(x)->(y) " + "pattern f implements I {y=Identity(x);return (y=y)} " + "pattern g implements J {y=Identity(x);return (y=y)} pattern U=f|g", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} pattern U=f|missing", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} pattern U=f|f", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a=(f|missing)(x);return a.y}", + "interface I(x)->(y) pattern f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a=f(x);a=Identity(a.y);return a}", + "pattern P(x){return misspelled(x)}", + "pattern P(x){return Identity(x)} pattern P(x){return Identity(x)}", + "pattern P(x:const int){return Identity(x)}", + "pattern P(x,a:const float){return LeakyRelu(x,alpha=a)}", + "pattern P(x:float[N]){y:float[x]=Identity(x);return y}", + "pattern P(x:float[]=[1.0]){return Identity(x)}", + }; + for (const auto* text : invalid) { + SCOPED_TRACE(text); + EXPECT_THROW(PatternLibrary{text}, OnnxRuntimeException); + } + std::string too_many = R"( + interface I(x)->(y) + pattern f implements I {return (y=x)} + pattern g implements I {return (y=x)} + pattern U=f|g + pattern P(x) { + )"; + for (int i = 0; i < 7; ++i) + too_many += "a" + std::to_string(i) + "=U(" + (i == 0 ? "x" : "a" + std::to_string(i - 1) + ".y") + ");"; + too_many += "return Identity(a6.y)}"; + EXPECT_THROW(PatternLibrary{too_many}, OnnxRuntimeException); + PatternLibrary valid("pattern P(x){return Identity(x)}"); + EXPECT_THROW(valid.GetPattern("missing"), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, ConstTensorQualifierAndTypedReaders) { + for (int scenario = 0; scenario < 5; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* c = scenario == 1 ? b.MakeInput({1}, 0.0f, 1.0f) + : b.MakeInitializer({1}, {2.0f}); + b.AddNode("Add", {x, c}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + if (scenario == 2) { + auto inputs = test.graph.GetInputsIncludingInitializers(); + inputs.push_back(c); + test.graph.SetInputs(inputs); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + if (scenario == 3 || scenario == 4) { + const auto* original = test.graph.GetInitializer(c->Name(), false); + ASSERT_NE(original, nullptr); + auto tensor = *original; + if (scenario == 3) { + tensor.clear_raw_data(); + tensor.clear_float_data(); + } else { + tensor.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + } + test.graph.RemoveInitializedTensor(c->Name()); + test.graph.AddInitializedTensor(tensor); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + InlinedVector values{99.0f}; + bool available = true; + const auto status = ReadConstantTensor(test.graph, c, values, available); + if (scenario == 3) { + EXPECT_FALSE(status.IsOK()); + EXPECT_EQ(values, (InlinedVector{99.0f})); + } else { + ASSERT_STATUS_OK(status); + EXPECT_EQ(available, scenario == 0); + EXPECT_EQ(values, (InlinedVector{scenario == 0 ? 2.0f : 99.0f})); + } + InlinedVector wrong_type; + ASSERT_STATUS_OK(ReadConstantTensor(test.graph, c, wrong_type, available)); + EXPECT_FALSE(available); + Statistics stats; + Options options; + options.match_only = true; + Rule rule{"Const", "(x,c:const float[])=>return Add(x,c)", "(x,c)=>Add(x,c)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 1 || scenario == 2 ? 0u : 1u); + } + TestGraph test; + auto* value = test.builder.MakeInitializer({2}, {-1, 8}); + InlinedVector values; + bool available; + ASSERT_STATUS_OK(ReadConstantTensor(test.graph, value, values, available)); + EXPECT_TRUE(available); + EXPECT_EQ(values, (InlinedVector{-1, 8})); + EXPECT_THROW(Pattern("(x:const const float[])=>Identity(x)"), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, ConstQualifierIsSpecificToItsTypeAlternative) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = scenario == 0 ? b.MakeInitializer({1}, {1.0f}) + : b.MakeInput({scenario == 1 ? 1 : 2}, 0.0f, 1.0f); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Options options; + options.match_only = true; + Statistics stats; + Rule rule{"ConstAlternative", "(x:const float[1] | float[2])=>return Identity(x)", "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 1 ? 0u : 1u); + } + Pattern optional("(x:const float[] | None)=>return Identity(x)"); + EXPECT_TRUE(optional.IsOptional(optional.NamedValue("x"))); +} + +TEST(PatternRewriterTests, PatternCompositionLiteralConstantsAreRemapped) { + const PatternLibrary library(R"( + interface Shift(x) -> (y) + pattern shift implements Shift { y=Add(x,3.0); return (y=y) } + pattern P(x) { p=shift(x); return p.y } + )"); + for (float constant : {3.0f, 4.0f}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + b.AddNode("Add", {x, b.MakeScalarInitializer(constant)}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Options options; + options.match_only = true; + Statistics stats; + Rule rule{"Literal", library.GetPattern("P"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, constant == 3.0f ? 1u : 0u); + } +} + +TEST(PatternRewriterTests, UnifiedPatternsDirectNestedCallsAndIdentityEntries) { + const PatternLibrary library(R"( + pattern entry(x) { c=wrapper(x); return Identity(c.y) } + pattern wrapper(x) { p=leaf(x); return (y=p.y) } + pattern leaf(x:float[N]) { a=Identity(x); return (y=a) } + pattern identity(x) { return (y=x) } + pattern mixed(x) { a=Identity(x); return (original=x,y=a) } + pattern use_identity(x) { p=identity(x); return Identity(p.y) } + pattern implicit(x) { return Identity(x) } + pattern use_implicit(x) { p=implicit(x); return Identity(p.output) } + pattern use_mixed(x) { p=mixed(x); return Add(p.original,p.y) } + )"); + for (const char* name : {"identity", "mixed"}) { + SCOPED_TRACE(name); + EXPECT_THROW(library.GetPattern(name), OnnxRuntimeException); + } + for (const char* entry : {"entry", "use_identity", "use_implicit", "use_mixed"}) { + SCOPED_TRACE(entry); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + if (std::string_view(entry) != "use_identity") b.AddNode("Identity", {x}, {a}); + if (std::string_view(entry) == "use_mixed") + b.AddNode("Add", {x, a}, {b.MakeOutput()}); + else + b.AddNode("Identity", {std::string_view(entry) == "use_identity" ? x : a}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + if (std::string_view(entry) == "entry") { + const auto wrapper = match.instance("c"); + EXPECT_EQ(wrapper.definition(), "wrapper"); + EXPECT_EQ(wrapper.instance("p").definition(), "leaf"); + EXPECT_EQ(wrapper.valueOf("x"), x); + EXPECT_EQ(wrapper.valueOf("y"), a); + EXPECT_EQ(std::get(wrapper.instance("p").dimensionOf("N")), 2); + EXPECT_EQ(wrapper.instance("p").valueOf("y"), a); + EXPECT_TRUE(wrapper.alternative().empty()); + EXPECT_THROW(wrapper.alternativeView(), OnnxRuntimeException); + } else if (std::string_view(entry) == "use_identity") { + EXPECT_EQ(match.instance("p").valueOf("y"), x); + } + return true; + }; + Rule rule{"UnifiedCall", library.GetPattern(entry), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + Options options; + options.max_iterations = 1; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + EXPECT_NO_THROW(library.GetPattern("leaf")); + EXPECT_NO_THROW(library.GetPattern("wrapper")); +} + +TEST(PatternRewriterTests, UnifiedPatternsIndependentOutputsAreCallableNotEntries) { + const PatternLibrary library(R"( + pattern branches(x,z) { a=Identity(x); b=Neg(z); return (left=a,right=b) } + pattern parent(x,z) { p=branches(x,z); return Add(p.left,p.right) } + pattern branch_alias=branches + pattern aliased_parent(x,z) { p=branch_alias(x,z); return Add(p.left,p.right) } + )"); + for (const char* entry : {"branches", "branch_alias"}) { + try { + library.GetPattern(entry); + FAIL() << "An independent-output definition must not serve as a standalone entry."; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("cannot be a rule entry"), std::string::npos); + } + } + for (const char* entry : {"parent", "aliased_parent"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* z = b.MakeInput({2}, 0.0f, 1.0f); + auto* left = b.MakeIntermediate(); + auto* right = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {left}); + b.AddNode("Neg", {z}, {right}); + b.AddNode("Add", {left, right}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.instance("p").valueOf("left"), left); + EXPECT_EQ(match.instance("p").valueOf("right"), right); + if (std::string_view(entry) == "aliased_parent") { + EXPECT_EQ(match.instance("p").alternative(), "branches"); + EXPECT_EQ(match.instance("p").alternativeView().definition(), "branches"); + EXPECT_TRUE(match.instance("p").alternativeView().alternative().empty()); + } + return true; + }; + Rule rule{"IndependentOutputs", library.GetPattern(entry), "(x,z)=>Add(x,z)", check, {{kOnnxDomain, 17, 17}}}; + Options options; + options.match_only = true; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + } +} + +namespace { +struct NestedSelectionView { + std::string_view definition; + std::string_view member; + std::unique_ptr selected; + const NodeArg* y; + explicit NestedSelectionView(BindingView view) + : definition(view.definition()), member(view.alternative()), selected(member.empty() ? nullptr : std::make_unique(view.alternativeView())), y(view.valueOf("y")) {} +}; +} // namespace + +TEST(PatternRewriterTests, UnifiedPatternsNestedAlternativesTagsRetryAndIsolation) { + for (const char* callee : {"Outer", "(Alias | Inner)"}) { + const PatternLibrary library(std::string(R"( + interface Unary(x:float[N]) -> (y:float[N]) + pattern first implements Unary { a=Identity(x); return (y=a) } + pattern second implements Unary { b=Identity(x); return (y=b) } + pattern Inner=first|second + pattern Alias=Inner + pattern Outer=Alias|Inner + pattern P(x,z) { p=)") + callee + + "(x); q=" + callee + "(z); return Add(p.y,q.y) }"); + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* z = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + b.AddNode("Identity", {z}, {c}); + b.AddNode("Add", {a, c}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector attempts; + auto path = [](BindingView view) { + std::string result; + while (!view.alternative().empty()) { + if (!result.empty()) result += '/'; + result += view.alternative(); + view = view.alternativeView(); + } + return result; + }; + auto check = [&](const Graph&, Match& match) { + const auto p = match.instance("p"), q = match.instance("q"); + const auto pp = path(p), qp = path(q); + attempts.push_back(pp + ":" + qp); + const NestedSelectionView typed(p); + EXPECT_EQ(typed.y, a); + EXPECT_NE(typed.selected, nullptr); + if (typed.selected) EXPECT_EQ(typed.selected->y, a); + EXPECT_EQ(p.valueOf("x"), x); + EXPECT_EQ(q.valueOf("x"), z); + EXPECT_THROW(match.attributeOf("leaked"), OnnxRuntimeException); + if (pp == "Inner/second" && qp == "Alias/Inner/first") { + EXPECT_THROW(p.valueOf("a"), OnnxRuntimeException); + EXPECT_EQ(p.alternativeView().alternativeView().valueOf("b"), a); + EXPECT_EQ(q.alternativeView().alternativeView().alternativeView().valueOf("a"), c); + return true; + } + ONNX_NAMESPACE::AttributeProto leaked; + leaked.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + leaked.set_i(123); + match.update("leaked", std::move(leaked)); + return match.Reject("Try the next nested selection."); + }; + Rule rule{"Nested", library.GetPattern("P"), "(x,z)=>Add(x,z)", check, {{kOnnxDomain, 17, 17}}}; + Options options; + options.indexed = indexed; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_EQ(attempts.size(), 13u); + EXPECT_EQ(attempts.front(), "Alias/Inner/first:Alias/Inner/first"); + EXPECT_EQ(attempts[1], "Alias/Inner/first:Alias/Inner/second"); + EXPECT_EQ(attempts.back(), "Inner/second:Alias/Inner/first"); + } + } +} + +TEST(PatternRewriterTests, UnifiedPatternsAlternativesCanBeEntriesOrCallPlainBodies) { + const PatternLibrary library(R"( + pattern positive(x:float[]) { y=Relu(x); return y } + pattern negative(x:float[]) { y=Neg(x); return y } + pattern Forms=positive|negative + pattern Alias=Forms + pattern Caller(x) { p=Alias(x); return Identity(p.y) } + )"); + for (bool indexed : {false, true}) { + for (const char* entry : {"Forms", "Alias", "Caller"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + if (std::string_view(entry) == "Caller") { + auto* y = b.MakeIntermediate(); + b.AddNode("Neg", {x}, {y}); + b.AddNode("Identity", {y}, {b.MakeOutput()}); + } else { + b.AddNode("Neg", {x}, {b.MakeOutput()}); + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + auto view = std::string_view(entry) == "Caller" ? match.instance("p") : match.bindings(); + if (std::string_view(entry) != "Forms") { + EXPECT_EQ(view.alternative(), "Forms"); + view = view.alternativeView(); + } + EXPECT_EQ(view.alternative(), "negative"); + EXPECT_EQ(view.alternativeView().definition(), "negative"); + EXPECT_TRUE(view.alternativeView().interfaceName().empty()); + return true; + }; + Rule rule{"EntryAlternative", library.GetPattern(entry), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Options options; + options.indexed = indexed; + options.max_iterations = 1; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + } + } +} + +TEST(PatternRewriterTests, UnifiedPatternsEntryValidityChecksEveryExpansion) { + const PatternLibrary library(R"( + interface Unary(x) -> (y) + pattern pass implements Unary { return (y=x) } + pattern op implements Unary { y=Identity(x); return (y=y) } + pattern Forms=op|pass + pattern Alias=Forms + pattern Caller(x) { p=Alias(x); return Identity(p.y) } + )"); + EXPECT_THROW(library.GetPattern("Forms"), OnnxRuntimeException); + EXPECT_THROW(library.GetPattern("Alias"), OnnxRuntimeException); + EXPECT_THROW(library.GetPattern("Unary"), OnnxRuntimeException); + EXPECT_NO_THROW(library.GetPattern("Caller")); +} + +TEST(PatternRewriterTests, UnifiedPatternsPublicAttributeInputsAndScopedAliases) { + const PatternLibrary library(R"( + pattern transpose(x:float[N,M],axes:ints) { y=Transpose(x,perm=axes); return y } + pattern nested(x,perm:ints) { a=transpose(x,perm); return (y=a.y) } + pattern Caller(x,axes:ints) { p=nested(x,axes); return Identity(p.y) } + )"); + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, 0.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Transpose", {x}, {y}).AddAttribute("perm", std::vector{1, 0}); + b.AddNode("Identity", {y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.instance("p").attributeOf("perm").ints(0), 1); + EXPECT_EQ(match.instance("p").instance("a").attributeOf("axes").ints(1), 0); + auto attribute = match.attributeOf("axes"); + EXPECT_THROW(match.update("p.a.axes", attribute), OnnxRuntimeException); + EXPECT_THROW(match.update("p.perm", attribute), OnnxRuntimeException); + EXPECT_THROW(match.update("axes", attribute), OnnxRuntimeException); + return true; + }; + Rule rule{"AttributePort", library.GetPattern("Caller"), "(x,p.a.axes)=>Transpose(x,perm=p.a.axes)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); +} + +TEST(PatternRewriterTests, UnifiedPatternsRejectCyclesCompatibilityAndObsoleteDeclarations) { + const char* invalid[] = { + "fragment old(x) { return Identity(x) }", + "union Old=a|b", + "pattern A=B pattern B=A", + "pattern A=A", + "pattern A=B pattern B(x){p=A(x);return p.y}", + "pattern A(x){p=B(x);return p.y} pattern B=A", + "interface I(x)->(y) pattern A=I", + "pattern A=missing", + "pattern a(x){return (y=x)} pattern b(z){return (y=z)} pattern C=a|b", + "pattern a(x){return (y=x)} pattern b(x){return (z=x)} pattern C=a|b", + "pattern a(x:float[2]){return (y=x)} pattern b(x:float[3]){return (y=x)} pattern C=a|b", + "pattern a(x){y:float[2]=Identity(x);return y} " + "pattern b(x){y:float[3]=Identity(x);return y} pattern C=a|b", + "interface I(x)->(y) pattern a implements I {return (y=x)} " + "pattern b(x){return (y=x)} pattern C=a|b", + "pattern a(x){return (y=x)} pattern B=a pattern C(x){p=(a|B|a)(x);return Identity(p.y)}", + "pattern a(x){return (y=x)} pattern B=a pattern C(x){p=(a|missing)(x);return Identity(p.y)}", + "pattern a(x){return (y=x)} pattern b(x){return (z=x)} pattern C(x){p=(a|b)(x);return Identity(p.y)}", + "pattern a(x,axes:ints){y=Transpose(x,perm=axes);return y} " + "pattern b(x){z=Identity(x);p=a(x,z);return p.y}", + "pattern a(x){return (left=x,right=x)} pattern b(x){p=a(x);q=a(x);return p.left,q.left}", + }; + for (const auto* text : invalid) { + SCOPED_TRACE(text); + EXPECT_THROW(PatternLibrary{text}, OnnxRuntimeException); + } + std::string text = "pattern base(x){return (y=x)} pattern other=base pattern A0=base|other "; + for (int i = 1; i < 6; ++i) { + const auto previous = "A" + std::to_string(i - 1); + text += "pattern B" + std::to_string(i) + "=" + previous + " pattern A" + std::to_string(i) + + "=" + previous + "|B" + std::to_string(i) + " "; + } + EXPECT_NO_THROW(PatternLibrary{text}); + text += "pattern too_many=A5|base"; + EXPECT_THROW(PatternLibrary{text}, OnnxRuntimeException); + EXPECT_THROW(Pattern("(x)=>return x"), OnnxRuntimeException); + EXPECT_THROW(Pattern("(x,z)=>{a=Identity(x);b=Identity(z);return a,b}"), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, UnifiedPatternsExpansionAndAnnotationThresholds) { + EXPECT_NO_THROW(PatternLibrary(R"( + pattern a(x:float[]) { y=Identity(x); return y } + pattern b(x:float[...]) { y=Neg(x); return y } + pattern Forms=a|b + pattern P(x) { p=(Forms)(x); return Identity(p.y) } + )") + .GetPattern("P")); + std::string body = "pattern node(x){y=Identity(x);return y} pattern P(x){"; + for (int i = 0; i < 64; ++i) + body += "a" + std::to_string(i) + "=node(" + (i == 0 ? "x" : "a" + std::to_string(i - 1) + ".y") + ");"; + EXPECT_NO_THROW(PatternLibrary(body + "return a63.y}").GetPattern("P")); + EXPECT_THROW(PatternLibrary(body + "a64=node(a63.y);return a64.y}"), OnnxRuntimeException); + std::string types; + for (int i = 0; i < 64; ++i) { + if (!types.empty()) types += '|'; + types += "float[" + std::to_string(i) + "]"; + } + // Inferred output ports must not reapply and square the same disjunction. + EXPECT_NO_THROW(PatternLibrary("pattern P(x){y:" + types + "=Identity(x);return y}").GetPattern("P")); + EXPECT_THROW(PatternLibrary("pattern P(x){y:" + types + "|float[64]=Identity(x);return y}"), + OnnxRuntimeException); +} + +TEST(PatternRewriterTests, PortableFloatLiteralsValidateWholeToken) { + for (const char* literal : {"+1.25", "-1.25", "1.25e+2", ".125", "1."}) { + SCOPED_TRACE(literal); + EXPECT_NO_THROW(Pattern(std::string("(x)=>return Add(x,") + literal + ")")); + EXPECT_NO_THROW(Replacement(std::string("(x)=>LeakyRelu(x,alpha=") + literal + ")")); + } + for (const char* literal : {"1e999", "1e", "1.2.3", "1e+", "1.0-2.0"}) { + SCOPED_TRACE(literal); + EXPECT_THROW(Pattern(std::string("(x)=>return Add(x,") + literal + ")"), OnnxRuntimeException); + EXPECT_THROW(Replacement(std::string("(x)=>LeakyRelu(x,alpha=") + literal + ")"), OnnxRuntimeException); + } +} + +TEST(PatternRewriterTests, AliasedReturnsAreCallableButNotDuplicateEntryOutputs) { + const PatternLibrary library(R"( + pattern IdentityAlias(x) { return (y=x) } + pattern Duplicate(x) { y=Identity(x); return (a=y,b=y) } + pattern NestedDuplicate(x) { y=Identity(x); p=IdentityAlias(y); return (a=y,b=p.y) } + pattern Caller(x) { p=Duplicate(x); return Add(p.a,p.b) } + pattern NestedCaller(x) { p=NestedDuplicate(x); return Add(p.a,p.b) } + pattern Forms=Duplicate|NestedDuplicate + )"); + for (const char* name : {"Duplicate", "NestedDuplicate", "Forms"}) { + SCOPED_TRACE(name); + try { + library.GetPattern(name); + FAIL() << "Aliased entry returns must not create duplicate replacement producers."; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("duplicate returned value"), std::string::npos); + } + } + for (const char* name : {"Caller", "NestedCaller"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, 0.0f, 1.0f); + auto* y = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {y}); + b.AddNode("Add", {y, y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"Aliases", library.GetPattern(name), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + Options options; + options.max_iterations = 1; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, ZeroNodeCompositionHasCumulativeCallBudget) { + std::string text = "pattern P0(x){return (y=x)} "; + for (int i = 1; i <= 7; ++i) { + const auto previous = "P" + std::to_string(i - 1); + text += "pattern P" + std::to_string(i) + "(x){a=" + previous + "(x);b=" + previous + "(a.y);return (y=b.y)} "; + } + // P7 contains 254 expanded calls; these wrappers reach exactly 256. + text += "pattern P8(x){a=P7(x);return (y=a.y)} pattern P9(x){a=P8(x);return (y=a.y)} "; + EXPECT_NO_THROW(PatternLibrary{text}); + const std::string overflow = text + "pattern P10(x){a=P9(x);return (y=a.y)}"; + try { + PatternLibrary library(overflow); + FAIL() << "Zero-node metadata expansion must be bounded before copying."; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("expanded call limit of 256"), std::string::npos); + } + EXPECT_THROW(PatternLibrary(text + "pattern P10(x){a=P7(x);b=P7(a.y);return (y=b.y)}"), + OnnxRuntimeException); +} + +TEST(PatternRewriterTests, InterfacePortsCannotBecomeAttributes) { + for (const char* text : { + "interface I(x,alpha:float)->(y) pattern P implements I {y=LeakyRelu(x,alpha=alpha);return y}", + "interface I(x)->(y,z) pattern P implements I:(alpha:float)->" + "{a=LeakyRelu(x,alpha=alpha);return (y=alpha,z=a)}", + "pattern Attr(x,alpha:float){y=LeakyRelu(x,alpha=alpha);return y} " + "interface I(x,alpha:float)->(y) pattern P implements I {a=Attr(x,alpha);return (y=a.y)}", + "interface I(x,alpha:float)->(y) " + "pattern A implements I {y=Mul(x,alpha);return y} " + "pattern B implements I {y=LeakyRelu(x,alpha=alpha);return y} pattern Forms=A|B"}) { + SCOPED_TRACE(text); + try { + PatternLibrary library(text); + FAIL() << "Interface ports must remain tensor bindings."; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("must remain a tensor"), std::string::npos); + } + } + EXPECT_NO_THROW(PatternLibrary(R"( + interface I(x) -> (y) + pattern P implements I: (alpha:float) -> { y=LeakyRelu(x,alpha=alpha);return y } + pattern C(x) { p=P(x); return Identity(p.y) } + )") + .GetPattern("C")); +} + +namespace { +class CommaDecimalPoint : public std::numpunct { + protected: + char do_decimal_point() const override { return ','; } +}; + +class ScopedGlobalLocale { + public: + explicit ScopedGlobalLocale(const std::locale& locale) : previous_(std::locale::global(locale)) {} + ~ScopedGlobalLocale() { std::locale::global(previous_); } + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ScopedGlobalLocale); + + private: + std::locale previous_; +}; +} // namespace + +TEST(PatternRewriterTests, GroupSignaturesAreIndependentOfGlobalLocale) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool second_constant : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* constant = b.MakeInitializer({2}, second_constant ? std::vector{1.0f, 5.5f} + : std::vector{1.5f, 5.0f}); + b.AddNode("Add", {x, constant}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector callbacks; + Rules rules; + rules.emplace_back("First", "(x)=>return Add(x,[1.5,5.0])", "(x)=>Identity(x)", [&](const Graph&, Match&) {callbacks.push_back(0); return true; }, InlinedVector{{kOnnxDomain, 17, 17}}); + rules.emplace_back("Second", "(x)=>return Add(x,[1.0,5.5])", "(x)=>Identity(x)", [&](const Graph&, Match&) {callbacks.push_back(1); return true; }, InlinedVector{{kOnnxDomain, 17, 17}}); + Statistics stats; + { + const ScopedGlobalLocale locale{std::locale(std::locale::classic(), new CommaDecimalPoint)}; + ASSERT_EQ(std::use_facet>(std::locale()).decimal_point(), ','); + ASSERT_STATUS_OK(PatternRewriter(std::move(rules), {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + } + ASSERT_EQ(callbacks.size(), 1u); + EXPECT_EQ(callbacks[0], second_constant ? 1 : 0); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } +} + +TEST(PatternRewriterTests, UnannotatedInterfaceOutputAliasesMustBePresent) { + const PatternLibrary library(R"( + interface I(v:float | None) -> (y) + pattern P implements I { return (y=v) } + pattern Alias=P + pattern Wrapped implements I { p=Alias(v);return (y=p.y) } + pattern Entry(x,v:float | None) { p=P(v);return Clip(x,p.y) } + pattern Nested(x,v:float | None) { p=Wrapped(v);return Clip(x,p.y) } + )"); + for (const char* entry : {"Entry", "Nested"}) { + for (bool indexed : {false, true}) { + for (int scenario = 0; scenario < 3; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + std::vector inputs{x}; + const NodeArg* minimum = nullptr; + if (scenario == 1) inputs.push_back(b.MakeOptionalTensor()); + if (scenario == 2) { + auto* value = b.MakeScalarInitializer(0.0f); + inputs.push_back(value); + minimum = value; + } + b.AddNode("Clip", inputs, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_NE(match.instance("p").valueOf("y"), nullptr); + EXPECT_EQ(match.instance("p").valueOf("y"), minimum); + return true; + }; + Rule rule{"RequiredOutput", library.GetPattern(entry), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(checks, scenario == 2 ? 1 : 0); + EXPECT_EQ(stats.matches, scenario == 2 ? 1u : 0u); + } + } + } +} + +TEST(PatternRewriterTests, InterleavedGroupsDispatchCallbacksInRegistrationOrder) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool throw_late : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* y = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + auto& root = b.AddNode("Add", {a, y}, {b.MakeOutput()}); + const NodeIndex index = root.Index(); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector callbacks; + InlinedVector rejected; + auto reject = [&](const Graph&, Match& match) { + callbacks.push_back(0); + ONNX_NAMESPACE::AttributeProto extra; + extra.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + extra.set_i(99); + match.update("candidate_only", std::move(extra)); + return match.Reject("first rejects"); + }; + auto accept = [&](const Graph&, Match& match) { + callbacks.push_back(1); + EXPECT_THROW(match.attributeOf("candidate_only"), OnnxRuntimeException); + return true; + }; + auto later = [&](const Graph&, Match& match) { + callbacks.push_back(2); + EXPECT_THROW(match.attributeOf("candidate_only"), OnnxRuntimeException); + ORT_ENFORCE(!throw_late, "lower priority callback"); + return true; + }; + Rule first{"First", "(x,y)=>return Add(x,y)", "(x,y)=>Add(x,y)", reject, {{kOnnxDomain, 17, 17}}}; + Rule middle{"Middle", "(x,y)=>{a=Identity(x);return Add(a,y)}", "(x,y)=>Add(x,y)", accept, {{kOnnxDomain, 17, 17}}}; + Rule last{"Last", "(x,y)=>return Add(x,y)", "(x,y)=>Add(x,y)", later, {{kOnnxDomain, 17, 17}}}; + Options options{.match_only = match_only, .indexed = indexed, .max_iterations = 1}; + options.on_failure = [&](const MatchFailure& failure) { + if (failure.root == index && failure.stage == FailureStage::CheckAndBind) rejected.push_back(failure.rule); + }; + Statistics stats; + PatternRewriter rewriter({first, middle, last}, options); + if (match_only && throw_late) { + EXPECT_THROW((void)rewriter.Run(test.graph, stats), OnnxRuntimeException); + } else { + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + EXPECT_EQ(stats.matches, match_only ? 2u : 1u); + } + ASSERT_EQ(callbacks.size(), match_only ? 3u : 2u); + EXPECT_EQ(callbacks[0], 0); + EXPECT_EQ(callbacks[1], 1); + if (match_only) EXPECT_EQ(callbacks[2], 2); + ASSERT_EQ(rejected.size(), 1u); + EXPECT_EQ(rejected[0], "First"); + // Indexed mode shares First/Last's traversal even when other groups intervene. + EXPECT_EQ(stats.structural_checks, indexed ? 2u : (match_only ? 6u : 5u)); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } +} + +TEST(PatternRewriterTests, AnnotationPrioritySurvivesReorderedAndAliasedCalls) { + for (int repeat = 0; repeat < 16; ++repeat) { + const PatternLibrary library(R"( + pattern F(x:float[M,_] | float[_,M], z:float[_,M] | float[M,_]) { + y=CastLike(x,z);return y + } + pattern Reordered(v,w) { p=F(w,v);return p.y } + pattern Aliased(v) { p=F(v,v);return p.y } + pattern Alias=F + pattern Nested(v,w) { p=Alias(w,v);return p.y } + )"); + for (const char* entry : {"F", "Reordered", "Aliased", "Nested"}) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* v = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* w = std::string_view(entry) == "Aliased" ? v : b.MakeInput({2, 3}, -1.0f, 1.0f); + b.AddNode("CastLike", {w, v}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + const auto view = std::string_view(entry) == "F" ? match.bindings() : match.instance("p"); + EXPECT_EQ(std::get(view.dimensionOf("M")), 2); + return true; + }; + Rule rule{"Priority", library.GetPattern(entry), std::string_view(entry) == "F" ? "(x)=>Identity(x)" : "(v)=>Identity(v)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(checks, 1); + } + } + } +} + +TEST(PatternRewriterTests, ComposedInterfaceTensorBindingsCannotBecomeAttributes) { + for (bool annotated : {false, true}) { + const std::string prefix = annotated + ? "interface I(a:float[])->(y:float[]) " + : "interface I(a)->(y) "; + const std::string definitions = prefix + R"( + pattern Pass implements I { return (y=a) } + pattern Alias=Pass + pattern Wrapped implements I { p=Alias(a);return (y=p.y) } + pattern Attr(x,alpha:float) { y=LeakyRelu(x,alpha=alpha);return y } + )"; + for (const char* call : {"Pass", "Alias", "Wrapped"}) { + for (const char* body : { + "p=CALL(a);return LeakyRelu(x,alpha=a)", + "p=CALL(a);q=Attr(x,p.y);return q.y", + "p=CALL(a);q=Attr(x,p.a);return q.y", + "q=Attr(x,a);p=CALL(a);return q.y"}) { + std::string text = body; + text.replace(text.find("CALL"), 4, call); + try { + const PatternLibrary library(definitions + "pattern Entry(x,a){" + text + "}"); + FAIL() << "Accepted interface tensor/attribute conflict: " << text; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("Declared tensor binding"), std::string::npos) << error.what(); + } + } + } + } +} + +TEST(PatternRewriterTests, ReplacementDefaultsAreVisibleToChainedRules) { + for (bool indexed : {false, true}) { + for (bool explicit_alpha : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {a}); + auto* output = b.MakeOutput(); + b.AddNode("LeakyRelu", {a}, {output}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + size_t checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_FLOAT_EQ(match.attributeOf("alpha").f(), explicit_alpha ? 0.02f : 0.01f); + EXPECT_FLOAT_EQ(match.nodes.back()->GetAttributes().at("alpha").f(), explicit_alpha ? 0.02f : 0.01f); + return true; + }; + Rules rules{ + {"RemoveIdentity", "(x)=>{a=Identity(x);return LeakyRelu(a,alpha=0.01)}", explicit_alpha ? "(x)=>LeakyRelu(x,alpha=0.02)" : "(x)=>LeakyRelu(x)", {}, {{kOnnxDomain, 17, 17}}}, + {"UsePRelu", "(x)=>return LeakyRelu(x,alpha=alpha)", explicit_alpha ? "(x)=>PRelu(x,0.02)" : "(x)=>PRelu(x,0.01)", check, {{kOnnxDomain, 17, 17}}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter(rules, {.indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 2u); + EXPECT_EQ(checks, 1u); + EXPECT_FALSE(stats.iteration_limit_reached); + ASSERT_NE(test.graph.GetProducerNode(output->Name()), nullptr); + EXPECT_EQ(test.graph.GetProducerNode(output->Name())->OpType(), "PRelu"); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, ForwardDefinitionNestingHasExplicitBoundary) { + for (bool reverse : {false, true}) { + for (bool aliases : {false, true}) { + for (int depth : {64, 65, 256}) { + std::string text; + for (int index = 0; index < depth; ++index) { + const int i = reverse ? depth - 1 - index : index; + const std::string name = "P" + std::to_string(i); + const std::string next = "P" + std::to_string(i + 1); + if (i == depth - 1) + text += "pattern " + name + "(x){return (y=x)} "; + else if (i > 0 && aliases) + text += "pattern " + name + "=" + next + " "; + else + text += "pattern " + name + "(x){p=" + next + "(x);" + + (i == 0 ? "y=Identity(p.y);return y" : "return (y=p.y)") + "} "; + } + if (depth == 64) { + const PatternLibrary library(text); + EXPECT_EQ(library.GetPattern("P0").Nodes().size(), 1u); + } else { + try { + const PatternLibrary library(text); + FAIL() << "Accepted depth " << depth; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("dependency nesting limit of 64"), std::string::npos) + << error.what(); + } + } + } + } + } +} + +TEST(PatternRewriterTests, ComposedTensorAnnotationsCannotBecomeAttributes) { + const std::string definitions = R"( + pattern Pass(a:float) { return (y=a) } + pattern Alias=Pass + pattern Wrapped(a) { p=Alias(a);return (y=p.y) } + pattern Attr(x,alpha:float) { y=LeakyRelu(x,alpha=alpha);return y } + )"; + for (const char* call : {"Pass", "Alias", "Wrapped"}) { + for (const char* body : { + "p=CALL(a);return LeakyRelu(x,alpha=a)", + "p=CALL(a);q=Attr(x,p.y);return q.y", + "q=Attr(x,a);p=CALL(a);return q.y"}) { + std::string text = body; + text.replace(text.find("CALL"), 4, call); + for (bool indexed : {false, true}) { + SCOPED_TRACE(text); + try { + const PatternLibrary library(definitions + "pattern Entry(x,a){" + text + "}"); + Rule rule{"Invalid", library.GetPattern("Entry"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + const PatternRewriter rewriter({rule}, {.match_only = true, .indexed = indexed}); + FAIL() << "Accepted tensor-to-attribute conversion"; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("Composed tensor annotation"), std::string::npos) + << error.what(); + } + } + } + } + const PatternLibrary library(definitions + R"( + pattern Entry(x,a:float) { p=Attr(x,a);return p.y } + pattern Nested(x,a:float) { p=Entry(x,a);return p.y } + )"); + for (const char* entry : {"Entry", "Nested"}) { + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + b.AddNode("LeakyRelu", {x}, {b.MakeOutput()}).AddAttribute("alpha", 0.25f); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_FLOAT_EQ(match.attributeOf("a").f(), 0.25f); + const auto child = match.instance("p"); + EXPECT_FLOAT_EQ(child.attributeOf(std::string_view(entry) == "Entry" ? "alpha" : "a").f(), 0.25f); + return true; + }; + Rule rule{"Valid", library.GetPattern(entry), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(checks, 1); + } + } +} + +TEST(PatternRewriterTests, RepeatedDimensionsRequireKnownEqualEvidence) { + using Dim = std::variant; + struct Case { + Dim first, second; + bool equal; + }; + const Case cases[] = { + {std::string{}, std::string{}, false}, {std::string{}, int64_t{2}, false}, {int64_t{2}, std::string{}, false}, {int64_t{2}, int64_t{2}, true}, {int64_t{2}, int64_t{3}, false}, {std::string("batch"), std::string("batch"), true}, {std::string("batch"), std::string("other"), false}, {std::string("batch"), int64_t{2}, false}, {int64_t{2}, std::string("batch"), false}, {std::string{}, std::string("batch"), false}, {std::string("batch"), std::string{}, false}}; + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool same_tensor : {false, true}) { + for (const auto& c : cases) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeSymbolicInput({c.first, same_tensor ? c.second : Dim{std::string{}}}); + auto* y = same_tensor ? nullptr : b.MakeSymbolicInput({c.second, std::string{}}); + if (same_tensor) + b.AddNode("Identity", {x}, {b.MakeOutput()}); + else + b.AddNode("CastLike", {x, y}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern(same_tensor ? "(x:float[M,M])=>return Identity(x)" + : "(x:float[M,K],y:float[M,N])=>return CastLike(x,y)"); + rule.replacement = Replacement("(x)=>Identity(x)"); + int checks = 0; + rule.check_and_bind = [&](const Graph&, Match& match) { + ++checks; + if (const auto* integer = std::get_if(&c.first)) + EXPECT_EQ(std::get(match.dimensionOf("M")), *integer); + else + EXPECT_EQ(std::get(match.dimensionOf("M")), std::get(c.first)); + if (!same_tensor) { + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("K"))); + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("N"))); + } + return true; + }; + InlinedVector failures; + Options options{.match_only = match_only, .indexed = indexed, .max_iterations = 1}; + options.on_failure = [&](const MatchFailure& failure) { + if (failure.stage == FailureStage::TypeShape) failures.emplace_back(failure.reason); + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, options).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, c.equal ? 1u : 0u); + EXPECT_EQ(stats.rewrites, c.equal && !match_only ? 1u : 0u); + EXPECT_EQ(checks, c.equal ? 1 : 0); + if (!c.equal) { + ASSERT_FALSE(failures.empty()); + EXPECT_NE(failures.front().find("dimension 'M'"), std::string::npos); + const bool unknown = c.first == Dim{std::string{}} || c.second == Dim{std::string{}}; + EXPECT_NE(failures.front().find(unknown ? "Unknown occurrence" : "Conflicting bindings"), std::string::npos); + } + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } + } +} + +TEST(PatternRewriterTests, SingletonAndSelectedAlternativeDimensionsAllowUnknowns) { + for (bool indexed : {false, true}) { + for (const char* annotation : { + "float[M,N]", "float[_,_]", "float[M,_] | float[_,M]", + "float[M,M] | float[M,_]", "float[M,M] | float[_,N]"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeSymbolicInput({std::string{}, std::string{}}); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:" + std::string(annotation) + ")=>return Identity(x)"); + int checks = 0; + rule.check_and_bind = [&](const Graph&, Match& match) { + ++checks; + if (std::string_view(annotation).find('M') != std::string::npos) + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("M"))); + if (std::string_view(annotation).find('N') != std::string::npos) + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("N"))); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(checks, 1); + } + // The first alternative tentatively binds M, but fails unknown equality. + // The successful alternative must neither inherit that binding nor its occurrence count. + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeSymbolicInput({int64_t{2}, std::string{}}); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto rule = IdentityPair(); + rule.pattern = Pattern("(x:float[M,M] | float[_,M])=>return Identity(x)"); + rule.check_and_bind = [](const Graph&, Match& match) { + EXPECT_TRUE(std::holds_alternative(match.dimensionOf("M"))); + return true; + }; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + TestGraph pair; + auto& pair_builder = pair.builder; + auto* unknown = pair_builder.MakeSymbolicInput({std::string{}}); + auto* known = pair_builder.MakeInput({2}, -1.0f, 1.0f); + pair_builder.AddNode("CastLike", {unknown, known}, {pair_builder.MakeOutput()}); + pair_builder.SetGraphOutputs(); + ASSERT_STATUS_OK(pair.graph.Resolve()); + rule.pattern = Pattern("(x:float[M] | float[_],y:float[M])=>return CastLike(x,y)"); + rule.check_and_bind = [](const Graph&, Match& match) { + EXPECT_EQ(std::get(match.dimensionOf("M")), 2); + return true; + }; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(pair.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + } +} + +TEST(PatternRewriterTests, ComposedDimensionOccurrencesAreHygienicAndConjunctive) { + const PatternLibrary library(R"( + interface Capture(a:float[M]) -> (y) + interface Equality(a:float[M]) -> (y:float[M]) + pattern Pass implements Capture { return (y=a) } + pattern Equal implements Equality { return (y=a) } + pattern Alias=Pass + pattern Wrapped(a) { p=Alias(a);return (y=p.y) } + pattern Single(x) { p=Wrapped(x);return Identity(p.y) } + pattern Independent(x,y) { p=Wrapped(x);q=Wrapped(y);return CastLike(p.y,q.y) } + pattern Repeated(x) { p=Equal(x);return Identity(p.y) } + )"); + for (bool indexed : {false, true}) { + for (const char* entry : {"Single", "Independent", "Repeated"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeSymbolicInput({std::string{}}); + if (std::string_view(entry) == "Independent") { + auto* y = b.MakeSymbolicInput({std::string{}}); + b.AddNode("CastLike", {x, y}, {b.MakeOutput()}); + } else { + b.AddNode("Identity", {x}, {b.MakeOutput()}); + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + auto view = match.instance("p").instance("p"); + EXPECT_TRUE(std::holds_alternative(view.dimensionOf("M"))); + if (std::string_view(entry) == "Independent") + EXPECT_TRUE(std::holds_alternative(match.instance("q").instance("p").dimensionOf("M"))); + return true; + }; + Rule rule{"Composition", library.GetPattern(entry), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + const bool expected = std::string_view(entry) != "Repeated"; + EXPECT_EQ(stats.matches, expected ? 1u : 0u); + EXPECT_EQ(checks, expected ? 1 : 0); + } + } +} + +TEST(PatternRewriterTests, DimensionEqualityRejectionRetriesGraphAlternatives) { + const PatternLibrary library(R"( + interface I(x) -> (y) + pattern Strict implements I { y:float[M,M]=Identity(x);return y } + pattern Loose implements I { y:float[M,N]=Identity(x);return y } + pattern Wild implements I { y:float[_,_]=Identity(x);return y } + pattern Forms=Strict | Loose | Wild + pattern Entry(x) { p=Forms(x);return p.y } + )"); + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeSymbolicInput({std::string{}, std::string{}}); + auto* output = b.MakeOutput(); + b.AddNode("Identity", {x}, {output}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + // Identity shape inference may supply symbolic output dimensions. The + // candidate metadata here deliberately has two anonymous dimensions. + output->SetShape(*x->Shape()); + InlinedVector checks; + auto check = [&](const Graph&, Match& match) { + const auto view = match.instance("p"); + checks.emplace_back(view.alternative()); + if (view.alternative() == "Loose") { + EXPECT_TRUE(std::holds_alternative(view.dimensionOf("M"))); + EXPECT_TRUE(std::holds_alternative(view.dimensionOf("N"))); + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attribute.set_i(1); + match.update("only_loose", attribute); + return match.Reject("try wild"); + } + EXPECT_EQ(view.alternative(), "Wild"); + EXPECT_THROW(match.attributeOf("only_loose"), OnnxRuntimeException); + return true; + }; + Rule rule{"Forms", library.GetPattern("Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + ASSERT_EQ(checks.size(), 2u); + EXPECT_EQ(checks[0], "Loose"); + EXPECT_EQ(checks[1], "Wild"); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } +} + +TEST(PatternRewriterTests, PositionalSignaturesCallsAndOptionalScaling) { + const PatternLibrary library(R"( + pattern plain(g:float[M,N],u:float[M,N]) -> (product:float[M,N]) { + s=Sigmoid(g);a=Mul(g,s);product=Mul(a,u);return product + } + pattern scaled(g:float[M,N],u:float[M,N],alpha:float[...]) -> (product:float[M,N]) { + v=Mul(g,alpha);s=Sigmoid(v);a=Mul(g,s);product=Mul(a,u);return product + } + pattern explicit(g,u,alpha:float[...] | None) { + product=plain(g,u) | scaled(g,u,alpha);return product + } + pattern wildcard(g,u) { + product=plain(g,u) | scaled(g,u,_);return product + } + pattern factored(g,u,alpha:float[...] | None) { + v=(Mul(g,alpha) | g);s=Sigmoid(v);a=Mul(g,s);product=Mul(a,u);return product + } + )"); + for (const char* entry : {"explicit", "wildcard", "factored"}) { + for (bool indexed : {false, true}) { + for (bool scaled : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* g = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* u = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* alpha = scaled ? b.MakeScalarInitializer(0.5f) : nullptr; + auto* v = g; + if (scaled) { + v = b.MakeIntermediate(); + b.AddNode("Mul", {g, alpha}, {v}); + } + auto* s = b.MakeIntermediate(); + auto* a = b.MakeIntermediate(); + auto* output = b.MakeOutput(); + b.AddNode("Sigmoid", {v}, {s}); + b.AddNode("Mul", {g, s}, {a}); + b.AddNode("Mul", {a, u}, {output}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_EQ(match.valueOf("product"), output); + if (std::string_view(entry) != "wildcard") { + EXPECT_EQ(match.optionalValueOf("alpha"), alpha); + EXPECT_EQ(match.valueOf("alpha"), alpha); + } else { + EXPECT_THROW(match.optionalValueOf("alpha"), OnnxRuntimeException); + EXPECT_THROW(match.optionalValueOf("_"), OnnxRuntimeException); + } + EXPECT_THROW(match.optionalValueOf("unknown"), OnnxRuntimeException); + if (std::string_view(entry) != "factored") { + auto view = match.instance("product"); + EXPECT_EQ(view.alternative(), scaled ? "scaled" : "plain"); + EXPECT_EQ(view.alternativeView().valueOf("product"), output); + if (scaled) EXPECT_EQ(view.valueOf("alpha"), alpha); + } else { + EXPECT_EQ(match.valueOf("v"), v); + EXPECT_EQ(match.instance("v").alternative(), scaled ? "Mul" : "g"); + } + return true; + }; + Rule rule{"Scaled", library.GetPattern(entry), "(g,u)=>Mul(g,u)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(checks, 1); + EXPECT_EQ(stats.matches, 1u); + } + } + } +} + +TEST(PatternRewriterTests, PositionalMultiOutputAndLegacyInstanceNamespaces) { + const PatternLibrary library(R"( + pattern Pair(x:float[4]) -> (first:float[2],second:float[2]) { + left,right=Split(x,axis=0);return right,left + } + pattern Alias=Pair + pattern Entry(x) { a,b=Alias(x);return Sub(a,b) } + pattern Legacy(x) { p=Pair(x);return Sub(p.first,p.second) } + )"); + for (const char* entry : {"Entry", "Legacy"}) { + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* left = b.MakeIntermediate(); + auto* right = b.MakeIntermediate(); + b.AddNode("Split", {x}, {left, right}).AddAttribute("axis", int64_t{0}); + b.AddNode("Sub", {right, left}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + if (std::string_view(entry) == "Entry") { + EXPECT_EQ(match.valueOf("a"), right); + EXPECT_EQ(match.valueOf("b"), left); + EXPECT_EQ(match.instance("a").valueOf("first"), right); + EXPECT_THROW(match.instance("b"), OnnxRuntimeException); + } else { + EXPECT_THROW(match.valueOf("p"), OnnxRuntimeException); + EXPECT_THROW(match.optionalValueOf("p"), OnnxRuntimeException); + EXPECT_EQ(match.instance("p").valueOf("first"), right); + } + return true; + }; + Rule rule{"Pair", library.GetPattern(entry), "(x)=>{a,b=Split(x,axis=0);return Sub(b,a)}", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}).Run(test.graph, stats)); + EXPECT_EQ(checks, 1); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } +} + +TEST(PatternRewriterTests, ExpressionIdentityPriorityAndBindingIsolation) { + const PatternLibrary library(R"( + pattern Long(x,a) -> (y) { y=Mul(x,a);return y } + pattern Short(x) -> (y) { return x } + pattern Entry(x,a:float[...] | None) { + v=Long(x,a) | Short(x);return Add(v,x) + } + pattern PreferShort(x,a:float[...] | None) { + v=Short(x) | Long(x,a);return Add(v,x) + } + pattern Direct(x) { v=x;return Add(v,x) } + )"); + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool prefer_short : {false, true}) { + for (bool reject_long : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* alpha = b.MakeScalarInitializer(0.5f); + auto* mul = b.MakeIntermediate(); + b.AddNode("Mul", {x, alpha}, {mul}); + b.AddNode("Add", {mul, x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector calls; + auto check = [&](const Graph&, Match& match) { + const std::string selected(match.instance("v").alternative()); + calls.push_back(selected); + if (selected == "Long") { + EXPECT_EQ(match.optionalValueOf("a"), alpha); + if (reject_long) return match.Reject("try identity"); + } + return true; + }; + Rule rule{"Priority", library.GetPattern(prefer_short ? "PreferShort" : "Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + // Short aliases v and x, so Add(v,x) cannot match this graph's distinct operands. + ASSERT_EQ(calls.size(), 1u); + EXPECT_EQ(calls.front(), "Long"); + EXPECT_EQ(stats.matches, reject_long ? 0u : 1u); + } + } + } + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* other = b.MakeInput({2}, -1.0f, 1.0f); + b.AddNode("Add", {x, other}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"Alias", library.GetPattern("Direct"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 0u); + } +} + +TEST(PatternRewriterTests, ExpressionAlternativesRetryOverlappingLongAndIdentity) { + const PatternLibrary library(R"( + pattern Scale(x,a) -> (y) { y=Mul(x,a);return y } + pattern Pass(x) -> (y) { return x } + pattern Entry(x,a:float[...] | None) { y=Scale(x,a) | Pass(x);return Sigmoid(y) } + pattern ShortFirst(x,a:float[...] | None) { y=Pass(x) | Scale(x,a);return Sigmoid(y) } + )"); + for (bool indexed : {false, true}) { + for (bool short_first : {false, true}) { + for (bool reject_first : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* alpha = b.MakeScalarInitializer(0.5f); + auto* mul = b.MakeIntermediate(); + b.AddNode("Mul", {x, alpha}, {mul}); + auto* output = b.MakeOutput(); + b.AddNode("Sigmoid", {mul}, {output}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector calls; + auto check = [&](const Graph&, Match& match) { + const auto selected = match.instance("y").alternative(); + calls.emplace_back(selected); + if (calls.size() == 1 && reject_first) { + ONNX_NAMESPACE::AttributeProto attr; + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attr.set_i(1); + match.update("only_first", attr); + return match.Reject("try next"); + } + EXPECT_THROW(match.attributeOf("only_first"), OnnxRuntimeException); + EXPECT_EQ(match.valueOf("y"), mul); + EXPECT_EQ(match.valueOf("x"), selected == "Scale" ? x : mul); + EXPECT_EQ(match.optionalValueOf("a"), selected == "Scale" ? alpha : nullptr); + EXPECT_EQ(match.instance("y").alternativeView().valueOf("y"), mul); + return true; + }; + Rule rule{"Overlap", library.GetPattern(short_first ? "ShortFirst" : "Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed, .max_iterations = 1}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + ASSERT_EQ(calls.size(), reject_first ? 2u : 1u); + EXPECT_EQ(calls.front(), short_first ? "Pass" : "Scale"); + if (reject_first) EXPECT_EQ(calls.back(), short_first ? "Scale" : "Pass"); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } +} + +TEST(PatternRewriterTests, OptionalBranchCaptureElsewhereAndRequiredSlotSafety) { + const PatternLibrary library(R"( + pattern Entry(x,alpha:float[...] | None) { + v=(Mul(x,alpha) | x);s=Sigmoid(v);return Add(s,alpha) + } + pattern Optional(x,alpha:float[...] | None) { y=Identity(x);return y } + pattern Elsewhere(x,alpha:float[...] | None) { + v=(Mul(x,alpha) | x);s=Sigmoid(v);return CastLike(s,alpha) + } + pattern RequiredResult(a:float | None) -> (y) { return a } + pattern RequiredAtCall(x,a:float | None) { v=RequiredResult(a);return Clip(x,v) } + )"); + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* alpha = b.MakeScalarInitializer(0.25f); + auto* s = b.MakeIntermediate(); + b.AddNode("Sigmoid", {x}, {s}); + b.AddNode("Add", {s, alpha}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_EQ(match.optionalValueOf("alpha"), alpha); + return true; + }; + Rule rule{"BoundElsewhere", library.GetPattern("Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(checks, 1); + TestGraph missing; + auto& mb = missing.builder; + auto* input = mb.MakeInput({2}, -1.0f, 1.0f); + mb.AddNode("Identity", {input}, {mb.MakeOutput()}); + mb.SetGraphOutputs(); + ASSERT_STATUS_OK(missing.graph.Resolve()); + Rule invalid{"AbsentReplacement", library.GetPattern("Optional"), "(alpha)=>Identity(alpha)", {}, {{kOnnxDomain, 17, 17}}}; + const auto status = PatternRewriter({invalid}, {.indexed = indexed}).Run(missing.graph, stats); + EXPECT_FALSE(status.IsOK()); + EXPECT_EQ(stats.rewrites, 0u); + EXPECT_EQ(missing.graph.NumberOfNodes(), 1); + ASSERT_STATUS_OK(missing.graph.Resolve()); + for (bool present : {false, true}) { + TestGraph clip; + auto& cb = clip.builder; + auto* clip_x = cb.MakeInput({2}, -1.0f, 1.0f); + if (present) + cb.AddNode("Clip", {clip_x, cb.MakeScalarInitializer(0.0f)}, {cb.MakeOutput()}); + else + cb.AddNode("Clip", {clip_x}, {cb.MakeOutput()}); + cb.SetGraphOutputs(); + ASSERT_STATUS_OK(clip.graph.Resolve()); + Rule required_result{"RequiredResult", library.GetPattern("RequiredAtCall"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({required_result}, {.match_only = true, .indexed = indexed}).Run(clip.graph, stats)); + EXPECT_EQ(stats.matches, present ? 1u : 0u); + } + for (bool floating : {false, true}) { + TestGraph elsewhere; + auto& eb = elsewhere.builder; + auto* elsewhere_x = eb.MakeInput({2}, -1.0f, 1.0f); + auto* a = floating ? eb.MakeScalarInitializer(0.25f) : eb.MakeScalarInitializer(1); + auto* sigmoid = eb.MakeIntermediate(); + eb.AddNode("Sigmoid", {elsewhere_x}, {sigmoid}); + eb.AddNode("CastLike", {sigmoid, a}, {eb.MakeOutput()}); + eb.SetGraphOutputs(); + ASSERT_STATUS_OK(elsewhere.graph.Resolve()); + Rule typed{"OptionalType", library.GetPattern("Elsewhere"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({typed}, {.match_only = true, .indexed = indexed}).Run(elsewhere.graph, stats)); + EXPECT_EQ(stats.matches, floating ? 1u : 0u); + } + } +} + +TEST(PatternRewriterTests, PositionalExpressionValidationAndLimits) { + const std::string definitions = R"( + pattern One(x:float[2]) -> (y:float[2]) { y=Identity(x);return y } + pattern Different(x:float[2],a:float[...]) -> (z:int64[2]) { z=Cast(x,to=7);return z } + pattern Two(x) -> (a,b) { a,b=Split(x,axis=0);return a,b } + pattern Alias(x) { return x } + pattern Attribute(x,a:float) { return LeakyRelu(x,alpha=a) } + )"; + for (const char* entry : { + "pattern E(x) { a,b=One(x);return Add(a,b) }", + "pattern E(x) { a,b,c=Two(x);return Add(a,b) }", + "pattern E(x) { a,a=Two(x);return a }", + "pattern E(x) { a,x=Two(x);return a }", + "pattern E(x) { p=Two(x);return Identity(p) }", + "pattern E(x) { p=Two(x);p=Identity(x);return p }", + "pattern E(x) { p=Two(x);v=p;return Identity(v) }", + "pattern E(x) { a,b=Two(x);b=One(x);return Add(a,b) }", + "pattern E(x) { p=Two(x);return LeakyRelu(x,alpha=p) }", + "pattern E(x) { p:float[2]=Two(x);return Add(p.a,p.b) }", + "pattern E(x,a) { y=One(x) | Different(x,a);return y }", + "pattern E(x) { y=One(x) | Two(x);return Identity(y) }", + "pattern E(x,a,b) { alpha=Add(a,b);v=(Mul(x,alpha)|x);return Sigmoid(v) }", + "pattern E(x,a) { v=(Mul(x,a)|x);return Sigmoid(v) }", + "pattern E(x) { y=One(x,x);return y }", + "pattern E(x) { y=Attribute(x,_);return y }", + "pattern E(x) { y=Alias(missing);return Identity(y) }", + "pattern E(x) { y=x | missing;return Identity(y) }", + "pattern E(x) -> (y:float[2] | None) { y=Identity(x);return y }"}) { + SCOPED_TRACE(entry); + EXPECT_THROW((void)PatternLibrary(definitions + entry), OnnxRuntimeException); + } + EXPECT_NO_THROW((void)PatternLibrary(R"( + pattern E(x,a,b) { + alpha=Add(a,b);v=(Mul(x,alpha)|x);s=Sigmoid(v);return Add(s,alpha) + } + )")); + std::string text = "pattern E(x) {"; + for (int i = 0; i < 7; ++i) { + const auto input = i == 0 ? std::string("x") : "v" + std::to_string(i - 1); + text += "v" + std::to_string(i) + "=(Identity(" + input + ") | " + input + ");"; + } + text += "return Sigmoid(v6)}"; + EXPECT_THROW((void)PatternLibrary(text), OnnxRuntimeException); +} + +TEST(PatternRewriterTests, WildcardActualsAreIndependentRealConstrainedValues) { + const PatternLibrary library(R"( + pattern Double(x:float[2]) -> (y) { y=Add(x,x);return y } + pattern Independent() { a=Double(_);b=Double(_);return Mul(a,b) } + pattern Constant(x,a:const float) -> (y) { y=Mul(x,a);return y } + pattern WithConstant(x) { y=Constant(x,_);return y } + pattern ClipMin(x,a:float | None) -> (y) { y=Clip(x,a);return y } + pattern WithMinimum(x) { y=ClipMin(x,_);return y } + )"); + for (bool indexed : {false, true}) { + for (bool same : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* z = b.MakeInput({2}, -1.0f, 1.0f); + auto* a = b.MakeIntermediate(); + auto* c = b.MakeIntermediate(); + b.AddNode("Add", {x, same ? x : z}, {a}); + b.AddNode("Add", {z, z}, {c}); + b.AddNode("Mul", {a, c}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_EQ(match.instance("a").valueOf("x"), x); + EXPECT_EQ(match.instance("b").valueOf("x"), z); + EXPECT_THROW(match.valueOf("_"), OnnxRuntimeException); + return true; + }; + Rule rule{"Independent", library.GetPattern("Independent"), "()=>Identity(0.0)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(checks, same ? 1 : 0); + } + for (int scenario = 0; scenario < 4; ++scenario) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* alpha = scenario == 0 ? b.MakeScalarInitializer(0.5f) + : scenario == 1 ? b.MakeInput({}, 0.5f, 0.5f) + : scenario == 2 ? b.MakeInitializer({1}, {0.5f}) + : b.MakeScalarInitializer(1); + if (scenario == 3) { + auto* integer = b.MakeInput({2}, 1, 2); + b.AddNode("Mul", {integer, alpha}, {b.MakeOutput()}); + } else { + b.AddNode("Mul", {x, alpha}, {b.MakeOutput()}); + } + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"Constant", library.GetPattern("WithConstant"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, scenario == 0 ? 1u : 0u); + } + for (bool present : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + std::vector inputs{x}; + if (present) inputs.push_back(b.MakeScalarInitializer(0.0f)); + b.AddNode("Clip", inputs, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"RealWildcard", library.GetPattern("WithMinimum"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, present ? 1u : 0u); + } + } +} + +TEST(PatternRewriterTests, NestedMultiOutputExpressionsPreservePositionalBindings) { + const PatternLibrary library(R"( + pattern Pair(x:float[4]) -> (first:float[2],second:float[2]) { + a,b=Split(x,axis=0);return b,a + } + pattern Added(x:float[4],alpha:float) -> (p:float[2],q:float[2]) { + z=Add(x,alpha);a,b=Pair(z);return a,b + } + pattern Entry(x,alpha:float | None) { + a,b=(Pair(x) | Added(x,alpha));return Sub(a,b) + } + )"); + for (bool indexed : {false, true}) { + for (bool added : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({4}, -1.0f, 1.0f); + auto* alpha = added ? b.MakeScalarInitializer(0.25f) : nullptr; + auto* z = x; + if (added) { + z = b.MakeIntermediate(); + b.AddNode("Add", {x, alpha}, {z}); + } + auto* left = b.MakeIntermediate(); + auto* right = b.MakeIntermediate(); + b.AddNode("Split", {z}, {left, right}).AddAttribute("axis", int64_t{0}); + b.AddNode("Sub", {right, left}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector attempts; + auto check = [&](const Graph&, Match& match) { + auto view = match.instance("a"); + attempts.emplace_back(view.alternative()); + if (added && view.alternative() == "Pair") return match.Reject("Prefer the included Add."); + EXPECT_EQ(match.valueOf("a"), right); + EXPECT_EQ(match.valueOf("b"), left); + EXPECT_EQ(match.optionalValueOf("alpha"), alpha); + if (added) { + auto nested = view.instance("a"); + EXPECT_EQ(nested.valueOf("first"), right); + EXPECT_EQ(nested.valueOf("second"), left); + } + return true; + }; + Rule rule{"NestedPair", library.GetPattern("Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(attempts, added ? InlinedVector({"Pair", "Added"}) + : InlinedVector({"Pair"})); + } + } +} + +TEST(PatternRewriterTests, PositionalSelectedDimensionsAndTensorAliasContracts) { + const PatternLibrary library(R"( + pattern Repeated(x:float[M,M]) -> (y) { return x } + pattern Single(x:float[M,N]) -> (y) { return x } + pattern Entry(x) { y=Repeated(x) | Single(x);return Identity(y) } + pattern Signature(x:float[M]) -> (y:float[M]) { y=Identity(x);return y } + )"); + for (bool indexed : {false, true}) { + for (bool known : {false, true}) { + TestGraph test; + auto& b = test.builder; + using Dim = std::variant; + const Dim dim = known ? Dim{int64_t{2}} : Dim{std::string{}}; + auto* x = b.MakeSymbolicInput({dim, dim}); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector attempts; + auto check = [&](const Graph&, Match& match) { + const auto view = match.instance("y"); + attempts.emplace_back(view.alternative()); + if (known) + EXPECT_EQ(std::get(view.dimensionOf("M")), 2); + else + EXPECT_TRUE(std::holds_alternative(view.dimensionOf("M"))); + return true; + }; + Rule rule{"SelectedDimensions", library.GetPattern("Entry"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(attempts, known ? InlinedVector({"Repeated", "Single"}) + : InlinedVector({"Single"})); + TestGraph signature; + auto& sb = signature.builder; + auto* input = sb.MakeSymbolicInput({dim}); + sb.AddNode("Identity", {input}, {sb.MakeOutput()}); + sb.SetGraphOutputs(); + ASSERT_STATUS_OK(signature.graph.Resolve()); + Rule signed_rule{"Signature", library.GetPattern("Signature"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + ASSERT_STATUS_OK(PatternRewriter({signed_rule}, {.match_only = true, .indexed = indexed}).Run(signature.graph, stats)); + EXPECT_EQ(stats.matches, known ? 1u : 0u); + } + } + for (const char* source : { + "pattern P(x,a:float){v=a;return LeakyRelu(x,alpha=v)}", + "pattern P(x,a:float){y=LeakyRelu(x,alpha=a);v=a;return Identity(y)}", + "pattern F(a)->(y){return a} pattern P(x,a){p=F(a);return LeakyRelu(x,alpha=a)}"}) { + SCOPED_TRACE(source); + EXPECT_THROW((void)PatternLibrary(source), OnnxRuntimeException); + } + EXPECT_NO_THROW((void)PatternLibrary(R"( + pattern P(x) { alias=x;return Add(alias,x) } + pattern Strings(x) { y=test::Op(x,label=")|(");return y } + )")); +} + +TEST(PatternRewriterTests, ExpressionPlanAndNestingBounds) { + const PatternLibrary grouped(R"(pattern E(x){y=(Identity(x)|x)|x;return Sigmoid(y)})"); + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* identity = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {identity}); + b.AddNode("Sigmoid", {identity}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + InlinedVector selections; + auto check = [&](const Graph&, Match& match) { + auto view = match.instance("y"); + if (view.alternative() == "") view = view.alternativeView(); + selections.emplace_back(view.alternative()); + EXPECT_THROW(view.alternativeView(), OnnxRuntimeException); + EXPECT_EQ(match.valueOf("y"), identity); + return true; + }; + Rule rule{"Grouped", grouped.GetPattern("E"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(selections, (InlinedVector{"Identity", "x", "x"})); + EXPECT_EQ(stats.matches, 3u); + } + for (int count : {6, 7}) { + std::string source = "pattern E(x){"; + for (int i = 0; i < count; ++i) { + const auto input = i == 0 ? std::string("x") : "v" + std::to_string(i - 1); + source += "v" + std::to_string(i) + "=(Identity(" + input + ") | " + input + ");"; + } + source += "return Sigmoid(v" + std::to_string(count - 1) + ")}"; + if (count == 6) { + const PatternLibrary library(source); + for (bool indexed : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + auto* value = x; + for (int i = 0; i < count; ++i) { + auto* output = b.MakeIntermediate(); + b.AddNode("Identity", {value}, {output}); + value = output; + } + b.AddNode("Sigmoid", {value}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + Rule rule{"AllPlans", library.GetPattern("E"), "(x)=>Identity(x)", {}, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 64u); + } + } else { + EXPECT_THROW((void)PatternLibrary(source), OnnxRuntimeException); + } + } + std::string malformed = "pattern E(x){y=(missing)"; + for (int i = 0; i < 1000; ++i) malformed += "()"; + malformed += ";return Identity(y)}"; + EXPECT_THROW((void)PatternLibrary(malformed), OnnxRuntimeException); + for (int nesting : {63, 64, 1000}) { + const std::string source = "pattern E(x){y=" + std::string(nesting, '(') + + "Identity(x)" + std::string(nesting, ')') + ";return y}"; + if (nesting == 63) + EXPECT_NO_THROW((void)PatternLibrary(source)); + else + EXPECT_THROW((void)PatternLibrary(source), OnnxRuntimeException); + } + for (int count : {32, 33}) { + std::string source; + for (int i = 0; i < count; ++i) { + source += "pattern F" + std::to_string(i) + "(x){y=("; + source += i + 1 == count ? "Identity(x)" : "F" + std::to_string(i + 1) + "(x)"; + source += ");return y}"; + } + if (count == 32) + EXPECT_NO_THROW((void)PatternLibrary(source)); + else + EXPECT_THROW((void)PatternLibrary(source), OnnxRuntimeException); + } +} + +TEST(PatternRewriterTests, InferredExpressionOutputsAllowUnconstrainedBranches) { + for (bool alias_first : {false, true}) { + const std::string text = std::string(R"( + pattern Fixed(x:float[2]) { y=)") + + (alias_first ? "x | Identity(x)" : "Identity(x) | x") + R"(;return y } + pattern Symbolic(x:float[M]) { y=)" + + (alias_first ? "x | Identity(x)" : "Identity(x) | x") + R"(;return y } + pattern FixedEntry(x) { y=Fixed(x);return Sigmoid(y) } + pattern Nested(x) { y=Symbolic(x);return y } + pattern SymbolicEntry(x) { y=Nested(x);return Sigmoid(y) } + )"; + const PatternLibrary library(text); + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (const char* entry : {"FixedEntry", "SymbolicEntry"}) { + TestGraph test; + auto& b = test.builder; + auto* x = std::string_view(entry) == "FixedEntry" + ? b.MakeInput({2}, -1.0f, 1.0f) + : b.MakeSymbolicInput({std::string{}}); + auto* identity = b.MakeIntermediate(); + b.AddNode("Identity", {x}, {identity}); + b.AddNode("Sigmoid", {identity}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + EXPECT_EQ(match.valueOf("y"), identity); + if (std::string_view(entry) == "SymbolicEntry") + EXPECT_TRUE(std::holds_alternative( + match.instance("y").instance("y").dimensionOf("M"))); + return true; + }; + Rule rule{"Inferred", library.GetPattern(entry), "(x)=>Sigmoid(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + EXPECT_EQ(checks, match_only ? 2 : 1); + EXPECT_EQ(stats.matches, match_only ? 2u : 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } + } +} + +TEST(PatternRewriterTests, DeclaredOutputsCannotShadowAliasesOrPositionalCallResults) { + const std::string copy = "pattern Copy(x)->(y){return x} "; + for (const char* body : {"a=x;return Identity(a)", "a=Copy(x);return Identity(a)", + "a,b=Pair(x);return Identity(a)"}) { + const std::string text = copy + + "pattern Pair(x)->(first,second){first=x;second=x;return first,second} " + "pattern E(x:float[2])->(a){" + + body + "}"; + try { + const PatternLibrary library(text); + FAIL() << "Accepted an output name that shadows an alias: " << body; + } catch (const OnnxRuntimeException& error) { + EXPECT_NE(std::string(error.what()).find("Output port shadows a different local value"), std::string::npos) + << body << ": " << error.what(); + } + } + for (const char* body : {"a=x;return a", "a=Copy(x);return a"}) { + const PatternLibrary library(copy + "pattern F(x:float[2])->(a){" + body + + "} " + "pattern E(x){a=F(x);return Identity(a)}"); + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2}, -1.0f, 1.0f); + b.AddNode("Identity", {x}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + auto check = [&](const Graph&, Match& match) { + EXPECT_EQ(match.valueOf("a"), x); + EXPECT_EQ(match.instance("a").valueOf("a"), x); + if (std::string_view(body).find("Copy") != std::string_view::npos) + EXPECT_EQ(match.instance("a").instance("a").valueOf("y"), x); + return true; + }; + Rule rule{"SameHandle", library.GetPattern("E"), "(x)=>Identity(x)", check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + EXPECT_EQ(stats.matches, 1u); + } + } + } +} + +TEST(PatternRewriterTests, OrdinaryOutputSignaturesPreserveTypedAttributeInputs) { + const PatternLibrary library(R"( + pattern Int(x,axis:int)->(y){y=Softmax(x,axis=axis);return y} + pattern Ints(x,axes:ints)->(y){y=Transpose(x,perm=axes);return y} + pattern Float(x,alpha:float)->(y){y=LeakyRelu(x,alpha=alpha);return y} + pattern NestedInt(x,axis:int)->(y){y=Int(x,axis);return y} + pattern NestedInts(x,axes:ints)->(y){y=Ints(x,axes);return y} + pattern NestedFloat(x,alpha:float)->(y){y=Float(x,alpha);return y} + )"); + for (bool indexed : {false, true}) { + for (bool match_only : {false, true}) { + for (bool nested : {false, true}) { + for (const char* kind : {"Int", "Ints", "Float"}) { + TestGraph test; + auto& b = test.builder; + auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); + const std::string_view type(kind); + auto& node = b.AddNode(type == "Int" ? "Softmax" : type == "Ints" ? "Transpose" + : "LeakyRelu", + {x}, {b.MakeOutput()}); + if (type == "Int") + node.AddAttribute("axis", int64_t{1}); + else if (type == "Ints") + node.AddAttribute("perm", std::vector{1, 0}); + else + node.AddAttribute("alpha", 0.25f); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + int checks = 0; + auto check = [&](const Graph&, Match& match) { + ++checks; + const auto view = nested ? match.instance("y") : match.bindings(); + if (type == "Int") { + EXPECT_EQ(match.attributeOf("axis").i(), 1); + EXPECT_EQ(view.attributeOf("axis").i(), 1); + } else if (type == "Ints") { + EXPECT_EQ(match.attributeOf("axes").ints_size(), 2); + EXPECT_EQ(view.attributeOf("axes").ints(0), 1); + } else { + EXPECT_FLOAT_EQ(match.attributeOf("alpha").f(), 0.25f); + EXPECT_FLOAT_EQ(view.attributeOf("alpha").f(), 0.25f); + } + return true; + }; + const std::string entry = std::string(nested ? "Nested" : "") + kind; + const char* recipe = type == "Int" ? "(x,axis)=>Softmax(x,axis=axis)" + : type == "Ints" ? "(x,axes)=>Transpose(x,perm=axes)" + : "(x,alpha)=>LeakyRelu(x,alpha=alpha)"; + Rule rule{"Attribute", library.GetPattern(entry), recipe, check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = match_only, .indexed = indexed, .max_iterations = 1}) + .Run(test.graph, stats)); + EXPECT_EQ(checks, 1); + EXPECT_EQ(stats.matches, 1u); + EXPECT_EQ(stats.rewrites, match_only ? 0u : 1u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } + } + } + } + for (const char* text : { + "interface I(x,axis:int)->(y)", + "pattern P(x)->(y:ints){y=Identity(x);return y}", + "pattern P(x,a:int)->(y){y=Softmax(x,axis=a);return y} " + "pattern E(x){y=P(x,_);return y}"}) { + SCOPED_TRACE(text); + EXPECT_THROW((void)PatternLibrary(text), OnnxRuntimeException); + } +} + +TEST(PatternRewriterTests, FactoredOptionalScalingNumericalRewrite) { + const PatternLibrary library(R"( + pattern activation(g:float[M,N],u:float[M,N],alpha:float[...] | None) -> (product:float[M,N]) { + scaled=(Mul(g,alpha) | g);s=Sigmoid(scaled);a=Mul(g,s);product=Mul(a,u);return product + } + pattern Fuse(g,u,alpha:float[...] | None) { + product=activation(g,u,alpha);return Identity(product) + } + )"); + for (bool indexed : {false, true}) { + for (int variant = 0; variant < 3; ++variant) { + TestGraph test; + auto& b = test.builder; + auto* g = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* u = b.MakeInput({2, 3}, -1.0f, 1.0f); + auto* alpha = variant == 0 ? nullptr + : variant == 1 ? b.MakeScalarInitializer(0.5f) + : b.MakeInitializer({1}, {0.5f}); + auto* scaled = g; + if (alpha) { + scaled = b.MakeIntermediate(); + b.AddNode("Mul", {g, alpha}, {scaled}); + } + auto* sigmoid = b.MakeIntermediate(); + auto* activated = b.MakeIntermediate(); + auto* product = b.MakeIntermediate(); + b.AddNode("Sigmoid", {scaled}, {sigmoid}); + b.AddNode("Mul", {g, sigmoid}, {activated}); + b.AddNode("Mul", {activated, u}, {product}); + b.AddNode("Identity", {product}, {b.MakeOutput()}); + b.SetGraphOutputs(); + ASSERT_STATUS_OK(test.graph.Resolve()); + const std::string original = test.model.ToProto().SerializeAsString(); + auto check = [alpha](const Graph&, Match& match) { + EXPECT_EQ(match.optionalValueOf("alpha"), alpha); + EXPECT_EQ(match.instance("product").optionalValueOf("alpha"), alpha); + return true; + }; + const Replacement plain(R"((g,u)=>{s=Sigmoid(g);a=Mul(g,s);return Mul(a,u)})"); + const Replacement with_alpha(R"((g,u,alpha)=>{v=Mul(g,alpha);s=Sigmoid(v);a=Mul(g,s);return Mul(a,u)})"); + auto replacement = [plain, with_alpha](const Graph&, const Match& match) { + return match.optionalValueOf("alpha") ? with_alpha : plain; + }; + Rule rule{"Factored", library.GetPattern("Fuse"), replacement, check, {{kOnnxDomain, 17, 17}}}; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}, {.indexed = indexed, .max_iterations = 1}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 1u); + EXPECT_EQ(Count(test.graph, "Identity"), 0u); + ASSERT_STATUS_OK(test.graph.Resolve()); + const std::string rewritten = test.model.ToProto().SerializeAsString(); + auto execute = [&](const std::string& serialized, std::vector& outputs) { + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + InferenceSessionWrapper session{options, GetEnvironment()}; + std::istringstream stream(serialized); + ORT_RETURN_IF_ERROR(session.Load(stream)); + ORT_RETURN_IF_ERROR(session.Initialize()); + return session.Run(RunOptions{}, b.feeds_, b.output_names_, &outputs); + }; + std::vector before, after; + ASSERT_STATUS_OK(execute(original, before)); + ASSERT_STATUS_OK(execute(rewritten, after)); + ASSERT_EQ(before.size(), 1u); + ASSERT_EQ(after.size(), 1u); + const auto& expected = before[0].Get(); + const auto& actual = after[0].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + for (int64_t i = 0; i < expected.Shape().Size(); ++i) + EXPECT_NEAR(expected.Data()[i], actual.Data()[i], 1e-6f); + } + } +} + +} // namespace test +} // namespace onnxruntime